Skip to content

Commit

Permalink
Optimize listing runs for Python API and CLI (#1430)
Browse files Browse the repository at this point in the history
  • Loading branch information
r4victor authored Jul 24, 2024
1 parent b45cbf8 commit a96b925
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 10 deletions.
4 changes: 3 additions & 1 deletion src/dstack/_internal/server/schemas/runs.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
from typing import List, Optional
from uuid import UUID

from pydantic import Field

from dstack._internal.core.models.common import CoreModel
from dstack._internal.core.models.instances import SSHKey
from dstack._internal.core.models.profiles import Profile
Expand All @@ -15,7 +17,7 @@ class ListRunsRequest(CoreModel):
only_active: bool = False
prev_submitted_at: Optional[datetime]
prev_run_id: Optional[UUID]
limit: int = 1000
limit: int = Field(100, ge=0, le=100)
ascending: bool = False


Expand Down
22 changes: 15 additions & 7 deletions src/dstack/api/_public/runs.py
Original file line number Diff line number Diff line change
Expand Up @@ -189,6 +189,7 @@ def logs(
start_time=next_start_time,
end_time=None,
descending=False,
limit=100,
diagnose=diagnose,
),
)
Expand Down Expand Up @@ -438,6 +439,7 @@ def get_plan(
regions=regions,
instance_types=instance_types,
spot_policy=spot_policy,
retry=None,
retry_policy=retry_policy,
max_duration=max_duration,
max_price=max_price,
Expand Down Expand Up @@ -503,13 +505,19 @@ def list(self, all: bool = False) -> List[Run]:
Returns:
list of runs
"""
runs = self._api_client.runs.list(project_name=self._project, repo_id=None)
if not all:
active = [run for run in runs if not run.status.is_finished()]
if active:
runs = active
else:
runs = runs[:1] # the most recent finished run
# Return only one page of latest runs (<=100). Returning all the pages may be costly.
# TODO: Consider introducing `since` filter with a reasonable default.
only_active = not all
runs = self._api_client.runs.list(
project_name=self._project,
repo_id=None,
only_active=only_active,
)
if only_active and len(runs) == 0:
runs = self._api_client.runs.list(
project_name=self._project,
repo_id=None,
)[:1]
return [self._model_to_run(run) for run in runs]

def get(self, run_name: str) -> Optional[Run]:
Expand Down
25 changes: 23 additions & 2 deletions src/dstack/api/server/_runs.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
from datetime import datetime
from typing import List, Optional
from uuid import UUID

from pydantic import parse_obj_as

Expand All @@ -25,8 +27,27 @@


class RunsAPIClient(APIClientGroup):
def list(self, project_name: Optional[str], repo_id: Optional[str]) -> List[Run]:
body = ListRunsRequest(project_name=project_name, repo_id=repo_id)
def list(
self,
project_name: Optional[str],
repo_id: Optional[str],
username: Optional[str] = None,
only_active: bool = False,
prev_submitted_at: Optional[datetime] = None,
prev_run_id: Optional[UUID] = None,
limit: int = 100,
ascending: bool = False,
) -> List[Run]:
body = ListRunsRequest(
project_name=project_name,
repo_id=repo_id,
username=username,
only_active=only_active,
prev_submitted_at=prev_submitted_at,
prev_run_id=prev_run_id,
limit=limit,
ascending=ascending,
)
resp = self._request("/api/runs/list", body=body.json())
return parse_obj_as(List[Run.__response__], resp.json())

Expand Down

0 comments on commit a96b925

Please sign in to comment.