Coverage for apps/lmm_service_manager.py: 55%
165 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Manage the LMM services and their URLs.
3"""
5import sys
6import logging
7import asyncio
9from http import HTTPStatus
11from aiohttp import TCPConnector
12from aiohttp import ClientSession
13from aiohttp import ClientTimeout
14from aiohttp.client_exceptions import ClientConnectorError
16from typing import Optional
17from typing import List
18from typing import Tuple
19from typing import Dict
21from kubernetes_asyncio.config.config_exception import ConfigException
22from kubernetes_asyncio.client.exceptions import ApiException
24# Local relative imports
25sys.path.append("..") # noqa: E402
26sys.path.append("../..") # noqa: E402
28from console_utils import setup_logging
29import k8s_utils
30from k8s_utils import K8sService
31from k8s_utils import K8sContainer
32from k8s_utils import get_k8s_services
33from k8s_utils import ServiceNotFoundError
35from streamwise_apps import STREAMWISE_APPS
36from streamwise_apps import VLLM_SERVICES
39SERVICE_LONG_TIMEOUT = ClientTimeout(
40 connect=10.0,
41 sock_connect=10.0,
42 sock_read=10 * 60.0,
43)
45STATUS_TIMEOUT = ClientTimeout(
46 total=0.5
47)
50class LMMServiceManager:
51 """
52 Manage the LMM services and their URLs.
53 """
55 def __init__(
56 self,
57 app_name: str,
58 k8s_cluster: Optional[str] = None
59 ) -> None:
60 """Initialize the service manager."""
61 self.app_name = app_name
62 self.k8s_cluster = k8s_cluster
63 self.services: Dict[str, K8sService] = {}
64 self.running = True
65 self.logger = self._get_logger()
67 connector = TCPConnector(
68 limit=100,
69 limit_per_host=10,
70 use_dns_cache=True,
71 force_close=True,
72 ssl=k8s_utils.VERIFY_SSL)
73 self.session: Optional[ClientSession] = ClientSession(
74 connector=connector,
75 timeout=SERVICE_LONG_TIMEOUT)
77 def _get_logger(self) -> logging.Logger:
78 """Get the logger for the service manager."""
79 logger = setup_logging(
80 path=f"/tmp/{self.app_name}",
81 file_name="service_manager.log",
82 level=logging.INFO,
83 use_global=True)
84 return logger
86 async def stop(self) -> None:
87 """Stop the dependents."""
88 self.running = False
89 if self.session:
90 await self.session.close()
91 self.session = None
93 async def init_k8s_services(self) -> None:
94 """Initialize the K8s services."""
95 await self.set_services()
97 async def update_container_status(
98 self,
99 service_name: str,
100 container: K8sContainer
101 ) -> int:
102 """Update the status of the container."""
103 url = container.get_url()
104 status, is_busy, gpu_model, num_gpus = await self.get_container_status(service_name, url)
105 container.status = status
106 container.busy = is_busy
107 container.set_gpu_model(gpu_model)
108 container.set_num_gpus(num_gpus)
109 return 1 # for counting
111 async def set_services(self) -> None:
112 """Set the K8s services and their containers asynchronously."""
113 try:
114 services = await get_k8s_services(self.k8s_cluster)
116 tasks = [
117 self.update_container_status(service_name, container)
118 for service_name, service in services.items()
119 for container in service.containers
120 ]
121 await asyncio.gather(*tasks)
123 self.services = services
124 except ConfigException as config_ex:
125 self.logger.error(f"Cannot load K8s services due to config error: {config_ex}")
126 except Exception as ex:
127 self.logger.error(f"Cannot load K8s services [{type(ex)}]: {ex}")
129 async def update_service_status(self) -> None:
130 """Update the status of all services and their containers asynchronously."""
131 self.logger.debug("Updating service status...")
133 # Update if containers/services change
134 await self.update_services()
136 tasks = [
137 self.update_container_status(service_name, container)
138 for service_name, service in self.services.items()
139 for container in service.containers
140 ]
142 num_containers = sum(await asyncio.gather(*tasks))
143 self.logger.debug(f"Services updated with {num_containers} containers.")
145 async def update_services(self) -> None:
146 """Update the K8s services and their containers asynchronously."""
147 try:
148 new_services = await get_k8s_services(self.k8s_cluster)
149 for service_name, new_service in new_services.items():
150 if service_name not in self.services:
151 logging.info(f"Service added: {service_name}")
152 else:
153 old_service = self.services[service_name]
154 old_container_names = [c.name for c in old_service.containers]
155 new_container_names = [c.name for c in new_service.containers]
156 if set(old_container_names) != set(new_container_names):
157 logging.info(f"Containers for {service_name} changed: {new_container_names}")
158 # Update in botch case for URLs, etc
159 self.services[service_name] = new_service
160 for service_name in list(self.services.keys()):
161 if service_name not in new_services:
162 logging.info(f"Service removed: {service_name}")
163 del self.services[service_name]
164 except ConfigException as config_ex:
165 self.logger.error(f"Cannot update K8s services due to config error: {config_ex}")
166 except ApiException as api_ex:
167 self.logger.error(f"Cannot update K8s services due to API error: {api_ex}")
168 except Exception as ex:
169 self.logger.error(f"Cannot update K8s services [{type(ex)}]: {ex}")
171 async def get_container_status(
172 self,
173 service_name: str,
174 url: Optional[str]
175 ) -> Tuple[Optional[str], bool, Optional[str], int]:
176 """ Get the status of a service asynchronously from its health endpoint. """
177 status = None
178 is_busy = False
179 gpu_model = None
180 num_gpus = 0
181 if not url:
182 return status, is_busy, gpu_model, num_gpus
183 if self.session is None:
184 return status, is_busy, gpu_model, num_gpus
185 try:
186 timeout = STATUS_TIMEOUT
187 health_url = f"{url}/health"
188 async with self.session.get(health_url, timeout=timeout) as response:
189 if response.status == HTTPStatus.OK:
190 if service_name in VLLM_SERVICES:
191 status = "ok" # vLLM does not return JSON
192 else:
193 try:
194 response_data = await response.json()
195 if service_name in ("streamwise") or service_name in STREAMWISE_APPS:
196 status = "ok" # No model reporting
197 elif service_name in response_data:
198 service_data = response_data[service_name]
199 is_busy = service_data.get("running", False)
200 # Services that can run multiple requests concurrently
201 if service_name in VLLM_SERVICES or service_name in (
202 "podcasttranscript",
203 "slidetranscript",
204 ):
205 is_busy = False
206 status = service_data.get("status", None)
207 gpu_model = service_data.get("gpu", None)
208 num_gpus = int(service_data.get("world_size", -1))
209 else:
210 status = "?"
211 except Exception as ex:
212 if service_name in VLLM_SERVICES:
213 status = "ok" # vLLM does not return JSON
214 else:
215 self.logger.error(f"Cannot parse JSON from {service_name} at {health_url}: {ex}")
216 status = "x"
217 else:
218 status = f"error {response.status}"
219 except TimeoutError:
220 self.logger.error(f"Timeout connecting to '{service_name}' at {health_url}.")
221 status = "timeout"
222 except ClientConnectorError:
223 self.logger.error(f"Cannot connect to '{service_name}' at {health_url}.")
224 except Exception as ex:
225 self.logger.error(f"Cannot connect to '{service_name}' at {health_url}: {type(ex)} {ex}")
226 status = "x"
227 return status, is_busy, gpu_model, num_gpus
229 async def start_updater(
230 self,
231 interval_seconds: float = 1.0
232 ) -> None:
233 """Start a background task to periodically update the status of all services and their containers."""
234 async def update_loop() -> None:
235 while self.running:
236 await self.update_service_status()
237 await asyncio.sleep(interval_seconds)
238 asyncio.create_task(update_loop())
239 self.logger.info("Started service status update.")
241 def print_service_status(self) -> list:
242 """Print the current status of all services and their containers."""
243 results = []
244 for service_name, service in self.services.items():
245 for container in service.containers:
246 results.append([
247 service_name,
248 container.get_url(),
249 container.status,
250 container.get_gpu_model(),
251 container.get_num_gpus() if container.get_num_gpus() >= 0 else "-"
252 ])
253 return results
255 async def warmup_services(self) -> None:
256 """Warmup all services by sending a warmup request."""
257 self.logger.info("Warming up services...")
258 # TODO implement proper warmup
260 def get_service_url(
261 self,
262 service_name: str,
263 exclude_busy: bool = True,
264 ) -> str:
265 """Get the URL of a service, optionally excluding busy containers."""
266 urls = self.get_service_urls(
267 service_name,
268 exclude_busy=exclude_busy)
269 return urls[0]
271 def get_service_urls(
272 self,
273 service_name: str,
274 exclude_busy: bool = True,
275 ) -> List[str]:
276 """Get the URLs of a service, optionally excluding busy containers."""
277 if service_name not in self.services:
278 raise ServiceNotFoundError(service_name)
279 service = self.services[service_name]
280 best_containers = service.get_best_containers(exclude_busy=exclude_busy)
281 if best_containers is None:
282 return []
283 return [url for c in best_containers if (url := c.get_url()) is not None]
285 def get_num_services(self) -> int:
286 return len(self.services)
288 def items(self) -> List[Tuple[str, K8sService]]:
289 return list(self.services.items())