From 4f6db4a661dcc88d8afc37665071336818ca0b12 Mon Sep 17 00:00:00 2001
From: Victor Skvortsov <vds003@gmail.com>
Date: Wed, 24 Jul 2024 11:57:52 +0500
Subject: [PATCH] Optimize listing runs for Python API and CLI

---
 src/dstack/_internal/server/schemas/runs.py |  4 +++-
 src/dstack/api/_public/runs.py              | 22 ++++++++++++------
 src/dstack/api/server/_runs.py              | 25 +++++++++++++++++++--
 3 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/src/dstack/_internal/server/schemas/runs.py b/src/dstack/_internal/server/schemas/runs.py
index eec4cc322..16bd28bbd 100644
--- a/src/dstack/_internal/server/schemas/runs.py
+++ b/src/dstack/_internal/server/schemas/runs.py
@@ -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
@@ -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
 
 
diff --git a/src/dstack/api/_public/runs.py b/src/dstack/api/_public/runs.py
index 2341b8936..3f3ccf79d 100644
--- a/src/dstack/api/_public/runs.py
+++ b/src/dstack/api/_public/runs.py
@@ -189,6 +189,7 @@ def logs(
                         start_time=next_start_time,
                         end_time=None,
                         descending=False,
+                        limit=100,
                         diagnose=diagnose,
                     ),
                 )
@@ -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,
@@ -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]:
diff --git a/src/dstack/api/server/_runs.py b/src/dstack/api/server/_runs.py
index e189c5164..a97bee504 100644
--- a/src/dstack/api/server/_runs.py
+++ b/src/dstack/api/server/_runs.py
@@ -1,4 +1,6 @@
+from datetime import datetime
 from typing import List, Optional
+from uuid import UUID
 
 from pydantic import parse_obj_as
 
@@ -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())