Coverage for streamwise/pod_manager.py: 78%
241 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2API interface to add a pod for a specified container in the Kubernetes cluster.
3"""
5import os
6import sys
7import logging
8import random
9import string
10import secrets
12from http import HTTPStatus
14from quart import jsonify
16from typing import Optional
17from typing import List
18from typing import Tuple
20from kubernetes_asyncio.client import CoreV1Api
21from kubernetes_asyncio.client import ApiClient
22from kubernetes_asyncio.client import CustomObjectsApi
24from kubernetes_asyncio.client import V1Affinity
25from kubernetes_asyncio.client import V1Toleration
26from kubernetes_asyncio.client import V1Container
27from kubernetes_asyncio.client import V1ContainerPort
28from kubernetes_asyncio.client import V1Pod
29from kubernetes_asyncio.client import V1PodSpec
30from kubernetes_asyncio.client import V1ResourceRequirements
31from kubernetes_asyncio.client import V1LocalObjectReference
32from kubernetes_asyncio.client import V1VolumeMount
33from kubernetes_asyncio.client import V1Volume
34from kubernetes_asyncio.client import V1EnvVarSource
35from kubernetes_asyncio.client import V1EnvVar
36from kubernetes_asyncio.client import V1SecretKeySelector
37from kubernetes_asyncio.client import V1CSIVolumeSource
38from kubernetes_asyncio.client import V1EmptyDirVolumeSource
39from kubernetes_asyncio.client import V1ObjectMeta
40from kubernetes_asyncio.client import V1DeleteOptions
42from kubernetes_asyncio.client import V1NodeAffinity
43from kubernetes_asyncio.client import V1NodeSelectorRequirement
44from kubernetes_asyncio.client import V1NodeSelectorTerm
45from kubernetes_asyncio.client import V1NodeSelector
47from kubernetes_asyncio.client import V1Service
48from kubernetes_asyncio.client import V1ServiceSpec
49from kubernetes_asyncio.client import V1ServicePort
51from kubernetes_asyncio.client.exceptions import ApiException
53from service_account_manager import get_streamwiseapp_service_account
54from service_account_manager import get_streamwise_service_account
56sys.path.append("..")
57from quart_utils import QuartReturn
58from quart_utils import get_docker_image
60from k8s_utils import load_k8s_config
62from streamwise_apps import STREAMWISE_APPS
63from streamwise_apps import VLLM_SERVICES
66def get_tls_cert_settings() -> Tuple[V1VolumeMount, V1Volume]:
67 """Get the TLS certificate volume mount and volume using the Secrets Store CSI Driver.
69 Fetches the certificate from Azure Key Vault (via SecretProviderClass 'streamwise-tls')
70 and mounts it at /certs so the run_httpserver.bash entrypoint auto-detects it for HTTPS.
71 Requires: deployment/k8s/tls-secret-provider.yaml applied first.
72 """
73 volume_mount = V1VolumeMount(
74 name="tls-csi",
75 mount_path="/certs",
76 read_only=True
77 )
78 volume = V1Volume(
79 name="tls-csi",
80 csi=V1CSIVolumeSource(
81 driver="secrets-store.csi.k8s.io",
82 read_only=True,
83 volume_attributes={"secretProviderClass": "streamwise-tls"}
84 )
85 )
86 return volume_mount, volume
89async def tls_cert_volume_exists(namespace: str, k8s_cluster: Optional[str] = None) -> bool:
90 """Check if the SecretProviderClass for TLS certificates exists in the namespace.
92 Returns True only if the SecretProviderClass 'streamwise-tls' custom resource is present,
93 meaning the Secrets Store CSI Driver is configured and the TLS volume can be mounted.
95 Args:
96 namespace: Kubernetes namespace to check for the SecretProviderClass.
97 k8s_cluster: Kubernetes cluster context name, or None for the default context.
98 """
99 try:
100 await load_k8s_config(k8s_cluster)
101 async with ApiClient() as api_client:
102 custom_api = CustomObjectsApi(api_client)
103 await custom_api.get_namespaced_custom_object(
104 group="secrets-store.csi.x-k8s.io",
105 version="v1",
106 namespace=namespace,
107 plural="secretproviderclasses",
108 name="streamwise-tls"
109 )
110 return True
111 except ApiException as api_ex:
112 if api_ex.status == 404:
113 logging.info("SecretProviderClass 'streamwise-tls' not found; skipping TLS volume mount.")
114 else:
115 logging.warning("Error checking TLS cert volume: %s", api_ex.reason)
116 return False
117 except Exception as ex:
118 logging.warning("Error checking TLS cert volume: %s", ex)
119 return False
122def get_vllm_settings() -> Tuple[List[V1VolumeMount], List[V1Volume]]:
123 """Get the volume mounts and volumes for VLLM-based containers."""
124 volume_mounts = [
125 V1VolumeMount(
126 mount_path="/root/.cache/huggingface",
127 name="hf-cache"
128 ),
129 V1VolumeMount(
130 name="shm",
131 mount_path="/dev/shm"
132 )
133 ]
134 volumes = [
135 V1Volume(
136 name="hf-cache",
137 empty_dir=V1EmptyDirVolumeSource()
138 ),
139 V1Volume(
140 name="shm",
141 empty_dir=V1EmptyDirVolumeSource(
142 medium="Memory",
143 size_limit="1Gi") # Original 64Mi
144 )
145 ]
146 return volume_mounts, volumes
149def get_gemma_settings(num_gpus: int) -> Tuple[List[str], List[V1VolumeMount], List[V1Volume]]:
150 """vLLM parameters for Gemma."""
151 args = [
152 "--model", "google/gemma-3-27b-it",
153 "--tensor-parallel-size", str(num_gpus),
154 "--guided-decoding-backend", "xgrammar"
155 ]
156 volume_mounts, volumes = get_vllm_settings()
157 return args, volume_mounts, volumes
160def get_llama32_settings(num_gpus: int) -> Tuple[List[str], List[V1VolumeMount], List[V1Volume]]:
161 """vLLM parameters for Llama."""
162 args = [
163 "--model", "meta-llama/Llama-3.2-90B-Vision",
164 "--tensor-parallel-size", str(num_gpus),
165 "--guided-decoding-backend", "xgrammar",
166 "--enable-prefix-caching",
167 "--max-model-len", "8192",
168 "--max-num-seq", "128"
169 ]
170 """
171 --enable-auto-tool-choice \
172 --tool-call-parser pythonic \
173 --chat-template examples/tool_chat_template_llama3.2_pythonic.jinja \
174 --trust-remote-code \
175 --limit-mm-per-prompt "image=1"
176 """
177 volume_mounts, volumes = get_vllm_settings()
178 return args, volume_mounts, volumes
181def get_whisper_settings() -> Tuple[List[str], List[V1VolumeMount], List[V1Volume]]:
182 args = [
183 "--model", "openai/whisper-large-v3",
184 ]
185 volume_mounts, volumes = get_vllm_settings()
186 return args, volume_mounts, volumes
189async def add_service_ip(
190 k8s_api: CoreV1Api,
191 pod_name: str,
192 target_port: int,
193 lb_rg: str,
194 lb_ip: str,
195 lb_port: int,
196 protocol: str = "TCP",
197 namespace: str = "default",
198) -> bool:
199 """Add a LoadBalancer service to expose the pod."""
200 service_name = f"{pod_name}-svc"
202 logging.info(
203 f"Creating load balancer for '{service_name}' for pod {pod_name}:{target_port} "
204 f"on {lb_ip}:{lb_port} in resource group '{lb_rg}'.")
206 service_body = V1Service(
207 metadata=V1ObjectMeta(
208 name=service_name,
209 annotations={
210 "service.beta.kubernetes.io/azure-load-balancer-resource-group": lb_rg
211 }
212 ),
213 spec=V1ServiceSpec(
214 type="LoadBalancer",
215 load_balancer_ip=lb_ip,
216 selector={
217 "app": pod_name
218 },
219 ports=[
220 V1ServicePort(
221 port=lb_port,
222 target_port=target_port,
223 protocol=protocol,
224 )
225 ]
226 )
227 )
228 service_response = await k8s_api.create_namespaced_service(
229 namespace=namespace,
230 body=service_body)
231 if service_response is None:
232 return False
233 return True
236def get_gpu_type_affinity(gpu_type: Optional[str]) -> List[str]:
237 if gpu_type is None or gpu_type == "N/A" or gpu_type == "Any":
238 return []
239 if gpu_type == "a+":
240 return get_gpu_type_affinity("a100") + get_gpu_type_affinity("h100") + get_gpu_type_affinity("h200")
241 if gpu_type == "h+":
242 return get_gpu_type_affinity("h100") + get_gpu_type_affinity("h200")
243 if gpu_type == "a100":
244 return [
245 "NVIDIA-A100-SXM4-40GB",
246 "NVIDIA-A100-SXM4-80GB",
247 "NVIDIA-A100-PCIe-40GB",
248 "NVIDIA-A100-PCIe-80GB",
249 "NVIDIA-A100-80GB-PCIe",
250 ]
251 if gpu_type == "h100":
252 return [
253 "NVIDIA-H100-SXM5-80GB",
254 "NVIDIA-H100-PCIe-80GB",
255 "NVIDIA-H100-NVL",
256 "NVIDIA-H100-80GB-HBM3",
257 "NVIDIA-H100",
258 ]
259 if gpu_type == "h200":
260 return [
261 "NVIDIA-H200-SXM5-141GB",
262 "NVIDIA-H200"
263 ]
264 if gpu_type == "gb200":
265 return [
266 "NVIDIA-GB200-NVL",
267 "NVIDIA-GB200-SXM6-192GB",
268 "NVIDIA-GB200",
269 ]
270 if gpu_type == "gb300":
271 return [
272 "NVIDIA-GB300-NVL",
273 "NVIDIA-GB300",
274 ]
275 if gpu_type == "v100":
276 return [
277 "Tesla-V100-PCIE-16GB",
278 "Tesla-V100-SXM2-16GB",
279 "Tesla-V100-SXM2-32GB"
280 ]
281 return []
284# Valid NVIDIA MIG profiles for A100 and H100 GPUs.
285# Each profile name maps to the Kubernetes resource suffix (nvidia.com/mig-<profile>).
286MIG_PROFILES = {
287 # A100 40 GB profiles
288 "1g.5gb",
289 "2g.10gb",
290 "3g.20gb",
291 "4g.20gb",
292 "7g.40gb",
293 # A100 80 GB / H100 80 GB profiles
294 "1g.10gb",
295 "2g.20gb",
296 "3g.40gb",
297 "4g.40gb",
298 "7g.80gb",
299}
302def get_mig_resource_name(mig_profile: str) -> Optional[str]:
303 """Return the Kubernetes resource name for a given MIG profile, or None if invalid."""
304 if mig_profile in MIG_PROFILES:
305 return f"nvidia.com/mig-{mig_profile}"
306 return None
309def get_container_port(container_name: str) -> int:
310 """Get the default container port for a given container name, with special cases for certain containers."""
311 if container_name in VLLM_SERVICES:
312 return 8000
313 if container_name == "streamwise":
314 return 18181
315 if container_name in STREAMWISE_APPS:
316 return 18080
317 return 8080
320def generate_custom_random() -> str:
321 """Generate a custom random string similar to Kubernetes random suffix generation.
322 Mimics K8s random suffix generation.
323 """
324 part1 = secrets.token_hex(4)
325 part2 = ''.join(random.choices(string.ascii_lowercase + string.digits, k=5))
326 return f"{part1}-{part2}"
329async def add_pod(
330 container_name: Optional[str],
331 cpu: int = 2,
332 memory_gib: int = 4,
333 ephemeral_storage_gib: int = 16,
334 gpu: int = 0,
335 gpu_type: Optional[str] = None,
336 mig_profile: Optional[str] = None,
337 tag: Optional[str] = None,
338 lb_rg: Optional[str] = None,
339 lb_ip: Optional[str] = None,
340 lb_port: Optional[int] = None,
341 namespace: str = "default",
342 k8s_cluster: Optional[str] = None,
343 use_https: bool = False
344) -> QuartReturn:
345 """
346 API interface to add a pod for the specified container.
347 If this does not work, try "deployment/helm/deploy.sh" to deploy namespace, etc.
349 When *mig_profile* is provided (e.g. ``"1g.5gb"``), the pod requests a MIG slice instead of a whole GPU.
350 The *gpu* parameter then represents the number of MIG instances requested (usually 1).
351 MIG is only supported on A100 and H100 GPUs.
352 Set *gpu_type* accordingly (``"a100"`` or ``"h+"``) so the pod is scheduled on a MIG-capable node.
353 """
354 if not container_name:
355 return jsonify({"error": "Missing required parameter 'container_name'"}), HTTPStatus.BAD_REQUEST
357 # When no GPU resources are requested, ignore any provided MIG profile to keep API behavior consistent.
358 if not gpu or gpu <= 0:
359 mig_profile = None
360 elif mig_profile and not get_mig_resource_name(mig_profile):
361 return jsonify({"error": f"Invalid MIG profile '{mig_profile}'"}), HTTPStatus.BAD_REQUEST
363 logging.info(
364 f"Adding pod for '{container_name}' with {cpu} CPU, {memory_gib} GiB memory, "
365 f"{ephemeral_storage_gib} GiB storage, {gpu} GPU(s) of type '{gpu_type}'"
366 + (f" MIG profile '{mig_profile}'" if mig_profile else "") + ". "
367 f"Load balancer RG: '{lb_rg}', IP: '{lb_ip}', port: '{lb_port}'.")
369 # Resources
370 resource_request = {
371 "cpu": cpu,
372 "memory": f"{memory_gib}Gi",
373 }
374 resource_limit = {
375 "cpu": cpu * 2,
376 "memory": f"{memory_gib * 2}Gi",
377 }
378 if ephemeral_storage_gib and ephemeral_storage_gib > 0:
379 resource_request["ephemeral-storage"] = f"{ephemeral_storage_gib}Gi"
380 resource_limit["ephemeral-storage"] = f"{ephemeral_storage_gib * 2}Gi"
381 if gpu and gpu > 0:
382 if mig_profile:
383 # Request a MIG slice: nvidia.com/mig-<profile> (e.g. nvidia.com/mig-1g.5gb)
384 mig_resource = get_mig_resource_name(mig_profile)
385 assert mig_resource is not None # already validated above
386 resource_request[mig_resource] = gpu
387 resource_limit[mig_resource] = gpu
388 else:
389 resource_request["nvidia.com/gpu"] = gpu
390 resource_limit["nvidia.com/gpu"] = gpu
392 image_url = await get_docker_image(container_name, tag=tag)
393 if image_url is None:
394 return jsonify({"error": f"Invalid container '{container_name}'"}), HTTPStatus.BAD_REQUEST
396 container_ports = [
397 V1ContainerPort(
398 container_port=get_container_port(container_name),
399 protocol="TCP"
400 )
401 ]
402 # TODO set probes and other container settings
404 env_vars = [
405 V1EnvVar(
406 name="HUGGING_FACE_HUB_TOKEN",
407 value_from=V1EnvVarSource(
408 secret_key_ref=V1SecretKeySelector(
409 name="hf-token",
410 key="token")))
411 ]
412 if gpu < 1:
413 env_vars.append(V1EnvVar(
414 name="NVIDIA_VISIBLE_DEVICES",
415 value="none"))
416 if container_name in "streamwise" or container_name in STREAMWISE_APPS:
417 env_vars.append(V1EnvVar(
418 name="LB_RESOURCE_GROUP",
419 value=os.getenv("LB_RESOURCE_GROUP", "resource_group")))
420 env_vars.append(V1EnvVar(
421 name="LB_IP_ADDRESS",
422 value=os.getenv("LB_IP_ADDRESS", "1.2.3.4")))
424 node_selector = {
425 "kubernetes.io/os": "linux",
426 "kubernetes.io/arch": "amd64"
427 }
428 node_affinity = None
429 gpu_type_affinity = get_gpu_type_affinity(gpu_type)
430 if gpu_type_affinity:
431 node_affinity = V1NodeAffinity(
432 required_during_scheduling_ignored_during_execution=V1NodeSelector(
433 node_selector_terms=[
434 V1NodeSelectorTerm(
435 match_expressions=[
436 V1NodeSelectorRequirement(
437 key="nvidia.com/gpu.product",
438 operator="In",
439 values=gpu_type_affinity
440 ),
441 ]
442 )
443 ]
444 )
445 )
447 # AKS adds the taint to Spot VMs, so we need to add a toleration to allow scheduling on Spot nodes
448 tolerations = []
449 if gpu and gpu > 0:
450 tolerations.append(
451 V1Toleration(
452 key="kubernetes.azure.com/scalesetpriority",
453 operator="Equal",
454 value="spot",
455 effect="NoSchedule"
456 )
457 )
459 # vLLM specific settings
460 args = None
461 volume_mounts: List[V1VolumeMount] = []
462 volumes: List[V1Volume] = []
463 if container_name == "gemma":
464 if not gpu or gpu <= 0:
465 return jsonify({"error": "Gemma requires at least one GPU"}), HTTPStatus.BAD_REQUEST
466 args, vllm_mounts, vllm_volumes = get_gemma_settings(gpu)
467 volume_mounts.extend(vllm_mounts)
468 volumes.extend(vllm_volumes)
469 elif container_name == "llama32":
470 if not gpu or gpu <= 0:
471 return jsonify({"error": "Llama 3.2 requires at least one GPU"}), HTTPStatus.BAD_REQUEST
472 args, vllm_mounts, vllm_volumes = get_llama32_settings(gpu)
473 volume_mounts.extend(vllm_mounts)
474 volumes.extend(vllm_volumes)
475 elif container_name == "whisper":
476 args, vllm_mounts, vllm_volumes = get_whisper_settings()
477 volume_mounts.extend(vllm_mounts)
478 volumes.extend(vllm_volumes)
480 # Mount TLS certificate from Azure Key Vault (Secrets Store CSI Driver) at /certs so
481 # the run_httpserver.bash entrypoint auto-enables HTTPS. Only mounted when HTTPS is
482 # enabled and the SecretProviderClass exists in the cluster.
483 if use_https and await tls_cert_volume_exists(namespace, k8s_cluster):
484 tls_mount, tls_volume = get_tls_cert_settings()
485 volume_mounts.append(tls_mount)
486 volumes.append(tls_volume)
488 containers = [
489 V1Container(
490 name=container_name,
491 image=image_url,
492 args=args,
493 env=env_vars,
494 ports=container_ports,
495 volume_mounts=volume_mounts,
496 resources=V1ResourceRequirements(
497 requests={k: str(v) for k, v in resource_request.items()},
498 limits={k: str(v) for k, v in resource_limit.items()}
499 ),
500 )
501 ]
502 image_pull_secrets = [
503 # Key for the Azure Container Registry (ACR)
504 V1LocalObjectReference(name="acr-secret")
505 ]
507 random_suffix = generate_custom_random()
508 pod_name = f"{container_name}-{random_suffix}"
509 pod_labels = {"app": pod_name}
511 pod_spec = V1PodSpec(
512 containers=containers,
513 volumes=volumes,
514 image_pull_secrets=image_pull_secrets,
515 node_selector=node_selector,
516 affinity=V1Affinity(node_affinity=node_affinity),
517 tolerations=tolerations
518 )
519 pod_metadata = V1ObjectMeta(name=pod_name, labels=pod_labels)
521 if container_name == "streamwise":
522 pod_spec.service_account_name = await get_streamwise_service_account(
523 k8s_cluster=k8s_cluster,
524 namespace=namespace)
525 if container_name in STREAMWISE_APPS:
526 pod_spec.service_account_name = await get_streamwiseapp_service_account(
527 k8s_cluster=k8s_cluster,
528 namespace=namespace)
530 await load_k8s_config(k8s_cluster)
531 async with ApiClient() as api_client:
532 k8s_api = CoreV1Api(api_client)
533 pod_response = await k8s_api.create_namespaced_pod(
534 namespace=namespace, # This needs to be created using deployment/helm/deploy.sh
535 body=V1Pod(
536 metadata=pod_metadata,
537 spec=pod_spec
538 )
539 )
540 if pod_response is None:
541 return jsonify({
542 "error": f"Failed to create pod for {container_name}"
543 }), HTTPStatus.INTERNAL_SERVER_ERROR
545 # Define the services to expose the port
546 if lb_rg:
547 success = await add_service_ip(
548 k8s_api=k8s_api,
549 pod_name=pod_name,
550 target_port=get_container_port(container_name),
551 lb_rg=lb_rg,
552 lb_ip=lb_ip or "",
553 lb_port=lb_port or 8080,
554 namespace=namespace
555 )
556 if not success:
557 return jsonify({
558 "error": f"Failed to create service for {container_name}"
559 }), HTTPStatus.INTERNAL_SERVER_ERROR
561 return jsonify({
562 "message": "Pod creation requested",
563 "pod_name": pod_name,
564 "container_name": container_name,
565 "image_url": image_url,
566 "resource_request": resource_request,
567 **({"mig_profile": mig_profile} if mig_profile else {}),
568 }), HTTPStatus.OK
571async def remove_pod(
572 pod_name: str,
573 namespace: str = "default",
574 k8s_cluster: Optional[str] = None
575) -> QuartReturn:
576 """API interface to remove a pod by name."""
577 if not pod_name:
578 return jsonify({"error": "Pod name is required"}), HTTPStatus.BAD_REQUEST
580 await load_k8s_config(k8s_cluster)
581 async with ApiClient() as api_client:
582 k8s_api = CoreV1Api(api_client)
583 try:
584 # kubectl delete pod <pod_name> -n rtgen --grace-period=0 --force
585 del_pod_response = await k8s_api.delete_namespaced_pod(
586 name=pod_name,
587 namespace=namespace,
588 grace_period_seconds=0,
589 propagation_policy="Foreground",
590 body=V1DeleteOptions())
591 if del_pod_response is None:
592 return jsonify({"error": f"Cannot remove pod {pod_name}"}), HTTPStatus.INTERNAL_SERVER_ERROR
593 except ApiException as api_ex:
594 if api_ex.status == HTTPStatus.NOT_FOUND:
595 logging.error(f"Pod {pod_name} not found for removal.")
596 return jsonify({"error": f"Pod {pod_name} not found"}), HTTPStatus.NOT_FOUND
597 else:
598 logging.error(f"Error removing pod {pod_name}: {api_ex.reason}.")
599 return jsonify({"error": api_ex.reason}), api_ex.status
600 except Exception as ex:
601 logging.error(f"Error removing pod {pod_name} [{type(ex)}]: {ex}.")
602 return jsonify({"error": str(ex)}), HTTPStatus.INTERNAL_SERVER_ERROR
604 # Remove associated public address (if it exists)
605 service_name = f"{pod_name}-svc"
606 try:
607 await k8s_api.delete_namespaced_service(
608 name=service_name,
609 namespace=namespace)
610 except ApiException as api_ex:
611 if api_ex.status == HTTPStatus.NOT_FOUND:
612 logging.warning(f"Service {service_name} not found for removal.")
613 else:
614 logging.error(f"Error removing service {service_name}: {api_ex.reason}.")
615 except Exception as ex:
616 logging.error(f"Error removing service {service_name}: {ex}.")
618 return jsonify({"message": f"Pod {pod_name} removed successfully"}), HTTPStatus.OK