Merge remote-tracking branch 'origin/master' into pr_kubernetes_offers

Author: un-def

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

@@ -182,6 +182,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:
@@ -256,6 +257,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(
@@ -1112,6 +1114,7 @@ def _create_job_pod(
     run_spec: RunSpec,
     job_spec: JobSpec,
     volumes: list[Volume],
+    requirements: Requirements,
     authorized_keys: list[str],
 ) -> None:
     node_affinity: Optional[client.V1NodeAffinity] = None
@@ -1120,7 +1123,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
     resource_requests = ResourceRequests.from_resources_spec(resources_spec)
     resource_limits = ResourceLimits.from_resources_spec(resources_spec)
     gpu_resource: Optional[AnyKubernetesGPUResource] = None

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/core/services/repos.py

@@ -2,7 +2,7 @@
 from contextlib import suppress
 from pathlib import Path
 from tempfile import NamedTemporaryFile
-from typing import Optional
+from typing import Optional, cast
 
 import git
 import git.cmd
@@ -205,8 +205,12 @@ def _get_repo_default_branch(url: str, env: dict[str, str]) -> Optional[str]:
     # See: https://github.com/git/git/commit/3d4355712b9fe77a96ad4ad877d92dc7ff6e0874
     # See: https://gist.github.com/ChrisTollefson/ab9c0a5d1dd4dd615217345c6936a307
     _git = git.cmd.Git()(c="credential.helper=")
+    # Type cast is required since GitPython 3.1.51 where Git.ls_remote() was implemented as
+    # an actual method wrapping Git.execute() but no proper @overload signatures were added.
+    # Our call is translated to:
+    # Git.execute(..., with_extended_output=False, as_process=False, stdout_as_string=True) -> str
+    output = cast(str, _git.ls_remote("--symref", url, "HEAD", env=env))
     # output example: "ref: refs/heads/dev\tHEAD\n545344f77c0df78367085952a97fc3a058eb4c65\tHEAD"
-    output: str = _git.ls_remote("--symref", url, "HEAD", env=env)
     for line in output.splitlines():
         # line format: `<oid> TAB <ref> LF`
         oid, _, ref = line.partition("\t")

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,