dstack 0.19.18__py3-none-any.whl → 0.19.20__py3-none-any.whl

dstack/_internal/server/routers/backends.py

@@ -27,7 +27,10 @@ from dstack._internal.server.services.config import (
     get_backend_config_yaml,
     update_backend_config_yaml,
 )
-from dstack._internal.server.utils.routers import get_base_api_additional_responses
+from dstack._internal.server.utils.routers import (
+    CustomORJSONResponse,
+    get_base_api_additional_responses,
+)
 
 root_router = APIRouter(
     prefix="/api/backends",
@@ -41,35 +44,37 @@ project_router = APIRouter(
 )
 
 
-@root_router.post("/list_types")
-async def list_backend_types() -> List[BackendType]:
-    return dstack._internal.core.backends.configurators.list_available_backend_types()
+@root_router.post("/list_types", response_model=List[BackendType])
+async def list_backend_types():
+    return CustomORJSONResponse(
+        dstack._internal.core.backends.configurators.list_available_backend_types()
+    )
 
 
-@project_router.post("/create")
+@project_router.post("/create", response_model=AnyBackendConfigWithCreds)
 async def create_backend(
     body: AnyBackendConfigWithCreds,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectAdmin()),
-) -> AnyBackendConfigWithCreds:
+):
     _, project = user_project
     config = await backends.create_backend(session=session, project=project, config=body)
     if settings.SERVER_CONFIG_ENABLED:
         await ServerConfigManager().sync_config(session=session)
-    return config
+    return CustomORJSONResponse(config)
 
 
-@project_router.post("/update")
+@project_router.post("/update", response_model=AnyBackendConfigWithCreds)
 async def update_backend(
     body: AnyBackendConfigWithCreds,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectAdmin()),
-) -> AnyBackendConfigWithCreds:
+):
     _, project = user_project
     config = await backends.update_backend(session=session, project=project, config=body)
     if settings.SERVER_CONFIG_ENABLED:
         await ServerConfigManager().sync_config(session=session)
-    return config
+    return CustomORJSONResponse(config)
 
 
 @project_router.post("/delete")
@@ -86,16 +91,16 @@ async def delete_backends(
         await ServerConfigManager().sync_config(session=session)
 
 
-@project_router.post("/{backend_name}/config_info")
+@project_router.post("/{backend_name}/config_info", response_model=AnyBackendConfigWithCreds)
 async def get_backend_config_info(
     backend_name: BackendType,
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectAdmin()),
-) -> AnyBackendConfigWithCreds:
+):
     _, project = user_project
     config = await backends.get_backend_config(project=project, backend_type=backend_name)
     if config is None:
         raise ResourceNotExistsError()
-    return config
+    return CustomORJSONResponse(config)
 
 
 @project_router.post("/create_yaml")
@@ -126,10 +131,12 @@ async def update_backend_yaml(
     )
 
 
-@project_router.post("/{backend_name}/get_yaml")
+@project_router.post("/{backend_name}/get_yaml", response_model=BackendInfoYAML)
 async def get_backend_yaml(
     backend_name: BackendType,
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectAdmin()),
-) -> BackendInfoYAML:
+):
     _, project = user_project
-    return await get_backend_config_yaml(project=project, backend_type=backend_name)
+    return CustomORJSONResponse(
+        await get_backend_config_yaml(project=project, backend_type=backend_name)
+    )

dstack/_internal/server/routers/files.py

@@ -12,6 +12,7 @@ from dstack._internal.server.security.permissions import Authenticated
 from dstack._internal.server.services import files
 from dstack._internal.server.settings import SERVER_CODE_UPLOAD_LIMIT
 from dstack._internal.server.utils.routers import (
+    CustomORJSONResponse,
     get_base_api_additional_responses,
     get_request_size,
 )
@@ -24,12 +25,12 @@ router = APIRouter(
 )
 
 
