[Feature]: Support dstack offer #2142 (#2540)

peterschmidt85 · web-flow · commit b0e159bb1f8e · 2025-04-22T12:55:24.000+02:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -32,7 +32,7 @@ dependencies = [
     "python-multipart>=0.0.16",
     "filelock",
     "psutil",
-    "gpuhunt>=0.1.2,<0.2.0",
+    "gpuhunt>=0.1.3,<0.2.0",
     "argcomplete>=3.5.0",
 ]
 
diff --git a/src/dstack/_internal/cli/commands/offer.py b/src/dstack/_internal/cli/commands/offer.py
@@ -0,0 +1,116 @@
+import argparse
+import contextlib
+import json
+from pathlib import Path
+
+from dstack._internal.cli.commands import APIBaseCommand
+from dstack._internal.cli.services.configurators.run import (
+    BaseRunConfigurator,
+)
+from dstack._internal.cli.utils.common import console
+from dstack._internal.cli.utils.run import print_run_plan
+from dstack._internal.core.models.configurations import (
+    ApplyConfigurationType,
+    TaskConfiguration,
+)
+from dstack._internal.core.models.runs import RunSpec
+from dstack.api.utils import load_profile
+
+
+class OfferConfigurator(BaseRunConfigurator):
+    # TODO: The command currently uses `BaseRunConfigurator` to register arguments.
+    #   This includes --env, --retry-policy, and other arguments that are unnecessary for this command.
+    #   Eventually, we should introduce a base `OfferConfigurator` that doesn't include those arguments—
+    #   `BaseRunConfigurator` will inherit from `OfferConfigurator`.
+    #
+    #   Additionally, it should have its own type: `ApplyConfigurationType.OFFER`.
+    TYPE = ApplyConfigurationType.TASK
+
+    @classmethod
+    def register_args(
+        cls,
+        parser: argparse.ArgumentParser,
+    ):
+        super().register_args(parser, default_max_offers=50)
+
+
+# TODO: Support aggregated offers
+# TODO: Add tests
+class OfferCommand(APIBaseCommand):
+    NAME = "offer"
+    DESCRIPTION = "List offers"
+
+    def _register(self):
+        super()._register()
+        self._parser.add_argument(
+            "--format",
+            choices=["plain", "json"],
+            default="plain",
+            help="Output format (default: plain)",
+        )
+        self._parser.add_argument(
+            "--json",
+            action="store_const",
+            const="json",
+            dest="format",
+            help="Output in JSON format (equivalent to --format json)",
+        )
+        OfferConfigurator.register_args(self._parser)
+
+    def _command(self, args: argparse.Namespace):
+        super()._command(args)
+        conf = TaskConfiguration(commands=[":"])
+
+        configurator = OfferConfigurator(api_client=self.api)
+        configurator.apply_args(conf, args, [])
+        profile = load_profile(Path.cwd(), profile_name=args.profile)
+
+        run_spec = RunSpec(
+            configuration=conf,
+            ssh_key_pub="(dummy)",
+            profile=profile,
+        )
+        if args.format == "plain":
+            status = console.status("Getting offers...")
+        else:
+            status = contextlib.nullcontext()
+        with status:
+            run_plan = self.api.client.runs.get_plan(
+                self.api.project,
+                run_spec,
+                max_offers=args.max_offers,
+            )
+
+        job_plan = run_plan.job_plans[0]
+
+        if args.format == "json":
+            output = {
+                "project": run_plan.project_name,
+                "user": run_plan.user,
+                "resources": job_plan.job_spec.requirements.resources.dict(),
+                "max_price": (job_plan.job_spec.requirements.max_price),
+                "spot": run_spec.configuration.spot_policy,
+                "reservation": run_plan.run_spec.configuration.reservation,
+                "offers": [],
+                "total_offers": job_plan.total_offers,
+            }
+
+            for offer in job_plan.offers:
+                output["offers"].append(
+                    {
+                        "backend": (
+                            "ssh" if offer.backend.value == "remote" else offer.backend.value
+                        ),
+                        "region": offer.region,
+                        "instance_type": offer.instance.name,
+                        "resources": offer.instance.resources.dict(),
+                        "spot": offer.instance.resources.spot,
+                        "price": float(offer.price),
+                        "availability": offer.availability.value,
+                    }
+                )
+
+            print(json.dumps(output, indent=2))
+            return
+        else:
+            print_run_plan(run_plan, include_run_properties=False)
diff --git a/src/dstack/_internal/cli/main.py b/src/dstack/_internal/cli/main.py
@@ -14,6 +14,7 @@
 from dstack._internal.cli.commands.init import InitCommand
 from dstack._internal.cli.commands.logs import LogsCommand
 from dstack._internal.cli.commands.metrics import MetricsCommand
