[Feature]: Allow listing available key resources such as gpu, region, and backends #2142 (WIP)

Moved the /gpus API to a separate routing; also moved GPU-related logic from runs.

Author: peterschmidt85

src/dstack/_internal/cli/commands/gpu.py

@@ -76,7 +76,7 @@ def _command(self, args: argparse.Namespace):
             status = contextlib.nullcontext()
 
         with status:
-            gpu_response = self.api.client.runs.get_gpus(
+            gpu_response = self.api.client.gpus.get_gpus(
                 self.api.project,
                 run_spec,
                 group_by=args.group_by,

src/dstack/_internal/server/app.py

@@ -29,6 +29,7 @@
     files,
     fleets,
     gateways,
+    gpus,
     instances,
     logs,
     metrics,
@@ -204,6 +205,7 @@ def register_routes(app: FastAPI, ui: bool = True):
     app.include_router(repos.router)
     app.include_router(runs.root_router)
     app.include_router(runs.project_router)
+    app.include_router(gpus.project_router)
     app.include_router(metrics.router)
     app.include_router(logs.router)
     app.include_router(secrets.router)

src/dstack/_internal/server/routers/gpus.py

@@ -0,0 +1,29 @@
+from typing import Tuple
+
+from fastapi import APIRouter, Depends
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from dstack._internal.server.db import get_session
+from dstack._internal.server.models import ProjectModel, UserModel
+from dstack._internal.server.schemas.gpus import GetRunGpusRequest, RunGpusResponse
+from dstack._internal.server.security.permissions import ProjectMember
+from dstack._internal.server.services.gpus import get_run_gpus_grouped
+from dstack._internal.server.utils.routers import get_base_api_additional_responses
+
+project_router = APIRouter(
+    prefix="/api/project/{project_name}/gpus",
+    tags=["gpus"],
+    responses=get_base_api_additional_responses(),
+)
+
+
+@project_router.post("/list", response_model=RunGpusResponse, response_model_exclude_none=True)
+async def get_run_gpus(
+    body: GetRunGpusRequest,
+    session: AsyncSession = Depends(get_session),
+    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
+) -> RunGpusResponse:
+    _, project = user_project
+    return await get_run_gpus_grouped(
+        session=session, project=project, run_spec=body.run_spec, group_by=body.group_by
+    )

src/dstack/_internal/server/routers/runs.py

@@ -10,11 +10,9 @@
 from dstack._internal.server.schemas.runs import (
     ApplyRunPlanRequest,
     DeleteRunsRequest,
-    GetRunGpusRequest,
     GetRunPlanRequest,
     GetRunRequest,
     ListRunsRequest,
-    RunGpusResponse,
     StopRunsRequest,
     SubmitRunRequest,
 )
@@ -181,18 +179,6 @@ async def delete_runs(
     await runs.delete_runs(session=session, project=project, runs_names=body.runs_names)
 
 
-@project_router.post("/gpus", response_model=RunGpusResponse, response_model_exclude_none=True)
-async def get_run_gpus(
-    body: GetRunGpusRequest,
-    session: AsyncSession = Depends(get_session),
-    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
-) -> RunGpusResponse:
-    _, project = user_project
-    return await runs.get_run_gpus_grouped(
-        session=session, project=project, run_spec=body.run_spec, group_by=body.group_by
-    )
-
-
 # apply_plan replaces submit_run since it can create new runs.
 @project_router.post("/submit", deprecated=True)
 async def submit_run(

src/dstack/_internal/server/schemas/gpus.py

@@ -0,0 +1,64 @@
+from typing import List, Literal, Optional
+
+import gpuhunt
+from pydantic import Field
+
+from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.common import CoreModel
+from dstack._internal.core.models.instances import InstanceAvailability
+from dstack._internal.core.models.resources import Range
+from dstack._internal.core.models.runs import RunSpec
+
+
+class BackendGpu(CoreModel):
+    """GPU specification from a backend offer."""
+
+    name: str
+    memory_mib: int
+    vendor: gpuhunt.AcceleratorVendor
+    availability: InstanceAvailability
+    spot: bool
+    count: int
+    price: float
+
+
+class BackendGpus(CoreModel):
+    """Backend GPU specifications."""
+
+    backend_type: BackendType
+    gpus: List[BackendGpu]
+    regions: List[str]
+
+
+class GetRunGpusRequest(CoreModel):
+    """Request for getting run GPUs with optional grouping."""
+
+    run_spec: RunSpec
+    group_by: Optional[List[Literal["backend", "region", "count"]]] = Field(
+        default=None,
+        description="List of fields to group by. Valid values: 'backend', 'region', 'count'",
+    )
+
+
+class GpuGroup(CoreModel):
+    """GPU group that can handle all grouping scenarios."""
+
+    name: str
+    memory_mib: int
+    vendor: gpuhunt.AcceleratorVendor
+    availability: List[InstanceAvailability]
+    spot: List[Literal["spot", "on-demand"]]
+    count: Range[int]
+    price: Range[float]
+    backends: Optional[List[BackendType]] = None
+    backend: Optional[BackendType] = None
+    regions: Optional[List[str]] = None
+    region: Optional[str] = None
+
+
+class RunGpusResponse(CoreModel):
+    """Response containing GPU specifications."""
+
+    gpus: List[GpuGroup] = Field(
+        description="List of GPU specifications, grouped according to the group_by parameter"
+    )

src/dstack/_internal/server/schemas/runs.py

@@ -1,14 +1,10 @@
 from datetime import datetime
-from typing import Annotated, List, Literal, Optional
+from typing import Annotated, List, Optional
 from uuid import UUID
 
-import gpuhunt
 from pydantic import Field
 
-from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import CoreModel
-from dstack._internal.core.models.instances import InstanceAvailability
-from dstack._internal.core.models.resources import Range
 from dstack._internal.core.models.runs import ApplyRunPlanInput, RunSpec
 
 
@@ -68,57 +64,3 @@ class StopRunsRequest(CoreModel):
 
 class DeleteRunsRequest(CoreModel):
     runs_names: List[str]
-
-
-class BackendGpu(CoreModel):
-    """GPU specification from a backend offer."""
-
-    name: str
-    memory_mib: int
-    vendor: gpuhunt.AcceleratorVendor
-    availability: InstanceAvailability
-    spot: bool
-    count: int
-    price: float
-
-
-class BackendGpus(CoreModel):
-    """Backend GPU specifications."""
-
-    backend_type: BackendType
-    gpus: List[BackendGpu]
-    regions: List[str]
-
-
-class GetRunGpusRequest(CoreModel):
-    """Request for getting run GPUs with optional grouping."""
-
-    run_spec: RunSpec
-    group_by: Optional[List[Literal["backend", "region", "count"]]] = Field(
-        default=None,
-        description="List of fields to group by. Valid values: 'backend', 'region', 'count'",
-    )
-
-
-class GpuGroup(CoreModel):
-    """GPU group that can handle all grouping scenarios."""
-
-    name: str
-    memory_mib: int
-    vendor: gpuhunt.AcceleratorVendor
-    availability: List[InstanceAvailability]
-    spot: List[Literal["spot", "on-demand"]]
-    count: Range[int]
-    price: Range[float]
-    backends: Optional[List[BackendType]] = None
-    backend: Optional[BackendType] = None
-    regions: Optional[List[str]] = None
-    region: Optional[str] = None
-
-
-class RunGpusResponse(CoreModel):
-    """Response containing GPU specifications."""
-
-    gpus: List[GpuGroup] = Field(
-        description="List of GPU specifications, grouped according to the group_by parameter"
-    )