Coverage for src/lilbee/server/routes/placement.py: 100%

1"""Placement routes: inspect, preview, and (when enabled) apply GPU placement. 

2 

3Every route requires auth: even the reads run resolve_placement_plan, which 

4spawns subprocess device probes. Every route needs the token. Applying or 

5clearing placement restarts the shared fleet's moved roles, which is unsafe 

6across concurrent HTTP clients, so PUT/DELETE are refused by default. They are gated on the 

7``allow_http_placement`` flag (LILBEE_ALLOW_HTTP_PLACEMENT), which an operator 

8turns on for a single-client / owned deployment to get the same apply/clear 

9capability the CLI and TUI have. 

10""" 

11 

12from __future__ import annotations 

13 

14import asyncio 

15import json 

16import logging 

17 

18from litestar import delete, get, post, put 

19from litestar.exceptions import HTTPException 

20from litestar.response import Stream 

21 

22from lilbee.core.config import cfg 

23from lilbee.providers.base import ProviderError 

24from lilbee.providers.fleet.placement_spec import PlacementError 

25from lilbee.server import handlers 

26from lilbee.server.handlers.sse import SSE_MEDIA_TYPE 

27from lilbee.server.models import GpusResponse, PlacementResponse, PlacementSpecBody 

28 

29log = logging.getLogger(__name__) 

30 

31_HTTP_UNPROCESSABLE = 422 

32_HTTP_CONFLICT = 409 

33_HTTP_UNAVAILABLE = 503 

34# OSError is deliberately absent: on this path it means the device probe could 

35# not run (no nvidia-smi on PATH, no permission on the device node, a failed 

36# spawn), which is a host fault and not something the caller's spec can fix. 

37_INPUT_ERRORS = (PlacementError, ValueError) 

38_PROBE_FAILED_DETAIL = ( 

39 "Could not probe the GPUs. Check that the vendor tool (nvidia-smi or " 

40 "rocm-smi) is installed and this process may read the device." 

41) 

42_MISSING_SPEC_DETAIL = "spec is required to apply placement; send {} to DELETE for auto." 

43 

44 

45def _spec_json(body: PlacementSpecBody) -> str | None: 

46 """Serialize the optional spec dict to JSON, or return None when absent.""" 

47 return json.dumps(body.spec) if body.spec is not None else None 

48 

49 

50def _refused() -> HTTPException: 

51 from lilbee.app.placement import placement_refused_message 

52 

53 return HTTPException(status_code=_HTTP_CONFLICT, detail=placement_refused_message()) 

54 

55 

56@get("/api/placement") 

57async def placement_route() -> PlacementResponse: 

58 """Current effective placement.""" 

59 try: 

60 return await handlers.placement() 

61 except ProviderError as exc: 

62 raise HTTPException(status_code=_HTTP_UNAVAILABLE, detail=str(exc)) from exc 

63 

64 

65@post("/api/placement/preview", status_code=200) 

66async def placement_preview_route(data: PlacementSpecBody) -> PlacementResponse: 

67 """Preview a candidate spec (or auto when no spec). Requires auth: runs subprocess probes.""" 

68 try: 

69 return await handlers.placement_preview(_spec_json(data)) 

70 except ProviderError as exc: 

71 raise HTTPException(status_code=_HTTP_UNAVAILABLE, detail=str(exc)) from exc 

72 except OSError as exc: 

73 log.exception("GPU probe failed") 

74 raise HTTPException(status_code=_HTTP_UNAVAILABLE, detail=_PROBE_FAILED_DETAIL) from exc 

75 except _INPUT_ERRORS as exc: 

76 raise HTTPException(status_code=_HTTP_UNPROCESSABLE, detail=str(exc)) from exc 

77 

78 

79@put("/api/placement") 

80async def placement_set_route(data: PlacementSpecBody) -> PlacementResponse: 

81 """Apply a manual spec. Refused unless allow_http_placement is enabled.""" 

82 if not cfg.allow_http_placement: 

83 raise _refused() 

84 spec_json = _spec_json(data) 

85 if spec_json is None: 

86 raise HTTPException(status_code=_HTTP_UNPROCESSABLE, detail=_MISSING_SPEC_DETAIL) 

87 try: 

88 return await handlers.placement_set(spec_json) 

89 except ProviderError as exc: 

90 raise HTTPException(status_code=_HTTP_UNAVAILABLE, detail=str(exc)) from exc 

91 except OSError as exc: 

92 log.exception("GPU probe failed") 

93 raise HTTPException(status_code=_HTTP_UNAVAILABLE, detail=_PROBE_FAILED_DETAIL) from exc 

94 except _INPUT_ERRORS as exc: 

95 raise HTTPException(status_code=_HTTP_UNPROCESSABLE, detail=str(exc)) from exc 

96 

97 

98@delete("/api/placement", status_code=200) 

99async def placement_clear_route() -> PlacementResponse: 

100 """Clear placement (back to auto). Refused unless allow_http_placement is enabled.""" 

101 if not cfg.allow_http_placement: 

102 raise _refused() 

103 try: 

104 return await handlers.placement_clear() 

105 except ProviderError as exc: 

106 raise HTTPException(status_code=_HTTP_UNAVAILABLE, detail=str(exc)) from exc 

107 

108 

109@get("/api/gpus") 

110async def gpus_route() -> GpusResponse: 

111 """Detected GPUs with free/total VRAM, plus the host-level Intel util notice.""" 

112 try: 

113 return await handlers.gpus() 

114 except ProviderError as exc: 

115 raise HTTPException(status_code=_HTTP_UNAVAILABLE, detail=str(exc)) from exc 

116 

117 

118@get("/api/gpus/stream", media_type=SSE_MEDIA_TYPE) 

119async def gpu_stats_stream_route() -> Stream: 

120 """Live per-GPU utilization + free memory as SSE for the placement view.""" 

121 from lilbee.app.placement import get_placement 

122 

123 try: 

124 # get_placement runs resolve_placement_plan, which spawns subprocess 

125 # device probes that can wedge for the probe timeout on a broken driver; 

126 # offload so the probe never blocks the event loop. 

127 view = await asyncio.to_thread(get_placement) 

128 except ProviderError as exc: 

129 raise HTTPException(status_code=_HTTP_UNAVAILABLE, detail=str(exc)) from exc 

130 return Stream(handlers.gpu_stats_stream(view.gpus), media_type=SSE_MEDIA_TYPE)