Coverage for k8s_utils.py: 62%
290 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Utilities for Kubernetes.
3"""
4import logging
6from typing import Union
7from typing import List
8from typing import Dict
9from typing import Optional
10from typing import Any
12from kubernetes_asyncio import config as k8s_config
13from kubernetes_asyncio import client as k8s_client
14from kubernetes_asyncio.client import ApiClient
15from kubernetes_asyncio.client import CoreV1Api
18# URL scheme used when constructing outbound service URLs (http or https).
19# Call set_service_scheme("https") at startup to enable HTTPS for all
20# service URL lookups (K8sContainer.get_url, get_k8s_containers, etc.).
21SERVICE_SCHEME: str = "http"
23# Whether to verify SSL certificates when making outbound HTTPS requests.
24# Set to False when services use self-signed certificates (e.g. Key Vault
25# generated certs in AKS). Call set_verify_ssl(False) at startup alongside
26# set_service_scheme("https") to disable certificate verification.
27VERIFY_SSL: bool = True
30def set_service_scheme(scheme: str) -> None:
31 """Set the URL scheme used for K8s service URLs."""
32 global SERVICE_SCHEME
33 if scheme not in ("http", "https"):
34 raise ValueError(f"Invalid service scheme: {scheme!r}. Must be 'http' or 'https'.")
35 SERVICE_SCHEME = scheme
38def set_verify_ssl(verify: bool) -> None:
39 """Set whether to verify SSL certificates for outbound HTTPS requests."""
40 global VERIFY_SSL
41 VERIFY_SSL = verify
44class NoActiveContainerError(Exception):
45 """Exception for no active containers."""
46 def __init__(
47 self,
48 service_name: str,
49 containers: Optional[List['K8sContainer']] = None,
50 ) -> None:
51 super().__init__(f"No active container for '{service_name}'")
52 self.service_name = service_name
53 self.containers = containers
55 def __str__(self) -> str:
56 return f"No active container for '{self.service_name}'"
59class NoRunnableContainerError(Exception):
60 """Exception for no runnable containers."""
61 def __init__(self, service_name: str) -> None:
62 super().__init__(f"No runnable containers for '{service_name}'")
63 self.service_name = service_name
65 def __str__(self) -> str:
66 return f"No runnable containers for '{self.service_name}'"
69class ServiceNotFoundError(Exception):
70 """Exception for service not found."""
71 def __init__(self, service_name: str) -> None:
72 super().__init__(f"Service '{service_name}' not found")
73 self.service_name = service_name
75 def __str__(self) -> str:
76 return f"Service '{self.service_name}' not found"
79class K8sContainer:
80 """A Kubernetes container."""
81 def __init__(self, name: str) -> None:
82 self.name = name
83 self.ip: Optional[str] = None
84 self.port: Optional[int] = None
85 self.resources: Dict[str, str] = {}
86 self.gpu_model: Optional[str] = None
87 self.num_gpus: int = -1
88 self.status: Optional[str] = None
89 self.busy = False
91 def set_gpu_model(self, gpu_model: Optional[str]) -> None:
92 self.gpu_model = gpu_model
94 def set_num_gpus(self, num_gpus: int) -> None:
95 self.num_gpus = num_gpus
97 def get_gpu_model(self) -> Optional[str]:
98 return self.gpu_model
100 def get_num_gpus(self) -> int:
101 if self.resources and "nvidia.com/gpu" in self.resources:
102 return int(self.resources["nvidia.com/gpu"])
103 return self.num_gpus
105 def get_url(self) -> Optional[str]:
106 if self.ip is None or self.port is None:
107 return None
108 return f"{SERVICE_SCHEME}://{self.ip}:{self.port}"
110 def is_active(self) -> bool:
111 if self.status is None:
112 return False
113 if self.status != "ok":
114 return False
115 if self.get_url() is None:
116 return False
117 return True
119 def is_busy(self) -> bool:
120 if self.busy is None:
121 return False
122 return self.busy
124 def __str__(self) -> str:
125 return f"K8sContainer(name={self.name}, status={self.status}, busy={self.busy}, " + \
126 f"ip={self.ip}, port={self.port}, resources={self.resources})"
128 def __repr__(self) -> str:
129 return self.__str__()
132class K8sService:
133 """A Kubernetes service with multiple containers."""
135 def __init__(self, name: str) -> None:
136 self.name = name
137 self.containers: List[K8sContainer] = []
139 def add_container(self, container: K8sContainer) -> None:
140 self.containers.append(container)
142 def merge_service(self, other: 'K8sService') -> None:
143 self.containers.extend(other.containers)
145 def get_active_containers(self) -> List[K8sContainer]:
146 """Get the list of active containers for this service."""
147 if not self.containers:
148 raise NoActiveContainerError(self.name, self.containers)
149 active_containers = [c for c in self.containers if c.is_active()]
150 if not active_containers:
151 raise NoActiveContainerError(self.name, self.containers)
152 return active_containers
154 def get_runnable_containers(self) -> List[K8sContainer]:
155 """Get the list of runnable (not busy) containers for this service."""
156 active_containers = self.get_active_containers()
157 nonbusy_containers = [c for c in active_containers if not c.is_busy()]
158 if not nonbusy_containers:
159 raise NoRunnableContainerError(self.name)
160 return nonbusy_containers
162 def get_best_containers(
163 self,
164 excluded_containers: Optional[List[K8sContainer]] = None,
165 exclude_busy: bool = True,
166 ) -> Optional[List[K8sContainer]]:
167 """Get the best available containers for this service, optionally excluding some containers."""
168 if exclude_busy:
169 runnable_containers = self.get_runnable_containers()
170 else:
171 runnable_containers = self.get_active_containers()
172 if not runnable_containers:
173 return None
174 if excluded_containers:
175 runnable_containers = [
176 c for c in runnable_containers
177 if c not in excluded_containers
178 ]
180 sorted_containers = sorted(
181 runnable_containers,
182 key=lambda c: (
183 # Prefer Active, then GPU, then CPU, then memory
184 # TODO check and rank for H200, H100, A100,...
185 int(c.resources.get("nvidia.com/gpu", "0")), # GPU
186 int(c.resources.get("cpu", "0")), # CPU
187 parse_k8s_resource_quantity(c.resources.get("memory", "0")) # Memory
188 ),
189 reverse=True
190 )
191 return sorted_containers
193 def get_best_container(
194 self,
195 excluded_containers: Optional[List[K8sContainer]] = None,
196 exclude_busy: bool = True,
197 ) -> Optional[K8sContainer]:
198 best_containers = self.get_best_containers(
199 excluded_containers=excluded_containers,
200 exclude_busy=exclude_busy)
201 if not best_containers:
202 return None
203 return best_containers[0]
205 def __str__(self) -> str:
206 return f"K8sService(name={self.name}, containers={len(self.containers)})"
209def parse_mig_resources(resources: Optional[Dict[str, str]]) -> Dict[str, int]:
210 """Extract MIG partition resources from a Kubernetes resource dict.
212 Scans for keys matching ``nvidia.com/mig-<profile>`` and returns a mapping
213 of profile name to count, e.g. ``{"1g.5gb": 7, "2g.10gb": 3}``.
214 """
215 mig: Dict[str, int] = {}
216 if resources is None:
217 return mig
218 MIG_RESOURCE_PREFIX = "nvidia.com/mig-"
219 for key, value in resources.items():
220 if key.startswith(MIG_RESOURCE_PREFIX):
221 profile = key[len(MIG_RESOURCE_PREFIX):]
222 mig[profile] = int(value)
223 return mig
226def parse_k8s_resource_quantity(quantity: str) -> Union[int, float]:
227 """Parse Kubernetes resource quantity strings into numeric values."""
228 try:
229 if quantity.endswith("m"):
230 return float(quantity[:-1]) / 1000.0
231 if quantity.endswith("Ki"):
232 return float(quantity[:-2]) * 1024
233 if quantity.endswith("Mi"):
234 return float(quantity[:-2]) * 1024 * 1024
235 if quantity.endswith("Gi"):
236 return float(quantity[:-2]) * 1024 * 1024 * 1024
237 if "." in quantity:
238 return float(quantity)
239 return int(quantity)
240 except Exception as ex:
241 logging.info(f"Error parsing resource quantity '{quantity}': {ex}")
242 return 0
245async def load_k8s_config(
246 context_name: Optional[str] = None
247) -> None:
248 """Load Kubernetes configuration."""
249 if context_name == "unittest":
250 return # Skip loading config for unit tests
251 elif context_name == "incluster":
252 k8s_config.load_incluster_config() # Running in a pod
253 else:
254 await k8s_config.load_kube_config(context=context_name)
257def is_k8s_node_ready(node: k8s_client.V1Node) -> bool:
258 """Check if a Kubernetes node is ready."""
259 for condition in node.status.conditions:
260 if condition.type == "Ready":
261 return condition.status == "True"
262 return False
265async def get_k8s_nodes(
266 context_name: Optional[str] = None
267) -> List[Dict[str, Any]]:
268 """Get the list of Kubernetes nodes."""
269 await load_k8s_config(context_name)
270 async with ApiClient() as api_client:
271 k8s_api = CoreV1Api(api_client)
272 nodes = await k8s_api.list_node()
273 ret = []
274 for node in nodes.items:
275 node_name = node.metadata.name
276 is_ready = is_k8s_node_ready(node)
277 labels = node.metadata.labels
279 allocatable_resources = node.status.allocatable
280 capacity_resources = node.status.capacity
281 images = None
282 if node.status.images is not None:
283 images = [{
284 "names": image.names,
285 "size_bytes": image.size_bytes
286 } for image in node.status.images]
287 addresses = []
288 if node.status.addresses is not None:
289 for address in node.status.addresses:
290 addresses.append(address.to_dict())
292 gpu_model = "N/A"
293 if "nvidia.com/gpu.product" in labels:
294 gpu_model = labels.get("nvidia.com/gpu.product", "N/A")
295 elif "beta.kubernetes.io/instance-type" in labels:
296 instance_type = labels["beta.kubernetes.io/instance-type"]
297 # gpu_model = get_gpu_model_from_instance_type(instance_type)
298 gpu_model = instance_type
300 region = "N/A"
301 if "azure/region" in labels:
302 region = labels["azure/region"]
303 elif "topology.kubernetes.io/region" in labels:
304 region = labels["topology.kubernetes.io/region"]
306 resource_group = "N/A"
307 if "network-resourcegroup" in labels:
308 resource_group = labels.get("network-resourcegroup", "N/A")
310 mig_enabled = any(k.startswith("nvidia.com/mig-") for k in allocatable_resources)
312 # Collect per-profile MIG resource counts (capacity and allocatable).
313 # These map directly to the Kubernetes resource names (nvidia.com/mig-<profile>)
314 # and let operators diagnose scheduling errors such as
315 # "Insufficient nvidia.com/mig-1g.5gb".
316 mig_resources: Dict[str, Dict[str, int]] = {}
317 for resource_key in sorted(capacity_resources):
318 if resource_key.startswith("nvidia.com/mig-"):
319 profile = resource_key[len("nvidia.com/mig-"):]
320 mig_resources[profile] = {
321 "capacity": int(capacity_resources.get(resource_key, 0)),
322 "allocatable": int(allocatable_resources.get(resource_key, 0)),
323 }
325 info = {
326 "node_name": node_name,
327 "region": region,
328 "resource_group": resource_group,
329 "addresses": addresses,
330 "is_ready": is_ready,
331 "capacity_resources": {
332 "cpu": parse_k8s_resource_quantity(capacity_resources.get("cpu", "N/A")),
333 "memory": parse_k8s_resource_quantity(capacity_resources.get("memory", "N/A")),
334 "storage": parse_k8s_resource_quantity(capacity_resources.get("ephemeral-storage", "N/A")),
335 "gpu": capacity_resources.get("nvidia.com/gpu", "N/A"),
336 },
337 "allocatable_resources": {
338 "cpu": parse_k8s_resource_quantity(allocatable_resources.get("cpu", "N/A")),
339 "memory": parse_k8s_resource_quantity(allocatable_resources.get("memory", "N/A")),
340 "storage": parse_k8s_resource_quantity(allocatable_resources.get("ephemeral-storage", "N/A")),
341 "gpu": allocatable_resources.get("nvidia.com/gpu", "N/A"),
342 },
343 "architecture": node.status.node_info.architecture,
344 "kernel_version": node.status.node_info.kernel_version,
345 "os_image": node.status.node_info.os_image,
346 "creation_timestamp": node.metadata.creation_timestamp,
347 "labels": labels,
348 "images": images,
349 "gpu_model": gpu_model,
350 "mig_enabled": mig_enabled,
351 "mig_resources": mig_resources,
352 }
353 ret.append(info)
354 return ret
357async def get_k8s_pods(
358 context_name: Optional[str] = None
359) -> List[Dict[str, Any]]:
360 """Get the list of Kubernetes pods."""
361 ret = []
362 await load_k8s_config(context_name)
363 async with ApiClient() as api_client:
364 k8s_api = CoreV1Api(api_client)
365 pods = await k8s_api.list_pod_for_all_namespaces()
366 for pod in pods.items:
367 pod_name = pod.metadata.name
368 namespace = pod.metadata.namespace
369 pod_status = pod.status.phase
370 pod_ip = pod.status.pod_ip
371 node_name = pod.spec.node_name
372 url = "N/A"
373 for container in pod.spec.containers:
374 container_name = container.name
375 resources = container.resources.requests
376 cpu: float = 0.0
377 memory: float = 0.0
378 gpu: int = 0
379 mig_profile: Optional[str] = None
380 if resources is not None:
381 cpu = parse_k8s_resource_quantity(resources.get("cpu", "0"))
382 memory = parse_k8s_resource_quantity(resources.get("memory", "0"))
383 gpu = int(resources.get("nvidia.com/gpu", 0))
385 mig_resources = parse_mig_resources(resources)
386 if mig_resources:
387 mig_profile = next(iter(mig_resources))
388 gpu = mig_resources[mig_profile]
390 if container.ports:
391 for container_port in container.ports:
392 if pod_ip and container_port and container_port.container_port:
393 # Assumes one open port per container
394 url = f"{SERVICE_SCHEME}://{pod_ip}:{container_port.container_port}"
395 ret.append({
396 "namespace": namespace,
397 "pod_name": pod_name,
398 "status": pod_status,
399 "pod_ip": pod_ip,
400 "container_name": container_name,
401 "url": url,
402 "node": node_name,
403 "cpu": cpu,
404 "memory": memory,
405 "gpu": gpu,
406 "mig_profile": mig_profile,
407 })
408 return ret
411async def get_k8s_load_balancers(
412 context_name: Optional[str] = None
413) -> List[Dict[str, Any]]:
414 ret: List[Dict[str, Any]] = []
416 await load_k8s_config(context_name)
418 async with ApiClient() as api_client:
419 k8s_api = CoreV1Api(api_client)
420 services = await k8s_api.list_service_for_all_namespaces()
421 for service in services.items:
422 service_spec = service.spec
423 if service_spec.type == "LoadBalancer" and len(service_spec.ports) > 0:
424 external_ip = service_spec.load_balancer_ip
425 if service.status.load_balancer.ingress:
426 external_ip = service.status.load_balancer.ingress[0].ip
427 spec_port = service.spec.ports[0]
428 external_url = None
429 if external_ip and spec_port:
430 external_url = f"{SERVICE_SCHEME}://{external_ip}:{spec_port.port}"
431 lb_info = {
432 "namespace": service.metadata.namespace,
433 "svc_name": service.metadata.name,
434 "pod_name": service_spec.selector.get("app", "N/A"),
435 "cluster_ip": service_spec.cluster_ip,
436 "external_ip": external_ip,
437 "external_port": spec_port.port,
438 "external_url": external_url,
439 "cluster_port": spec_port.target_port,
440 "node_port": spec_port.node_port,
441 }
442 ret.append(lb_info)
443 return ret
446def get_svc_name_from_container_name(container_name: str) -> str:
447 """Get the service name from the container name."""
448 service_name = container_name.lower().replace("-", "_")
449 return service_name
452async def get_k8s_services_ns(
453 context_name: Optional[str] = None,
454 namespace: str = "rtgen"
455) -> Dict[str, K8sService]:
456 svc_map = {}
457 await load_k8s_config(context_name)
458 async with ApiClient() as api_client:
459 k8s_api = CoreV1Api(api_client)
460 pods = await k8s_api.list_namespaced_pod(namespace)
461 for pod in pods.items:
462 pod_name = pod.metadata.name
463 pod_ip = pod.status.pod_ip
464 logging.debug(f"Pod: {pod_name} IP: {pod_ip}.")
465 for container in pod.spec.containers:
466 resources = container.resources.requests
468 ret_container = K8sContainer(name=pod_name)
469 ret_container.ip = pod_ip
470 ret_container.resources = resources
472 if not container.ports:
473 continue
474 for container_port in container.ports:
475 svc_name = get_svc_name_from_container_name(container.name)
476 ret_container.port = container_port.container_port
478 if svc_name not in svc_map:
479 svc_map[svc_name] = K8sService(svc_name)
480 svc_map[svc_name].add_container(ret_container)
481 return svc_map
484async def get_k8s_services(
485 context_name: Optional[str] = None
486) -> Dict[str, K8sService]:
487 svc_map = {}
489 await load_k8s_config(context_name)
490 async with ApiClient() as api_client:
491 k8s_api = CoreV1Api(api_client)
492 namespaces = await k8s_api.list_namespace()
493 for ns in namespaces.items:
494 namespace = ns.metadata.name
495 if namespace.startswith("rtgen"):
496 svc_map_ns = await get_k8s_services_ns(context_name, namespace)
497 if svc_map_ns:
498 for service_name, svc in svc_map_ns.items():
499 if service_name not in svc_map:
500 svc_map[service_name] = svc
501 else:
502 svc_map[service_name].merge_service(svc)
503 return svc_map