From 2b8d84a9f745591ed488954ad6664c61b28a17cd Mon Sep 17 00:00:00 2001 From: mbeaulne Date: Wed, 26 Aug 2026 14:07:09 -0400 Subject: [PATCH] Adds a compute resource endpoint --- cloud_pipelines_backend/api_router.py | 19 ++ cloud_pipelines_backend/api_server_sql.py | 3 + cloud_pipelines_backend/compute_resources.py | 91 +++++++++ cloud_pipelines_backend/errors.py | 7 + tests/test_compute_resources.py | 188 +++++++++++++++++++ 5 files changed, 308 insertions(+) create mode 100644 cloud_pipelines_backend/compute_resources.py create mode 100644 tests/test_compute_resources.py diff --git a/cloud_pipelines_backend/api_router.py b/cloud_pipelines_backend/api_router.py index 0e8fa41..6455ea3 100644 --- a/cloud_pipelines_backend/api_router.py +++ b/cloud_pipelines_backend/api_router.py @@ -13,6 +13,7 @@ from . import api_server_sql from . import backend_types_sql from . import component_library_api_server as components_api +from . import compute_resources from . import database_ops from . import errors from .instrumentation import contextual_logging @@ -141,6 +142,19 @@ def handle_item_already_exists_error( content={"message": str(exc)}, ) + @app.exception_handler(errors.UnsupportedGpuError) + def handle_unsupported_gpu_error( + request: fastapi.Request, exc: errors.UnsupportedGpuError + ): + return fastapi.responses.JSONResponse( + status_code=fastapi.status.HTTP_422_UNPROCESSABLE_CONTENT, + content={ + "reason": "unsupported_gpu", + "detail": str(exc), + "unsupported_gpus": exc.unsupported_gpus, + }, + ) + @app.exception_handler(errors.ApiValidationError) def handle_api_validation_error( request: fastapi.Request, exc: errors.ApiValidationError @@ -340,6 +354,11 @@ def get_container_log( router.get("/api/pipeline_runs/", tags=["pipelineRuns"], **default_config)( inject_session_dependency(list_pipeline_runs_func) ) + router.get( + "/api/pipeline_runs/capabilities", + tags=["pipelineRuns"], + **default_config, + )(compute_resources.get_pipeline_run_capabilities) router.get("/api/pipeline_runs/{id}", tags=["pipelineRuns"], **default_config)( inject_session_dependency(pipeline_run_service.get) ) diff --git a/cloud_pipelines_backend/api_server_sql.py b/cloud_pipelines_backend/api_server_sql.py index cf71a16..5c29ba7 100644 --- a/cloud_pipelines_backend/api_server_sql.py +++ b/cloud_pipelines_backend/api_server_sql.py @@ -9,6 +9,7 @@ from . import backend_types_sql as bts from . import component_structures as structures +from . import compute_resources from . import errors from . import filter_query_sql @@ -120,6 +121,8 @@ def create( # TODO: Load and validate all components # TODO: Fetch missing components and populate component specs + compute_resources.validate_pipeline_gpu_resources(root_task) + pipeline_name = root_task.component_ref.spec.name with session.begin(): diff --git a/cloud_pipelines_backend/compute_resources.py b/cloud_pipelines_backend/compute_resources.py new file mode 100644 index 0000000..5048dcf --- /dev/null +++ b/cloud_pipelines_backend/compute_resources.py @@ -0,0 +1,91 @@ +import dataclasses +import json +from collections import abc +from typing import Literal + +from . import component_structures as structures +from . import errors +from .launchers import kubernetes_launchers + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class GpuResource: + id: str + display_name: str + status: Literal["supported", "deprecated"] + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class GetPipelineRunCapabilitiesResponse: + gpus: tuple[GpuResource, ...] + + +GPU_RESOURCES = ( + GpuResource( + id="NVIDIA-B300", + display_name="NVIDIA B300", + status="supported", + ), + GpuResource( + id="NVIDIA-H200", + display_name="NVIDIA H200", + status="deprecated", + ), +) + +_VALID_GPU_IDS = frozenset(gpu.id for gpu in GPU_RESOURCES) +_GPU_ANNOTATION_KEY = kubernetes_launchers.RESOURCES_ACCELERATORS_ANNOTATION_KEY + + +def get_pipeline_run_capabilities() -> GetPipelineRunCapabilitiesResponse: + return GetPipelineRunCapabilitiesResponse(gpus=GPU_RESOURCES) + + +def validate_pipeline_gpu_resources(root_task: structures.TaskSpec) -> None: + unsupported_gpus: set[str] = set() + + for task in _walk_tasks(root_task): + annotation_value = (task.annotations or {}).get(_GPU_ANNOTATION_KEY) + if annotation_value is None: + continue + + accelerators = _parse_accelerator_annotation(annotation_value) + unsupported_gpus.update(set(accelerators) - _VALID_GPU_IDS) + + if unsupported_gpus: + raise errors.UnsupportedGpuError(unsupported_gpus=sorted(unsupported_gpus)) + + +def _walk_tasks(root_task: structures.TaskSpec) -> abc.Iterator[structures.TaskSpec]: + yield root_task + + component_spec = root_task.component_ref.spec + if component_spec and isinstance( + component_spec.implementation, structures.GraphImplementation + ): + for child_task in component_spec.implementation.graph.tasks.values(): + yield from _walk_tasks(child_task) + + +def _parse_accelerator_annotation(annotation_value: object) -> abc.Mapping[str, object]: + if isinstance(annotation_value, str): + try: + accelerators = json.loads(annotation_value) + except json.JSONDecodeError: + gpu_id, separator, quantity = annotation_value.rpartition(":") + if not separator or not gpu_id or not quantity: + raise errors.ApiValidationError( + f"GPU resource annotation {_GPU_ANNOTATION_KEY!r} must be a JSON object or a SkyPilot ':' string." + ) + accelerators = {gpu_id: quantity} + else: + accelerators = annotation_value + + if not isinstance(accelerators, dict) or not all( + isinstance(gpu_id, str) for gpu_id in accelerators + ): + raise errors.ApiValidationError( + f"GPU resource annotation {_GPU_ANNOTATION_KEY!r} must be a JSON object with GPU identifiers as keys." + ) + + return accelerators diff --git a/cloud_pipelines_backend/errors.py b/cloud_pipelines_backend/errors.py index dee8652..a979572 100644 --- a/cloud_pipelines_backend/errors.py +++ b/cloud_pipelines_backend/errors.py @@ -26,3 +26,10 @@ class ApiValidationError(Exception): """Base for all filter/annotation validation errors -> 422.""" pass + + +class UnsupportedGpuError(ApiValidationError): + def __init__(self, *, unsupported_gpus: list[str]): + self.unsupported_gpus = unsupported_gpus + gpu_list = ", ".join(repr(gpu) for gpu in unsupported_gpus) + super().__init__(f"Unsupported GPU resource(s): {gpu_list}.") diff --git a/tests/test_compute_resources.py b/tests/test_compute_resources.py new file mode 100644 index 0000000..d40fca1 --- /dev/null +++ b/tests/test_compute_resources.py @@ -0,0 +1,188 @@ +import fastapi +import pytest +from fastapi.testclient import TestClient +from sqlalchemy import orm + +from cloud_pipelines_backend import api_router +from cloud_pipelines_backend import api_server_sql +from cloud_pipelines_backend import backend_types_sql as bts +from cloud_pipelines_backend import component_structures as structures +from cloud_pipelines_backend import compute_resources +from cloud_pipelines_backend import database_ops +from cloud_pipelines_backend import errors + +_GPU_ANNOTATION_KEY = "cloud-pipelines.net/launchers/generic/resources.accelerators" + + +def _make_container_task( + *, gpu_id: str | None = None, name: str = "test-task" +) -> structures.TaskSpec: + annotations = {_GPU_ANNOTATION_KEY: f'{{"{gpu_id}": "1"}}'} if gpu_id else None + return structures.TaskSpec( + component_ref=structures.ComponentReference( + spec=structures.ComponentSpec( + name=name, + implementation=structures.ContainerImplementation( + container=structures.ContainerSpec(image="test-image:latest") + ), + ) + ), + annotations=annotations, + ) + + +def _make_graph_task(*children: structures.TaskSpec) -> structures.TaskSpec: + return structures.TaskSpec( + component_ref=structures.ComponentReference( + spec=structures.ComponentSpec( + name="test-pipeline", + implementation=structures.GraphImplementation( + graph=structures.GraphSpec( + tasks={ + f"task-{index}": task for index, task in enumerate(children) + } + ) + ), + ) + ) + ) + + +def _make_test_client() -> tuple[TestClient, orm.sessionmaker]: + engine = database_ops.create_db_engine(database_uri="sqlite://") + bts._TableBase.metadata.create_all(engine) + session_factory = orm.sessionmaker(engine) + + def get_session(): + with session_factory() as session: + yield session + + def get_user_details() -> api_router.UserDetails: + return api_router.UserDetails( + name="test-user", + permissions=api_router.Permissions(read=True, write=True, admin=False), + ) + + app = fastapi.FastAPI() + api_router._setup_routes_internal( + app=app, + get_session=get_session, + user_details_getter=get_user_details, + ) + return TestClient(app), session_factory + + +def test_get_pipeline_run_capabilities_returns_supported_and_deprecated_gpus() -> None: + response = compute_resources.get_pipeline_run_capabilities() + + assert response.gpus == ( + compute_resources.GpuResource( + id="NVIDIA-B300", display_name="NVIDIA B300", status="supported" + ), + compute_resources.GpuResource( + id="NVIDIA-H200", display_name="NVIDIA H200", status="deprecated" + ), + ) + + +def test_deprecated_gpu_is_valid() -> None: + compute_resources.validate_pipeline_gpu_resources( + _make_container_task(gpu_id="NVIDIA-H200") + ) + + +def test_deprecated_gpu_in_skypilot_string_format_is_valid() -> None: + task = _make_container_task() + task.annotations = {_GPU_ANNOTATION_KEY: "NVIDIA-H200:1"} + + compute_resources.validate_pipeline_gpu_resources(task) + + +def test_unsupported_gpu_in_skypilot_string_format_is_rejected() -> None: + task = _make_container_task() + task.annotations = {_GPU_ANNOTATION_KEY: "NVIDIA-A100:1"} + + with pytest.raises(errors.UnsupportedGpuError) as exc_info: + compute_resources.validate_pipeline_gpu_resources(task) + + assert exc_info.value.unsupported_gpus == ["NVIDIA-A100"] + + +def test_malformed_skypilot_string_format_is_rejected() -> None: + task = _make_container_task() + task.annotations = {_GPU_ANNOTATION_KEY: "NVIDIA-H200"} + + with pytest.raises(errors.ApiValidationError, match="SkyPilot"): + compute_resources.validate_pipeline_gpu_resources(task) + + +def test_supported_gpu_is_valid_in_nested_task() -> None: + compute_resources.validate_pipeline_gpu_resources( + _make_graph_task(_make_container_task(gpu_id="NVIDIA-B300")) + ) + + +def test_unsupported_gpu_in_nested_task_is_rejected() -> None: + with pytest.raises(errors.UnsupportedGpuError) as exc_info: + compute_resources.validate_pipeline_gpu_resources( + _make_graph_task(_make_container_task(gpu_id="NVIDIA-A100")) + ) + + assert exc_info.value.unsupported_gpus == ["NVIDIA-A100"] + + +def test_rejected_gpu_does_not_create_pipeline_run() -> None: + engine = database_ops.create_db_engine(database_uri="sqlite://") + bts._TableBase.metadata.create_all(engine) + + with orm.Session(engine) as session: + with pytest.raises(errors.UnsupportedGpuError): + api_server_sql.PipelineRunsApiService_Sql().create( + session=session, + root_task=_make_container_task(gpu_id="NVIDIA-A100"), + ) + + assert session.query(bts.PipelineRun).count() == 0 + assert session.query(bts.ExecutionNode).count() == 0 + + +def test_pipeline_run_capabilities_api_response() -> None: + client, _ = _make_test_client() + + response = client.get("/api/pipeline_runs/capabilities") + + assert response.status_code == 200 + assert response.json() == { + "gpus": [ + { + "id": "NVIDIA-B300", + "display_name": "NVIDIA B300", + "status": "supported", + }, + { + "id": "NVIDIA-H200", + "display_name": "NVIDIA H200", + "status": "deprecated", + }, + ] + } + + +def test_pipeline_run_api_returns_structured_unsupported_gpu_error() -> None: + client, session_factory = _make_test_client() + root_task = _make_container_task(gpu_id="NVIDIA-A100") + + response = client.post( + "/api/pipeline_runs/", + json={"root_task": root_task.to_json_dict()}, + ) + + assert response.status_code == 422 + assert response.json() == { + "reason": "unsupported_gpu", + "detail": "Unsupported GPU resource(s): 'NVIDIA-A100'.", + "unsupported_gpus": ["NVIDIA-A100"], + } + with session_factory() as session: + assert session.query(bts.PipelineRun).count() == 0 + assert session.query(bts.ExecutionNode).count() == 0