Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion src/dstack/_internal/core/backends/base/compute.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,7 @@ def run_job(
project_ssh_public_key: str,
project_ssh_private_key: str,
volumes: List[Volume],
placement_group: Optional[PlacementGroup],
) -> JobProvisioningData:
"""
Launches a new instance for the job. It should return `JobProvisioningData` ASAP.
Expand Down Expand Up @@ -287,6 +288,7 @@ def run_job(
project_ssh_public_key: str,
project_ssh_private_key: str,
volumes: List[Volume],
placement_group: Optional[PlacementGroup],
) -> JobProvisioningData:
"""
The default `run_job()` implementation for all backends that support `create_instance()`.
Expand All @@ -303,7 +305,9 @@ def run_job(
)
instance_offer = instance_offer.copy()
self._restrict_instance_offer_az_to_volumes_az(instance_offer, volumes)
return self.create_instance(instance_offer, instance_config, placement_group=None)
return self.create_instance(
instance_offer, instance_config, placement_group=placement_group
)

def _restrict_instance_offer_az_to_volumes_az(
self,
Expand Down Expand Up @@ -335,6 +339,7 @@ def run_jobs(
instance_offer: InstanceOfferWithAvailability,
project_ssh_public_key: str,
project_ssh_private_key: str,
placement_group: Optional[PlacementGroup],
) -> ComputeGroupProvisioningData:
pass

Expand Down
2 changes: 2 additions & 0 deletions src/dstack/_internal/core/backends/kubernetes/compute.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@
Resources,
SSHConnectionParams,
)
from dstack._internal.core.models.placement import PlacementGroup
from dstack._internal.core.models.resources import CPUSpec, GPUSpec, Memory
from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run
from dstack._internal.core.models.volumes import Volume
Expand Down Expand Up @@ -131,6 +132,7 @@ def run_job(
project_ssh_public_key: str,
project_ssh_private_key: str,
volumes: list[Volume],
placement_group: Optional[PlacementGroup],
) -> JobProvisioningData:
instance_name = generate_unique_instance_name_for_job(run, job)
assert run.run_spec.ssh_key_pub is not None
Expand Down
1 change: 1 addition & 0 deletions src/dstack/_internal/core/backends/local/compute.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@ def run_job(
project_ssh_public_key: str,
project_ssh_private_key: str,
volumes: List[Volume],
placement_group: Optional[PlacementGroup],
) -> JobProvisioningData:
return JobProvisioningData(
backend=instance_offer.backend,
Expand Down
3 changes: 3 additions & 0 deletions src/dstack/_internal/core/backends/runpod/compute.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
InstanceOfferWithAvailability,
SSHKey,
)
from dstack._internal.core.models.placement import PlacementGroup
from dstack._internal.core.models.resources import Memory, Range
from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run
from dstack._internal.core.models.volumes import Volume, VolumeProvisioningData
Expand Down Expand Up @@ -109,6 +110,7 @@ def run_job(
project_ssh_public_key: str,
project_ssh_private_key: str,
volumes: List[Volume],
placement_group: Optional[PlacementGroup],
) -> JobProvisioningData:
assert run.run_spec.ssh_key_pub is not None
instance_config = InstanceConfiguration(
Expand Down Expand Up @@ -216,6 +218,7 @@ def run_jobs(
instance_offer: InstanceOfferWithAvailability,
project_ssh_public_key: str,
project_ssh_private_key: str,
placement_group: Optional[PlacementGroup],
) -> ComputeGroupProvisioningData:
master_job_configuration = job_configurations[0]
master_job = master_job_configuration.job
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ class {{ backend_name }}Compute(
project_ssh_public_key: str,
project_ssh_private_key: str,
volumes: List[Volume],
placement_group: Optional[PlacementGroup],
) -> JobProvisioningData:
# TODO: Implement if create_instance() is not implemented. Delete otherwise.
raise NotImplementedError()
Expand Down
2 changes: 2 additions & 0 deletions src/dstack/_internal/core/backends/vastai/compute.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
InstanceOfferWithAvailability,
InstanceRuntime,
)
from dstack._internal.core.models.placement import PlacementGroup
from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run
from dstack._internal.core.models.volumes import Volume
from dstack._internal.utils.logging import get_logger
Expand Down Expand Up @@ -82,6 +83,7 @@ def run_job(
project_ssh_public_key: str,
project_ssh_private_key: str,
volumes: List[Volume],
placement_group: Optional[PlacementGroup],
) -> JobProvisioningData:
instance_name = generate_unique_instance_name_for_job(
run, job, max_length=MAX_INSTANCE_NAME_LEN
Expand Down
Loading