-@router.post("/get_archive_by_hash")
+@router.post("/get_archive_by_hash", response_model=FileArchive)
 async def get_archive_by_hash(
     body: GetFileArchiveByHashRequest,
     session: Annotated[AsyncSession, Depends(get_session)],
     user: Annotated[UserModel, Depends(Authenticated())],
-) -> FileArchive:
+):
     archive = await files.get_archive_by_hash(
         session=session,
         user=user,
@@ -37,16 +38,16 @@ async def get_archive_by_hash(
     )
     if archive is None:
         raise ResourceNotExistsError()
-    return archive
+    return CustomORJSONResponse(archive)
 
 
-@router.post("/upload_archive")
+@router.post("/upload_archive", response_model=FileArchive)
 async def upload_archive(
     request: Request,
     file: UploadFile,
     session: Annotated[AsyncSession, Depends(get_session)],
     user: Annotated[UserModel, Depends(Authenticated())],
-) -> FileArchive:
+):
     request_size = get_request_size(request)
     if SERVER_CODE_UPLOAD_LIMIT > 0 and request_size > SERVER_CODE_UPLOAD_LIMIT:
         diff_size_fmt = sizeof_fmt(request_size)
@@ -64,4 +65,4 @@ async def upload_archive(
         user=user,
         file=file,
     )
-    return archive
+    return CustomORJSONResponse(archive)

dstack/_internal/server/routers/fleets.py

@@ -18,7 +18,10 @@ from dstack._internal.server.schemas.fleets import (
     ListFleetsRequest,
 )
 from dstack._internal.server.security.permissions import Authenticated, ProjectMember
-from dstack._internal.server.utils.routers import get_base_api_additional_responses
+from dstack._internal.server.utils.routers import (
+    CustomORJSONResponse,
+    get_base_api_additional_responses,
+)
 
 root_router = APIRouter(
     prefix="/api/fleets",
@@ -32,12 +35,12 @@ project_router = APIRouter(
 )
 
 
-@root_router.post("/list")
+@root_router.post("/list", response_model=List[Fleet])
 async def list_fleets(
     body: ListFleetsRequest,
     session: AsyncSession = Depends(get_session),
     user: UserModel = Depends(Authenticated()),
-) -> List[Fleet]:
+):
     """
     Returns all fleets and instances within them visible to user sorted by descending `created_at`.
     `project_name` and `only_active` can be specified as filters.
@@ -45,36 +48,40 @@ async def list_fleets(
     The results are paginated. To get the next page, pass `created_at` and `id` of
     the last fleet from the previous page as `prev_created_at` and `prev_id`.
     """
-    return await fleets_services.list_fleets(
-        session=session,
-        user=user,
-        project_name=body.project_name,
-        only_active=body.only_active,
-        prev_created_at=body.prev_created_at,
-        prev_id=body.prev_id,
-        limit=body.limit,
-        ascending=body.ascending,
+    return CustomORJSONResponse(
+        await fleets_services.list_fleets(
+            session=session,
+            user=user,
+            project_name=body.project_name,
+            only_active=body.only_active,
+            prev_created_at=body.prev_created_at,
+            prev_id=body.prev_id,
+            limit=body.limit,
+            ascending=body.ascending,
+        )
     )
 
 
-@project_router.post("/list")
+@project_router.post("/list", response_model=List[Fleet])
 async def list_project_fleets(
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
-) -> List[Fleet]:
+):
     """
     Returns all fleets in the project.
     """
     _, project = user_project