+from dstack._internal.cli.commands.offer import OfferCommand
 from dstack._internal.cli.commands.ps import PsCommand
 from dstack._internal.cli.commands.server import ServerCommand
 from dstack._internal.cli.commands.stats import StatsCommand
@@ -65,6 +66,7 @@ def main():
     FleetCommand.register(subparsers)
     GatewayCommand.register(subparsers)
     InitCommand.register(subparsers)
+    OfferCommand.register(subparsers)
     LogsCommand.register(subparsers)
     MetricsCommand.register(subparsers)
     PsCommand.register(subparsers)
diff --git a/src/dstack/_internal/cli/services/configurators/run.py b/src/dstack/_internal/cli/services/configurators/run.py
@@ -92,7 +92,7 @@ def apply_configuration(
                 profile=profile,
             )
 
-        print_run_plan(run_plan, offers_limit=configurator_args.max_offers)
+        print_run_plan(run_plan, max_offers=configurator_args.max_offers)
 
         confirm_message = "Submit a new run?"
         stop_run_name = None
@@ -274,7 +274,7 @@ def delete_configuration(
         console.print(f"Run [code]{conf.name}[/] deleted")
 
     @classmethod
-    def register_args(cls, parser: argparse.ArgumentParser):
+    def register_args(cls, parser: argparse.ArgumentParser, default_max_offers: int = 3):
         configuration_group = parser.add_argument_group(f"{cls.TYPE.value} Options")
         configuration_group.add_argument(
             "-n",
@@ -286,7 +286,7 @@ def register_args(cls, parser: argparse.ArgumentParser):
             "--max-offers",
             help="Number of offers to show in the run plan",
             type=int,
-            default=3,
+            default=default_max_offers,
         )
         cls.register_env_args(configuration_group)
         configuration_group.add_argument(
diff --git a/src/dstack/_internal/cli/utils/run.py b/src/dstack/_internal/cli/utils/run.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Optional, Union
 
 from rich.markup import escape
 from rich.table import Table
@@ -24,7 +24,9 @@
 from dstack.api import Run
 
 
-def print_run_plan(run_plan: RunPlan, offers_limit: int = 3):
+def print_run_plan(
+    run_plan: RunPlan, max_offers: Optional[int] = None, include_run_properties: bool = True
+):
     job_plan = run_plan.job_plans[0]
 
     props = Table(box=None, show_header=False)
@@ -39,29 +41,30 @@ def print_run_plan(run_plan: RunPlan, offers_limit: int = 3):
         if job_plan.job_spec.max_duration
         else "-"
     )
-    inactivity_duration = None
-    if isinstance(run_plan.run_spec.configuration, DevEnvironmentConfiguration):
-        inactivity_duration = "-"
-        if isinstance(run_plan.run_spec.configuration.inactivity_duration, int):
-            inactivity_duration = format_pretty_duration(
-                run_plan.run_spec.configuration.inactivity_duration
-            )
-    if job_plan.job_spec.retry is None:
-        retry = "-"
-    else:
-        retry = escape(job_plan.job_spec.retry.pretty_format())
-
-    profile = run_plan.run_spec.merged_profile
-    creation_policy = profile.creation_policy
-    # FIXME: This assumes the default idle_duration is the same for client and server.
-    # If the server changes idle_duration, old clients will see incorrect value.
-    termination_policy, termination_idle_time = get_termination(
-        profile, DEFAULT_RUN_TERMINATION_IDLE_TIME
-    )
-    if termination_policy == TerminationPolicy.DONT_DESTROY:
-        idle_duration = "-"
-    else:
-        idle_duration = format_pretty_duration(termination_idle_time)
+    if include_run_properties:
+        inactivity_duration = None
+        if isinstance(run_plan.run_spec.configuration, DevEnvironmentConfiguration):
+            inactivity_duration = "-"
+            if isinstance(run_plan.run_spec.configuration.inactivity_duration, int):
+                inactivity_duration = format_pretty_duration(
+                    run_plan.run_spec.configuration.inactivity_duration
+                )
+        if job_plan.job_spec.retry is None:
+            retry = "-"
+        else:
+            retry = escape(job_plan.job_spec.retry.pretty_format())
+
+        profile = run_plan.run_spec.merged_profile
+        creation_policy = profile.creation_policy
+        # FIXME: This assumes the default idle_duration is the same for client and server.
+        # If the server changes idle_duration, old clients will see incorrect value.
+        termination_policy, termination_idle_time = get_termination(
+            profile, DEFAULT_RUN_TERMINATION_IDLE_TIME
+        )
+        if termination_policy == TerminationPolicy.DONT_DESTROY:
+            idle_duration = "-"
+        else:
+            idle_duration = format_pretty_duration(termination_idle_time)
 
     if req.spot is None:
         spot_policy = "auto"
@@ -75,30 +78,32 @@ def th(s: str) -> str:
 
     props.add_row(th("Project"), run_plan.project_name)
     props.add_row(th("User"), run_plan.user)
-    props.add_row(th("Configuration"), run_plan.run_spec.configuration_path)
-    props.add_row(th("Type"), run_plan.run_spec.configuration.type)
+    if include_run_properties:
+        props.add_row(th("Configuration"), run_plan.run_spec.configuration_path)
+        props.add_row(th("Type"), run_plan.run_spec.configuration.type)
     props.add_row(th("Resources"), pretty_req)
-    props.add_row(th("Max price"), max_price)
-    props.add_row(th("Max duration"), max_duration)
-    if inactivity_duration is not None:  # None means n/a
-        props.add_row(th("Inactivity duration"), inactivity_duration)
     props.add_row(th("Spot policy"), spot_policy)
-    props.add_row(th("Retry policy"), retry)
-    props.add_row(th("Creation policy"), creation_policy)
-    props.add_row(th("Idle duration"), idle_duration)
+    props.add_row(th("Max price"), max_price)
+    if include_run_properties:
+        props.add_row(th("Retry policy"), retry)
+        props.add_row(th("Creation policy"), creation_policy)
+        props.add_row(th("Idle duration"), idle_duration)
+        props.add_row(th("Max duration"), max_duration)
+        if inactivity_duration is not None:  # None means n/a
+            props.add_row(th("Inactivity duration"), inactivity_duration)
     props.add_row(th("Reservation"), run_plan.run_spec.configuration.reservation or "-")
 
     offers = Table(box=None)
     offers.add_column("#")
     offers.add_column("BACKEND")
     offers.add_column("REGION")
-    offers.add_column("INSTANCE")
+    offers.add_column("INSTANCE TYPE")
     offers.add_column("RESOURCES")
     offers.add_column("SPOT")
     offers.add_column("PRICE")
     offers.add_column()
 
-    job_plan.offers = job_plan.offers[:offers_limit]
+    job_plan.offers = job_plan.offers[:max_offers] if max_offers else job_plan.offers
 
     for i, offer in enumerate(job_plan.offers, start=1):
         r = offer.instance.resources
diff --git a/src/dstack/_internal/server/routers/runs.py b/src/dstack/_internal/server/routers/runs.py
@@ -100,6 +100,7 @@ async def get_plan(
         project=project,
         user=user,
         run_spec=body.run_spec,
+        max_offers=body.max_offers,
     )
     return run_plan
 
diff --git a/src/dstack/_internal/server/schemas/runs.py b/src/dstack/_internal/server/schemas/runs.py
@@ -26,6 +26,9 @@ class GetRunRequest(CoreModel):
 
 class GetRunPlanRequest(CoreModel):
     run_spec: RunSpec
+    max_offers: Optional[int] = Field(
+        description="The maximum number of offers to return", ge=1, le=10000
+    )
 
 
 class SubmitRunRequest(CoreModel):
diff --git a/src/dstack/_internal/server/services/runs.py b/src/dstack/_internal/server/services/runs.py
@@ -92,6 +92,8 @@
     JobTerminationReason.FAILED_TO_START_DUE_TO_NO_CAPACITY,
 }
 
