Add pipeline run GPU capabilities and validation

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)
     )

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():

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 '<gpu-id>:<quantity>' 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

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}.")

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