From 3ac6c3382e688a0962c32f5270027fff96e61853 Mon Sep 17 00:00:00 2001 From: Silu Panda <31051721+SiluPanda@users.noreply.github.com> Date: Sun, 30 Aug 2026 11:12:16 -0700 Subject: [PATCH] Fix shared blocks fleet offer selection --- .../pipeline_tasks/jobs_submitted.py | 34 ++- .../_internal/server/services/instances.py | 29 +- .../_internal/server/services/offers.py | 100 ++++++- .../server/services/requirements/combine.py | 18 +- .../_internal/server/services/runs/plan.py | 110 +++++++- .../pipeline_tasks/test_submitted_jobs.py | 132 ++++++++- .../server/services/runs/test_plan.py | 262 +++++++++++++++++- .../_internal/server/services/test_offers.py | 108 +++++++- 8 files changed, 732 insertions(+), 61 deletions(-) diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py index 8923e5f061..b25f7176f1 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py @@ -116,8 +116,8 @@ from dstack._internal.server.services.locking import get_locker from dstack._internal.server.services.logging import fmt from dstack._internal.server.services.offers import ( + generate_shared_offer, get_instance_offer_with_restricted_az, - get_offers_by_requirements, ) from dstack._internal.server.services.pipelines import PipelineHinterProtocol from dstack._internal.server.services.placement import ( @@ -128,6 +128,7 @@ ) from dstack._internal.server.services.runs import run_model_to_run from dstack._internal.server.services.runs.plan import ( + _get_backend_offers_in_fleet, find_optimal_fleet_with_offers, get_instance_offers_in_fleet, get_run_candidate_fleet_models_filters, @@ -135,9 +136,6 @@ get_targeted_instance_offers, select_run_candidate_fleet_models_with_filters, ) -from dstack._internal.server.services.runs.spec import ( - check_run_spec_requires_instance_mounts, -) from dstack._internal.server.services.secrets import get_project_secrets_mapping from dstack._internal.server.services.volumes import volume_model_to_volume from dstack._internal.server.utils import tracing @@ -1738,8 +1736,11 @@ async def _promote_or_create_instance_models_for_provisioned_jobs( provisioned_job_model.used_instance_id = instance_model.id instance_models.append(instance_model) + runtime_offer = offer + if offer.blocks < offer.total_blocks: + runtime_offer = generate_shared_offer(offer, offer.blocks, offer.total_blocks) provisioned_job_model.job_runtime_data = _prepare_job_runtime_data( - offer, context.multinode + runtime_offer, context.multinode ).model_dump_json() events.emit( session, @@ -1817,8 +1818,8 @@ def _create_instance_model_for_job( price=offer.price, region=offer.region, volume_attachments=[], - total_blocks=1, - busy_blocks=1, + total_blocks=offer.total_blocks, + busy_blocks=offer.blocks, ) @@ -1848,8 +1849,8 @@ def _promote_placeholder_instance( instance_model.region = offer.region instance_model.termination_policy = termination_policy instance_model.termination_idle_time = termination_idle_time - instance_model.total_blocks = 1 - instance_model.busy_blocks = 1 + instance_model.total_blocks = offer.total_blocks + instance_model.busy_blocks = offer.blocks async def _process_volume_attachments( @@ -2405,18 +2406,15 @@ async def _provision_new_capacity( placement_group_models=placement_group_models, fleet_model=fleet_model, ) - multinode = requirements.multinode or is_multinode_job(job) - offers = await get_offers_by_requirements( + offers = await _get_backend_offers_in_fleet( project=project, - profile=profile, - requirements=requirements, - exclude_not_available=True, - multinode=multinode, - master_job_provisioning_data=master_job_provisioning_data, + fleet_model=fleet_model, + run_spec=run.run_spec, + job=job, volumes=volumes, - privileged=job.job_spec.privileged, - instance_mounts=check_run_spec_requires_instance_mounts(run.run_spec), placement_group=placement_group_model_to_placement_group_optional(placement_group_model), + exclude_not_available=True, + master_job_provisioning_data=master_job_provisioning_data, ) offers_iter = iter(offers) offers_tried = 0 diff --git a/src/dstack/_internal/server/services/instances.py b/src/dstack/_internal/server/services/instances.py index 2b763c93ad..3d0771759c 100644 --- a/src/dstack/_internal/server/services/instances.py +++ b/src/dstack/_internal/server/services/instances.py @@ -59,7 +59,7 @@ from dstack._internal.server.schemas.runner import InstanceHealthResponse, TaskStatus from dstack._internal.server.services import events from dstack._internal.server.services.logging import fmt -from dstack._internal.server.services.offers import generate_shared_offer +from dstack._internal.server.services.offers import get_matching_shared_offer from dstack._internal.server.services.projects import list_user_project_models from dstack._internal.server.services.runner.client import ShimClient from dstack._internal.utils import common as common_utils @@ -693,7 +693,6 @@ def get_shared_instances_with_offers( volumes: Optional[List[List[Volume]]] = None, ) -> list[tuple[InstanceModel, InstanceOfferWithAvailability]]: instances_with_offers: list[tuple[InstanceModel, InstanceOfferWithAvailability]] = [] - query_filter = requirements_to_query_filter(requirements) filtered_instances = filter_instances( instances=instances, profile=profile, @@ -711,18 +710,20 @@ def get_shared_instances_with_offers( continue total_blocks = common_utils.get_or_error(instance.total_blocks) idle_blocks = total_blocks - instance.busy_blocks - min_blocks = total_blocks if multinode else 1 - for blocks in range(min_blocks, total_blocks + 1): - shared_offer = generate_shared_offer(offer, blocks, total_blocks) - catalog_item = offer_to_catalog_item(shared_offer) - if gpuhunt.matches(catalog_item, query_filter): - if blocks <= idle_blocks: - shared_offer.availability = InstanceAvailability.IDLE - else: - shared_offer.availability = InstanceAvailability.BUSY - if shared_offer.availability == InstanceAvailability.IDLE or not idle_only: - instances_with_offers.append((instance, shared_offer)) - break + shared_offer = get_matching_shared_offer( + offer, + requirements=requirements, + blocks=total_blocks, + multinode=multinode, + ) + if shared_offer is None: + continue + if shared_offer.blocks <= idle_blocks: + shared_offer.availability = InstanceAvailability.IDLE + else: + shared_offer.availability = InstanceAvailability.BUSY + if shared_offer.availability == InstanceAvailability.IDLE or not idle_only: + instances_with_offers.append((instance, shared_offer)) return instances_with_offers diff --git a/src/dstack/_internal/server/services/offers.py b/src/dstack/_internal/server/services/offers.py index 8193ab0863..55e350175d 100644 --- a/src/dstack/_internal/server/services/offers.py +++ b/src/dstack/_internal/server/services/offers.py @@ -6,8 +6,13 @@ import gpuhunt from dstack._internal.core.backends.base.backend import Backend -from dstack._internal.core.backends.base.compute import ComputeWithPlacementGroupSupport +from dstack._internal.core.backends.base.compute import ( + ComputeWithCreateInstanceSupport, + ComputeWithPlacementGroupSupport, +) +from dstack._internal.core.backends.base.offers import filter_offers_by_requirements from dstack._internal.core.backends.features import ( + BACKENDS_WITH_CREATE_INSTANCE_SUPPORT, BACKENDS_WITH_INSTANCE_VOLUMES_SUPPORT, BACKENDS_WITH_MULTINODE_SUPPORT, BACKENDS_WITH_PRIVILEGED_SUPPORT, @@ -31,6 +36,7 @@ async def get_offers_by_requirements( project: ProjectModel, profile: Profile, requirements: Requirements, + shared_offer_requirements: Optional[Requirements] = None, exclude_not_available=False, multinode: bool = False, master_job_provisioning_data: Optional[JobProvisioningData] = None, @@ -90,6 +96,11 @@ async def get_offers_by_requirements( if backend_types is not None: backends = [b for b in backends if b.TYPE in backend_types or b.TYPE == BackendType.DSTACK] + if blocks != 1 and shared_offer_requirements is not None: + backends = [ + b for b in backends if isinstance(b.compute(), ComputeWithCreateInstanceSupport) + ] + offers = await backends_services.get_backend_offers( backends=backends, requirements=requirements, @@ -109,8 +120,16 @@ async def get_offers_by_requirements( volumes_locations=volumes_locations, ) + if blocks != 1 and shared_offer_requirements is not None: + offers = _filter_shared_block_supported_offers(offers) + if blocks != 1: - offers = _get_shareable_offers(offers, blocks) + offers = _get_shareable_offers( + offers, + blocks, + requirements=shared_offer_requirements, + multinode=multinode, + ) if max_offers is not None: offers = itertools.islice(offers, max_offers) @@ -159,6 +178,7 @@ def generate_shared_offer( ) -> InstanceOfferWithAvailability: full_resources = offer.instance.resources resources = Resources( + cpu_arch=full_resources.cpu_arch, cpus=full_resources.cpus // total_blocks * blocks, memory_mib=full_resources.memory_mib // total_blocks * blocks, gpus=full_resources.gpus[: len(full_resources.gpus) // total_blocks * blocks], @@ -180,6 +200,30 @@ def generate_shared_offer( ) +def get_matching_shared_offer( + offer: InstanceOfferWithAvailability, + requirements: Requirements, + blocks: Union[int, Literal["auto"]], + *, + multinode: bool = False, +) -> Optional[InstanceOfferWithAvailability]: + resources = offer.instance.resources + gpu_count = len(resources.gpus) + if gpu_count > 0 and resources.gpus[0].vendor == gpuhunt.AcceleratorVendor.GOOGLE: + # TPUs cannot be shared. + gpu_count = 1 + divisible, total_blocks = is_divisible_into_blocks(resources.cpus, gpu_count, blocks) + if not divisible: + return None + + min_blocks = total_blocks if multinode else 1 + shared_offers = ( + generate_shared_offer(offer, selected_blocks, total_blocks) + for selected_blocks in range(min_blocks, total_blocks + 1) + ) + return next(filter_offers_by_requirements(shared_offers, requirements), None) + + def get_instance_offer_with_restricted_az( instance_offer: InstanceOfferWithAvailability, master_job_provisioning_data: Optional[JobProvisioningData], @@ -254,20 +298,50 @@ def _filter_offers( def _get_shareable_offers( offers: Iterable[Tuple[Backend, InstanceOfferWithAvailability]], blocks: Union[int, Literal["auto"]], + *, + requirements: Optional[Requirements] = None, + multinode: bool = False, ) -> Iterator[Tuple[Backend, InstanceOfferWithAvailability]]: """ Yields offers that can be shared with `total_blocks` set. """ for backend, offer in offers: - resources = offer.instance.resources - cpu_count = resources.cpus - gpu_count = len(resources.gpus) - if gpu_count > 0 and resources.gpus[0].vendor == gpuhunt.AcceleratorVendor.GOOGLE: - # TPUs cannot be shared - gpu_count = 1 - divisible, total_blocks = is_divisible_into_blocks(cpu_count, gpu_count, blocks) - if not divisible: + if requirements is None: + resources = offer.instance.resources + cpu_count = resources.cpus + gpu_count = len(resources.gpus) + if gpu_count > 0 and resources.gpus[0].vendor == gpuhunt.AcceleratorVendor.GOOGLE: + # TPUs cannot be shared + gpu_count = 1 + divisible, total_blocks = is_divisible_into_blocks(cpu_count, gpu_count, blocks) + if not divisible: + continue + new_offer = offer.model_copy() + new_offer.total_blocks = total_blocks + yield (backend, new_offer) + continue + + shared_offer = get_matching_shared_offer( + offer, + requirements=requirements, + blocks=blocks, + multinode=multinode, + ) + if shared_offer is None: + continue + annotated_offer = offer.model_copy() + annotated_offer.blocks = shared_offer.blocks + annotated_offer.total_blocks = shared_offer.total_blocks + yield backend, annotated_offer + + +def _filter_shared_block_supported_offers( + offers: Iterable[Tuple[Backend, InstanceOfferWithAvailability]], +) -> Iterator[Tuple[Backend, InstanceOfferWithAvailability]]: + for backend, offer in offers: + if backend.TYPE == offer.backend: + yield backend, offer + continue + if offer.backend not in BACKENDS_WITH_CREATE_INSTANCE_SUPPORT: continue - new_offer = offer.model_copy() - new_offer.total_blocks = total_blocks - yield (backend, new_offer) + yield backend, offer diff --git a/src/dstack/_internal/server/services/requirements/combine.py b/src/dstack/_internal/server/services/requirements/combine.py index 36043b2117..24598befb3 100644 --- a/src/dstack/_internal/server/services/requirements/combine.py +++ b/src/dstack/_internal/server/services/requirements/combine.py @@ -62,10 +62,26 @@ def combine_fleet_and_run_profiles( def combine_fleet_and_run_requirements( fleet_requirements: Requirements, run_requirements: Requirements +) -> Optional[Requirements]: + try: + resources = _combine_resources(fleet_requirements.resources, run_requirements.resources) + except CombineError: + return None + return combine_fleet_and_run_requirements_with_resources( + fleet_requirements=fleet_requirements, + run_requirements=run_requirements, + resources=resources, + ) + + +def combine_fleet_and_run_requirements_with_resources( + fleet_requirements: Requirements, + run_requirements: Requirements, + resources: ResourcesSpec, ) -> Optional[Requirements]: try: return Requirements( - resources=_combine_resources(fleet_requirements.resources, run_requirements.resources), + resources=resources.model_copy(deep=True), max_price=_get_min_optional(fleet_requirements.max_price, run_requirements.max_price), spot=_combine_spot_optional(fleet_requirements.spot, run_requirements.spot), reservation=get_single_value_optional( diff --git a/src/dstack/_internal/server/services/runs/plan.py b/src/dstack/_internal/server/services/runs/plan.py index ae626b1c30..3f9429f574 100644 --- a/src/dstack/_internal/server/services/runs/plan.py +++ b/src/dstack/_internal/server/services/runs/plan.py @@ -18,7 +18,9 @@ InstanceOfferWithAvailability, InstanceStatus, ) +from dstack._internal.core.models.placement import PlacementGroup from dstack._internal.core.models.profiles import CreationPolicy, Profile +from dstack._internal.core.models.resources import ResourcesSpec from dstack._internal.core.models.runs import ( Job, JobPlan, @@ -64,6 +66,7 @@ from dstack._internal.server.services.requirements.combine import ( combine_fleet_and_run_profiles, combine_fleet_and_run_requirements, + combine_fleet_and_run_requirements_with_resources, ) from dstack._internal.server.services.runs.spec import ( check_run_spec_requires_instance_mounts, @@ -83,6 +86,42 @@ _PER_FLEET_MAX_OFFERS = 100 +def _is_cloud_blocks_fleet(fleet_spec: FleetSpec) -> bool: + return fleet_spec.configuration.ssh_config is None and fleet_spec.configuration.blocks != 1 + + +def _get_cloud_blocks_backend_resources( + fleet_resources: ResourcesSpec, + run_resources: ResourcesSpec, +) -> Optional[ResourcesSpec]: + resources = fleet_resources.model_copy(deep=True) + + run_cpu_arch = run_resources.cpu.arch + if run_cpu_arch is not None: + if resources.cpu.arch is not None and resources.cpu.arch != run_cpu_arch: + return None + resources.cpu.arch = run_cpu_arch + + run_shm_size = run_resources.shm_size + if run_shm_size is not None and ( + resources.shm_size is None or run_shm_size < resources.shm_size + ): + resources.shm_size = run_shm_size + + run_disk = run_resources.disk + if run_disk is None: + return resources + if resources.disk is None: + resources.disk = run_disk.model_copy(deep=True) + return resources + + disk_size = resources.disk.size.intersect(run_disk.size) + if disk_size is None: + return None + resources.disk.size = disk_size + return resources + + async def get_job_plans( session: AsyncSession, project: ProjectModel, @@ -526,14 +565,51 @@ def get_run_profile_and_requirements_in_fleet( if profile is None: raise ValueError("Cannot combine fleet profile") fleet_requirements = get_fleet_requirements(fleet_spec) - requirements = combine_fleet_and_run_requirements( - fleet_requirements, job.job_spec.requirements - ) + if not _is_cloud_blocks_fleet(fleet_spec): + requirements = combine_fleet_and_run_requirements( + fleet_requirements, job.job_spec.requirements + ) + else: + requirements = combine_fleet_and_run_requirements_with_resources( + fleet_requirements=fleet_requirements, + run_requirements=job.job_spec.requirements, + resources=job.job_spec.requirements.resources, + ) if requirements is None: raise ValueError("Cannot combine fleet requirements") return profile, requirements +def _get_backend_offer_requirements_in_fleet( + job: Job, + run_spec: RunSpec, + fleet_spec: FleetSpec, +) -> tuple[Profile, Requirements, Optional[Requirements]]: + profile, offer_requirements = get_run_profile_and_requirements_in_fleet( + job=job, + run_spec=run_spec, + fleet_spec=fleet_spec, + ) + if not _is_cloud_blocks_fleet(fleet_spec): + return profile, offer_requirements, None + + fleet_requirements = get_fleet_requirements(fleet_spec) + backend_resources = _get_cloud_blocks_backend_resources( + fleet_resources=fleet_requirements.resources, + run_resources=job.job_spec.requirements.resources, + ) + if backend_resources is None: + raise ValueError("Cannot combine fleet requirements") + backend_requirements = combine_fleet_and_run_requirements_with_resources( + fleet_requirements=fleet_requirements, + run_requirements=job.job_spec.requirements, + resources=backend_resources, + ) + if backend_requirements is None: + raise ValueError("Cannot combine fleet requirements") + return profile, backend_requirements, offer_requirements + + async def _select_candidate_fleet_models( session: AsyncSession, project: ProjectModel, @@ -766,6 +842,9 @@ async def _get_backend_offers_in_fleet( job: Job, volumes: Optional[list[list[Volume]]], fleet_spec: Optional[FleetSpec] = None, + placement_group: Optional[PlacementGroup] = None, + exclude_not_available: bool = False, + master_job_provisioning_data: Optional[JobProvisioningData] = None, max_offers: Optional[int] = None, full_offers: bool = False, unallocated_resources: bool = False, @@ -774,19 +853,22 @@ async def _get_backend_offers_in_fleet( fleet_spec = get_fleet_spec(fleet_model) try: check_can_create_new_cloud_instance_in_fleet(fleet_model, fleet_spec) - profile, requirements = get_run_profile_and_requirements_in_fleet( - job=job, - run_spec=run_spec, - fleet_spec=fleet_spec, + profile, requirements, shared_offer_requirements = ( + _get_backend_offer_requirements_in_fleet( + job=job, + run_spec=run_spec, + fleet_spec=fleet_spec, + ) ) except ValueError: backend_offers = [] else: # Master job offers must be in the same cluster as existing instances. - master_instance_provisioning_data = get_fleet_master_instance_provisioning_data( - fleet_model=fleet_model, - fleet_spec=fleet_spec, - ) + if master_job_provisioning_data is None: + master_job_provisioning_data = get_fleet_master_instance_provisioning_data( + fleet_model=fleet_model, + fleet_spec=fleet_spec, + ) # Handle multinode for old jobs that don't have requirements.multinode set. # TODO: Drop multinode param. multinode = requirements.multinode or is_multinode_job(job) @@ -794,11 +876,15 @@ async def _get_backend_offers_in_fleet( project=project, profile=profile, requirements=requirements, + shared_offer_requirements=shared_offer_requirements, + exclude_not_available=exclude_not_available, multinode=multinode, - master_job_provisioning_data=master_instance_provisioning_data, + master_job_provisioning_data=master_job_provisioning_data, volumes=volumes, privileged=job.job_spec.privileged, instance_mounts=check_run_spec_requires_instance_mounts(run_spec), + placement_group=placement_group, + blocks=fleet_spec.configuration.blocks, max_offers=max_offers, full_offers=full_offers, unallocated_resources=unallocated_resources, diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py index e511241d7f..b3b651ceac 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py @@ -25,7 +25,7 @@ ) from dstack._internal.core.models.envs import Env from dstack._internal.core.models.fleets import FleetNodesSpec, InstanceGroupPlacement -from dstack._internal.core.models.instances import InstanceStatus +from dstack._internal.core.models.instances import InstanceOfferWithAvailability, InstanceStatus from dstack._internal.core.models.placement import PlacementGroup from dstack._internal.core.models.profiles import ( CreationPolicy, @@ -466,6 +466,136 @@ async def test_provisions_new_capacity_for_assigned_job_with_placeholder( ) assert len(res.scalars().all()) == 1 + async def test_provisions_shared_block_offer_for_new_capacity_in_blocks_fleet( + self, test_db, session: AsyncSession, worker: JobSubmittedWorker + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + fleet_spec = get_fleet_spec() + fleet_spec.configuration.nodes = FleetNodesSpec(min=0, target=0, max=2) + fleet_spec.configuration.blocks = "auto" + fleet_spec.configuration.resources = ResourcesSpec( + cpu=CPUSpec.parse("4..8"), + memory=Range[Memory](min=Memory.parse("8GB"), max=Memory.parse("32GB")), + gpu=None, + disk="100GB..", + ) + fleet = await create_fleet(session=session, project=project, spec=fleet_spec) + run_spec = get_run_spec( + repo_id=repo.name, + configuration=TaskConfiguration( + image="debian", + commands=["echo"], + resources=ResourcesSpec( + cpu=CPUSpec.parse("arm:1"), + memory=Range[Memory](min=Memory.parse("2GB"), max=None), + gpu=None, + disk="200GB..", + ), + ), + ) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_spec=run_spec, + fleet=fleet, + ) + job = await create_job(session=session, run=run) + + offer = get_instance_offer_with_availability( + backend=BackendType.AWS, + cpu_count=4, + memory_gib=8, + disk_gib=300, + ) + offer.instance.resources.cpu_arch = gpuhunt.CPUArchitecture.ARM + with patch("dstack._internal.server.services.backends.get_project_backends") as m: + backend_mock = Mock() + compute_mock = Mock(spec=ComputeMockSpec) + backend_mock.TYPE = BackendType.AWS + backend_mock.compute.return_value = compute_mock + m.return_value = [backend_mock] + compute_mock.get_offers.return_value = [offer] + compute_mock.run_job.return_value = get_job_provisioning_data( + dockerized=True, + backend=BackendType.AWS, + ) + + await _process_job(session=session, worker=worker, job_model=job) + + job = await _get_job(session, job.id) + assert job.status == JobStatus.SUBMITTED + assert job.instance is not None + placeholder_id = job.instance.id + + await _process_job(session=session, worker=worker, job_model=job) + + job = await _get_job(session, job.id) + assert job.status == JobStatus.PROVISIONING + assert job.instance is not None and job.instance.id == placeholder_id + assert job.instance.status == InstanceStatus.PROVISIONING + instance_offer = validate_json_extra_ignore( + InstanceOfferWithAvailability, get_or_error(job.instance.offer) + ) + assert instance_offer.blocks == 1 + assert instance_offer.total_blocks == 4 + assert instance_offer.instance.resources.cpu_arch == gpuhunt.CPUArchitecture.ARM + assert instance_offer.instance.resources.cpus == 4 + assert instance_offer.instance.resources.memory_mib == 8 * 1024 + assert instance_offer.instance.resources.disk.size_mib == 300 * 1024 + assert job.instance.total_blocks == 4 + assert job.instance.busy_blocks == 1 + runtime = validate_json_extra_ignore(JobRuntimeData, get_or_error(job.job_runtime_data)) + assert runtime.network_mode == NetworkMode.BRIDGE + assert runtime.offer is not None + assert runtime.offer.blocks == 1 + assert runtime.offer.total_blocks == 4 + assert runtime.offer.instance.resources.cpu_arch == gpuhunt.CPUArchitecture.ARM + assert runtime.offer.instance.resources.cpus == 1 + assert runtime.offer.instance.resources.memory_mib == 2 * 1024 + assert runtime.offer.instance.resources.disk.size_mib == 300 * 1024 + run_job_offer = compute_mock.run_job.call_args.args[2] + assert run_job_offer.blocks == 1 + assert run_job_offer.total_blocks == 4 + assert run_job_offer.instance.resources.cpu_arch == gpuhunt.CPUArchitecture.ARM + assert run_job_offer.instance.resources.cpus == 4 + assert run_job_offer.instance.resources.memory_mib == 8 * 1024 + assert run_job_offer.instance.resources.disk.size_mib == 300 * 1024 + run_job_requirements = compute_mock.run_job.call_args.args[7] + assert run_job_requirements.resources.cpu.arch == gpuhunt.CPUArchitecture.ARM + assert run_job_requirements.resources.cpu.count == Range[int](min=1, max=1) + assert run_job_requirements.resources.memory == Range[Memory]( + min=Memory.parse("2GB"), + max=None, + ) + assert run_job_requirements.resources.disk.size == Range[Memory]( + min=Memory.parse("200GB"), + max=None, + ) + + job.instance.status = InstanceStatus.IDLE + await session.commit() + + second_run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_spec=run_spec, + fleet=fleet, + ) + second_job = await create_job(session=session, run=second_run) + await _process_job(session=session, worker=worker, job_model=second_job) + + second_job = await _get_job(session, second_job.id) + assert second_job.status == JobStatus.SUBMITTED + assert second_job.instance is not None and second_job.instance.id == placeholder_id + await session.refresh(second_job.instance) + assert second_job.instance.busy_blocks == 2 + async def test_multinode_master_reuses_placeholder_when_provisioning_falls_back_to_run_job( self, test_db, session: AsyncSession, worker: JobSubmittedWorker ): diff --git a/src/tests/_internal/server/services/runs/test_plan.py b/src/tests/_internal/server/services/runs/test_plan.py index e55ff6de67..b6bc751585 100644 --- a/src/tests/_internal/server/services/runs/test_plan.py +++ b/src/tests/_internal/server/services/runs/test_plan.py @@ -1,6 +1,7 @@ import copy -from unittest.mock import AsyncMock, Mock +from unittest.mock import AsyncMock, Mock, patch +import gpuhunt import pytest from sqlalchemy.ext.asyncio import AsyncSession @@ -24,6 +25,7 @@ ) from dstack._internal.core.models.resources import CPUSpec, GPUSpec, Memory, Range, ResourcesSpec from dstack._internal.server.services.jobs import get_jobs_from_run_spec +from dstack._internal.server.services.offers import get_matching_shared_offer from dstack._internal.server.services.projects import get_project_model_by_name from dstack._internal.server.services.runs import get_plan from dstack._internal.server.services.runs.plan import ( @@ -32,9 +34,11 @@ _get_backend_offers_in_fleet, get_backend_offers_in_run_candidate_fleets, get_job_plans, + get_run_profile_and_requirements_in_fleet, get_targeted_instance_offers, ) from dstack._internal.server.testing.common import ( + ComputeMockSpec, create_export, create_fleet, create_instance, @@ -46,6 +50,7 @@ get_job_provisioning_data, get_remote_connection_info, get_run_spec, + get_ssh_fleet_configuration, ) pytestmark = pytest.mark.usefixtures("image_config_mock") @@ -807,6 +812,35 @@ async def test_skips_backend_offers_when_instances_specified( class TestGetBackendOffersInFleet: + @pytest.mark.asyncio + async def test_ssh_blocks_fleet_keeps_old_requirement_combination_semantics(self) -> None: + fleet_spec = get_fleet_spec(get_ssh_fleet_configuration(blocks="auto")) + fleet_spec.configuration.resources = ResourcesSpec( + cpu=CPUSpec.parse("4..8"), + memory=Range[Memory](min=Memory.parse("8GB"), max=Memory.parse("32GB")), + gpu=None, + ) + run_spec = get_run_spec( + repo_id="repo", + configuration=TaskConfiguration( + image="debian", + commands=["echo"], + resources=ResourcesSpec( + cpu=CPUSpec.parse("1"), + memory=Range[Memory](min=Memory.parse("2GB"), max=None), + gpu=None, + ), + ), + ) + jobs = await get_jobs_from_run_spec(run_spec=run_spec, secrets={}, replica_num=0) + + with pytest.raises(ValueError, match="Cannot combine fleet requirements"): + get_run_profile_and_requirements_in_fleet( + job=jobs[0], + run_spec=run_spec, + fleet_spec=fleet_spec, + ) + @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) async def test_keeps_unconstrained_offers_for_non_empty_cluster_fleet_without_elected_master( @@ -853,3 +887,229 @@ async def test_keeps_unconstrained_offers_for_non_empty_cluster_fleet_without_el get_offers_by_requirements_mock.await_args.kwargs["master_job_provisioning_data"] is None ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_returns_single_block_offer_when_new_capacity_can_share_growing_blocks_fleet( + self, test_db, session: AsyncSession + ) -> None: + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + repo = await create_repo(session=session, project_id=project.id) + fleet_spec = get_fleet_spec() + fleet_spec.configuration.nodes = FleetNodesSpec(min=0, target=0, max=2) + fleet_spec.configuration.blocks = "auto" + fleet_spec.configuration.resources = ResourcesSpec( + cpu=CPUSpec.parse("4..8"), + memory=Range[Memory](min=Memory.parse("8GB"), max=Memory.parse("32GB")), + gpu=None, + disk="100GB..", + ) + fleet = await create_fleet(session=session, project=project, spec=fleet_spec) + run_spec = get_run_spec( + repo_id=repo.name, + configuration=TaskConfiguration( + image="debian", + commands=["echo"], + resources=ResourcesSpec( + cpu=CPUSpec.parse("arm:1"), + memory=Range[Memory](min=Memory.parse("2GB"), max=None), + gpu=None, + disk="200GB..", + ), + ), + ) + jobs = await get_jobs_from_run_spec(run_spec=run_spec, secrets={}, replica_num=0) + offer = get_instance_offer_with_availability(cpu_count=4, memory_gib=8, disk_gib=300) + offer.instance.resources.cpu_arch = gpuhunt.CPUArchitecture.ARM + with ( + patch("dstack._internal.server.services.backends.get_project_backends") as m, + patch( + "dstack._internal.server.services.offers.get_matching_shared_offer", + wraps=get_matching_shared_offer, + ) as matching_shared_offer_mock, + ): + backend = Mock() + backend.TYPE = BackendType.AWS + compute_mock = Mock(spec=ComputeMockSpec) + backend.compute.return_value = compute_mock + compute_mock.get_offers.return_value = [offer] + m.return_value = [backend] + + offers = await _get_backend_offers_in_fleet( + project=project, + fleet_model=fleet, + run_spec=run_spec, + job=jobs[0], + volumes=None, + ) + + assert [ + (selected_backend, selected_offer.blocks) + for selected_backend, selected_offer in offers + ] == [(backend, 1)] + assert offers[0][1].total_blocks == 4 + assert offers[0][1].instance.resources.cpus == 4 + assert offers[0][1].instance.resources.memory_mib == 8 * 1024 + queried_requirements = compute_mock.get_offers.call_args.args[0] + shared_offer_requirements = matching_shared_offer_mock.call_args.kwargs["requirements"] + assert queried_requirements.resources.cpu.count == Range[int](min=4, max=8) + assert queried_requirements.resources.cpu.arch == gpuhunt.CPUArchitecture.ARM + assert queried_requirements.resources.memory == Range[Memory]( + min=Memory.parse("8GB"), + max=Memory.parse("32GB"), + ) + assert queried_requirements.resources.disk.size == Range[Memory]( + min=Memory.parse("200GB"), + max=None, + ) + assert shared_offer_requirements.resources.cpu.count == Range[int](min=1, max=1) + assert shared_offer_requirements.resources.cpu.arch == gpuhunt.CPUArchitecture.ARM + assert shared_offer_requirements.resources.memory == Range[Memory]( + min=Memory.parse("2GB"), + max=None, + ) + assert shared_offer_requirements.resources.disk.size == Range[Memory]( + min=Memory.parse("200GB"), + max=None, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_returns_multi_block_offer_when_new_capacity_needs_more_than_one_block( + self, test_db, session: AsyncSession + ) -> None: + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + repo = await create_repo(session=session, project_id=project.id) + fleet_spec = get_fleet_spec() + fleet_spec.configuration.nodes = FleetNodesSpec(min=0, target=0, max=2) + fleet_spec.configuration.blocks = "auto" + fleet_spec.configuration.resources = ResourcesSpec( + cpu=CPUSpec.parse("4..8"), + memory=Range[Memory](min=Memory.parse("8GB"), max=Memory.parse("32GB")), + gpu=None, + ) + fleet = await create_fleet(session=session, project=project, spec=fleet_spec) + run_spec = get_run_spec( + repo_id=repo.name, + configuration=TaskConfiguration( + image="debian", + commands=["echo"], + resources=ResourcesSpec( + cpu=CPUSpec.parse("2"), + memory=Range[Memory](min=Memory.parse("3GB"), max=None), + gpu=None, + ), + ), + ) + jobs = await get_jobs_from_run_spec(run_spec=run_spec, secrets={}, replica_num=0) + offer = get_instance_offer_with_availability(cpu_count=4, memory_gib=8) + with ( + patch("dstack._internal.server.services.backends.get_project_backends") as m, + patch( + "dstack._internal.server.services.offers.get_matching_shared_offer", + wraps=get_matching_shared_offer, + ) as matching_shared_offer_mock, + ): + backend = Mock() + backend.TYPE = BackendType.AWS + compute_mock = Mock(spec=ComputeMockSpec) + backend.compute.return_value = compute_mock + compute_mock.get_offers.return_value = [offer] + m.return_value = [backend] + + offers = await _get_backend_offers_in_fleet( + project=project, + fleet_model=fleet, + run_spec=run_spec, + job=jobs[0], + volumes=None, + ) + + assert [ + (selected_backend, selected_offer.blocks) + for selected_backend, selected_offer in offers + ] == [(backend, 2)] + assert offers[0][1].total_blocks == 4 + assert offers[0][1].instance.resources.cpus == 4 + assert offers[0][1].instance.resources.memory_mib == 8 * 1024 + queried_requirements = compute_mock.get_offers.call_args.args[0] + shared_offer_requirements = matching_shared_offer_mock.call_args.kwargs["requirements"] + assert queried_requirements.resources.cpu.count == Range[int](min=4, max=8) + assert queried_requirements.resources.memory == Range[Memory]( + min=Memory.parse("8GB"), + max=Memory.parse("32GB"), + ) + assert shared_offer_requirements.resources.cpu.count == Range[int](min=2, max=2) + assert shared_offer_requirements.resources.memory == Range[Memory]( + min=Memory.parse("3GB"), + max=None, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_excludes_non_create_backends_for_cloud_blocks_fleet( + self, test_db, session: AsyncSession + ) -> None: + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + repo = await create_repo(session=session, project_id=project.id) + fleet_spec = get_fleet_spec() + fleet_spec.configuration.nodes = FleetNodesSpec(min=0, target=0, max=2) + fleet_spec.configuration.blocks = "auto" + fleet_spec.configuration.resources = ResourcesSpec( + cpu=CPUSpec.parse("4..8"), + memory=Range[Memory](min=Memory.parse("8GB"), max=Memory.parse("32GB")), + gpu=None, + ) + fleet = await create_fleet(session=session, project=project, spec=fleet_spec) + run_spec = get_run_spec( + repo_id=repo.name, + configuration=TaskConfiguration( + image="debian", + commands=["echo"], + resources=ResourcesSpec( + cpu=CPUSpec.parse("1"), + memory=Range[Memory](min=Memory.parse("2GB"), max=None), + gpu=None, + ), + ), + ) + jobs = await get_jobs_from_run_spec(run_spec=run_spec, secrets={}, replica_num=0) + + with patch("dstack._internal.server.services.backends.get_project_backends") as m: + aws_backend = Mock() + aws_backend.TYPE = BackendType.AWS + aws_compute_mock = Mock(spec=ComputeMockSpec) + aws_backend.compute.return_value = aws_compute_mock + aws_offer = get_instance_offer_with_availability( + backend=BackendType.AWS, cpu_count=4, memory_gib=8 + ) + aws_compute_mock.get_offers.return_value = [aws_offer] + + kubernetes_backend = Mock() + kubernetes_backend.TYPE = BackendType.KUBERNETES + kubernetes_compute_mock = Mock() + kubernetes_backend.compute.return_value = kubernetes_compute_mock + kubernetes_offer = get_instance_offer_with_availability( + backend=BackendType.KUBERNETES, + region="", + cpu_count=4, + memory_gib=8, + ) + kubernetes_compute_mock.get_offers.return_value = [kubernetes_offer] + + m.return_value = [aws_backend, kubernetes_backend] + + offers = await _get_backend_offers_in_fleet( + project=project, + fleet_model=fleet, + run_spec=run_spec, + job=jobs[0], + volumes=None, + ) + + assert [backend.TYPE for backend, _ in offers] == [BackendType.AWS] + aws_compute_mock.get_offers.assert_called_once() + kubernetes_compute_mock.get_offers.assert_not_called() diff --git a/src/tests/_internal/server/services/test_offers.py b/src/tests/_internal/server/services/test_offers.py index 017d6aa813..61c8466560 100644 --- a/src/tests/_internal/server/services/test_offers.py +++ b/src/tests/_internal/server/services/test_offers.py @@ -4,10 +4,11 @@ from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.profiles import Profile -from dstack._internal.core.models.resources import ResourcesSpec +from dstack._internal.core.models.resources import CPUSpec, Memory, Range, ResourcesSpec from dstack._internal.core.models.runs import Requirements from dstack._internal.server.services.offers import get_offers_by_requirements from dstack._internal.server.testing.common import ( + ComputeMockSpec, get_instance_offer_with_availability, get_kubernetes_volume_configuration, get_volume, @@ -211,3 +212,108 @@ async def test_returns_no_offers_for_multinode_instance_mounts_and_non_multinode ) m.assert_awaited_once() assert res == [] + + @pytest.mark.asyncio + async def test_filters_shared_offers_before_max_offers(self): + profile = Profile(name="test") + requirements = Requirements( + resources=ResourcesSpec( + cpu=CPUSpec.parse("4..8"), + memory=Range[Memory](min=Memory.parse("8GB"), max=Memory.parse("32GB")), + gpu=None, + ) + ) + shared_offer_requirements = Requirements( + resources=ResourcesSpec( + cpu=CPUSpec.parse("2"), + memory=Range[Memory](min=Memory.parse("3GB"), max=None), + gpu=None, + ) + ) + with patch("dstack._internal.server.services.backends.get_project_backends") as m: + backend_mock = Mock() + backend_mock.TYPE = BackendType.SEEWEB + backend_mock.compute.return_value = Mock(spec=ComputeMockSpec) + nonmatching_offer = get_instance_offer_with_availability( + backend=BackendType.SEEWEB, + instance_type="nonmatching", + cpu_count=6, + memory_gib=8, + ) + matching_offer = get_instance_offer_with_availability( + backend=BackendType.SEEWEB, + instance_type="matching", + cpu_count=4, + memory_gib=8, + ) + backend_mock.compute.return_value.get_offers.return_value = [ + nonmatching_offer, + matching_offer, + ] + m.return_value = [backend_mock] + res = await get_offers_by_requirements( + project=Mock(), + profile=profile, + requirements=requirements, + shared_offer_requirements=shared_offer_requirements, + blocks="auto", + max_offers=1, + ) + + assert [ + (backend, offer.instance.name, offer.blocks, offer.total_blocks) + for backend, offer in res + ] == [(backend_mock, "matching", 2, 4)] + + @pytest.mark.asyncio + async def test_filters_non_create_dstack_wrapper_offers_before_max_offers(self): + profile = Profile(name="test") + requirements = Requirements( + resources=ResourcesSpec( + cpu=CPUSpec.parse("4..8"), + memory=Range[Memory](min=Memory.parse("8GB"), max=Memory.parse("32GB")), + gpu=None, + ) + ) + shared_offer_requirements = Requirements( + resources=ResourcesSpec( + cpu=CPUSpec.parse("2"), + memory=Range[Memory](min=Memory.parse("3GB"), max=None), + gpu=None, + ) + ) + with patch("dstack._internal.server.services.backends.get_project_backends") as m: + backend_mock = Mock() + backend_mock.TYPE = BackendType.DSTACK + backend_mock.compute.return_value = Mock(spec=ComputeMockSpec) + container_offer = get_instance_offer_with_availability( + backend=BackendType.KUBERNETES, + region="", + instance_type="container-offer", + cpu_count=4, + memory_gib=8, + ) + vm_offer = get_instance_offer_with_availability( + backend=BackendType.SEEWEB, + instance_type="vm-offer", + cpu_count=4, + memory_gib=8, + ) + backend_mock.compute.return_value.get_offers.return_value = [ + container_offer, + vm_offer, + ] + m.return_value = [backend_mock] + res = await get_offers_by_requirements( + project=Mock(), + profile=profile, + requirements=requirements, + shared_offer_requirements=shared_offer_requirements, + blocks="auto", + max_offers=1, + ) + + assert [ + (backend, offer.backend, offer.instance.name, offer.blocks, offer.total_blocks) + for backend, offer in res + ] == [(backend_mock, BackendType.SEEWEB, "vm-offer", 2, 4)]