From 4e5068bc1dcb31ebe6be31d7ed612c93d6bac8ed Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Mon, 20 Jul 2026 13:38:29 +0000 Subject: [PATCH] CLI: rework `offer` command * Use unconstrained ResourcesSpec by default * Drop _some_ irrelevant options * Drop OfferConfigurator * Rework `print_run_plan` -- move the offers table to a reusable function, drop offer-specific customization -- the `offer` command now uses its own ad-hoc reimplementation rather than parametrized catch-all-use-cases `print_run_plan`. --- src/dstack/_internal/cli/commands/offer.py | 317 +++++++++++------- src/dstack/_internal/cli/utils/offers.py | 49 +++ src/dstack/_internal/cli/utils/run.py | 114 ++----- .../_internal/cli/commands/test_offer.py | 66 ++++ src/tests/_internal/cli/utils/test_offer.py | 91 ----- 5 files changed, 337 insertions(+), 300 deletions(-) create mode 100644 src/dstack/_internal/cli/utils/offers.py create mode 100644 src/tests/_internal/cli/commands/test_offer.py delete mode 100644 src/tests/_internal/cli/utils/test_offer.py diff --git a/src/dstack/_internal/cli/commands/offer.py b/src/dstack/_internal/cli/commands/offer.py index 9ca84ed718..0781e80b96 100644 --- a/src/dstack/_internal/cli/commands/offer.py +++ b/src/dstack/_internal/cli/commands/offer.py @@ -1,53 +1,27 @@ import argparse +from contextlib import nullcontext from pathlib import Path -from typing import List, Literal, cast +from typing import List, Literal, Optional, cast + +from rich.table import Table from dstack._internal.cli.commands import APIBaseCommand -from dstack._internal.cli.services.configurators.run import ( - BaseRunConfigurator, +from dstack._internal.cli.models.offers import OfferCommandOutput, OfferRequirements +from dstack._internal.cli.services.profile import ( + apply_profile_args, + load_profile_from_args, + register_profile_args, ) -from dstack._internal.cli.services.profile import load_profile_from_args, register_profile_args -from dstack._internal.cli.services.resources import register_resources_args -from dstack._internal.cli.utils.common import console +from dstack._internal.cli.services.resources import apply_resources_args, register_resources_args +from dstack._internal.cli.utils.common import NO_OFFERS_WARNING, console from dstack._internal.cli.utils.gpu import print_gpu_json, print_gpu_table -from dstack._internal.cli.utils.run import print_offers_json, print_run_plan +from dstack._internal.cli.utils.offers import print_offers_table from dstack._internal.core.errors import CLIError -from dstack._internal.core.models.configurations import ApplyConfigurationType, TaskConfiguration -from dstack._internal.core.models.gpus import GpuGroup -from dstack._internal.core.models.runs import RunSpec - - -class OfferConfigurator(BaseRunConfigurator): - TYPE = ApplyConfigurationType.TASK - - @classmethod - def register_args(cls, parser: argparse.ArgumentParser): - configuration_group = parser.add_argument_group(f"{cls.TYPE.value} Options") - parser.add_argument( - "--group-by", - action="append", - help=( - "Group results by fields ([code]gpu[/code], [code]backend[/code], [code]region[/code], [code]count[/code]). " - "Optional, but if used, must include [code]gpu[/code]. " - "The use of [code]region[/code] also requires [code]backend[/code]. " - "Can be repeated or comma-separated (e.g. [code]--group-by gpu,backend[/code])." - ), - ) - configuration_group.add_argument( - "-n", - "--name", - dest="run_name", - help="The name of the run. If not specified, a random name is assigned", - ) - configuration_group.add_argument( - "--max-offers", - help="Number of offers to show in the run plan", - type=int, - default=50, - ) - cls.register_env_args(configuration_group) - register_resources_args(configuration_group) - register_profile_args(parser) +from dstack._internal.core.models.configurations import TaskConfiguration +from dstack._internal.core.models.instances import InstanceOfferWithAvailability +from dstack._internal.core.models.profiles import SpotPolicy +from dstack._internal.core.models.resources import ResourcesSpec +from dstack._internal.core.models.runs import Requirements, RunPlan, RunSpec, get_policy_map class OfferCommand(APIBaseCommand): @@ -69,84 +43,199 @@ def _register(self): dest="format", help="Output in JSON format (equivalent to --format json)", ) - OfferConfigurator.register_args(self._parser) + self._parser.add_argument( + "--group-by", + action="append", + help=( + "Group results by fields ([code]gpu[/code], [code]backend[/code], [code]region[/code], [code]count[/code]). " + "Optional, but if used, must include [code]gpu[/code]. " + "The use of [code]region[/code] also requires [code]backend[/code]. " + "Can be repeated or comma-separated (e.g. [code]--group-by gpu,backend[/code])." + ), + ) + self._parser.add_argument( + "--max-offers", + help="Number of offers to show", + type=int, + default=50, + ) + resources_group = self._parser.add_argument_group("Resources") + register_resources_args(resources_group) + # TODO: register only relevant options + register_profile_args(self._parser) def _command(self, args: argparse.Namespace): super()._command(args) - # Set image and user so that the server (a) does not default gpu.vendor - # to nvidia — `dstack offer` should show all vendors, and (b) does not - # attempt to pull image config from the Docker registry. - conf = TaskConfiguration(commands=[":"], image="scratch", user="root") - - configurator = OfferConfigurator(api_client=self.api) - configurator.apply_args(conf, args) - profile = load_profile_from_args(args=args, repo_dir=Path.cwd()) - - run_spec = RunSpec( - configuration=conf, - profile=profile, + group_by = _process_group_by_args(args.group_by or []) + with console.status("Getting offers...") if args.format == "plain" else nullcontext(): + if group_by: + self._list_gpus(args, group_by) + else: + self._list_offers(args) + + def _list_offers(self, args: argparse.Namespace) -> None: + run_spec = _get_run_spec(args) + run_plan = self.api.client.runs.get_plan( + project_name=self.api.project, + run_spec=run_spec, + max_offers=args.max_offers, ) + job_plan = run_plan.job_plans[0] + if args.format == "plain": + _print_offers_header( + project_name=self.api.project, + requirements=job_plan.job_spec.requirements, + ) + _print_offers_table( + offers=job_plan.offers, + total_offers=job_plan.total_offers, + max_price=job_plan.max_price, + show_fleet_hint=run_spec.merged_profile.fleets is None, + ) + elif args.format == "json": + _print_offers_json(run_plan) + else: + raise NotImplementedError(args.format) + + def _list_gpus(self, args: argparse.Namespace, group_by: list[str]) -> None: + run_spec = _get_run_spec(args) + gpus = self.api.client.gpus.list_gpus( + project_name=self.api.project, + run_spec=run_spec, + group_by=[g for g in group_by if g != "gpu"], + ) + if args.format == "plain": + print_gpu_table(gpus, run_spec, group_by, self.api.project) + elif args.format == "json": + print_gpu_json( + gpus, + run_spec, + cast(List[Literal["gpu", "backend", "region", "count"]], group_by), + self.api.project, + ) + else: + raise NotImplementedError(args.format) - if args.group_by: - args.group_by = self._process_group_by_args(args.group_by) - if args.group_by and "gpu" not in args.group_by: - group_values = ", ".join(args.group_by) - raise CLIError(f"Cannot group by '{group_values}' without also grouping by 'gpu'") +def _process_group_by_args(group_by_args: List[str]) -> List[str]: + valid_choices = {"gpu", "backend", "region", "count"} + processed = [] - if args.format == "plain": - with console.status("Getting offers..."): - if args.group_by: - gpus = self._list_gpus(args, run_spec) - print_gpu_table(gpus, run_spec, args.group_by, self.api.project) - else: - run_plan = self.api.client.runs.get_plan( - self.api.project, - run_spec, - max_offers=args.max_offers, - ) - print_run_plan( - run_plan, - include_run_properties=False, - show_offer_fleet_hint=run_spec.merged_profile.fleets is None, - ) - else: - if args.group_by: - gpus = self._list_gpus(args, run_spec) - print_gpu_json( - gpus, - run_spec, - cast(List[Literal["gpu", "backend", "region", "count"]], args.group_by), - self.api.project, - ) + for arg in group_by_args: + values = [v.strip() for v in arg.split(",") if v.strip()] + for value in values: + if value in valid_choices: + processed.append(value) else: - run_plan = self.api.client.runs.get_plan( - self.api.project, - run_spec, - max_offers=args.max_offers, + raise CLIError( + f"Invalid group-by value: '{value}'. Valid choices are: {', '.join(sorted(valid_choices))}" ) - print_offers_json(run_plan, run_spec) - - def _process_group_by_args(self, group_by_args: List[str]) -> List[str]: - valid_choices = {"gpu", "backend", "region", "count"} - processed = [] - - for arg in group_by_args: - values = [v.strip() for v in arg.split(",") if v.strip()] - for value in values: - if value in valid_choices: - processed.append(value) - else: - raise CLIError( - f"Invalid group-by value: '{value}'. Valid choices are: {', '.join(sorted(valid_choices))}" - ) - - return processed - - def _list_gpus(self, args: argparse.Namespace, run_spec: RunSpec) -> List[GpuGroup]: - group_by = [g for g in args.group_by if g != "gpu"] or None - return self.api.client.gpus.list_gpus( - self.api.project, - run_spec, - group_by=group_by, + + if processed and "gpu" not in processed: + group_values = ", ".join(processed) + raise CLIError(f"Cannot group by '{group_values}' without also grouping by 'gpu'") + + return processed + + +def _get_run_spec(args: argparse.Namespace) -> RunSpec: + # Set image and user so that the server (a) does not default gpu.vendor + # to nvidia — `dstack offer` should show all vendors, and (b) does not + # attempt to pull image config from the Docker registry. + conf = TaskConfiguration( + resources=ResourcesSpec.unconstrained(), + commands=[":"], + image="scratch", + user="root", + ) + apply_resources_args(args, conf) + apply_profile_args(args, conf) + profile = load_profile_from_args(args=args, repo_dir=Path.cwd()) + return RunSpec( + configuration=conf, + profile=profile, + ) + + +def _print_offers_header( + project_name: str, + requirements: Requirements, +): + def th(s: str) -> str: + return f"[bold]{s}[/bold]" + + props = Table(box=None, show_header=False) + props.add_column(no_wrap=True) # key + props.add_column() # value + + pretty_req = requirements.pretty_format(resources_only=True) + max_price = ( + f"${requirements.max_price:3f}".rstrip("0").rstrip(".") + if requirements.max_price + else "off" + ) + + if requirements.spot is None: + spot_policy = "auto" + elif requirements.spot: + spot_policy = "spot" + else: + spot_policy = "on-demand" + + props.add_row(th("Project"), project_name) + props.add_row(th("Resources"), pretty_req) + props.add_row(th("Spot policy"), spot_policy) + props.add_row(th("Max price"), max_price) + console.print(props) + console.print() + + +_FLEET_HINT = ( + "Hint: Existing fleets are ignored, and all available offers are shown." + " To filter by fleet, pass --fleet NAME." +) + + +def _print_offers_table( + offers: list[InstanceOfferWithAvailability], + total_offers: int, + max_price: Optional[float], + show_fleet_hint: bool, +): + if len(offers) > 0: + show_fleet_hint_before_table = ( + show_fleet_hint and total_offers <= len(offers) and len(offers) < 3 + ) + show_fleet_hint_after_table = show_fleet_hint and not show_fleet_hint_before_table + if show_fleet_hint_before_table: + console.print(f"[secondary]{_FLEET_HINT}[/]") + console.print() + print_offers_table( + offers=offers, + total_offers=total_offers, + max_price=max_price or 0.0, + mute_tail_rows=False, ) + console.print() + if show_fleet_hint_after_table: + console.print(f"[secondary]{_FLEET_HINT}[/]") + else: + console.print(NO_OFFERS_WARNING) + + +def _print_offers_json(run_plan: RunPlan): + job_plan = run_plan.job_plans[0] + requirements = OfferRequirements( + resources=job_plan.job_spec.requirements.resources, + max_price=job_plan.job_spec.requirements.max_price, + spot=get_policy_map(run_plan.run_spec.configuration.spot_policy, default=SpotPolicy.AUTO), + reservation=run_plan.run_spec.configuration.reservation, + ) + output = OfferCommandOutput( + project=run_plan.project_name, + user=run_plan.user, + requirements=requirements, + offers=job_plan.offers, + total_offers=job_plan.total_offers, + ) + print(output.json()) diff --git a/src/dstack/_internal/cli/utils/offers.py b/src/dstack/_internal/cli/utils/offers.py new file mode 100644 index 0000000000..bfb75213fe --- /dev/null +++ b/src/dstack/_internal/cli/utils/offers.py @@ -0,0 +1,49 @@ +import shutil +from collections.abc import Sequence + +from rich.table import Table + +from dstack._internal.cli.utils.common import console, format_backend, format_instance_availability +from dstack._internal.core.models.instances import InstanceOfferWithAvailability + + +def print_offers_table( + offers: Sequence[InstanceOfferWithAvailability], + total_offers: int, + max_price: float, + mute_tail_rows: bool, +): + table = Table(box=None, expand=shutil.get_terminal_size(fallback=(120, 40)).columns <= 110) + table.add_column("#") + table.add_column("BACKEND", style="grey58", ratio=2) + table.add_column("RESOURCES", ratio=4) + table.add_column("INSTANCE TYPE", style="grey58", no_wrap=True, ratio=2) + table.add_column("PRICE", style="grey58", ratio=1) + table.add_column() + + for i, offer in enumerate(offers, start=1): + r = offer.instance.resources + + instance = offer.instance.name + if offer.total_blocks > 1: + instance += f" ({offer.blocks}/{offer.total_blocks})" + table.add_row( + f"{i}", + format_backend(offer.backend, offer.region), + r.pretty_format(include_spot=True), + instance, + f"${offer.price:.4f}".rstrip("0").rstrip("."), + format_instance_availability(offer.availability), + style=None if i == 1 or not mute_tail_rows else "secondary", + ) + if total_offers > len(offers): + table.add_row("", "...", style="secondary") + + if len(offers) > 0: + console.print(table) + if total_offers > len(offers): + console.print( + f"[secondary] Shown {len(offers)} of {total_offers} offers, " + f"${max_price:3f}".rstrip("0").rstrip(".") + + "max[/]" + ) diff --git a/src/dstack/_internal/cli/utils/run.py b/src/dstack/_internal/cli/utils/run.py index dab2d264f0..d034acb203 100644 --- a/src/dstack/_internal/cli/utils/run.py +++ b/src/dstack/_internal/cli/utils/run.py @@ -5,7 +5,6 @@ from rich.markup import escape from rich.table import Table -from dstack._internal.cli.models.offers import OfferCommandOutput, OfferRequirements from dstack._internal.cli.models.runs import PsCommandOutput from dstack._internal.cli.utils.common import ( NO_FLEETS_WARNING, @@ -13,8 +12,8 @@ add_row_from_dict, console, format_backend, - format_instance_availability, ) +from dstack._internal.cli.utils.offers import print_offers_table from dstack._internal.core.models.configurations import DevEnvironmentConfiguration from dstack._internal.core.models.instances import ( InstanceOfferWithAvailability, @@ -23,7 +22,6 @@ from dstack._internal.core.models.profiles import ( DEFAULT_RUN_TERMINATION_IDLE_TIME, CreationPolicy, - SpotPolicy, TerminationPolicy, ) from dstack._internal.core.models.runs import ( @@ -35,7 +33,6 @@ ProbeSpec, RunPlan, RunStatus, - get_policy_map, ) from dstack._internal.core.models.runs import ( Run as CoreRun, @@ -56,34 +53,6 @@ class RunWaitStatus(str, Enum): WAITING_FOR_SCHEDULE = "waiting for schedule" -_OFFER_FLEET_HINT = ( - "Hint: Existing fleets are ignored, and all available offers are shown." - " To filter by fleet, pass --fleet NAME." -) - - -def print_offers_json(run_plan: RunPlan, run_spec): - """Print offers information in JSON format.""" - job_plan = run_plan.job_plans[0] - - requirements = OfferRequirements( - resources=job_plan.job_spec.requirements.resources, - max_price=job_plan.job_spec.requirements.max_price, - spot=get_policy_map(run_spec.configuration.spot_policy, default=SpotPolicy.AUTO), - reservation=run_plan.run_spec.configuration.reservation, - ) - - output = OfferCommandOutput( - project=run_plan.project_name, - user=run_plan.user, - requirements=requirements, - offers=job_plan.offers, - total_offers=job_plan.total_offers, - ) - - print(output.json()) - - def print_runs_json(project: str, runs: List[Run]) -> None: """Print runs information in JSON format.""" output = PsCommandOutput( @@ -96,10 +65,8 @@ def print_runs_json(project: str, runs: List[Run]) -> None: def print_run_plan( run_plan: RunPlan, max_offers: Optional[int] = None, - include_run_properties: bool = True, no_fleets: bool = False, verbose: bool = False, - show_offer_fleet_hint: bool = False, extra_properties: Optional[Dict[str, str]] = None, ): run_spec = run_plan.get_effective_run_spec() @@ -153,78 +120,35 @@ def th(s: str) -> str: props.add_row(th("Project"), run_plan.project_name) props.add_row(th("User"), run_plan.user) - if include_run_properties: - configuration_type = run_spec.configuration.type - if run_spec.configuration.type == "task": - configuration_type += f" (nodes={run_spec.configuration.nodes})" - props.add_row(th("Type"), configuration_type) + configuration_type = run_spec.configuration.type + if run_spec.configuration.type == "task": + configuration_type += f" (nodes={run_spec.configuration.nodes})" + props.add_row(th("Type"), configuration_type) props.add_row(th("Resources"), pretty_req) props.add_row(th("Spot policy"), spot_policy) props.add_row(th("Max price"), max_price) - if include_run_properties: - props.add_row(th("Retry policy"), retry) - if verbose or creation_policy != CreationPolicy.REUSE_OR_CREATE: - 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: # only set for dev-environment - props.add_row(th("Inactivity duration"), inactivity_duration) + props.add_row(th("Retry policy"), retry) + if verbose or creation_policy != CreationPolicy.REUSE_OR_CREATE: + 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: # only set for dev-environment + props.add_row(th("Inactivity duration"), inactivity_duration) if verbose or run_spec.configuration.reservation: props.add_row(th("Reservation"), run_spec.configuration.reservation or "no") for key, value in (extra_properties or {}).items(): props.add_row(th(key), value) - - offers = Table(box=None, expand=shutil.get_terminal_size(fallback=(120, 40)).columns <= 110) - offers.add_column("#") - offers.add_column("BACKEND", style="grey58", ratio=2) - offers.add_column("RESOURCES", ratio=4) - offers.add_column("INSTANCE TYPE", style="grey58", no_wrap=True, ratio=2) - offers.add_column("PRICE", style="grey58", ratio=1) - offers.add_column() - - displayed_offers = job_plan.offers[:max_offers] if max_offers else job_plan.offers - - for i, offer in enumerate(displayed_offers, start=1): - r = offer.instance.resources - - instance = offer.instance.name - if offer.total_blocks > 1: - instance += f" ({offer.blocks}/{offer.total_blocks})" - offers.add_row( - f"{i}", - format_backend(offer.backend, offer.region), - r.pretty_format(include_spot=True), - instance, - f"${offer.price:.4f}".rstrip("0").rstrip("."), - format_instance_availability(offer.availability), - style=None if i == 1 or not include_run_properties else "secondary", - ) - if job_plan.total_offers > len(displayed_offers): - offers.add_row("", "...", style="secondary") - console.print(props) console.print() + + displayed_offers = job_plan.offers[:max_offers] if max_offers else job_plan.offers if len(displayed_offers) > 0: - show_offer_fleet_hint_before_table = ( - show_offer_fleet_hint - and job_plan.total_offers <= len(displayed_offers) - and len(displayed_offers) < 3 + print_offers_table( + offers=displayed_offers, + total_offers=job_plan.total_offers, + max_price=job_plan.max_price or 0.0, + mute_tail_rows=True, ) - show_offer_fleet_hint_after_table = ( - show_offer_fleet_hint and not show_offer_fleet_hint_before_table - ) - if show_offer_fleet_hint_before_table: - console.print(f"[secondary]{_OFFER_FLEET_HINT}[/]") - console.print() - console.print(offers) - if job_plan.total_offers > len(displayed_offers): - console.print( - f"[secondary] Shown {len(displayed_offers)} of {job_plan.total_offers} offers, " - f"${job_plan.max_price:3f}".rstrip("0").rstrip(".") - + "max[/]" - ) - if show_offer_fleet_hint_after_table: - console.print(f"[secondary]{_OFFER_FLEET_HINT}[/]") console.print() else: console.print(NO_FLEETS_WARNING if no_fleets else NO_OFFERS_WARNING) diff --git a/src/tests/_internal/cli/commands/test_offer.py b/src/tests/_internal/cli/commands/test_offer.py new file mode 100644 index 0000000000..7dce6d6a48 --- /dev/null +++ b/src/tests/_internal/cli/commands/test_offer.py @@ -0,0 +1,66 @@ +from dstack._internal.cli.commands.offer import _print_offers_table +from dstack._internal.cli.utils.common import console +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.instances import ( + InstanceAvailability, + InstanceOfferWithAvailability, + InstanceType, + Resources, +) + +_FLEET_HINT = ( + "Hint: Existing fleets are ignored, and all available offers are shown." + " To filter by fleet, pass --fleet NAME." +) +_FLEET_HINT_START = "Hint: Existing fleets are ignored" + + +def _get_offer(index: int) -> InstanceOfferWithAvailability: + return InstanceOfferWithAvailability( + backend=BackendType.AWS, + instance=InstanceType( + name=f"instance-{index}", + resources=Resources(cpus=2, memory_mib=8192, spot=False, gpus=[]), + ), + region="us-east-1", + price=float(index), + availability=InstanceAvailability.AVAILABLE, + ) + + +def _get_max_price(offers: list[InstanceOfferWithAvailability]) -> float: + return max((offer.price for offer in offers), default=0.0) + + +class TestPrintOffersTableFleetHint: + def test_prints_hint_before_short_offer_table(self): + offers = [_get_offer(1), _get_offer(2)] + + with console.capture() as capture: + _print_offers_table( + offers=offers, + total_offers=2, + max_price=_get_max_price(offers), + show_fleet_hint=True, + ) + + output = capture.get() + assert " ".join(_FLEET_HINT.split()) in " ".join(output.split()) + assert output.index(_FLEET_HINT_START) < output.index("1 aws (us-east-1)") + + def test_prints_hint_after_truncated_offer_table(self): + offers = [_get_offer(index) for index in range(1, 4)] + + with console.capture() as capture: + _print_offers_table( + offers=offers, + total_offers=10, + max_price=_get_max_price(offers), + show_fleet_hint=True, + ) + + output = capture.get() + shown_footer = "Shown 3 of 10 offers, $3max" + assert shown_footer in output + assert " ".join(_FLEET_HINT.split()) in " ".join(output.split()) + assert output.index(shown_footer) < output.index(_FLEET_HINT_START) diff --git a/src/tests/_internal/cli/utils/test_offer.py b/src/tests/_internal/cli/utils/test_offer.py deleted file mode 100644 index 9a9a4dfb16..0000000000 --- a/src/tests/_internal/cli/utils/test_offer.py +++ /dev/null @@ -1,91 +0,0 @@ -import pytest - -from dstack._internal.cli.utils.common import console -from dstack._internal.cli.utils.run import print_run_plan -from dstack._internal.core.models.backends.base import BackendType -from dstack._internal.core.models.common import ApplyAction -from dstack._internal.core.models.instances import ( - InstanceAvailability, - InstanceOfferWithAvailability, - InstanceType, - Resources, -) -from dstack._internal.core.models.runs import JobPlan, RunPlan -from dstack._internal.server.services.jobs import get_jobs_from_run_spec -from dstack._internal.server.testing.common import get_run_spec - -_OFFER_FLEET_HINT = ( - "Hint: Existing fleets are ignored, and all available offers are shown." - " To filter by fleet, pass --fleet NAME." -) -_OFFER_FLEET_HINT_START = "Hint: Existing fleets are ignored" - - -def _get_offer(index: int) -> InstanceOfferWithAvailability: - return InstanceOfferWithAvailability( - backend=BackendType.AWS, - instance=InstanceType( - name=f"instance-{index}", - resources=Resources(cpus=2, memory_mib=8192, spot=False, gpus=[]), - ), - region="us-east-1", - price=float(index), - availability=InstanceAvailability.AVAILABLE, - ) - - -async def _get_run_plan( - *, offers: list[InstanceOfferWithAvailability], total_offers: int -) -> RunPlan: - run_spec = get_run_spec(repo_id="test-repo") - job = (await get_jobs_from_run_spec(run_spec=run_spec, secrets={}, replica_num=0))[0] - return RunPlan( - project_name="test-project", - user="test-user", - run_spec=run_spec, - effective_run_spec=run_spec, - job_plans=[ - JobPlan( - job_spec=job.job_spec, - offers=offers, - total_offers=total_offers, - max_price=max((offer.price for offer in offers), default=None), - ) - ], - action=ApplyAction.CREATE, - ) - - -class TestPrintRunPlanOfferHint: - @pytest.mark.asyncio - async def test_prints_hint_before_short_offer_table(self): - run_plan = await _get_run_plan(offers=[_get_offer(1), _get_offer(2)], total_offers=2) - - with console.capture() as capture: - print_run_plan( - run_plan, - include_run_properties=False, - show_offer_fleet_hint=True, - ) - - output = capture.get() - assert " ".join(_OFFER_FLEET_HINT.split()) in " ".join(output.split()) - assert output.index(_OFFER_FLEET_HINT_START) < output.index("1 aws (us-east-1)") - - @pytest.mark.asyncio - async def test_prints_hint_after_truncated_offer_table(self): - offers = [_get_offer(index) for index in range(1, 4)] - run_plan = await _get_run_plan(offers=offers, total_offers=10) - - with console.capture() as capture: - print_run_plan( - run_plan, - include_run_properties=False, - show_offer_fleet_hint=True, - ) - - output = capture.get() - shown_footer = "Shown 3 of 10 offers, $3max" - assert shown_footer in output - assert " ".join(_OFFER_FLEET_HINT.split()) in " ".join(output.split()) - assert output.index(shown_footer) < output.index(_OFFER_FLEET_HINT_START)