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 

11 

12from lilbee.app.placement import ( 

13 PlacementView, 

14 get_placement, 

15 preview_placement, 

16 set_placement, 

17) 

18from lilbee.cli import theme 

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

20from lilbee.cli.helpers import json_output 

21from lilbee.core.config import cfg 

22from lilbee.providers.base import ProviderError 

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

24 

25_PLACEMENT_ERRORS = (PlacementError, ProviderError, OSError) 

26 

27placement_app = typer.Typer( 

28 name="placement", 

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

30 no_args_is_help=True, 

31) 

32 

33_GIB = 1024**3 

34 

35 

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

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

38 if spec is None: 

39 return None 

40 if spec == "-": 

41 raw = sys.stdin.read() 

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

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

44 else: 

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

46 return PlacementSpec.from_json(raw) 

47 

48 

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

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

51 try: 

52 view = action() 

53 except _PLACEMENT_ERRORS as exc: 

54 if cfg.json_mode: 

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

56 else: 

57 console.print(f"[{theme.ERROR}]{exc}[/{theme.ERROR}]") 

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

59 if cfg.json_mode: 

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

61 from lilbee.server.models import PlacementResponse 

62 

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

64 else: 

65 _render_view(view) 

66 

67 

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

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

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

71 table = Table(title=title) 

72 table.add_column("GPU") 

73 table.add_column("Name") 

74 table.add_column("Free / Total") 

75 table.add_column("Roles") 

76 

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

78 for role_view in view.roles: 

79 for idx in role_view.devices: 

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

81 

82 for g in view.gpus: 

83 free_gib = g.free_bytes / _GIB 

84 total_gib = g.total_bytes / _GIB 

85 table.add_row( 

86 g.label, 

87 g.name or "(unnamed)", 

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

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

90 ) 

91 console.print(table) 

92 

93 for role_view in view.roles: 

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

95 console.print( 

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

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

98 ) 

99 

100 if view.co_tenants: 

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

102 console.print( 

103 f" [{theme.MUTED}]{names}: share memory, one loaded at a time[/{theme.MUTED}]" 

104 ) 

105 

106 for role in view.unplaceable: 

107 console.print(f" [{theme.ERROR}]{role.value}: does not fit, no server[/{theme.ERROR}]") 

108 

109 for skipped in view.skipped_not_installed: 

110 console.print( 

111 f" [{theme.WARNING}]{skipped.role.value}: {skipped.model} not downloaded, " 

112 f"pull it to place it[/{theme.WARNING}]" 

113 ) 

114 

115 if view.rejected_spec_json: 

116 console.print( 

117 f" [{theme.WARNING}]a saved placement does not fit this hardware and is being " 

118 f"ignored; run 'lilbee placement clear' or set a new one[/{theme.WARNING}]" 

119 ) 

120 

121 

122@placement_app.command("show") 

123def show( 

124 data_dir: Path | None = data_dir_option, 

125 use_global: bool = global_option, 

126) -> None: 

127 """Show the current effective placement.""" 

128 apply_overrides(data_dir=data_dir, use_global=use_global) 

129 _guard(get_placement) 

130 

131 

132@placement_app.command("preview") 

133def preview( 

134 spec: str | None = typer.Option( 

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

136 ), 

137 data_dir: Path | None = data_dir_option, 

138 use_global: bool = global_option, 

139) -> None: 

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

141 apply_overrides(data_dir=data_dir, use_global=use_global) 

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

143 

144 

145@placement_app.command("set") 

146def set_cmd( 

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

148 data_dir: Path | None = data_dir_option, 

149 use_global: bool = global_option, 

150) -> None: 

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

152 apply_overrides(data_dir=data_dir, use_global=use_global) 

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

154 

155 

156@placement_app.command("clear") 

157def clear( 

158 data_dir: Path | None = data_dir_option, 

159 use_global: bool = global_option, 

160) -> None: 

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

162 apply_overrides(data_dir=data_dir, use_global=use_global) 

163 _guard(lambda: set_placement(None))