From e201a27a64d7abad87e481b773137fa223d6ee1a Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Thu, 9 Jul 2026 13:53:07 +0000 Subject: [PATCH] Pass effective Requirements to Compute's run_job/run_jobs Most backend don't need it -- they provision machines/containers based on offers, which are already filtered against the effective Requirements, but both Slurm and Kubernetes backends completely ignore provided offers and rely on Requirements instead, which they previously took from the JobSpec, which is generated from the Run configuration only. Fixes: https://github.com/dstackai/dstack/issues/4019 --- src/dstack/_internal/core/backends/base/compute.py | 3 +++ src/dstack/_internal/core/backends/kubernetes/compute.py | 5 ++++- src/dstack/_internal/core/backends/runpod/compute.py | 2 ++ src/dstack/_internal/core/backends/slurm/compute.py | 7 ++++++- .../_internal/core/backends/template/compute.py.jinja | 1 + src/dstack/_internal/core/backends/vastai/compute.py | 1 + .../server/background/pipeline_tasks/jobs_submitted.py | 2 ++ 7 files changed, 19 insertions(+), 2 deletions(-) diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py index 477e9747e0..7cb7c84b35 100644 --- a/src/dstack/_internal/core/backends/base/compute.py +++ b/src/dstack/_internal/core/backends/base/compute.py @@ -128,6 +128,7 @@ def run_job( project_ssh_private_key: str, volumes: List[Volume], placement_group: Optional[PlacementGroup], + requirements: Requirements, ) -> JobProvisioningData: """ Launches a new instance for the job. It should return `JobProvisioningData` ASAP. @@ -307,6 +308,7 @@ def run_job( project_ssh_private_key: str, volumes: List[Volume], placement_group: Optional[PlacementGroup], + requirements: Requirements, ) -> JobProvisioningData: """ The default `run_job()` implementation for all backends that support `create_instance()`. @@ -358,6 +360,7 @@ def run_jobs( project_ssh_public_key: str, project_ssh_private_key: str, placement_group: Optional[PlacementGroup], + requirements: Requirements, ) -> ComputeGroupProvisioningData: pass diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py index a9572fa336..43a9cce9c1 100644 --- a/src/dstack/_internal/core/backends/kubernetes/compute.py +++ b/src/dstack/_internal/core/backends/kubernetes/compute.py @@ -173,6 +173,7 @@ def run_job( project_ssh_private_key: str, volumes: list[Volume], placement_group: Optional[PlacementGroup], + requirements: Requirements, ) -> JobProvisioningData: cluster = self.region_cluster_map.get(instance_offer.region) if cluster is None: @@ -247,6 +248,7 @@ def run_job( run_spec=run.run_spec, job_spec=job.job_spec, volumes=volumes, + requirements=requirements, authorized_keys=authorized_keys, ) exit_stack.callback( @@ -1098,6 +1100,7 @@ def _create_job_pod( run_spec: RunSpec, job_spec: JobSpec, volumes: list[Volume], + requirements: Requirements, authorized_keys: list[str], ) -> None: resources_requests: dict[str, str] = {} @@ -1108,7 +1111,7 @@ def _create_job_pod( volume_mounts: list[client.V1VolumeMount] = [] env_vars: list[client.V1EnvVar] = [] - resources_spec = job_spec.requirements.resources + resources_spec = requirements.resources assert isinstance(resources_spec.cpu, CPUSpec) if (cpu_min := resources_spec.cpu.count.min) is not None: resources_requests["cpu"] = str(cpu_min) diff --git a/src/dstack/_internal/core/backends/runpod/compute.py b/src/dstack/_internal/core/backends/runpod/compute.py index 7158c6950b..bafdb91b89 100644 --- a/src/dstack/_internal/core/backends/runpod/compute.py +++ b/src/dstack/_internal/core/backends/runpod/compute.py @@ -130,6 +130,7 @@ def run_job( project_ssh_private_key: str, volumes: List[Volume], placement_group: Optional[PlacementGroup], + requirements: Requirements, ) -> JobProvisioningData: assert run.run_spec.ssh_key_pub is not None instance_config = InstanceConfiguration( @@ -245,6 +246,7 @@ def run_jobs( project_ssh_public_key: str, project_ssh_private_key: str, placement_group: Optional[PlacementGroup], + requirements: Requirements, ) -> ComputeGroupProvisioningData: master_job_configuration = job_configurations[0] master_job = master_job_configuration.job diff --git a/src/dstack/_internal/core/backends/slurm/compute.py b/src/dstack/_internal/core/backends/slurm/compute.py index e3389272b5..9d7b40c539 100644 --- a/src/dstack/_internal/core/backends/slurm/compute.py +++ b/src/dstack/_internal/core/backends/slurm/compute.py @@ -127,12 +127,14 @@ def run_job( project_ssh_private_key: str, volumes: list[Volume], placement_group: Optional[PlacementGroup], + requirements: Requirements, ) -> JobProvisioningData: compute_provisioning_data = self._run_slurm_job( run=run, job=job, instance_offer=instance_offer, project_ssh_public_key=project_ssh_public_key, + requirements=requirements, ) return compute_provisioning_data.job_provisioning_datas[0] @@ -144,6 +146,7 @@ def run_jobs( project_ssh_public_key: str, project_ssh_private_key: str, placement_group: Optional[PlacementGroup], + requirements: Requirements, ) -> ComputeGroupProvisioningData: master_job = job_configurations[0].job return self._run_slurm_job( @@ -151,6 +154,7 @@ def run_jobs( job=master_job, instance_offer=instance_offer, project_ssh_public_key=project_ssh_public_key, + requirements=requirements, ) def terminate_instance( @@ -172,6 +176,7 @@ def _run_slurm_job( job: Job, instance_offer: InstanceOfferWithAvailability, project_ssh_public_key: str, + requirements: Requirements, ) -> ComputeGroupProvisioningData: if job.job_spec.registry_auth is not None: self._skip_offer_cache.add(run, job, instance_offer) @@ -196,7 +201,7 @@ def _run_slurm_job( authorized_keys = [project_ssh_public_key.strip(), run.run_spec.ssh_key_pub.strip()] node_count = job.job_spec.jobs_per_replica - resources_spec = job.job_spec.requirements.resources + resources_spec = requirements.resources requested_resources = get_requested_resources_from_resources_spec(resources_spec) partitions = _get_cluster_partitions(cluster, requested_resources) diff --git a/src/dstack/_internal/core/backends/template/compute.py.jinja b/src/dstack/_internal/core/backends/template/compute.py.jinja index cb4c4a8b09..046e429ea2 100644 --- a/src/dstack/_internal/core/backends/template/compute.py.jinja +++ b/src/dstack/_internal/core/backends/template/compute.py.jinja @@ -84,6 +84,7 @@ class {{ backend_name }}Compute( project_ssh_private_key: str, volumes: List[Volume], placement_group: Optional[PlacementGroup], + requirements: Requirements, ) -> JobProvisioningData: # TODO: Implement if create_instance() is not implemented. Delete otherwise. raise NotImplementedError() diff --git a/src/dstack/_internal/core/backends/vastai/compute.py b/src/dstack/_internal/core/backends/vastai/compute.py index 915d98f105..fc1a957078 100644 --- a/src/dstack/_internal/core/backends/vastai/compute.py +++ b/src/dstack/_internal/core/backends/vastai/compute.py @@ -118,6 +118,7 @@ def run_job( project_ssh_private_key: str, volumes: List[Volume], placement_group: Optional[PlacementGroup], + requirements: Requirements, ) -> JobProvisioningData: instance_name = generate_unique_instance_name_for_job( run, job, max_length=MAX_INSTANCE_NAME_LEN 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 35f613833c..0ae9bb4432 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py @@ -2292,6 +2292,7 @@ async def _provision_new_capacity( project_ssh_public_key, project_ssh_private_key, placement_group_model_to_placement_group_optional(placement_group_model), + requirements, ) return _ProvisionNewCapacityResult( provisioning_data=compute_group_provisioning_data, @@ -2315,6 +2316,7 @@ async def _provision_new_capacity( project_ssh_private_key, offer_volumes, placement_group_model_to_placement_group_optional(placement_group_model), + requirements, ) return _ProvisionNewCapacityResult( provisioning_data=job_provisioning_data,