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

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

dstack/_internal/server/routers/projects.py

@@ -23,7 +23,10 @@ from dstack._internal.server.security.permissions import (
     ProjectMemberOrPublicAccess,
 )
 from dstack._internal.server.services import projects
-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/projects",
@@ -32,30 +35,34 @@ router = APIRouter(
 )
 
 
-@router.post("/list")
+@router.post("/list", response_model=List[Project])
 async def list_projects(
     session: AsyncSession = Depends(get_session),
     user: UserModel = Depends(Authenticated()),
-) -> List[Project]:
+):
     """
     Returns all projects visible to user sorted by descending `created_at`.
 
     `members` and `backends` are always empty - call `/api/projects/{project_name}/get` to retrieve them.
     """
-    return await projects.list_user_accessible_projects(session=session, user=user)
+    return CustomORJSONResponse(
+        await projects.list_user_accessible_projects(session=session, user=user)
+    )
 
 
-@router.post("/create")
+@router.post("/create", response_model=Project)
 async def create_project(
     body: CreateProjectRequest,
     session: AsyncSession = Depends(get_session),
     user: UserModel = Depends(Authenticated()),
-) -> Project:
-    return await projects.create_project(
-        session=session,
-        user=user,
-        project_name=body.project_name,
-        is_public=body.is_public,
+):
+    return CustomORJSONResponse(
+        await projects.create_project(
+            session=session,
+            user=user,
+            project_name=body.project_name,
+            is_public=body.is_public,
+        )
     )
 
 
@@ -72,23 +79,24 @@ async def delete_projects(
     )
 
 
-@router.post("/{project_name}/get")
+@router.post("/{project_name}/get", response_model=Project)
 async def get_project(
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMemberOrPublicAccess()),
-) -> Project:
+):
     _, project = user_project
-    return projects.project_model_to_project(project)
+    return CustomORJSONResponse(projects.project_model_to_project(project))
 
 
 @router.post(
     "/{project_name}/set_members",
+    response_model=Project,
 )
 async def set_project_members(
     body: SetProjectMembersRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectManager()),
-) -> Project:
+):
     user, project = user_project
     await projects.set_project_members(
         session=session,
@@ -97,17 +105,18 @@ async def set_project_members(
         members=body.members,
     )
     await session.refresh(project)
-    return projects.project_model_to_project(project)
+    return CustomORJSONResponse(projects.project_model_to_project(project))
 
 
 @router.post(
     "/{project_name}/add_members",
+    response_model=Project,
 )
 async def add_project_members(
     body: AddProjectMemberRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectManagerOrPublicProject()),
-) -> Project:
+):
     user, project = user_project
     await projects.add_project_members(
         session=session,
@@ -116,17 +125,18 @@ async def add_project_members(
         members=body.members,
     )
     await session.refresh(project)
-    return projects.project_model_to_project(project)
+    return CustomORJSONResponse(projects.project_model_to_project(project))
 
 
 @router.post(
     "/{project_name}/remove_members",
+    response_model=Project,
 )
 async def remove_project_members(
     body: RemoveProjectMemberRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectManagerOrSelfLeave()),
-) -> Project:
+):
     user, project = user_project
     await projects.remove_project_members(
         session=session,
@@ -135,17 +145,18 @@ async def remove_project_members(
         usernames=body.usernames,
     )
     await session.refresh(project)
-    return projects.project_model_to_project(project)
+    return CustomORJSONResponse(projects.project_model_to_project(project))
 
 
 @router.post(
     "/{project_name}/update",
+    response_model=Project,
 )
 async def update_project(
     body: UpdateProjectRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectAdmin()),
-) -> Project:
+):
     user, project = user_project
     await projects.update_project(
         session=session,
@@ -154,4 +165,4 @@ async def update_project(
         is_public=body.is_public,
     )
     await session.refresh(project)
-    return projects.project_model_to_project(project)
+    return CustomORJSONResponse(projects.project_model_to_project(project))

dstack/_internal/server/routers/repos.py

@@ -16,6 +16,7 @@ from dstack._internal.server.security.permissions import ProjectMember
 from dstack._internal.server.services import repos
 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,
 )
@@ -28,21 +29,21 @@ router = APIRouter(
 )
 
 
-@router.post("/list")
+@router.post("/list", response_model=List[RepoHead])
 async def list_repos(
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
-) -> List[RepoHead]:
+):
     _, project = user_project
-    return await repos.list_repos(session=session, project=project)
+    return CustomORJSONResponse(await repos.list_repos(session=session, project=project))
 
 
-@router.post("/get")
+@router.post("/get", response_model=RepoHeadWithCreds)
 async def get_repo(
     body: GetRepoRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
-) -> RepoHeadWithCreds:
+):
     user, project = user_project
     repo = await repos.get_repo(
         session=session,
@@ -53,7 +54,7 @@ async def get_repo(
     )
     if repo is None:
         raise ResourceNotExistsError()
-    return repo
+    return CustomORJSONResponse(repo)
 
 
 @router.post("/init")