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

dstack/_internal/server/routers/runs.py

@@ -18,7 +18,10 @@ from dstack._internal.server.schemas.runs import (
 )
 from dstack._internal.server.security.permissions import Authenticated, ProjectMember
 from dstack._internal.server.services import runs
-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/runs",
@@ -32,12 +35,15 @@ project_router = APIRouter(
 )
 
 
-@root_router.post("/list")
+@root_router.post(
+    "/list",
+    response_model=List[Run],
+)
 async def list_runs(
     body: ListRunsRequest,
     session: AsyncSession = Depends(get_session),
     user: UserModel = Depends(Authenticated()),
-) -> List[Run]:
+):
     """
     Returns all runs visible to user sorted by descending `submitted_at`.
     `project_name`, `repo_id`, `username`, and `only_active` can be specified as filters.
@@ -47,26 +53,33 @@ async def list_runs(
     The results are paginated. To get the next page, pass `submitted_at` and `id` of
     the last run from the previous page as `prev_submitted_at` and `prev_run_id`.
     """
-    return await runs.list_user_runs(
-        session=session,
-        user=user,
-        project_name=body.project_name,
-        repo_id=body.repo_id,
-        username=body.username,
-        only_active=body.only_active,
-        prev_submitted_at=body.prev_submitted_at,
-        prev_run_id=body.prev_run_id,
-        limit=body.limit,
-        ascending=body.ascending,
+    return CustomORJSONResponse(
+        await runs.list_user_runs(
+            session=session,
+            user=user,
+            project_name=body.project_name,
+            repo_id=body.repo_id,
+            username=body.username,
+            only_active=body.only_active,
+            include_jobs=body.include_jobs,
+            job_submissions_limit=body.job_submissions_limit,
+            prev_submitted_at=body.prev_submitted_at,
+            prev_run_id=body.prev_run_id,
+            limit=body.limit,
+            ascending=body.ascending,
+        )
     )
 
 
-@project_router.post("/get")
+@project_router.post(
+    "/get",
+    response_model=Run,
+)
 async def get_run(
     body: GetRunRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
-) -> Run:
+):
     """
     Returns a run given `run_name` or `id`.
     If given `run_name`, does not return deleted runs.
@@ -81,15 +94,18 @@ async def get_run(
     )
     if run is None:
         raise ResourceNotExistsError("Run not found")
-    return run
+    return CustomORJSONResponse(run)
 
 
-@project_router.post("/get_plan")
+@project_router.post(
+    "/get_plan",
+    response_model=RunPlan,
+)
 async def get_plan(
     body: GetRunPlanRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
-) -> RunPlan:
+):
     """
     Returns a run plan for the given run spec.
     This is an optional step before calling `/apply`.
@@ -102,15 +118,18 @@ async def get_plan(
         run_spec=body.run_spec,
         max_offers=body.max_offers,
     )
-    return run_plan
+    return CustomORJSONResponse(run_plan)
 
 
-@project_router.post("/apply")
+@project_router.post(
+    "/apply",
+    response_model=Run,
+)
 async def apply_plan(
     body: ApplyRunPlanRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
-) -> Run:
+):
     """
     Creates a new run or updates an existing run.
     Errors if the expected current resource from the plan does not match the current resource.
@@ -118,12 +137,14 @@ async def apply_plan(
     If the existing run is active and cannot be updated, it must be stopped first.
     """
     user, project = user_project
-    return await runs.apply_plan(
-        session=session,
-        user=user,
-        project=project,
-        plan=body.plan,
-        force=body.force,
+    return CustomORJSONResponse(
+        await runs.apply_plan(
+            session=session,
+            user=user,
+            project=project,
+            plan=body.plan,
+            force=body.force,
+        )
     )
 
 

dstack/_internal/server/routers/secrets.py

@@ -14,6 +14,7 @@ from dstack._internal.server.schemas.secrets import (
 )
 from dstack._internal.server.security.permissions import ProjectAdmin
 from dstack._internal.server.services import secrets as secrets_services
+from dstack._internal.server.utils.routers import CustomORJSONResponse
 
 router = APIRouter(
     prefix="/api/project/{project_name}/secrets",
@@ -21,24 +22,26 @@ router = APIRouter(
 )
 
 
-@router.post("/list")
+@router.post("/list", response_model=List[Secret])
 async def list_secrets(
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectAdmin()),
-) -> List[Secret]:
+):
     _, project = user_project
