Coverage for streamwise/streamwise.py: 68%
539 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2StreamWise Cluster Manager.
3HTTP server that forwards the requests to each sub-module.
4"""
6from __future__ import annotations
8import sys
9import os
10import json
11import logging
12import asyncio
13import traceback
14import argparse
15import random
17from typing import List
18from typing import Dict
19from typing import Optional
20from typing import Any
21from typing import Union
23from kubernetes_asyncio.client.exceptions import ApiException
25from quart import Quart
26from quart import request
27from quart import jsonify
28from quart import render_template
30from http import HTTPStatus
32import file_manager
33import http_session_manager
34import pod_manager
35import node_manager
36import job_manager
37import allocator_bridge
38from container_config import get_minimum_service_container_specs
40from service_manager import get_services
41from service_manager import get_service_timestamps
43sys.path.append("..")
44from console_utils import setup_logging
45from streamwise_apps import STREAMWISE_APPS
47from quart_utils import QuartReturn
48from quart_utils import json_pretty_filter
49from quart_utils import format_datetime
50from quart_utils import get_friendly_container_name
51from quart_utils import get_friendly_pod_name
52from quart_utils import get_class_emoji
53from quart_utils import get_docker_image
54from quart_utils import is_rtgen_container
55from quart_utils import format_bytes
56from quart_utils import format_url
57from quart_utils import format_gpu_model
58from quart_utils import format_gpu_model_mig
59from quart_utils import get_aspect_ratio
60from quart_utils import get_file_type_emoji
61from quart_utils import get_content_type_emoji
62from quart_utils import get_friendly_region_name
64from k8s_utils import get_k8s_nodes
65from k8s_utils import get_k8s_pods
66from k8s_utils import get_k8s_load_balancers
69# Quart/Flask app configuration
70HOST = "0.0.0.0"
71PORT = 18181
72TMP_DIR = "/tmp"
73LOG_FILE_NAME = "streamwise.log"
74app = Quart(__name__)
75route = app.route
76template_filter = app.template_filter
79# Kubernetes cluster configuration
80# K8S_CLUSTER = "incluster" # If running in a pod
81K8S_CLUSTER = None # Use default context
83k8s_cluster = K8S_CLUSTER
85# This needs to be created using deployment/helm/deploy.sh
86NAMESPACE = "rtgen"
88# Set to True when the server is started with --certfile (HTTPS mode)
89use_https: bool = False
92# Template filters
93@template_filter("get_friendly_container_name")
94async def get_friendly_container_name_template(container_name: str) -> str:
95 return await get_friendly_container_name(container_name)
98@template_filter("get_friendly_pod_name")
99async def get_friendly_pod_name_template(pod_name: str) -> str:
100 return await get_friendly_pod_name(pod_name)
103@template_filter("get_class_emoji")
104async def get_class_emoji_template(container_name: str) -> str:
105 return await get_class_emoji(container_name)
108@template_filter("json_pretty")
109def json_pretty_filter_template(value: str, max_len: int = 128) -> str:
110 return json_pretty_filter(value, max_len=max_len)
113@template_filter("format_datetime")
114def format_datetime_template(value: int) -> str:
115 return format_datetime(value)
118@template_filter("format_bytes")
119def format_bytes_template(memory: int) -> str:
120 return format_bytes(memory)
123@template_filter("format_gpu_model")
124def format_gpu_model_template(gpu_model: str) -> Optional[str]:
125 return format_gpu_model(gpu_model)
128@template_filter("format_gpu_model_mig")
129def format_gpu_model_mig_template(
130 gpu_model: str,
131 mig_profile: Optional[str] = None
132) -> Optional[str]:
133 return format_gpu_model_mig(gpu_model, mig_profile)
136@template_filter("format_url")
137def format_url_template(url: Optional[str]) -> Optional[str]:
138 return format_url(url)
141@template_filter("get_aspect_ratio")
142def get_aspect_ratio_template(ratio: float) -> str:
143 return get_aspect_ratio(ratio)
146@template_filter("get_file_type_emoji")
147def get_file_type_emoji_template(file_type: str) -> str:
148 return get_file_type_emoji(file_type)
151@template_filter("get_content_type_emoji")
152def get_content_type_emoji_template(content_type: str) -> str:
153 return get_content_type_emoji(content_type)
156@template_filter("get_friendly_region_name")
157def get_friendly_region_name_template(region: str) -> str:
158 return get_friendly_region_name(region)
161@template_filter("is_rtgen_container")
162async def is_rtgen_container_template(container_name: str) -> bool:
163 return await is_rtgen_container(container_name)
166@template_filter('get_docker_image')
167async def get_docker_image_template(container_name: str) -> Optional[str]:
168 return await get_docker_image(container_name)
171# Setup and cleanup
172@app.before_serving
173async def startup() -> None:
174 """Initialize sessions before the server starts."""
175 await http_session_manager.startup()
178@app.after_serving
179async def shutdown() -> None:
180 """Cleanup tasks after server stops."""
181 await http_session_manager.shutdown()
184@app.errorhandler(HTTPStatus.INTERNAL_SERVER_ERROR)
185async def internal_error(ex: Exception) -> QuartReturn:
186 """Handle internal server errors and display a user-friendly error page."""
187 tb = traceback.format_exc()
188 error_message = getattr(ex, "description", str(ex))
189 logging.error(f"Internal error: {error_message}")
190 return await render_template(
191 "error.html",
192 error_message=str(ex),
193 exception=ex,
194 traceback=tb,
195 ), HTTPStatus.INTERNAL_SERVER_ERROR
198# HTTP routes
199@route("/", methods=["GET"])
200async def index() -> QuartReturn:
201 """Main index page showing all services and nodes."""
202 svcs = []
203 nodes = []
204 pods = []
205 lbs = []
207 try:
208 svcs = await get_services(
209 namespace=NAMESPACE,
210 k8s_cluster=k8s_cluster)
211 nodes = await get_k8s_nodes(k8s_cluster)
212 pods = await get_k8s_pods(k8s_cluster)
213 lbs = await get_k8s_load_balancers(k8s_cluster)
214 except Exception as ex:
215 logging.error(f"Error fetching index data: {ex}: {traceback.format_exc()}")
217 for svc in svcs:
218 pod_name = svc.get("pod_name")
219 lb = await get_lb_pod(pod_name)
220 if lb:
221 svc["load_balancer"] = await get_lb_pod(pod_name)
223 app_svcs = [svc for svc in svcs if svc.get("container_name") in STREAMWISE_APPS]
224 _system_containers = set(STREAMWISE_APPS) | {"streamwise"}
225 wrapper_svcs = [svc for svc in svcs if svc.get("container_name") not in _system_containers]
227 return await render_template(
228 "index.html",
229 k8s_cluster=k8s_cluster if k8s_cluster else "default",
230 svcs=svcs,
231 app_svcs=app_svcs,
232 wrapper_svcs=wrapper_svcs,
233 nodes=nodes,
234 pods=pods,
235 lbs=lbs)
238@route("/health", methods=["GET"])
239async def health() -> QuartReturn:
240 """Get health status."""
241 health = {
242 "status": "ok",
243 "k8s_cluster": k8s_cluster,
244 }
245 return jsonify(health), HTTPStatus.OK
248async def get_lb_pod(pod_name: Optional[str]) -> Optional[Dict[str, Any]]:
249 """Get load balancer info for a pod."""
250 if not pod_name:
251 return None
252 lbs = await get_k8s_load_balancers(k8s_cluster)
253 for lb in lbs:
254 if lb["pod_name"] == pod_name:
255 return lb
256 return None
259@route("/service/<service_name>", methods=["GET"])
260async def service_info(service_name: str) -> str:
261 """Display information about a specific service."""
262 services = await get_services(
263 container_name_filter=service_name,
264 details=True,
265 namespace=NAMESPACE,
266 k8s_cluster=k8s_cluster)
268 ret_svcs = [
269 svc
270 for svc in services
271 if svc.get("container_name") == service_name
272 ]
273 first_svc = ret_svcs[0] if ret_svcs else None
275 # Add load balancer info if available
276 lb = None
277 lbs = await get_k8s_load_balancers(k8s_cluster)
278 for svc in ret_svcs:
279 lb = next((
280 lb
281 for lb in lbs
282 if lb["pod_name"] == svc["pod_name"]), None)
283 if lb:
284 svc["load_balancer"] = lb
286 return await render_template(
287 "service.html",
288 service_name=service_name,
289 svcs=ret_svcs,
290 svc=first_svc,
291 lb=lb)
294@route("/service/<service_name>/<container_ip>", methods=["GET"])
295async def container_info(
296 service_name: str,
297 container_ip: str
298) -> str:
299 """Display information about a specific container instance of a service."""
300 services = await get_services(
301 container_name_filter=service_name,
302 details=True,
303 namespace=NAMESPACE,
304 k8s_cluster=k8s_cluster)
305 ret_svcs = [
306 svc
307 for svc in services
308 if svc.get("container_name") == service_name
309 and svc.get("pod_ip") == container_ip
310 ]
311 first_svc = ret_svcs[0] if ret_svcs else None
312 return await render_template(
313 "service.html",
314 service_name=service_name,
315 svcs=ret_svcs,
316 svc=first_svc)
319@route("/service/<service_name>/timeline", methods=["GET"])
320async def service_timelines(service_name: str) -> QuartReturn:
321 """Display timeline information for a specific service."""
322 services = await get_services(
323 container_name_filter=service_name,
324 details=False,
325 namespace=NAMESPACE,
326 k8s_cluster=k8s_cluster)
328 async def fetch_all_timestamps() -> List[Dict]:
329 tasks = [
330 get_service_timestamps(
331 svc["pod_name"],
332 svc["container_name"],
333 svc["url"]
334 )
335 for svc in services if svc.get("url") and svc["url"] != "N/A"
336 ]
337 results = await asyncio.gather(*tasks, return_exceptions=True)
339 timestamps = []
340 for result in results:
341 if isinstance(result, Exception):
342 logging.error(f"Error fetching timestamps: {result}")
343 elif isinstance(result, list) and result:
344 timestamps.extend(result)
345 return timestamps
347 try:
348 # TODO doesn't seem to retrieve from all instances
349 timestamps = await fetch_all_timestamps()
350 return await render_template(
351 "service_timeline.html",
352 service_name=service_name,
353 timestamps=timestamps)
354 except Exception as ex:
355 logging.error(f"Error fetching timestamps for {service_name}: {ex}")
356 return jsonify({"error": str(ex)}), HTTPStatus.INTERNAL_SERVER_ERROR
359@route("/service/timeline", methods=["GET"])
360async def services_timelines() -> QuartReturn:
361 """Display timeline information for all services."""
362 services = await get_services(
363 details=False,
364 namespace=NAMESPACE,
365 k8s_cluster=k8s_cluster)
367 async def fetch_all_timestamps() -> List[Dict]:
368 tasks = []
369 for service in services:
370 if service["url"] != "N/A":
371 task = get_service_timestamps(
372 service["pod_name"],
373 service["container_name"],
374 service["url"]
375 )
376 tasks.append(task)
377 ret = []
378 for result in await asyncio.gather(*tasks, return_exceptions=True):
379 if isinstance(result, list) and result:
380 ret.extend(result)
381 return ret
383 try:
384 timestamps = await fetch_all_timestamps()
385 return await render_template(
386 "service_timeline.html",
387 service_name="All Services",
388 timestamps=timestamps)
389 except Exception as ex:
390 logging.error(f"Error fetching timestamp: {ex}.")
391 return jsonify({"error": str(ex)}), HTTPStatus.INTERNAL_SERVER_ERROR
394@route("/node/<node_name>", methods=["GET"])
395async def node_info(node_name: str) -> QuartReturn:
396 """Display information about a specific node."""
397 return await node_manager.node_info(
398 node_name,
399 k8s_cluster=k8s_cluster)
402@route("/nodes", methods=["GET"])
403async def nodes_info() -> QuartReturn:
404 """Display information about all nodes."""
405 return await node_manager.nodes_info(k8s_cluster=k8s_cluster)
408@route("/audio_waveform/<container_ip>/<container_port>/<file_name>", methods=["GET"])
409async def get_audio_waveform(container_ip: str, container_port: int, file_name: str) -> QuartReturn:
410 """Generate and return a waveform PNG image for a WAV audio file."""
411 return await file_manager.get_audio_waveform(container_ip, container_port, file_name)
414@route("/video/<container_ip>/<container_port>/<file_name>", methods=["GET"])
415async def get_video_info(container_ip: str, container_port: int, file_name: str) -> QuartReturn:
416 """Get video file information from a container."""
417 return await file_manager.get_video_info(container_ip, container_port, file_name)
420@route("/files", methods=["GET"])
421async def list_files() -> QuartReturn:
422 """List files in the TMP_DIR directory."""
423 return await file_manager.list_files(TMP_DIR)
426@route("/file/<file_name>", methods=["GET"])
427async def download_local_file(file_name: str) -> QuartReturn:
428 """Download a file."""
429 return await file_manager.download_local_file(
430 TMP_DIR,
431 file_name)
434@route("/file_download/<container_ip>/<container_port>/<file_name>", methods=["GET"])
435async def download_service_file(
436 container_ip: str,
437 container_port: int,
438 file_name: str
439) -> QuartReturn:
440 """Download a file from a container."""
441 return await file_manager.download_service_file(
442 container_ip,
443 container_port,
444 file_name)
447@route("/file_view/<service_name>/<container_ip>/<container_port>/<file_name>", methods=["GET"])
448async def file_view(
449 service_name: str,
450 container_ip: str,
451 container_port: int,
452 file_name: str
453) -> QuartReturn:
454 """View the contents of a file from a container."""
455 return await file_manager.file_view(
456 service_name,
457 container_ip,
458 container_port,
459 file_name)
462@route("/file_stream/<container_ip>/<container_port>/<file_name>")
463async def file_stream(
464 container_ip: str,
465 container_port: int,
466 file_name: str
467) -> QuartReturn:
468 """Stream a file from a container."""
469 return await file_manager.file_stream(
470 container_ip,
471 container_port,
472 file_name)
475@route("/job/", methods=["GET"])
476async def submit_job() -> str:
477 """Render the job submission form."""
478 svcs = []
479 try:
480 svcs = await get_services(
481 namespace=NAMESPACE,
482 k8s_cluster=k8s_cluster)
483 except Exception as ex:
484 logging.exception("Error fetching services for /job/: %s", ex)
485 return await render_template(
486 "submit_job.html",
487 svcs=svcs)
490@route("/job/<container_ip>/<container_port>", methods=["GET"])
491async def submit_job_container(
492 container_ip: str,
493 container_port: int
494) -> str:
495 """Render the job submission form for a specific container."""
496 svcs = []
497 try:
498 svcs = await get_services(
499 namespace=NAMESPACE,
500 k8s_cluster=k8s_cluster)
501 except Exception as ex:
502 logging.exception("Error fetching services for /job/%s/%s: %s",
503 container_ip, container_port, ex)
504 return await render_template(
505 "submit_job.html",
506 svcs=svcs,
507 container_ip=container_ip,
508 container_port=container_port)
511@route("/api/job/<service_name>/<container_ip>/<container_port>", methods=["POST"])
512async def api_submit_job(
513 service_name: str,
514 container_ip: str,
515 container_port: int
516) -> QuartReturn:
517 """API interface to submit a job to the specified service."""
518 return await job_manager.submit_job(
519 service_name,
520 container_ip,
521 container_port)
524@route("/pod", methods=["GET"])
525async def add_pods() -> str:
526 """Render the add pod form."""
527 lb_rg = os.getenv("LB_RESOURCE_GROUP", "resource_group")
528 lb_ip = os.getenv("LB_IP_ADDRESS", "1.2.3.4")
529 return await render_template(
530 "add_pod.html",
531 lb_rg=lb_rg,
532 lb_ip=lb_ip,
533 random_lb_port=8080 + random.randint(0, 920),
534 )
537@route("/pod/<service_name>", methods=["GET"])
538async def add_pod(service_name: str) -> str:
539 """Render the add pod form for a specific service."""
540 lb_rg = os.getenv("LB_RESOURCE_GROUP", "resource_group")
541 lb_ip = os.getenv("LB_IP_ADDRESS", "1.2.3.4")
542 return await render_template(
543 "add_pod.html",
544 service_name=service_name,
545 lb_rg=lb_rg,
546 lb_ip=lb_ip,
547 random_lb_port=8080 + random.randint(0, 920),
548 )
551@route("/auto_deploy", methods=["GET"])
552async def auto_deploy_page() -> str:
553 """Render the standalone auto-deploy page. (TODO: enable customization after auto-deploy plan is generated)"""
554 return await render_template("auto_deploy.html")
557@route("/api/pod/<pod_name>", methods=["DELETE"])
558async def api_remove_pod(pod_name: str) -> QuartReturn:
559 """API interface to remove a pod by name."""
560 namespace = request.args.get("namespace")
561 if not namespace:
562 return jsonify({"error": "Namespace is required"}), HTTPStatus.BAD_REQUEST
563 return await pod_manager.remove_pod(
564 pod_name,
565 namespace=namespace,
566 k8s_cluster=k8s_cluster)
569@route("/api/pods/wrappers", methods=["DELETE"])
570async def api_delete_all_wrappers() -> QuartReturn:
571 """Delete all wrapper pods (non-app, non-system pods) in the namespace."""
572 svcs = await get_services(namespace=NAMESPACE, k8s_cluster=k8s_cluster)
573 # Exclude app pods and the streamwise management pod itself
574 excluded = set(STREAMWISE_APPS) | {"streamwise"}
575 wrapper_pods = [
576 svc["pod_name"] for svc in svcs
577 if svc.get("container_name") not in excluded and svc.get("pod_name")
578 ]
579 deleted = 0
580 errors: list[str] = []
581 for pod_name in wrapper_pods:
582 try:
583 await pod_manager.remove_pod(
584 pod_name, namespace=NAMESPACE, k8s_cluster=k8s_cluster)
585 deleted += 1
586 except Exception as e:
587 errors.append(f"{pod_name}: {e}")
588 result: dict[str, object] = {"deleted": deleted, "total": len(wrapper_pods)}
589 if errors:
590 result["errors"] = errors
591 return jsonify(result), HTTPStatus.OK
594@route("/api/services", methods=["GET"])
595async def api_get_services() -> QuartReturn:
596 """API interface to get the list of services."""
597 services = await get_services(
598 namespace=NAMESPACE,
599 k8s_cluster=k8s_cluster)
600 return jsonify(services), HTTPStatus.OK
603@route("/api/nodes", methods=["GET"])
604async def api_get_nodes() -> QuartReturn:
605 """API interface to get the list of nodes."""
606 nodes = await get_k8s_nodes(k8s_cluster)
607 return jsonify(nodes), HTTPStatus.OK
610def parse_gpu_info(
611 gpu_info: Optional[Union[int, str]]
612) -> tuple[int, Optional[str]]:
613 num_gpus = 1
614 mig_profile = None
615 if isinstance(gpu_info, int):
616 num_gpus = gpu_info
617 elif isinstance(gpu_info, str):
618 num_gpus = 1
619 mig_profile = gpu_info
620 else:
621 num_gpus = 0
622 return num_gpus, mig_profile
625@route("/api/service", methods=["POST"])
626async def api_add_service(
627 max_gpus: int = 1
628) -> QuartReturn:
629 """API interface to add pods for all services."""
630 try:
631 container_specs = get_minimum_service_container_specs(max_gpus=max_gpus)
633 for container_name, spec in container_specs.items():
634 num_gpus, mig_profile = parse_gpu_info(spec.gpu)
635 await pod_manager.add_pod(
636 container_name,
637 spec.cpu,
638 spec.memory_gib,
639 ephemeral_storage_gib=spec.ephemeral_storage_gib,
640 gpu=num_gpus,
641 mig_profile=mig_profile,
642 namespace=NAMESPACE,
643 k8s_cluster=k8s_cluster)
644 return jsonify({"message": "Services added successfully"}), HTTPStatus.OK
645 except ApiException as api_ex:
646 body = json.loads(api_ex.body) if api_ex.body else {}
647 message = body.get("message", "No message")
648 if message == "namespaces \"rtgen\" not found":
649 message += ".\nRun: 'kubectl create namespace rtgen'"
650 logging.error(f"K8s API error adding services: {message}.")
651 return jsonify({"error": message}), HTTPStatus.INTERNAL_SERVER_ERROR
652 except Exception as ex:
653 logging.error(f"Error adding services: {ex}.")
654 return jsonify({"error": str(ex)}), HTTPStatus.INTERNAL_SERVER_ERROR
657@route("/api/apps", methods=["POST"])
658async def api_add_apps() -> QuartReturn:
659 """API interface to add pods for all applications."""
660 try:
661 lb_rg = os.getenv("LB_RESOURCE_GROUP")
662 lb_ip = os.getenv("LB_IP_ADDRESS")
663 # CPU, memory GiB, ephemeral storage GiB, GPU count
664 container_dict: dict[str, tuple[int, int, int, int]] = {
665 "streamcast": (1, 4, 4, 0),
666 "streampersona": (1, 4, 4, 0),
667 "streamchat": (1, 4, 4, 0),
668 "streamshort": (1, 4, 4, 0),
669 "streammovie": (1, 4, 4, 0),
670 "streamanimate": (1, 4, 4, 0),
671 "streamlecture": (1, 4, 4, 0),
672 "streamdub": (1, 4, 4, 0),
673 "streamedit": (1, 4, 4, 0),
674 }
675 lb_ports = random.sample(range(8080, 9000), len(container_dict))
676 for (container_name, (cpu, mem_gib, storage_gib, gpu)), lb_port in zip(container_dict.items(), lb_ports):
677 await pod_manager.add_pod(
678 container_name,
679 cpu,
680 mem_gib,
681 ephemeral_storage_gib=storage_gib,
682 gpu=gpu,
683 lb_rg=lb_rg,
684 lb_ip=lb_ip,
685 lb_port=lb_port,
686 namespace=NAMESPACE,
687 k8s_cluster=k8s_cluster)
688 return jsonify({"message": "Applications added successfully"}), HTTPStatus.OK
689 except ApiException as api_ex:
690 body = json.loads(api_ex.body) if api_ex.body else {}
691 message = body.get("message", "No message")
692 if message == "namespaces \"rtgen\" not found":
693 message += ".\nRun: 'kubectl create namespace rtgen'"
694 logging.error(f"K8s API error adding applications: {message}.")
695 return jsonify({"error": message}), HTTPStatus.INTERNAL_SERVER_ERROR
696 except Exception as ex:
697 logging.error(f"Error adding applications: {ex}.")
698 return jsonify({"error": str(ex)}), HTTPStatus.INTERNAL_SERVER_ERROR
701@route("/api/pod", methods=["POST"])
702async def api_add_pod() -> QuartReturn:
703 """API interface to add a pod for the specified container."""
704 form = await request.form
705 container_name = form.get("container_name")
706 cpu = int(form.get("cpu", 2))
707 memory_gib = int(form.get("memory", 4))
708 ephemeral_storage_gib = int(form.get("ephemeralStorage", 16))
709 gpu = int(form.get("gpu", 0))
710 gpu_type = form.get("gpu_type")
711 mig_profile = form.get("mig_profile", "").strip() or None
712 tag = form.get("tag", "").strip() or None
713 lb_rg = form.get("lb_rg")
714 lb_ip = form.get("lb_ip")
715 lb_port = form.get("lb_port")
716 try:
717 return await pod_manager.add_pod(
718 container_name=container_name,
719 cpu=cpu,
720 memory_gib=memory_gib,
721 ephemeral_storage_gib=ephemeral_storage_gib,
722 gpu=gpu,
723 gpu_type=gpu_type,
724 mig_profile=mig_profile,
725 tag=tag,
726 lb_rg=lb_rg,
727 lb_ip=lb_ip,
728 lb_port=int(lb_port) if lb_port else None,
729 namespace=NAMESPACE,
730 k8s_cluster=k8s_cluster,
731 use_https=use_https,
732 )
733 except ApiException as api_ex:
734 body = json.loads(api_ex.body) if api_ex.body else {}
735 message = body.get("message", "No message")
736 if message == "namespaces \"rtgen\" not found":
737 message += ".\nRun: 'kubectl create namespace rtgen'"
738 logging.error(f"K8s API error adding services: {message}.")
739 return jsonify({"error": message}), HTTPStatus.INTERNAL_SERVER_ERROR
740 except Exception as ex:
741 logging.error(f"Error adding pod for {container_name}: {ex}.")
742 traceback.print_exc()
743 return jsonify({"error": str(ex)}), HTTPStatus.INTERNAL_SERVER_ERROR
746@route("/api/auto_deploy", methods=["POST"])
747async def api_auto_deploy() -> QuartReturn:
748 """Run the model allocator to produce an optimized deployment plan.
750 Expects JSON body:
751 {
752 "gpu_budget": {"A100": 8, "H100": 0, ...},
753 "workflow": "streamcast"
754 }
756 Returns the deployment plan with estimated metrics and per-container specs.
757 """
758 try:
759 data = await request.get_json()
760 if not data:
761 return jsonify({"error": "Request body must be JSON"}), HTTPStatus.BAD_REQUEST
763 gpu_budget = data.get("gpu_budget")
764 workflow_name = data.get("workflow")
766 if not gpu_budget or not isinstance(gpu_budget, dict):
767 return jsonify({"error": "Missing or invalid 'gpu_budget' field"}), HTTPStatus.BAD_REQUEST
768 for gpu_type_name, count in gpu_budget.items():
769 if isinstance(count, bool) or not isinstance(count, int) or count < 0:
770 return (
771 jsonify(
772 {
773 "error": (
774 "Invalid 'gpu_budget' field: each GPU type count must be a "
775 "non-negative integer"
776 )
777 }
778 ),
779 HTTPStatus.BAD_REQUEST,
780 )
781 if not workflow_name or not isinstance(workflow_name, str):
782 return jsonify({"error": "Missing or invalid 'workflow' field"}), HTTPStatus.BAD_REQUEST
784 plan = await asyncio.to_thread(
785 allocator_bridge.run_allocator,
786 gpu_budget=gpu_budget,
787 workflow_name=workflow_name,
788 )
789 result_json = allocator_bridge.deployment_plan_to_json(plan)
791 # Enrich specs with friendly names from services.json and uppercase GPU types
792 for spec in result_json.get("specs", []):
793 spec["friendly_name"] = await get_friendly_container_name(spec["container_name"])
794 if spec.get("gpu_type"):
795 spec["gpu_type"] = spec["gpu_type"].upper()
797 return jsonify(result_json), HTTPStatus.OK
799 except ValueError as ve:
800 return jsonify({"error": str(ve)}), HTTPStatus.BAD_REQUEST
801 except AssertionError as ae:
802 msg = str(ae) if str(ae) else (
803 "GPU budget too small. Each GPU type must have at least 8 GPUs "
804 "(one full server). Use a single GPU type with 8+ GPUs, or "
805 "ensure each type has at least 8."
806 )
807 return jsonify({"error": msg}), HTTPStatus.BAD_REQUEST
808 except Exception as ex:
809 logging.exception("Error in auto_deploy: %s", ex)
810 return jsonify({"error": str(ex)}), HTTPStatus.INTERNAL_SERVER_ERROR
813@route("/api/auto_deploy/confirm", methods=["POST"])
814async def api_auto_deploy_confirm() -> QuartReturn:
815 """Execute a deployment plan produced by /api/auto_deploy.
817 Expects JSON body:
818 {
819 "specs": [...],
820 "workflow": "streamcast" (optional: also deploys the application container)
821 }
823 Deploys all model wrapper containers in the plan, plus the application
824 container if a workflow name is provided.
825 """
826 try:
827 data = await request.get_json()
828 if not data:
829 return jsonify({"error": "Request body must be JSON"}), HTTPStatus.BAD_REQUEST
831 specs = data.get("specs")
832 if not specs or not isinstance(specs, list):
833 return jsonify({"error": "Missing or invalid 'specs' field"}), HTTPStatus.BAD_REQUEST
835 workflow = data.get("workflow")
837 deployed: List[str] = []
838 errors: List[str] = []
840 for spec in specs:
841 container_name = spec.get("container_name")
842 if not container_name:
843 errors.append("Spec missing 'container_name'")
844 continue
846 try:
847 add_pod_result = await pod_manager.add_pod(
848 container_name=container_name,
849 cpu=int(spec.get("cpu", 4)),
850 memory_gib=int(spec.get("memory_gib", 16)),
851 ephemeral_storage_gib=int(spec.get("ephemeral_storage_gib", 16)),
852 gpu=int(spec.get("gpu", 0)),
853 gpu_type=spec.get("gpu_type"),
854 mig_profile=spec.get("mig_profile"),
855 namespace=NAMESPACE,
856 k8s_cluster=k8s_cluster,
857 )
859 status_code = HTTPStatus.OK
860 if isinstance(add_pod_result, tuple) and len(add_pod_result) >= 2:
861 status_value = add_pod_result[1]
862 if isinstance(status_value, HTTPStatus):
863 status_code = status_value
864 elif isinstance(status_value, int):
865 status_code = HTTPStatus(status_value)
867 if status_code >= HTTPStatus.BAD_REQUEST:
868 msg = f"Failed to deploy '{container_name}' (status={int(status_code)})"
869 logging.error(msg)
870 errors.append(msg)
871 else:
872 deployed.append(container_name)
873 except Exception as pod_ex:
874 msg = f"Failed to deploy '{container_name}': {pod_ex}"
875 logging.error(msg)
876 errors.append(msg)
878 # Also deploy the application container if workflow is specified
879 if workflow and workflow in STREAMWISE_APPS:
880 try:
881 add_pod_result = await pod_manager.add_pod(
882 container_name=workflow,
883 cpu=4,
884 memory_gib=16,
885 ephemeral_storage_gib=16,
886 gpu=0,
887 gpu_type=None,
888 mig_profile=None,
889 namespace=NAMESPACE,
890 k8s_cluster=k8s_cluster,
891 )
892 status_code = HTTPStatus.OK
893 if isinstance(add_pod_result, tuple) and len(add_pod_result) >= 2:
894 status_value = add_pod_result[1]
895 if isinstance(status_value, HTTPStatus):
896 status_code = status_value
897 elif isinstance(status_value, int):
898 status_code = HTTPStatus(status_value)
899 if status_code >= HTTPStatus.BAD_REQUEST:
900 msg = f"Failed to deploy app '{workflow}' (status={int(status_code)})"
901 logging.error(msg)
902 errors.append(msg)
903 else:
904 deployed.append(workflow)
905 except Exception as app_ex:
906 msg = f"Failed to deploy app '{workflow}': {app_ex}"
907 logging.error(msg)
908 errors.append(msg)
910 total_deployed = len(deployed)
911 total_specs = len(specs) + (1 if workflow and workflow in STREAMWISE_APPS else 0)
912 status = HTTPStatus.OK if not errors else HTTPStatus.MULTI_STATUS
913 return jsonify({
914 "deployed": deployed,
915 "errors": errors,
916 "message": f"Deployed {total_deployed}/{total_specs} containers.",
917 }), status
919 except Exception as ex:
920 logging.exception("Error in auto_deploy/confirm: %s", ex)
921 return jsonify({"error": str(ex)}), HTTPStatus.INTERNAL_SERVER_ERROR
924@route("/api/auto_deploy/workflows", methods=["GET"])
925async def api_auto_deploy_workflows() -> QuartReturn:
926 """Return available workflows and GPU types for the auto-deploy UI."""
927 return jsonify({
928 "workflows": allocator_bridge.get_available_workflows(),
929 "gpu_types": allocator_bridge.get_available_gpu_types(),
930 }), HTTPStatus.OK
933@route("/api/auto_deploy/cluster_gpus", methods=["GET"])
934async def api_auto_deploy_cluster_gpus() -> QuartReturn:
935 """Return aggregated GPU counts by type from the current cluster.
937 Inspects all ready nodes and sums up allocatable GPUs grouped by the
938 nvidia.com/gpu.product label (mapped to canonical names like A100, H100, etc.).
939 """
940 try:
941 nodes = await get_k8s_nodes(k8s_cluster)
942 gpu_counts: dict[str, int] = {}
943 for node in nodes:
944 if not node.get("is_ready"):
945 continue
946 gpu_model = node.get("gpu_model", "N/A")
947 if gpu_model == "N/A":
948 continue
949 gpu_count = node.get("allocatable_resources", {}).get("gpu", 0)
950 if isinstance(gpu_count, str):
951 try:
952 gpu_count = int(gpu_count)
953 except ValueError:
954 continue
955 if gpu_count <= 0:
956 continue
957 # Map gpu_model label to canonical type name
958 canonical = _gpu_label_to_canonical(gpu_model)
959 gpu_counts[canonical] = gpu_counts.get(canonical, 0) + gpu_count
960 return jsonify({"gpu_budget": gpu_counts}), HTTPStatus.OK
961 except Exception as ex:
962 logging.exception("Error in cluster_gpus: %s", ex)
963 return jsonify({"error": str(ex)}), HTTPStatus.INTERNAL_SERVER_ERROR
966def _gpu_label_to_canonical(gpu_model: str) -> str:
967 """Map a GPU product label to a canonical type name for the allocator."""
968 model_upper = gpu_model.upper()
969 if "H100" in model_upper:
970 return "H100"
971 elif "H200" in model_upper:
972 return "H200"
973 elif "A100" in model_upper:
974 return "A100"
975 elif "GB200" in model_upper:
976 return "GB200"
977 elif "GB300" in model_upper:
978 return "GB300"
979 elif "V100" in model_upper:
980 return "V100"
981 elif "A10" in model_upper:
982 return "A10"
983 # Fallback: return as-is
984 return gpu_model
987@route("/api/node/<node_name>", methods=["DELETE"])
988async def api_remove_node(node_name: str) -> QuartReturn:
989 return await node_manager.remove_node(
990 node_name,
991 k8s_cluster=k8s_cluster)
994if __name__ == "__main__":
995 setup_logging(
996 path=TMP_DIR,
997 file_name=LOG_FILE_NAME)
999 parser = argparse.ArgumentParser(description="StreamWise Cluster Manager")
1000 parser.add_argument("--k8s_cluster", type=str, default=K8S_CLUSTER, help="Kubernetes cluster context name")
1001 parser.add_argument("--host", type=str, default=HOST, help="Host to bind the server to")
1002 parser.add_argument("--port", type=int, default=PORT, help="Port to bind the server to")
1003 parser.add_argument("--certfile", type=str, default=None, help="Path to SSL certificate file for HTTPS")
1004 parser.add_argument("--keyfile", type=str, default=None, help="Path to SSL private key file for HTTPS")
1005 parser.add_argument("--use-https", action="store_true", default=False,
1006 help="Use HTTPS for outbound service connections (health checks, file fetches, jobs)")
1007 args = parser.parse_args()
1009 k8s_cluster = args.k8s_cluster
1010 host = args.host
1011 port = args.port
1012 use_https = bool(args.certfile)
1014 if args.use_https:
1015 http_session_manager.set_service_scheme("https")
1016 http_session_manager.set_verify_ssl(False)
1018 try:
1019 scheme = "https" if args.certfile else "http"
1020 logging.info(f"Starting on {scheme}://{host}:{port} for K8S cluster '{k8s_cluster}'.")
1021 app.run(
1022 host=host,
1023 port=port,
1024 certfile=args.certfile,
1025 keyfile=args.keyfile,
1026 # threaded=True,
1027 # debug=True,
1028 )
1029 except OSError as os_err:
1030 logging.error(f"OS error starting: {os_err}")
1031 except Exception as ex:
1032 logging.error(f"Error starting: {ex}")