From e7b4e6f628bd98db4fa921163a006e5174c661c7 Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Thu, 6 Aug 2026 09:25:19 +0000 Subject: [PATCH] Add `for_offers_only` parameter to `/runs/get_plan` API method --- src/dstack/_internal/cli/commands/offer.py | 1 + .../_internal/cli/services/presets/create.py | 4 +- .../_internal/server/compatibility/runs.py | 18 +++ src/dstack/_internal/server/routers/runs.py | 11 +- src/dstack/_internal/server/schemas/runs.py | 6 + .../server/services/runs/__init__.py | 2 + .../_internal/server/services/runs/plan.py | 27 +--- src/dstack/api/_public/runs.py | 4 + src/dstack/api/server/_runs.py | 2 + .../server/compatibility/test_runs.py | 81 ++++++++++++ .../_internal/server/routers/test_runs.py | 123 ++++++++++++++---- .../server/services/runs/test_plan.py | 3 + 12 files changed, 231 insertions(+), 51 deletions(-) create mode 100644 src/tests/_internal/server/compatibility/test_runs.py diff --git a/src/dstack/_internal/cli/commands/offer.py b/src/dstack/_internal/cli/commands/offer.py index 5d697c539..6e52f4cfc 100644 --- a/src/dstack/_internal/cli/commands/offer.py +++ b/src/dstack/_internal/cli/commands/offer.py @@ -91,6 +91,7 @@ def _list_offers(self, args: argparse.Namespace) -> None: max_offers=args.max_offers, full_offers=args.full_offers, unallocated_resources=args.unallocated, + for_offers_only=True, ) job_plan = run_plan.job_plans[0] if args.format == "plain": diff --git a/src/dstack/_internal/cli/services/presets/create.py b/src/dstack/_internal/cli/services/presets/create.py index 4ca4d3665..c5ac21316 100644 --- a/src/dstack/_internal/cli/services/presets/create.py +++ b/src/dstack/_internal/cli/services/presets/create.py @@ -808,7 +808,9 @@ def _print_fleet_offers(api: Client, allowed_fleets: tuple[str, ...]) -> None: offer_configuration.fleets = list(allowed_fleets) run_spec = RunSpec(configuration=offer_configuration, profile=None) with console.status("Getting offers..."): - run_plan = api.client.runs.get_plan(api.project, run_spec, max_offers=10) + run_plan = api.client.runs.get_plan( + api.project, run_spec, max_offers=10, for_offers_only=True + ) props = Table(box=None, show_header=False) props.add_column(no_wrap=True) props.add_column() diff --git a/src/dstack/_internal/server/compatibility/runs.py b/src/dstack/_internal/server/compatibility/runs.py index 752f5f784..39a0521ca 100644 --- a/src/dstack/_internal/server/compatibility/runs.py +++ b/src/dstack/_internal/server/compatibility/runs.py @@ -52,3 +52,21 @@ def patch_run_spec(run_spec: RunSpec, client_version: Optional[Version]) -> None and isinstance(run_spec.configuration.gateway, EntityReference) ): run_spec.configuration.gateway = run_spec.configuration.gateway.format() + + +def is_run_plan_for_offers_only( + run_spec: RunSpec, for_offers_only: bool, client_version: Optional[Version] +) -> bool: + """ + Clients < 0.21.0 don't support `for_offers_only` argument and rely on a magic configuration + that triggers "offer collection only" path. + + TODO: Drop once clients < 0.21.0 are no longer supported. + + NOTE: A real task with `commands == [":"]` would also match this special `dstack offer` path. + """ + if for_offers_only: + return True + if client_version is not None and client_version < Version("0.21.0"): + return run_spec.configuration.type == "task" and run_spec.configuration.commands == [":"] + return False diff --git a/src/dstack/_internal/server/routers/runs.py b/src/dstack/_internal/server/routers/runs.py index c2434e146..a2b68f490 100644 --- a/src/dstack/_internal/server/routers/runs.py +++ b/src/dstack/_internal/server/routers/runs.py @@ -6,7 +6,11 @@ from dstack._internal.core.errors import ResourceNotExistsError from dstack._internal.core.models.runs import Run, RunPlan -from dstack._internal.server.compatibility.runs import patch_run, patch_run_plan +from dstack._internal.server.compatibility.runs import ( + is_run_plan_for_offers_only, + patch_run, + patch_run_plan, +) from dstack._internal.server.db import get_session from dstack._internal.server.models import ProjectModel, UserModel from dstack._internal.server.schemas.runs import ( @@ -133,6 +137,10 @@ async def get_plan( user, project = user_project if not user.ssh_public_key and not body.run_spec.ssh_key_pub: await users.refresh_ssh_key(session=session, actor=user) + # TODO: Use body.for_offers_only directly once clients < 0.21.0 are no longer supported + for_offers_only = is_run_plan_for_offers_only( + run_spec=body.run_spec, for_offers_only=body.for_offers_only, client_version=client_version + ) run_plan = await runs.get_plan( session=session, project=project, @@ -142,6 +150,7 @@ async def get_plan( full_offers=body.full_offers, unallocated_resources=body.unallocated_resources, legacy_repo_dir=legacy_repo_dir, + for_offers_only=for_offers_only, ) patch_run_plan(run_plan, client_version) return CustomJSONResponse(run_plan) diff --git a/src/dstack/_internal/server/schemas/runs.py b/src/dstack/_internal/server/schemas/runs.py index 09359db30..fcfca906f 100644 --- a/src/dstack/_internal/server/schemas/runs.py +++ b/src/dstack/_internal/server/schemas/runs.py @@ -52,6 +52,12 @@ class GetRunPlanRequest(CoreModel): bool, Field(description="Subtract allocated resources to return only unallocated resources"), ] = False + for_offers_only: Annotated[ + bool, + Field( + description="Set to True if the run plan is requested for offer collection only, not a real run submission" + ), + ] = False class SubmitRunRequest(CoreModel): diff --git a/src/dstack/_internal/server/services/runs/__init__.py b/src/dstack/_internal/server/services/runs/__init__.py index 02b72c981..5270de31b 100644 --- a/src/dstack/_internal/server/services/runs/__init__.py +++ b/src/dstack/_internal/server/services/runs/__init__.py @@ -534,6 +534,7 @@ async def get_plan( max_offers: Optional[int], full_offers: bool, unallocated_resources: bool, + for_offers_only: bool, legacy_repo_dir: bool = False, ) -> RunPlan: effective_run_spec = RunSpec.model_validate(run_spec.model_dump()) @@ -574,6 +575,7 @@ async def get_plan( max_offers=max_offers, full_offers=full_offers, unallocated_resources=unallocated_resources, + for_offers_only=for_offers_only, ) run_plan = RunPlan( project_name=project.name, diff --git a/src/dstack/_internal/server/services/runs/plan.py b/src/dstack/_internal/server/services/runs/plan.py index 5da7a7fe7..8aecfafe7 100644 --- a/src/dstack/_internal/server/services/runs/plan.py +++ b/src/dstack/_internal/server/services/runs/plan.py @@ -89,14 +89,14 @@ async def get_job_plans( max_offers: Optional[int], full_offers: bool, unallocated_resources: bool, + for_offers_only: bool, ) -> list[JobPlan]: """ Returns job plans for the given run spec. Normal run planning (`dstack apply`) selects the best fleet candidate for each planned job and builds offers from that path. `dstack offer` without `--group-by` uses the same - `/runs/get_plan` API, but its synthetic run spec is detected by - `_should_select_best_fleet_candidate()`. In that case, planning skips + `/runs/get_plan` API but with `for_offers_only=True`. In that case, planning skips best-fleet-candidate selection and collects offers directly: global offers when no fleets are specified, or offers from the selected fleets when `--fleet` is used. @@ -119,7 +119,7 @@ async def get_job_plans( job_num=0, ) - if _should_select_best_fleet_candidate(run_spec) and run_spec.merged_profile.instances is None: + if not for_offers_only and run_spec.merged_profile.instances is None: candidate_fleet_models = await _select_candidate_fleet_models( session=session, project=project, @@ -1016,27 +1016,6 @@ def _get_job_plan( ) -def _should_select_best_fleet_candidate(run_spec: RunSpec) -> bool: - """ - Returns ``True`` for normal run planning and ``False`` for `dstack offer` without - `--group-by`. - - Both `dstack apply` and `dstack offer` without `--group-by` call `/runs/get_plan`. The - current way to recognize `dstack offer` without `--group-by` is the synthetic task spec - that the CLI sends with `type == "task"` and `commands == [":"]`. - TODO: Replace this command-shape hack with an explicit request/API signal for - `dstack offer` without `--group-by`. - - When this function returns ``False``, the planner skips best-fleet-candidate selection - and goes directly to the special `dstack offer` collection path: - global offers when no fleets are specified, or offers from the selected fleets when - `--fleet` is used. - - A real task with `commands == [":"]` would also match this special `dstack offer` path. - """ - return not (run_spec.configuration.type == "task" and run_spec.configuration.commands == [":"]) - - def _get_offers_from_instances( instances: list[InstanceModel], ) -> list[tuple[InstanceModel, InstanceOfferWithAvailability]]: diff --git a/src/dstack/api/_public/runs.py b/src/dstack/api/_public/runs.py index 09c24f028..3fe274c98 100644 --- a/src/dstack/api/_public/runs.py +++ b/src/dstack/api/_public/runs.py @@ -478,6 +478,7 @@ def get_run_plan( max_offers: Optional[int] = None, full_offers: bool = False, unallocated_resources: bool = False, + for_offers_only: bool = False, ) -> RunPlan: """ Get a run plan. @@ -498,6 +499,8 @@ def get_run_plan( full_offers: Return full offers not adjusted by requirements. unallocated_resources: Subtract allocated resources to return only unallocated resources. + for_offers_only: Set to True if the run plan is requested for offer collection only, + not a real run submission. Returns: Run plan. @@ -554,6 +557,7 @@ def get_run_plan( max_offers=max_offers, full_offers=full_offers, unallocated_resources=unallocated_resources, + for_offers_only=for_offers_only, ) return run_plan diff --git a/src/dstack/api/server/_runs.py b/src/dstack/api/server/_runs.py index 3c40fa47c..9790c3987 100644 --- a/src/dstack/api/server/_runs.py +++ b/src/dstack/api/server/_runs.py @@ -77,12 +77,14 @@ def get_plan( max_offers: Optional[int] = None, full_offers: bool = False, unallocated_resources: bool = False, + for_offers_only: bool = False, ) -> RunPlan: body = GetRunPlanRequest( run_spec=run_spec, max_offers=max_offers, full_offers=full_offers, unallocated_resources=unallocated_resources, + for_offers_only=for_offers_only, ) body = copy.deepcopy(body) patch_run_spec(body.run_spec) diff --git a/src/tests/_internal/server/compatibility/test_runs.py b/src/tests/_internal/server/compatibility/test_runs.py new file mode 100644 index 000000000..9cdcd074a --- /dev/null +++ b/src/tests/_internal/server/compatibility/test_runs.py @@ -0,0 +1,81 @@ +from typing import Optional + +import pytest +from packaging.version import Version + +from dstack._internal.core.models.configurations import ( + AnyRunConfiguration, + DevEnvironmentConfiguration, + TaskConfiguration, +) +from dstack._internal.server.compatibility.runs import is_run_plan_for_offers_only +from dstack._internal.server.testing.common import get_run_spec + +_OFFER_CLI_CONFIGURATION = TaskConfiguration(commands=[":"], image="scratch", user="root") + + +class TestIsRunPlanForOffersOnly: + @pytest.mark.parametrize( + ("configuration", "for_offers_only", "client_version", "expected"), + [ + pytest.param(_OFFER_CLI_CONFIGURATION, True, Version("0.21.0"), True, id="flag-set"), + pytest.param( + DevEnvironmentConfiguration(), + True, + Version("0.21.0"), + True, + id="flag-set-for-any-configuration", + ), + pytest.param( + _OFFER_CLI_CONFIGURATION, + False, + Version("0.20.30"), + True, + id="old-client-sends-offer-cli-configuration", + ), + pytest.param( + TaskConfiguration(commands=["echo"], image="scratch"), + False, + Version("0.20.30"), + False, + id="old-client-sends-regular-task", + ), + pytest.param( + DevEnvironmentConfiguration(), + False, + Version("0.20.30"), + False, + id="old-client-sends-configuration-without-commands", + ), + pytest.param( + _OFFER_CLI_CONFIGURATION, + False, + Version("0.21.0"), + False, + id="new-client-does-not-rely-on-offer-cli-configuration", + ), + pytest.param( + _OFFER_CLI_CONFIGURATION, + False, + None, + False, + id="dev-client-does-not-rely-on-offer-cli-configuration", + ), + ], + ) + def test_returns_expected( + self, + configuration: AnyRunConfiguration, + for_offers_only: bool, + client_version: Optional[Version], + expected: bool, + ) -> None: + run_spec = get_run_spec(repo_id="test-repo", configuration=configuration) + assert ( + is_run_plan_for_offers_only( + run_spec=run_spec, + for_offers_only=for_offers_only, + client_version=client_version, + ) + is expected + ) diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py index 277d6cca6..1d637888d 100644 --- a/src/tests/_internal/server/routers/test_runs.py +++ b/src/tests/_internal/server/routers/test_runs.py @@ -2075,13 +2075,10 @@ async def test_returns_run_plan_instance_volumes( @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) @pytest.mark.parametrize( - "configuration", + "for_offers_only", [ - pytest.param({"type": "dev-environment", "ide": "vscode"}, id="regular-configuration"), - pytest.param( - {"type": "task", "commands": [":"], "image": "scratch"}, - id="special-configuration-used-by-dstack-offer-cli-command", - ), + pytest.param(False, id="run-plan"), + pytest.param(True, id="offer-collection"), ], ) async def test_returns_run_plan_with_offer_from_imported_fleet( @@ -2089,7 +2086,7 @@ async def test_returns_run_plan_with_offer_from_imported_fleet( test_db, session: AsyncSession, client: AsyncClient, - configuration: dict, + for_offers_only: bool, ) -> None: importer_user = await create_user(session, global_role=GlobalRole.USER) exporter_project = await create_project(session, name="exporter-project") @@ -2121,8 +2118,8 @@ async def test_returns_run_plan_with_offer_from_imported_fleet( exported_fleets=[fleet], ) - run_spec = {"configuration": configuration} - body = {"run_spec": run_spec} + run_spec = {"configuration": {"type": "dev-environment", "ide": "vscode"}} + body = {"run_spec": run_spec, "for_offers_only": for_offers_only} response = await client.post( "/api/project/importer-project/runs/get_plan", headers=get_auth_headers(importer_user.token), @@ -2334,16 +2331,14 @@ async def test_returns_no_offers_if_imported_fleet_specified_without_project_pre @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) @pytest.mark.parametrize( - "configuration", + ("configuration", "for_offers_only"), [ - pytest.param({"type": "dev-environment"}, id="regular-configuration"), + pytest.param({"type": "dev-environment"}, False, id="run-plan"), + pytest.param({"type": "dev-environment"}, True, id="offer-collection"), pytest.param( - {"type": "task", "commands": [":"], "image": "scratch"}, - id="special-configuration-used-by-dstack-offer-cli-command", - ), - pytest.param( - {"type": "task", "commands": [":"], "image": "scratch", "fleets": ["test-fleet"]}, - id="special-configuration-used-by-dstack-offer-cli-command-with-fleets", # --fleet + {"type": "dev-environment", "fleets": ["test-fleet"]}, + True, + id="offer-collection-with-fleets", # `dstack offer --fleet` ), ], ) @@ -2353,6 +2348,7 @@ async def test_preserves_backend_specific_offer_order( session: AsyncSession, client: AsyncClient, configuration: dict, + for_offers_only: bool, ) -> None: user = await create_user(session=session, global_role=GlobalRole.USER) project = await create_project(session=session, owner=user) @@ -2372,7 +2368,7 @@ async def test_preserves_backend_specific_offer_order( run_spec = get_run_spec( repo_id=repo.name, configuration=parse_run_configuration(configuration) ) - body = {"run_spec": run_spec.model_dump()} + body = {"run_spec": run_spec.model_dump(), "for_offers_only": for_offers_only} backend_mock_aws = Mock() backend_mock_aws.TYPE = BackendType.AWS @@ -2443,7 +2439,7 @@ async def test_offer_cli_preserves_backend_specific_offer_order_across_fleets( fleets=["fleet-aws", "fleet-vastai"], ), ) - body = {"run_spec": run_spec.model_dump()} + body = {"run_spec": run_spec.model_dump(), "for_offers_only": True} backend_mock_aws = Mock() backend_mock_aws.TYPE = BackendType.AWS @@ -2544,7 +2540,7 @@ async def test_offer_cli_returns_offers_from_all_specified_fleets( response = await client.post( f"/api/project/{project.name}/runs/get_plan", headers=get_auth_headers(user.token), - json={"run_spec": run_spec.model_dump()}, + json={"run_spec": run_spec.model_dump(), "for_offers_only": True}, ) assert response.status_code == 200, response.json() @@ -2591,7 +2587,7 @@ async def test_offer_cli_deduplicates_identical_backend_offers_across_specified_ fleets=["fleet-a", "fleet-b"], ), ) - body = {"run_spec": run_spec.model_dump()} + body = {"run_spec": run_spec.model_dump(), "for_offers_only": True} with patch("dstack._internal.server.services.backends.get_project_backends") as m: backend_mock_aws = Mock() @@ -2679,7 +2675,7 @@ async def test_offer_cli_keeps_identical_existing_instances_from_specified_fleet response = await client.post( f"/api/project/{project.name}/runs/get_plan", headers=get_auth_headers(user.token), - json={"run_spec": run_spec.model_dump()}, + json={"run_spec": run_spec.model_dump(), "for_offers_only": True}, ) assert response.status_code == 200, response.json() @@ -2713,7 +2709,7 @@ async def test_offer_cli_without_fleet_keeps_global_offers( user="root", ), ) - body = {"run_spec": run_spec.model_dump()} + body = {"run_spec": run_spec.model_dump(), "for_offers_only": True} with patch("dstack._internal.server.services.backends.get_project_backends") as m: backend_mock_aws = Mock() backend_mock_aws.TYPE = BackendType.AWS @@ -2805,7 +2801,7 @@ async def test_offer_without_fleets_uses_global_offer_collection( response = await client.post( f"/api/project/{project.name}/runs/get_plan", headers=get_auth_headers(user.token), - json={"run_spec": run_spec.model_dump()}, + json={"run_spec": run_spec.model_dump(), "for_offers_only": True}, ) assert response.status_code == 200, response.json() @@ -2868,7 +2864,7 @@ async def test_offer_with_fleets_uses_selected_fleet_offer_collection( response = await client.post( f"/api/project/{project.name}/runs/get_plan", headers=get_auth_headers(user.token), - json={"run_spec": run_spec.model_dump()}, + json={"run_spec": run_spec.model_dump(), "for_offers_only": True}, ) assert response.status_code == 200, response.json() @@ -2944,6 +2940,83 @@ async def test_regular_run_plan_uses_best_fleet_candidate_selection( assert job_plan["total_offers"] == 1 assert job_plan["offers"][0]["price"] == 3.0 + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + @pytest.mark.parametrize( + ("client_version", "for_offers_only", "expected_offer_collection"), + [ + pytest.param(None, True, True, id="dev-client-with-flag"), + pytest.param("0.21.0", True, True, id="new-client-with-flag"), + pytest.param("0.20.30", None, True, id="old-client-without-flag"), + pytest.param("0.21.0", None, False, id="new-client-without-flag"), + pytest.param(None, None, False, id="dev-client-without-flag"), + ], + ) + async def test_collects_offers_only_if_requested_by_for_offers_only( + self, + test_db, + session: AsyncSession, + client: AsyncClient, + client_version: Optional[str], + for_offers_only: Optional[bool], + expected_offer_collection: bool, + ) -> None: + """ + Clients prior to 0.21.0 don't send `for_offers_only`, so the synthetic run spec that + `dstack offer` sends still triggers offer collection. For newer clients, the same run + spec is planned as a regular run unless `for_offers_only` is set. + """ + user = await create_user(session=session, global_role=GlobalRole.USER) + project = await create_project(session=session, owner=user) + await add_project_member( + session=session, + project=project, + user=user, + project_role=ProjectRole.USER, + ) + repo = await create_repo(session=session, project_id=project.id) + run_spec = get_run_spec( + repo_id=repo.name, + configuration=TaskConfiguration( + commands=[":"], + image="scratch", + user="root", + ), + ) + body: dict = {"run_spec": run_spec.model_dump()} + if for_offers_only is not None: + body["for_offers_only"] = for_offers_only + headers = get_auth_headers(user.token) + if client_version is not None: + headers["X-API-Version"] = client_version + offer = get_instance_offer_with_availability(price=1.0) + with ( + patch( + "dstack._internal.server.services.runs.plan.get_non_fleet_offers", + new=AsyncMock(return_value=([(Mock(), offer)], [])), + ) as get_non_fleet_offers_mock, + patch( + "dstack._internal.server.services.runs.plan._select_candidate_fleet_models", + new=AsyncMock(return_value=[Mock()]), + ) as select_candidate_fleet_models_mock, + patch( + "dstack._internal.server.services.runs.plan.find_optimal_fleet_with_offers", + new=AsyncMock(return_value=(Mock(), [(Mock(), offer)], [])), + ) as find_optimal_fleet_with_offers_mock, + ): + response = await client.post( + f"/api/project/{project.name}/runs/get_plan", + headers=headers, + json=body, + ) + + assert response.status_code == 200, response.json() + assert get_non_fleet_offers_mock.await_count == int(expected_offer_collection) + assert select_candidate_fleet_models_mock.await_count == int(not expected_offer_collection) + assert find_optimal_fleet_with_offers_mock.await_count == int( + not expected_offer_collection + ) + @pytest.mark.parametrize( ("client_version", "expected_availability"), [ diff --git a/src/tests/_internal/server/services/runs/test_plan.py b/src/tests/_internal/server/services/runs/test_plan.py index df457d33f..4f10ba01c 100644 --- a/src/tests/_internal/server/services/runs/test_plan.py +++ b/src/tests/_internal/server/services/runs/test_plan.py @@ -136,6 +136,7 @@ async def test_skips_backend_offers_by_creation_policy( max_offers=None, full_offers=False, unallocated_resources=False, + for_offers_only=False, ) find_optimal_fleet_with_offers_mock.assert_awaited_once() @@ -176,6 +177,7 @@ async def test_excludes_backend_offers_when_instances_specified( max_offers=None, full_offers=False, unallocated_resources=False, + for_offers_only=False, ) get_targeted_instance_offers_mock.assert_awaited_once() @@ -217,6 +219,7 @@ async def test_empty_dev_environment_with_fleet_does_not_use_targeted_instances( max_offers=None, full_offers=False, unallocated_resources=False, + for_offers_only=False, ) select_instances_mock.assert_not_awaited()