Pass effective Requirements to Compute's run_job/run_jobs

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
 

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)

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

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,13 +146,15 @@ 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(
             run=run,
             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)

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()

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

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,