-    return await secrets_services.list_secrets(
-        session=session,
-        project=project,
+    return CustomORJSONResponse(
+        await secrets_services.list_secrets(
+            session=session,
+            project=project,
+        )
     )
 
 
-@router.post("/get")
+@router.post("/get", response_model=Secret)
 async def get_secret(
     body: GetSecretRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectAdmin()),
-) -> Secret:
+):
     _, project = user_project
     secret = await secrets_services.get_secret(
         session=session,
@@ -47,21 +50,23 @@ async def get_secret(
     )
     if secret is None:
         raise ResourceNotExistsError()
-    return secret
+    return CustomORJSONResponse(secret)
 
 
-@router.post("/create_or_update")
+@router.post("/create_or_update", response_model=Secret)
 async def create_or_update_secret(
     body: CreateOrUpdateSecretRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectAdmin()),
-) -> Secret:
+):
     _, project = user_project
-    return await secrets_services.create_or_update_secret(
-        session=session,
-        project=project,
-        name=body.name,
-        value=body.value,
+    return CustomORJSONResponse(
+        await secrets_services.create_or_update_secret(
+            session=session,
+            project=project,
+            name=body.name,
+            value=body.value,
+        )
     )
 
 

dstack/_internal/server/routers/server.py

@@ -2,6 +2,7 @@ from fastapi import APIRouter
 
 from dstack._internal import settings
 from dstack._internal.core.models.server import ServerInfo
+from dstack._internal.server.utils.routers import CustomORJSONResponse
 
 router = APIRouter(
     prefix="/api/server",
@@ -9,8 +10,10 @@ router = APIRouter(
 )
 
 
-@router.post("/get_info")
-async def get_server_info() -> ServerInfo:
-    return ServerInfo(
-        server_version=settings.DSTACK_VERSION,
+@router.post("/get_info", response_model=ServerInfo)
+async def get_server_info():
+    return CustomORJSONResponse(
+        ServerInfo(
+            server_version=settings.DSTACK_VERSION,
+        )
     )

dstack/_internal/server/routers/users.py

@@ -16,7 +16,10 @@ from dstack._internal.server.schemas.users import (
 )
 from dstack._internal.server.security.permissions import Authenticated, GlobalAdmin
 from dstack._internal.server.services import users
-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/users",
@@ -25,41 +28,41 @@ router = APIRouter(
 )
 
 
-@router.post("/list")
+@router.post("/list", response_model=List[User])
 async def list_users(
     session: AsyncSession = Depends(get_session),
     user: UserModel = Depends(Authenticated()),
-) -> List[User]:
-    return await users.list_users_for_user(session=session, user=user)
+):
+    return CustomORJSONResponse(await users.list_users_for_user(session=session, user=user))
 
 
-@router.post("/get_my_user")
+@router.post("/get_my_user", response_model=User)
 async def get_my_user(
     user: UserModel = Depends(Authenticated()),
-) -> User:
-    return users.user_model_to_user(user)
+):
+    return CustomORJSONResponse(users.user_model_to_user(user))
 
 
-@router.post("/get_user")
+@router.post("/get_user", response_model=UserWithCreds)
 async def get_user(
     body: GetUserRequest,
     session: AsyncSession = Depends(get_session),
     user: UserModel = Depends(Authenticated()),
-) -> UserWithCreds:
+):
     res = await users.get_user_with_creds_by_name(
         session=session, current_user=user, username=body.username
     )
     if res is None:
         raise ResourceNotExistsError()
-    return res
+    return CustomORJSONResponse(res)
 
 
-@router.post("/create")
+@router.post("/create", response_model=User)
 async def create_user(
     body: CreateUserRequest,
     session: AsyncSession = Depends(get_session),
     user: UserModel = Depends(GlobalAdmin()),
-) -> User:
+):
     res = await users.create_user(
         session=session,
         username=body.username,
@@ -67,15 +70,15 @@ async def create_user(
         email=body.email,
         active=body.active,
     )
-    return users.user_model_to_user(res)
+    return CustomORJSONResponse(users.user_model_to_user(res))
 
 
-@router.post("/update")
+@router.post("/update", response_model=User)
 async def update_user(
     body: UpdateUserRequest,
     session: AsyncSession = Depends(get_session),
     user: UserModel = Depends(GlobalAdmin()),
-) -> User:
+):
     res = await users.update_user(
         session=session,
         username=body.username,
@@ -85,19 +88,19 @@ async def update_user(
     )
     if res is None:
         raise ResourceNotExistsError()
-    return users.user_model_to_user(res)
+    return CustomORJSONResponse(users.user_model_to_user(res))
 
 
-@router.post("/refresh_token")
+@router.post("/refresh_token", response_model=UserWithCreds)
 async def refresh_token(
     body: RefreshTokenRequest,
     session: AsyncSession = Depends(get_session),
     user: UserModel = Depends(Authenticated()),
-) -> UserWithCreds:
+):
     res = await users.refresh_user_token(session=session, user=user, username=body.username)
     if res is None:
         raise ResourceNotExistsError()
