Coverage for src/lilbee/cli/placement.py: 100%

1"""`lilbee placement` sub-app: inspect, preview, and set GPU placement.""" 

2 

3from __future__ import annotations 

4 

5import sys 

6from collections.abc import Callable 

7from pathlib import Path 

8 

9import typer 

10from rich.table import Table 

11from rich.text import Text 

12 

13from lilbee.app.placement import ( 

14 PlacementView, 

15 get_placement, 

16 preview_placement, 

17 set_placement, 

18) 

19from lilbee.cli import theme 

20from lilbee.cli.app import apply_overrides, console, data_dir_option, global_option 

21from lilbee.cli.helpers import json_output 

22from lilbee.core.config import cfg 

23from lilbee.providers.base import ProviderError 

24from lilbee.providers.fleet.placement_spec import PlacementError, PlacementSpec 

25from lilbee.runtime.console import styled 

26 

27_PLACEMENT_ERRORS = (PlacementError, ProviderError, OSError) 

28 

29placement_app = typer.Typer( 

30 name="placement", 

31 help="Inspect and override multi-GPU model placement.", 

32 no_args_is_help=True, 

33) 

34 

35_GIB = 1024**3 

36 

37 

38def _read_spec(spec: str | None) -> PlacementSpec | None: 

39 """Parse a spec from a file path or stdin ('-'); return None when omitted.""" 

40 if spec is None: 

41 return None 

42 if spec == "-": 

43 raw = sys.stdin.read() 

44 elif spec.lstrip().startswith("{"): 

45 raw = spec # inline JSON rather than a file path 

46 else: 

47 raw = Path(spec).read_text(encoding="utf-8") 

48 return PlacementSpec.from_json(raw) 

49 

50 

51def _guard(action: Callable[[], PlacementView]) -> None: 

52 """Run a placement action and render it, turning known failures into a clean exit.""" 

53 try: 

54 view = action() 

55 except _PLACEMENT_ERRORS as exc: 

56 if cfg.json_mode: 

57 json_output({"error": str(exc)}) 

58 else: 

59 console.print(Text(str(exc), style=theme.ERROR), soft_wrap=True) 

60 raise typer.Exit(code=1) from exc 

61 if cfg.json_mode: 

62 # The same canonical shape the HTTP and MCP surfaces return. 

63 from lilbee.server.models import PlacementResponse 

64 

65 json_output(PlacementResponse.from_view(view).model_dump(mode="json")) 

66 else: 

67 _render_view(view) 

68 

69 

70def _render_view(view: PlacementView) -> None: 

71 """Print a Rich table of GPU rows plus per-role and unplaceable lines.""" 

72 title = "Placement (manual)" if view.manual else "Placement (auto)" 

73 table = Table(title=title) 

74 table.add_column("GPU") 

75 table.add_column("Name") 

76 table.add_column("Free / Total") 

77 table.add_column("Roles") 

78 

79 placed: dict[int, list[str]] = {g.index: [] for g in view.gpus} 

80 for role_view in view.roles: 

81 for idx in role_view.devices: 

82 placed.setdefault(idx, []).append(role_view.role.value) 

83 

84 for g in view.gpus: 

85 free_gib = g.free_bytes / _GIB 

86 total_gib = g.total_bytes / _GIB 

87 table.add_row( 

88 g.label, 

89 g.name or "(unnamed)", 

90 f"{free_gib:.0f} / {total_gib:.0f} GiB", 

91 Text(", ".join(placed.get(g.index, [])) or "-"), 

92 ) 

93 console.print(table) 

94 

95 for role_view in view.roles: 

96 split_info = f" split={list(role_view.tensor_split)}" if role_view.tensor_split else "" 

97 console.print( 

98 Text( 

99 f" {role_view.role.value}: devices={list(role_view.devices)}" 

100 f" replicas={role_view.replicas}{split_info} {role_view.model}" 

101 ), 

102 soft_wrap=True, 

103 ) 

104 

105 if view.co_tenants: 

106 names = ", ".join(role.value for role in view.co_tenants) 

107 console.print(f" {names}: share memory, one loaded at a time", style=theme.MUTED) 

108 

109 for role in view.unplaceable: 

110 console.print(f" {role.value}: does not fit, no server", style=theme.ERROR) 

111 

112 for skipped in view.skipped_not_installed: 

113 console.print( 

114 Text.assemble( 

115 (f" {skipped.role.value}: ", theme.WARNING), 

116 skipped.model, 

117 (" not downloaded, pull it to place it", theme.WARNING), 

118 ), 

119 soft_wrap=True, 

120 ) 

121 

122 if view.rejected_spec_json: 

123 console.print( 

124 styled( 

125 " ", 

126 ( 

127 "a saved placement does not fit this hardware and is being " 

128 "ignored; run 'lilbee placement clear' or set a new one", 

129 theme.WARNING, 

130 ), 

131 ) 

132 ) 

133 

134 

135@placement_app.command("show") 

136def show( 

137 data_dir: Path | None = data_dir_option, 

138 use_global: bool = global_option, 

139) -> None: 

140 """Show the current effective placement.""" 

141 apply_overrides(data_dir=data_dir, use_global=use_global) 

142 _guard(get_placement) 

143 

144 

145@placement_app.command("preview") 

146def preview( 

147 spec: str | None = typer.Option( 

148 None, "--spec", help="Spec JSON file, or - for stdin; omit for auto." 

149 ), 

150 data_dir: Path | None = data_dir_option, 

151 use_global: bool = global_option, 

152) -> None: 

153 """Preview what a spec (or auto) would place, without applying it.""" 

154 apply_overrides(data_dir=data_dir, use_global=use_global) 

155 _guard(lambda: preview_placement(_read_spec(spec))) 

156 

157 

158@placement_app.command("set") 

159def set_cmd( 

160 spec: str = typer.Option(..., "--spec", help="Spec JSON file, or - for stdin."), 

161 data_dir: Path | None = data_dir_option, 

162 use_global: bool = global_option, 

163) -> None: 

164 """Validate, persist, and apply a manual placement spec.""" 

165 apply_overrides(data_dir=data_dir, use_global=use_global) 

166 _guard(lambda: set_placement(_read_spec(spec))) 

167 

168 

169@placement_app.command("clear") 

170def clear( 

171 data_dir: Path | None = data_dir_option, 

172 use_global: bool = global_option, 

173) -> None: 

174 """Clear the manual placement and return to automatic placement.""" 

175 apply_overrides(data_dir=data_dir, use_global=use_global) 

176 _guard(lambda: set_placement(None))