-    return await fleets_services.list_project_fleets(session=session, project=project)
+    return CustomORJSONResponse(
+        await fleets_services.list_project_fleets(session=session, project=project)
+    )
 
 
-@project_router.post("/get")
+@project_router.post("/get", response_model=Fleet)
 async def get_fleet(
     body: GetFleetRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
-) -> Fleet:
+):
     """
     Returns a fleet given `name` or `id`.
     If given `name`, does not return deleted fleets.
@@ -86,15 +93,15 @@ async def get_fleet(
     )
     if fleet is None:
         raise ResourceNotExistsError()
-    return fleet
+    return CustomORJSONResponse(fleet)
 
 
-@project_router.post("/get_plan")
+@project_router.post("/get_plan", response_model=FleetPlan)
 async def get_plan(
     body: GetFleetPlanRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
-) -> FleetPlan:
+):
     """
     Returns a fleet plan for the given fleet configuration.
     """
@@ -105,45 +112,49 @@ async def get_plan(
         user=user,
         spec=body.spec,
     )
-    return plan
+    return CustomORJSONResponse(plan)
 
 
-@project_router.post("/apply")
+@project_router.post("/apply", response_model=Fleet)
 async def apply_plan(
     body: ApplyFleetPlanRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
-) -> Fleet:
+):
     """
     Creates a new fleet or updates an existing fleet.
     Errors if the expected current resource from the plan does not match the current resource.
     Use `force: true` to apply even if the current resource does not match.
     """
     user, project = user_project
-    return await fleets_services.apply_plan(
-        session=session,
-        user=user,
-        project=project,
-        plan=body.plan,
-        force=body.force,
+    return CustomORJSONResponse(
+        await fleets_services.apply_plan(
+            session=session,
+            user=user,
+            project=project,
+            plan=body.plan,
+            force=body.force,
+        )
     )
 
 
-@project_router.post("/create")
+@project_router.post("/create", response_model=Fleet)
 async def create_fleet(
     body: CreateFleetRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
-) -> Fleet:
+):
     """
     Creates a fleet given a fleet configuration.
     """
     user, project = user_project
-    return await fleets_services.create_fleet(
-        session=session,
-        project=project,
-        user=user,
-        spec=body.spec,
+    return CustomORJSONResponse(
+        await fleets_services.create_fleet(
+            session=session,
+            project=project,
+            user=user,
+            spec=body.spec,
+        )
     )
 
 

dstack/_internal/server/routers/gateways.py

@@ -13,7 +13,10 @@ from dstack._internal.server.security.permissions import (
     ProjectAdmin,
     ProjectMemberOrPublicAccess,
 )
-from dstack._internal.server.utils.routers import get_base_api_additional_responses
+from dstack._internal.server.utils.routers import (
+    CustomORJSONResponse,
+    get_base_api_additional_responses,
+)
 
 router = APIRouter(
     prefix="/api/project/{project_name}/gateways",
@@ -22,40 +25,44 @@ router = APIRouter(
 )
 
 
-@router.post("/list")
+@router.post("/list", response_model=List[models.Gateway])
 async def list_gateways(
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMemberOrPublicAccess()),
-) -> List[models.Gateway]:
+):
     _, project = user_project
-    return await gateways.list_project_gateways(session=session, project=project)
+    return CustomORJSONResponse(
+        await gateways.list_project_gateways(session=session, project=project)
+    )
 
 
-@router.post("/get")
+@router.post("/get", response_model=models.Gateway)
 async def get_gateway(
     body: schemas.GetGatewayRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMemberOrPublicAccess()),
-) -> models.Gateway:
+):
     _, project = user_project
     gateway = await gateways.get_gateway_by_name(session=session, project=project, name=body.name)
     if gateway is None:
         raise ResourceNotExistsError()
-    return gateway
+    return CustomORJSONResponse(gateway)
 
 
-@router.post("/create")
+@router.post("/create", response_model=models.Gateway)
 async def create_gateway(
     body: schemas.CreateGatewayRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectAdmin()),
-) -> models.Gateway:
+):
     user, project = user_project
-    return await gateways.create_gateway(
-        session=session,
-        user=user,
-        project=project,
-        configuration=body.configuration,
+    return CustomORJSONResponse(
+        await gateways.create_gateway(
+            session=session,
+            user=user,
+            project=project,
+            configuration=body.configuration,
+        )
     )
 
 
@@ -83,13 +90,15 @@ async def set_default_gateway(
     await gateways.set_default_gateway(session=session, project=project, name=body.name)
 
 
-@router.post("/set_wildcard_domain")
+@router.post("/set_wildcard_domain", response_model=models.Gateway)
 async def set_gateway_wildcard_domain(
     body: schemas.SetWildcardDomainRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectAdmin()),
-) -> models.Gateway:
+):
     _, project = user_project
-    return await gateways.set_gateway_wildcard_domain(
-        session=session, project=project, name=body.name, wildcard_domain=body.wildcard_domain
+    return CustomORJSONResponse(
+        await gateways.set_gateway_wildcard_domain(
+            session=session, project=project, name=body.name, wildcard_domain=body.wildcard_domain
+        )
     )

dstack/_internal/server/routers/instances.py

@@ -9,7 +9,10 @@ from dstack._internal.server.db import get_session
 from dstack._internal.server.models import UserModel
 from dstack._internal.server.schemas.instances import ListInstancesRequest
 from dstack._internal.server.security.permissions import Authenticated
-from dstack._internal.server.utils.routers import get_base_api_additional_responses
+from dstack._internal.server.utils.routers import (
+    CustomORJSONResponse,
+    get_base_api_additional_responses,
+)
 
 root_router = APIRouter(
     prefix="/api/instances",
@@ -18,12 +21,12 @@ root_router = APIRouter(
 )
 
 
-@root_router.post("/list")
+@root_router.post("/list", response_model=List[Instance])
 async def list_instances(
     body: ListInstancesRequest,
     session: AsyncSession = Depends(get_session),
     user: UserModel = Depends(Authenticated()),
-) -> List[Instance]:
+):
     """
     Returns all instances visible to user sorted by descending `created_at`.
     `project_names` and `fleet_ids` can be specified as filters.