-    return users.user_model_to_user_with_creds(res)
+    return CustomORJSONResponse(users.user_model_to_user_with_creds(res))
 
 
 @router.post("/delete")

dstack/_internal/server/routers/volumes.py

@@ -15,7 +15,10 @@ from dstack._internal.server.schemas.volumes import (
     ListVolumesRequest,
 )
 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/volumes",
@@ -25,12 +28,12 @@ root_router = APIRouter(
 project_router = APIRouter(prefix="/api/project/{project_name}/volumes", tags=["volumes"])
 
 
-@root_router.post("/list")
+@root_router.post("/list", response_model=List[Volume])
 async def list_volumes(
     body: ListVolumesRequest,
     session: AsyncSession = Depends(get_session),
     user: UserModel = Depends(Authenticated()),
-) -> List[Volume]:
+):
     """
     Returns all volumes visible to user sorted by descending `created_at`.
     `project_name` and `only_active` can be specified as filters.
@@ -38,36 +41,40 @@ async def list_volumes(
     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 volumes_services.list_volumes(
-        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 volumes_services.list_volumes(
+            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[Volume])
 async def list_project_volumes(
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
-) -> List[Volume]:
+):
     """
     Returns all volumes in the project.
     """
     _, project = user_project
-    return await volumes_services.list_project_volumes(session=session, project=project)
+    return CustomORJSONResponse(
+        await volumes_services.list_project_volumes(session=session, project=project)
+    )
 
 
-@project_router.post("/get")
+@project_router.post("/get", response_model=Volume)
 async def get_volume(
     body: GetVolumeRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
-) -> Volume:
+):
     """
     Returns a volume given a volume name.
     """
@@ -77,24 +84,26 @@ async def get_volume(
     )
     if volume is None:
         raise ResourceNotExistsError()
-    return volume
+    return CustomORJSONResponse(volume)
 
 
-@project_router.post("/create")
+@project_router.post("/create", response_model=Volume)
 async def create_volume(
     body: CreateVolumeRequest,
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
-) -> Volume:
+):
     """
     Creates a volume given a volume configuration.
     """
     user, project = user_project
-    return await volumes_services.create_volume(
-        session=session,
-        project=project,
-        user=user,
-        configuration=body.configuration,
+    return CustomORJSONResponse(
+        await volumes_services.create_volume(
+            session=session,
+            project=project,
+            user=user,
+            configuration=body.configuration,
+        )
     )
 
 

dstack/_internal/server/schemas/logs.py

@@ -9,8 +9,8 @@ from dstack._internal.core.models.common import CoreModel
 class PollLogsRequest(CoreModel):
     run_name: str
     job_submission_id: UUID4
-    start_time: Optional[datetime]
-    end_time: Optional[datetime]
+    start_time: Optional[datetime] = None
+    end_time: Optional[datetime] = None
     descending: bool = False
     next_token: Optional[str] = None
     limit: int = Field(100, ge=0, le=1000)

dstack/_internal/server/schemas/runs.py

@@ -9,12 +9,24 @@ from dstack._internal.core.models.runs import ApplyRunPlanInput, RunSpec
 
 
 class ListRunsRequest(CoreModel):
-    project_name: Optional[str]
-    repo_id: Optional[str]
-    username: Optional[str]
+    project_name: Optional[str] = None
+    repo_id: Optional[str] = None
+    username: Optional[str] = None
     only_active: bool = False
-    prev_submitted_at: Optional[datetime]
-    prev_run_id: Optional[UUID]
+    include_jobs: bool = Field(
+        True,
+        description=("Whether to include `jobs` in the response"),
+    )
+    job_submissions_limit: Optional[int] = Field(
+        None,
+        ge=0,
+        description=(
+            "Limit number of job submissions returned per job to avoid large responses."
+            "Drops older job submissions. No effect with `include_jobs: false`"
+        ),
+    )
+    prev_submitted_at: Optional[datetime] = None
+    prev_run_id: Optional[UUID] = None
     limit: int = Field(100, ge=0, le=100)
     ascending: bool = False