Support placement groups for elastic fleets (#3282)

* Support placement groups for elastic fleets

* Test placement group creation

* Update run_job() signatures

* Fix non-master jobs stuck with waiting master job

* Simplify tests

* Fix run_jobs signature

Author: r4victor

src/dstack/_internal/core/backends/base/compute.py

@@ -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.
@@ -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()`.
@@ -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,
@@ -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
 

src/dstack/_internal/core/backends/kubernetes/compute.py

@@ -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
@@ -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

src/dstack/_internal/core/backends/local/compute.py

@@ -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,

src/dstack/_internal/core/backends/runpod/compute.py

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

src/dstack/_internal/core/backends/template/compute.py.jinja

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

src/dstack/_internal/core/backends/vastai/compute.py

@@ -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
@@ -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