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,