+DEFAULT_MAX_OFFERS = 50
+
 
 async def list_user_runs(
     session: AsyncSession,
@@ -275,6 +277,7 @@ async def get_plan(
     project: ProjectModel,
     user: UserModel,
     run_spec: RunSpec,
+    max_offers: Optional[int],
 ) -> RunPlan:
     _validate_run_spec_and_set_defaults(run_spec)
 
@@ -342,7 +345,7 @@ async def get_plan(
 
         job_plan = JobPlan(
             job_spec=job_spec,
-            offers=job_offers[:50],
+            offers=job_offers[: (max_offers or DEFAULT_MAX_OFFERS)],
             total_offers=len(job_offers),
             max_price=max((offer.price for offer in job_offers), default=None),
         )
diff --git a/src/dstack/api/server/_runs.py b/src/dstack/api/server/_runs.py
@@ -53,11 +53,13 @@ def get(self, project_name: str, run_name: str) -> Run:
         resp = self._request(f"/api/project/{project_name}/runs/get", body=json_body)
         return parse_obj_as(Run.__response__, resp.json())
 
-    def get_plan(self, project_name: str, run_spec: RunSpec) -> RunPlan:
-        body = GetRunPlanRequest(run_spec=run_spec)
+    def get_plan(
+        self, project_name: str, run_spec: RunSpec, max_offers: Optional[int] = None
+    ) -> RunPlan:
+        body = GetRunPlanRequest(run_spec=run_spec, max_offers=max_offers)
         resp = self._request(
             f"/api/project/{project_name}/runs/get_plan",
-            body=body.json(exclude=_get_run_spec_excludes(run_spec)),
+            body=body.json(exclude=_get_get_plan_excludes(body)),
         )
         return parse_obj_as(RunPlan.__response__, resp.json())
 
@@ -96,6 +98,17 @@ def _get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[Dict]:
     return None
 
 
+def _get_get_plan_excludes(request: GetRunPlanRequest) -> Optional[Dict]:
+    """
+    Excludes new fields when they are not set to keep
+    clients backward-compatibility with older servers.
+    """
+    run_spec_excludes = _get_run_spec_excludes(request.run_spec)
+    if request.max_offers is None:
+        run_spec_excludes["max_offers"] = True
+    return run_spec_excludes
+
+
 def _get_run_spec_excludes(run_spec: RunSpec) -> Optional[Dict]:
     """
     Returns `run_spec` exclude mapping to exclude certain fields from the request.

Original file line number	Diff line number	Diff line change
`@@ -32,7 +32,7 @@ dependencies = [`
`32`	`32`	`"python-multipart>=0.0.16",`
`33`	`33`	`"filelock",`
`34`	`34`	`"psutil",`
`35`		`- "gpuhunt>=0.1.2,<0.2.0",`
	`35`	`+ "gpuhunt>=0.1.3,<0.2.0",`
`36`	`36`	`"argcomplete>=3.5.0",`
`37`	`37`	`]`
`38`	`38`
Original file line number	Diff line number	Diff line change
`@@ -100,6 +100,7 @@ async def get_plan(`
`100`	`100`	`project=project,`
`101`	`101`	`user=user,`
`102`	`102`	`run_spec=body.run_spec,`
	`103`	`+ max_offers=body.max_offers,`
`103`	`104`	`)`
`104`	`105`	`return run_plan`
`105`	`106`