dstack 0.19.7__py3-none-any.whl → 0.19.9__py3-none-any.whl

dstack/_internal/cli/services/args.py

@@ -19,8 +19,8 @@ def port_mapping(v: str) -> PortMapping:
     return PortMapping.parse(v)
 
 
-def cpu_spec(v: str) -> resources.Range[int]:
-    return parse_obj_as(resources.Range[int], v)
+def cpu_spec(v: str) -> dict:
+    return resources.CPUSpec.parse(v)
 
 
 def memory_spec(v: str) -> resources.Range[resources.Memory]:

dstack/_internal/cli/services/configurators/run.py

@@ -6,9 +6,10 @@ from pathlib import Path
 from typing import Dict, List, Optional, Set, Tuple
 
 import gpuhunt
+from pydantic import parse_obj_as
 
 import dstack._internal.core.models.resources as resources
-from dstack._internal.cli.services.args import disk_spec, gpu_spec, port_mapping
+from dstack._internal.cli.services.args import cpu_spec, disk_spec, gpu_spec, port_mapping
 from dstack._internal.cli.services.configurators.base import (
     ApplyEnvVarsConfiguratorMixin,
     BaseApplyConfigurator,
@@ -39,6 +40,7 @@ from dstack._internal.core.models.configurations import (
     TaskConfiguration,
 )
 from dstack._internal.core.models.repos.base import Repo
+from dstack._internal.core.models.resources import CPUSpec
 from dstack._internal.core.models.runs import JobSubmission, JobTerminationReason, RunStatus
 from dstack._internal.core.services.configs import ConfigManager
 from dstack._internal.core.services.diff import diff_models
@@ -72,6 +74,7 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
     ):
         self.apply_args(conf, configurator_args, unknown_args)
         self.validate_gpu_vendor_and_image(conf)
+        self.validate_cpu_arch_and_image(conf)
         if repo is None:
             repo = self.api.repos.load(Path.cwd())
         config_manager = ConfigManager()
@@ -95,6 +98,8 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
         print_run_plan(run_plan, max_offers=configurator_args.max_offers)
 
         confirm_message = "Submit a new run?"
+        if conf.name:
+            confirm_message = f"Submit the run [code]{conf.name}[/]?"
         stop_run_name = None
         if run_plan.current_resource is not None:
             changed_fields = []
@@ -127,11 +132,6 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
                     f"Active run [code]{conf.name}[/] already exists and cannot be updated in-place."
                 )
                 confirm_message = "Stop and override the run?"
-            else:
-                console.print(f"Finished run [code]{conf.name}[/] already exists.")
-                confirm_message = "Override the run?"
-        elif conf.name:
-            confirm_message = f"Submit the run [code]{conf.name}[/]?"
 
         if not command_args.yes and not confirm_ask(confirm_message):
             console.print("\nExiting...")
@@ -289,6 +289,14 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
             default=default_max_offers,
         )
         cls.register_env_args(configuration_group)
+        configuration_group.add_argument(
+            "--cpu",
+            type=cpu_spec,
+            help="Request CPU for the run. "
+            "The format is [code]ARCH[/]:[code]COUNT[/] (all parts are optional)",
+            dest="cpu_spec",
+            metavar="SPEC",
+        )
         configuration_group.add_argument(
             "--gpu",
             type=gpu_spec,
@@ -310,6 +318,8 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
         apply_profile_args(args, conf)
         if args.run_name:
             conf.name = args.run_name
+        if args.cpu_spec:
+            conf.resources.cpu = resources.CPUSpec.parse_obj(args.cpu_spec)
         if args.gpu_spec:
             conf.resources.gpu = resources.GPUSpec.parse_obj(args.gpu_spec)
         if args.disk_spec:
@@ -342,7 +352,7 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
 
     def validate_gpu_vendor_and_image(self, conf: BaseRunConfiguration) -> None:
         """
-        Infers `resources.gpu.vendor` if not set, requires `image` if the vendor is AMD.
+        Infers and sets `resources.gpu.vendor` if not set, requires `image` if the vendor is AMD.
         """
         gpu_spec = conf.resources.gpu
         if gpu_spec is None:
@@ -400,6 +410,29 @@ class BaseRunConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
                 "`image` is required if `resources.gpu.vendor` is `tenstorrent`"
             )
 