@@ -31,14 +34,16 @@ async def list_instances(
     The results are paginated. To get the next page, pass `created_at` and `id` of
     the last instance from the previous page as `prev_created_at` and `prev_id`.
     """
-    return await instances.list_user_instances(
-        session=session,
-        user=user,
-        project_names=body.project_names,
-        fleet_ids=body.fleet_ids,
-        only_active=body.only_active,
-        prev_created_at=body.prev_created_at,
-        prev_id=body.prev_id,
-        limit=body.limit,
-        ascending=body.ascending,
+    return CustomORJSONResponse(
+        await instances.list_user_instances(
+            session=session,
+            user=user,
+            project_names=body.project_names,
+            fleet_ids=body.fleet_ids,
+            only_active=body.only_active,
+            prev_created_at=body.prev_created_at,
+            prev_id=body.prev_id,
+            limit=body.limit,
+            ascending=body.ascending,
+        )
     )

dstack/_internal/server/routers/logs.py

@@ -7,7 +7,10 @@ from dstack._internal.server.models import ProjectModel, UserModel
 from dstack._internal.server.schemas.logs import PollLogsRequest
 from dstack._internal.server.security.permissions import ProjectMember
 from dstack._internal.server.services import logs
-from dstack._internal.server.utils.routers import get_base_api_additional_responses
+from dstack._internal.server.utils.routers import (
+    CustomORJSONResponse,
+    get_base_api_additional_responses,
+)
 
 router = APIRouter(
     prefix="/api/project/{project_name}/logs",
@@ -18,13 +21,14 @@ router = APIRouter(
 
 @router.post(
     "/poll",
+    response_model=JobSubmissionLogs,
 )
 async def poll_logs(
     body: PollLogsRequest,
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
-) -> JobSubmissionLogs:
+):
     _, project = user_project
     # The runner guarantees logs have different timestamps if throughput < 1k logs / sec.
     # Otherwise, some logs with duplicated timestamps may be filtered out.
     # This limitation is imposed by cloud log services that support up to millisecond timestamp resolution.
-    return await logs.poll_logs_async(project=project, request=body)
+    return CustomORJSONResponse(await logs.poll_logs_async(project=project, request=body))

dstack/_internal/server/routers/metrics.py

@@ -11,7 +11,10 @@ from dstack._internal.server.models import ProjectModel, UserModel
 from dstack._internal.server.security.permissions import ProjectMember
 from dstack._internal.server.services import metrics
 from dstack._internal.server.services.jobs import get_run_job_model
-from dstack._internal.server.utils.routers import get_base_api_additional_responses
+from dstack._internal.server.utils.routers import (
+    CustomORJSONResponse,
+    get_base_api_additional_responses,
+)
 
 router = APIRouter(
     prefix="/api/project/{project_name}/metrics",
@@ -22,6 +25,7 @@ router = APIRouter(
 
 @router.get(
     "/job/{run_name}",
+    response_model=JobMetrics,
 )
 async def get_job_metrics(
     run_name: str,
@@ -32,7 +36,7 @@ async def get_job_metrics(
     before: Optional[datetime] = None,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
-) -> JobMetrics:
+):
     """
     Returns job-level metrics such as hardware utilization
     given `run_name`, `replica_num`, and `job_num`.
@@ -63,10 +67,12 @@ async def get_job_metrics(
     if job_model is None:
         raise ResourceNotExistsError("Found no job with given parameters")
 
-    return await metrics.get_job_metrics(
-        session=session,
-        job_model=job_model,
-        limit=limit,
-        after=after,
-        before=before,
+    return CustomORJSONResponse(
+        await metrics.get_job_metrics(
+            session=session,
+            job_model=job_model,
+            limit=limit,
+            after=after,
+            before=before,
+        )
     )