+    def validate_cpu_arch_and_image(self, conf: BaseRunConfiguration) -> None:
+        """
+        Infers `resources.cpu.arch` if not set, requires `image` if the architecture is ARM.
+        """
+        # TODO: Remove in 0.20. Use conf.resources.cpu directly
+        cpu_spec = parse_obj_as(CPUSpec, conf.resources.cpu)
+        arch = cpu_spec.arch
+        if arch is None:
+            gpu_spec = conf.resources.gpu
+            if (
+                gpu_spec is not None
+                and gpu_spec.vendor in [None, gpuhunt.AcceleratorVendor.NVIDIA]
+                and gpu_spec.name
+                and any(map(gpuhunt.is_nvidia_superchip, gpu_spec.name))
+            ):
+                arch = gpuhunt.CPUArchitecture.ARM
+            else:
+                arch = gpuhunt.CPUArchitecture.X86
+        # NOTE: We don't set the inferred resources.cpu.arch for compatibility with older servers.
+        # Servers with ARM support set the arch using the same logic.
+        if arch == gpuhunt.CPUArchitecture.ARM and conf.image is None:
+            raise ConfigurationError("`image` is required if `resources.cpu.arch` is `arm`")
+
 
 class RunWithPortsConfigurator(BaseRunConfigurator):
     @classmethod
@@ -524,7 +557,9 @@ def print_finished_message(run: Run):
         console.print("[code]Done[/]")
         return
 
-    termination_reason, termination_reason_message = _get_run_termination_reason(run)
+    termination_reason, termination_reason_message, exit_status = (
+        _get_run_termination_reason_and_exit_status(run)
+    )
     message = "Run failed due to unknown reason. Check CLI, server, and run logs."
     if run.status == RunStatus.TERMINATED:
         message = "Run terminated due to unknown reason. Check CLI, server, and run logs."
@@ -536,13 +571,15 @@ def print_finished_message(run: Run):
             "Check CLI and server logs for more details."
         )
     elif termination_reason is not None:
+        exit_status_details = f"Exit status: {exit_status}.\n" if exit_status else ""
         error_details = (
             f"Error: {termination_reason_message}\n" if termination_reason_message else ""
         )
         message = (
             f"Run failed with error code {termination_reason.name}.\n"
+            f"{exit_status_details}"
             f"{error_details}"
-            "Check CLI, server, and run logs for more details."
+            f"Check [bold]dstack logs -d {run.name}[/bold] for more details."
         )
     console.print(f"[error]{message}[/]")
 
@@ -553,14 +590,20 @@ def get_run_exit_code(run: Run) -> int:
     return 1
 
 
-def _get_run_termination_reason(run: Run) -> Tuple[Optional[JobTerminationReason], Optional[str]]:
+def _get_run_termination_reason_and_exit_status(
+    run: Run,
+) -> Tuple[Optional[JobTerminationReason], Optional[str], Optional[int]]:
     if len(run._run.jobs) == 0:
-        return None, None
+        return None, None, None
     job = run._run.jobs[0]
     if len(job.job_submissions) == 0:
-        return None, None
+        return None, None, None
     job_submission = job.job_submissions[0]
-    return job_submission.termination_reason, job_submission.termination_reason_message
+    return (
+        job_submission.termination_reason,
+        job_submission.termination_reason_message,
+        job_submission.exit_status,
+    )
 
 
 def _run_resubmitted(run: Run, current_job_submission: Optional[JobSubmission]) -> bool:

dstack/_internal/cli/utils/run.py

@@ -1,4 +1,4 @@
-import os
+import shutil
 from typing import Any, Dict, List, Optional, Union
 
 from rich.markup import escape
@@ -95,7 +95,7 @@ def print_run_plan(
             props.add_row(th("Inactivity duration"), inactivity_duration)
     props.add_row(th("Reservation"), run_spec.configuration.reservation or "-")
 
-    offers = Table(box=None, expand=os.get_terminal_size()[0] <= 110)
+    offers = Table(box=None, expand=shutil.get_terminal_size(fallback=(120, 40)).columns <= 110)
     offers.add_column("#")
     offers.add_column("BACKEND", style="grey58", ratio=2)
     offers.add_column("RESOURCES", ratio=4)
@@ -149,7 +149,7 @@ def print_run_plan(
 def get_runs_table(
     runs: List[Run], verbose: bool = False, format_date: DateFormatter = pretty_date
 ) -> Table:
-    table = Table(box=None, expand=os.get_terminal_size()[0] <= 110)
+    table = Table(box=None, expand=shutil.get_terminal_size(fallback=(120, 40)).columns <= 110)
     table.add_column("NAME", style="bold", no_wrap=True, ratio=2)
     table.add_column("BACKEND", style="grey58", ratio=2)
     table.add_column("RESOURCES", ratio=3 if not verbose else 2)
@@ -218,6 +218,11 @@ def _get_run_error(run: Run) -> str:
 
 
 def _get_job_error(job: Job) -> str:
-    if job.job_submissions[-1].termination_reason is None:
+    job_submission = job.job_submissions[-1]
+    termination_reason = job_submission.termination_reason
+    exit_status = job_submission.exit_status
+    if termination_reason is None:
         return ""
-    return job.job_submissions[-1].termination_reason.name
+    if exit_status:
+        return f"{termination_reason.name} {exit_status}"
+    return termination_reason.name

dstack/_internal/core/backends/aws/compute.py

@@ -159,6 +159,7 @@ class AWSCompute(
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         project_name = instance_config.project_name
         ec2_resource = self.session.resource("ec2", region_name=instance_offer.region)
@@ -248,7 +249,7 @@ class AWSCompute(
                         spot=instance_offer.instance.resources.spot,
                         subnet_id=subnet_id,
                         allocate_public_ip=allocate_public_ip,
-                        placement_group_name=instance_config.placement_group_name,
+                        placement_group_name=placement_group.name if placement_group else None,
                         enable_efa=enable_efa,
                         max_efa_interfaces=max_efa_interfaces,
                         reservation_id=instance_config.reservation,
@@ -291,6 +292,7 @@ class AWSCompute(
     def create_placement_group(
         self,
         placement_group: PlacementGroup,
+        master_instance_offer: InstanceOffer,
     ) -> PlacementGroupProvisioningData:
         ec2_client = self.session.client("ec2", region_name=placement_group.configuration.region)
         logger.debug("Creating placement group %s...", placement_group.name)
@@ -323,6 +325,16 @@ class AWSCompute(
                 raise e
         logger.debug("Deleted placement group %s", placement_group.name)
 
+    def is_suitable_placement_group(
+        self,
+        placement_group: PlacementGroup,
+        instance_offer: InstanceOffer,
+    ) -> bool:
+        return (
+            placement_group.configuration.backend == BackendType.AWS
+            and placement_group.configuration.region == instance_offer.region
+        )
+
     def create_gateway(
         self,
         configuration: GatewayComputeConfiguration,

dstack/_internal/core/backends/azure/compute.py

@@ -62,6 +62,7 @@ from dstack._internal.core.models.instances import (
     InstanceOfferWithAvailability,
     InstanceType,
 )
+from dstack._internal.core.models.placement import PlacementGroup
 from dstack._internal.core.models.resources import Memory, Range
 from dstack._internal.core.models.runs import JobProvisioningData, Requirements
 from dstack._internal.utils.logging import get_logger
@@ -109,6 +110,7 @@ class AzureCompute(
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         instance_name = generate_unique_instance_name(
             instance_config, max_length=azure_resources.MAX_RESOURCE_NAME_LEN
@@ -136,6 +138,10 @@ class AzureCompute(
             location=location,
         )
 
+        managed_identity_resource_group, managed_identity_name = parse_vm_managed_identity(
+            self.config.vm_managed_identity
+        )
+
         base_tags = {
             "owner": "dstack",
             "dstack_project": instance_config.project_name,
@@ -159,7 +165,8 @@ class AzureCompute(
             network_security_group=network_security_group,
             network=network,
             subnet=subnet,
-            managed_identity=None,
+            managed_identity_name=managed_identity_name,
+            managed_identity_resource_group=managed_identity_resource_group,
             image_reference=_get_image_ref(
                 compute_client=self._compute_client,
                 location=location,
@@ -255,7 +262,8 @@ class AzureCompute(
             network_security_group=network_security_group,
             network=network,
             subnet=subnet,
-            managed_identity=None,
+            managed_identity_name=None,
+            managed_identity_resource_group=None,
             image_reference=_get_gateway_image_ref(),
             vm_size="Standard_B1ms",
             instance_name=instance_name,
@@ -338,6 +346,21 @@ def get_resource_group_network_subnet_or_error(
     return resource_group, network_name, subnet_name
 
 
+def parse_vm_managed_identity(
+    vm_managed_identity: Optional[str],
+) -> Tuple[Optional[str], Optional[str]]:
+    if vm_managed_identity is None:
+        return None, None
+    try:
+        resource_group, managed_identity = vm_managed_identity.split("/")
+        return resource_group, managed_identity
+    except Exception:
+        raise ComputeError(
+            "`vm_managed_identity` specified in incorrect format."
+            " Supported format: 'managedIdentityResourceGroup/managedIdentityName'"
+        )
+
+
 def _parse_config_vpc_id(vpc_id: str) -> Tuple[str, str]:
     resource_group, network_name = vpc_id.split("/")
     return resource_group, network_name
@@ -466,7 +489,8 @@ def _launch_instance(
     network_security_group: str,
     network: str,
     subnet: str,
-    managed_identity: Optional[str],
+    managed_identity_name: Optional[str],
+    managed_identity_resource_group: Optional[str],
     image_reference: ImageReference,
     vm_size: str,
     instance_name: str,
@@ -488,6 +512,20 @@ def _launch_instance(
         public_ip_address_configuration = VirtualMachinePublicIPAddressConfiguration(
             name="public_ip_config",
         )
+    managed_identity = None
+    if managed_identity_name is not None:
+        if managed_identity_resource_group is None:
+            managed_identity_resource_group = resource_group
+        managed_identity = VirtualMachineIdentity(
+            type=ResourceIdentityType.USER_ASSIGNED,
+            user_assigned_identities={
+                azure_utils.get_managed_identity_id(
+                    subscription_id,
+                    managed_identity_resource_group,
+                    managed_identity_name,
+                ): UserAssignedIdentitiesValue(),
+            },
+        )
     try:
         poller = compute_client.virtual_machines.begin_create_or_update(
             resource_group,
@@ -552,16 +590,7 @@ def _launch_instance(
                 ),
                 priority="Spot" if spot else "Regular",
                 eviction_policy="Delete" if spot else None,
-                identity=None
-                if managed_identity is None
-                else VirtualMachineIdentity(
-                    type=ResourceIdentityType.USER_ASSIGNED,
-                    user_assigned_identities={
-                        azure_utils.get_managed_identity_id(
-                            subscription_id, resource_group, managed_identity
-                        ): UserAssignedIdentitiesValue()
-                    },
-                ),
+                identity=managed_identity,
                 user_data=base64.b64encode(user_data.encode()).decode(),
                 tags=tags,
             ),

dstack/_internal/core/backends/azure/configurator.py

@@ -4,6 +4,7 @@ from typing import List, Optional, Tuple
 
 import azure.core.exceptions
 from azure.core.credentials import TokenCredential
+from azure.mgmt import msi as msi_mgmt
 from azure.mgmt import network as network_mgmt
 from azure.mgmt import resource as resource_mgmt
 from azure.mgmt import subscription as subscription_mgmt
@@ -97,6 +98,7 @@ class AzureConfigurator(Configurator):
         self._check_config_locations(config)
         self._check_config_tags(config)
         self._check_config_resource_group(config=config, credential=credential)
+        self._check_config_vm_managed_identity(config=config, credential=credential)
         self._check_config_vpc(config=config, credential=credential)
 
     def create_backend(
@@ -260,6 +262,25 @@ class AzureConfigurator(Configurator):
                     except BackendError as e:
                         raise ServerClientError(e.args[0])
 
+    def _check_config_vm_managed_identity(
+        self, config: AzureBackendConfigWithCreds, credential: auth.AzureCredential
+    ):
+        try:
+            resource_group, identity_name = compute.parse_vm_managed_identity(
+                config.vm_managed_identity
+            )
+        except BackendError as e:
+            raise ServerClientError(e.args[0])
+        if resource_group is None or identity_name is None:
+            return
+        msi_client = msi_mgmt.ManagedServiceIdentityClient(credential, config.subscription_id)
+        try:
+            msi_client.user_assigned_identities.get(resource_group, identity_name)
+        except azure.core.exceptions.ResourceNotFoundError:
+            raise ServerClientError(
+                f"Managed identity {identity_name} not found in resource group {resource_group}"
+            )
+
     def _set_client_creds_tenant_id(
         self,
         creds: AzureClientCreds,

dstack/_internal/core/backends/azure/models.py

@@ -62,6 +62,15 @@ class AzureBackendConfig(CoreModel):
             )
         ),
     ] = None
+    vm_managed_identity: Annotated[
+        Optional[str],
+        Field(
+            description=(
+                "The managed identity to associate with provisioned VMs."
+                " Must have a format `managedIdentityResourceGroup/managedIdentityName`"
+            )
+        ),
+    ] = None
     tags: Annotated[
         Optional[Dict[str, str]],
         Field(description="The tags that will be assigned to resources created by `dstack`"),

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

@@ -6,7 +6,7 @@ import threading
 from abc import ABC, abstractmethod
 from functools import lru_cache
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Dict, List, Literal, Optional
 
 import git
 import requests
@@ -25,6 +25,7 @@ from dstack._internal.core.models.gateways import (
 )
 from dstack._internal.core.models.instances import (
     InstanceConfiguration,
+    InstanceOffer,
     InstanceOfferWithAvailability,
     SSHKey,
 )
@@ -44,6 +45,8 @@ logger = get_logger(__name__)
 DSTACK_SHIM_BINARY_NAME = "dstack-shim"
 DSTACK_RUNNER_BINARY_NAME = "dstack-runner"
 
+GoArchType = Literal["amd64", "arm64"]
+
 
 class Compute(ABC):
     """
@@ -144,6 +147,7 @@ class ComputeWithCreateInstanceSupport(ABC):
         self,
         instance_offer: InstanceOfferWithAvailability,
         instance_config: InstanceConfiguration,
+        placement_group: Optional[PlacementGroup],
     ) -> JobProvisioningData:
         """
         Launches a new instance. It should return `JobProvisioningData` ASAP.
@@ -176,7 +180,7 @@ class ComputeWithCreateInstanceSupport(ABC):
         )
         instance_offer = instance_offer.copy()
         self._restrict_instance_offer_az_to_volumes_az(instance_offer, volumes)
-        return self.create_instance(instance_offer, instance_config)
+        return self.create_instance(instance_offer, instance_config, placement_group=None)
 
     def _restrict_instance_offer_az_to_volumes_az(
         self,
@@ -225,9 +229,15 @@ class ComputeWithPlacementGroupSupport(ABC):
     def create_placement_group(
         self,
         placement_group: PlacementGroup,
+        master_instance_offer: InstanceOffer,
     ) -> PlacementGroupProvisioningData:
         """
         Creates a placement group.
+
+        Args:
+            placement_group: details about the placement group to be created
+            master_instance_offer: the first instance dstack will attempt to add
+                                   to the placement group
         """
         pass
 
@@ -242,10 +252,27 @@ class ComputeWithPlacementGroupSupport(ABC):
         """
         pass
 
+    @abstractmethod
+    def is_suitable_placement_group(
+        self,
+        placement_group: PlacementGroup,
+        instance_offer: InstanceOffer,
+    ) -> bool:
+        """
+        Checks if the instance offer can be provisioned in the placement group.
+
+        Should return immediately, without performing API calls.
+
+        Can be called with an offer originating from a different backend, because some backends
+        (BackendType.DSTACK) produce offers on behalf of other backends. Should return `False`
+        in that case.
+        """
+        pass
+
 
 class ComputeWithGatewaySupport(ABC):
     """
-    Must be subclassed and imlemented to support gateways.
+    Must be subclassed and implemented to support gateways.
     """
 
     @abstractmethod
@@ -418,6 +445,21 @@ def generate_unique_volume_name(
     )
 
 
+def generate_unique_placement_group_name(
+    project_name: str,
+    fleet_name: str,
+    max_length: int = _DEFAULT_MAX_RESOURCE_NAME_LEN,
+) -> str:
+    """
+    Generates a unique placement group name valid across all backends.
+    """
+    return generate_unique_backend_name(
+        resource_name=fleet_name,
+        project_name=project_name,
+        max_length=max_length,
+    )
+
+
 def generate_unique_backend_name(
     resource_name: str,
     project_name: Optional[str],
@@ -483,13 +525,14 @@ def get_shim_env(
     base_path: Optional[PathLike] = None,
     bin_path: Optional[PathLike] = None,
     backend_shim_env: Optional[Dict[str, str]] = None,
+    arch: Optional[str] = None,
 ) -> Dict[str, str]:
     log_level = "6"  # Trace
     envs = {
         "DSTACK_SHIM_HOME": get_dstack_working_dir(base_path),
         "DSTACK_SHIM_HTTP_PORT": str(DSTACK_SHIM_HTTP_PORT),
         "DSTACK_SHIM_LOG_LEVEL": log_level,
-        "DSTACK_RUNNER_DOWNLOAD_URL": get_dstack_runner_download_url(),
+        "DSTACK_RUNNER_DOWNLOAD_URL": get_dstack_runner_download_url(arch),
         "DSTACK_RUNNER_BINARY_PATH": get_dstack_runner_binary_path(bin_path),
         "DSTACK_RUNNER_HTTP_PORT": str(DSTACK_RUNNER_HTTP_PORT),
         "DSTACK_RUNNER_SSH_PORT": str(DSTACK_RUNNER_SSH_PORT),
@@ -509,16 +552,19 @@ def get_shim_commands(
     base_path: Optional[PathLike] = None,
     bin_path: Optional[PathLike] = None,
     backend_shim_env: Optional[Dict[str, str]] = None,
+    arch: Optional[str] = None,
 ) -> List[str]:
     commands = get_shim_pre_start_commands(
         base_path=base_path,
         bin_path=bin_path,
+        arch=arch,
     )
     shim_env = get_shim_env(
         authorized_keys=authorized_keys,
         base_path=base_path,
         bin_path=bin_path,
         backend_shim_env=backend_shim_env,
+        arch=arch,
     )
     for k, v in shim_env.items():
         commands += [f'export "{k}={v}"']
@@ -539,35 +585,63 @@ def get_dstack_runner_version() -> str:
     return version or "latest"
 
 
-def get_dstack_runner_download_url() -> str:
-    if url := os.environ.get("DSTACK_RUNNER_DOWNLOAD_URL"):
-        return url
-    build = get_dstack_runner_version()
-    if settings.DSTACK_VERSION is not None:
-        bucket = "dstack-runner-downloads"
-    else:
-        bucket = "dstack-runner-downloads-stgn"
-    return (
-        f"https://{bucket}.s3.eu-west-1.amazonaws.com/{build}/binaries/dstack-runner-linux-amd64"
-    )
-
-
-def get_dstack_shim_download_url() -> str:
-    if url := os.environ.get("DSTACK_SHIM_DOWNLOAD_URL"):
-        return url
-    build = get_dstack_runner_version()
-    if settings.DSTACK_VERSION is not None:
-        bucket = "dstack-runner-downloads"
-    else:
-        bucket = "dstack-runner-downloads-stgn"
-    return f"https://{bucket}.s3.eu-west-1.amazonaws.com/{build}/binaries/dstack-shim-linux-amd64"
+def normalize_arch(arch: Optional[str] = None) -> GoArchType:
+    """
+    Converts the given free-form architecture string to the Go GOARCH format.
+    Only 64-bit x86 and ARM are supported. If the word size is not specified (e.g., `x86`, `arm`),
+    64-bit is implied.
+    If the arch is not specified, falls back to `amd64`.
+    """
+    if not arch:
+        return "amd64"
+    arch_lower = arch.lower()
+    if "32" in arch_lower or arch_lower in ["i386", "i686"]:
+        raise ValueError(f"32-bit architectures are not supported: {arch}")
+    if arch_lower.startswith("x86") or arch_lower.startswith("amd"):
+        return "amd64"
+    if arch_lower.startswith("arm") or arch_lower.startswith("aarch"):
+        return "arm64"
+    raise ValueError(f"Unsupported architecture: {arch}")
+
+
+def get_dstack_runner_download_url(arch: Optional[str] = None) -> str:
+    url_template = os.environ.get("DSTACK_RUNNER_DOWNLOAD_URL")
+    if not url_template:
+        if settings.DSTACK_VERSION is not None:
+            bucket = "dstack-runner-downloads"
+        else:
+            bucket = "dstack-runner-downloads-stgn"
+        url_template = (
+            f"https://{bucket}.s3.eu-west-1.amazonaws.com"
+            "/{version}/binaries/dstack-runner-linux-{arch}"
+        )
+    version = get_dstack_runner_version()
+    arch = normalize_arch(arch)
+    return url_template.format(version=version, arch=arch)
+
+
+def get_dstack_shim_download_url(arch: Optional[str] = None) -> str:
+    url_template = os.environ.get("DSTACK_SHIM_DOWNLOAD_URL")
+    if not url_template:
+        if settings.DSTACK_VERSION is not None:
+            bucket = "dstack-runner-downloads"
+        else:
+            bucket = "dstack-runner-downloads-stgn"
+        url_template = (
+            f"https://{bucket}.s3.eu-west-1.amazonaws.com"
+            "/{version}/binaries/dstack-shim-linux-{arch}"
+        )
+    version = get_dstack_runner_version()
+    arch = normalize_arch(arch)
+    return url_template.format(version=version, arch=arch)
 
 
 def get_shim_pre_start_commands(
     base_path: Optional[PathLike] = None,
     bin_path: Optional[PathLike] = None,
+    arch: Optional[str] = None,
 ) -> List[str]:
-    url = get_dstack_shim_download_url()
+    url = get_dstack_shim_download_url(arch)
     dstack_shim_binary_path = get_dstack_shim_binary_path(bin_path)
     dstack_working_dir = get_dstack_working_dir(base_path)
     return [