Add ARM64 support (#2595)

* Support NVIDIA Grace superchips on Lambda
* Support ARM with SSH fleets

Closes: #2101

Author: un-def

docs/docs/concepts/dev-environments.md

@@ -175,6 +175,8 @@ name: vscode
 ide: vscode
 
 resources:
+  # 16 or more x86_64 cores
+  cpu: 16..
   # 200GB or more RAM
   memory: 200GB..
   # 4 GPUs from 40GB to 80GB
@@ -187,10 +189,16 @@ resources:
 
 </div>
 
+The `cpu` property also allows you to specify the CPU architecture, `x86` or `arm`. Examples:
+`x86:16` (16 x86-64 cores), `arm:8..` (at least 8 ARM64 cores).
+If the architecture is not specified, `dstack` tries to infer it from the `gpu` specification
+using `x86` as the fallback value.
+
 The `gpu` property allows specifying not only memory size but also GPU vendor, names
 and their quantity. Examples: `nvidia` (one NVIDIA GPU), `A100` (one A100), `A10G,A100` (either A10G or A100),
 `A100:80GB` (one A100 of 80GB), `A100:2` (two A100), `24GB..40GB:2` (two GPUs between 24GB and 40GB),
 `A100:40GB:2` (two A100 GPUs of 40GB).
+If the vendor is not specified, `dstack` tries to infer it from the GPU name using `nvidia` as the fallback value.
 
 ??? info "Google Cloud TPU"
     To use TPUs, specify its architecture via the `gpu` property.

docs/docs/concepts/services.md

@@ -316,6 +316,8 @@ commands:
 port: 8000
 
 resources:
+  # 16 or more x86_64 cores
+  cpu: 16..
   # 2 GPUs of 80GB
   gpu: 80GB:2
 
@@ -325,10 +327,16 @@ resources:
 
 </div>
 
+The `cpu` property also allows you to specify the CPU architecture, `x86` or `arm`. Examples:
+`x86:16` (16 x86-64 cores), `arm:8..` (at least 8 ARM64 cores).
+If the architecture is not specified, `dstack` tries to infer it from the `gpu` specification
+using `x86` as the fallback value.
+
 The `gpu` property allows specifying not only memory size but also GPU vendor, names
 and their quantity. Examples: `nvidia` (one NVIDIA GPU), `A100` (one A100), `A10G,A100` (either A10G or A100),
 `A100:80GB` (one A100 of 80GB), `A100:2` (two A100), `24GB..40GB:2` (two GPUs between 24GB and 40GB),
 `A100:40GB:2` (two A100 GPUs of 40GB).
+If the vendor is not specified, `dstack` tries to infer it from the GPU name using `nvidia` as the fallback value.
 
 ??? info "Google Cloud TPU"
     To use TPUs, specify its architecture via the `gpu` property.

docs/docs/concepts/tasks.md

@@ -192,6 +192,8 @@ commands:
   - python fine-tuning/qlora/train.py
   
 resources:
+  # 16 or more x86_64 cores
+  cpu: 16..
   # 200GB or more RAM
   memory: 200GB..
   # 4 GPUs from 40GB to 80GB
@@ -204,10 +206,16 @@ resources:
 
 </div>
 
+The `cpu` property also allows you to specify the CPU architecture, `x86` or `arm`. Examples:
+`x86:16` (16 x86-64 cores), `arm:8..` (at least 8 ARM64 cores).
+If the architecture is not specified, `dstack` tries to infer it from the `gpu` specification
+using `x86` as the fallback value.
+
 The `gpu` property allows specifying not only memory size but also GPU vendor, names
 and their quantity. Examples: `nvidia` (one NVIDIA GPU), `A100` (one A100), `A10G,A100` (either A10G or A100),
 `A100:80GB` (one A100 of 80GB), `A100:2` (two A100), `24GB..40GB:2` (two GPUs between 24GB and 40GB),
 `A100:40GB:2` (two A100 GPUs of 40GB).
+If the vendor is not specified, `dstack` tries to infer it from the GPU name using `nvidia` as the fallback value.
 
 ??? info "Google Cloud TPU"
     To use TPUs, specify its architecture via the `gpu` property.

docs/docs/reference/api/python/index.md

@@ -136,10 +136,21 @@ finally:
       show_root_toc_entry: false
       heading_level: 4
       item_id_mapping:
+        cpu: dstack.api.CPU
         gpu: dstack.api.GPU
         memory: dstack.api.Memory
         Range: dstack.api.Range
 
+### `dstack.api.CPU` { #dstack.api.CPU data-toc-label="CPU" }
+
+#SCHEMA# dstack.api.CPU
+    overrides:
+      show_root_heading: false
+      show_root_toc_entry: false
+      heading_level: 4
+      item_id_mapping:
+        Range: dstack.api.Range
+
 ### `dstack.api.GPU` { #dstack.api.GPU data-toc-label="GPU" }
 
 #SCHEMA# dstack.api.GPU

docs/docs/reference/dstack.yml/dev-environment.md

@@ -35,6 +35,14 @@ The `dev-environment` configuration type allows running [dev environments](../..
         required: true
       item_id_prefix: resources-
 
+#### `resources.cpu` { #resources-cpu data-toc-label="cpu" }
+
+#SCHEMA# dstack._internal.core.models.resources.CPUSpec
+    overrides:
+      show_root_heading: false
+      type:
+        required: true
+
 #### `resources.gpu` { #resources-gpu data-toc-label="gpu" }
 
 #SCHEMA# dstack._internal.core.models.resources.GPUSpec

docs/docs/reference/dstack.yml/fleet.md

@@ -46,15 +46,23 @@ The `fleet` configuration type allows creating and updating fleets.
         required: true
       item_id_prefix: resources-
 
-#### `resouces.gpu` { #resources-gpu data-toc-label="gpu" } 
+#### `resources.cpu` { #resources-cpu data-toc-label="cpu" }
+
+#SCHEMA# dstack._internal.core.models.resources.CPUSpec
+    overrides:
+      show_root_heading: false
+      type:
+        required: true
+
+#### `resources.gpu` { #resources-gpu data-toc-label="gpu" }
 
 #SCHEMA# dstack._internal.core.models.resources.GPUSpec
     overrides:
       show_root_heading: false
       type:
         required: true
 
-#### `resouces.disk` { #resources-disk data-toc-label="disk" }
+#### `resources.disk` { #resources-disk data-toc-label="disk" }
 
 #SCHEMA# dstack._internal.core.models.resources.DiskSpec
     overrides:

docs/docs/reference/dstack.yml/service.md

@@ -129,15 +129,23 @@ The `service` configuration type allows running [services](../../concepts/servic
         required: true
       item_id_prefix: resources-
 
-#### `resouces.gpu` { #resources-gpu data-toc-label="gpu" }
+#### `resources.cpu` { #resources-cpu data-toc-label="cpu" }
+
+#SCHEMA# dstack._internal.core.models.resources.CPUSpec
+    overrides:
+      show_root_heading: false
+      type:
+        required: true
+
+#### `resources.gpu` { #resources-gpu data-toc-label="gpu" }
 
 #SCHEMA# dstack._internal.core.models.resources.GPUSpec
     overrides:
       show_root_heading: false
       type:
         required: true
 
-#### `resouces.disk` { #resources-disk data-toc-label="disk" }
+#### `resources.disk` { #resources-disk data-toc-label="disk" }
 
 #SCHEMA# dstack._internal.core.models.resources.DiskSpec
     overrides:

docs/docs/reference/dstack.yml/task.md

@@ -35,15 +35,23 @@ The `task` configuration type allows running [tasks](../../concepts/tasks.md).
         required: true
       item_id_prefix: resources-
 
-#### `resouces.gpu` { #resources-gpu data-toc-label="gpu" }
+#### `resources.cpu` { #resources-cpu data-toc-label="cpu" }
+
+#SCHEMA# dstack._internal.core.models.resources.CPUSpec
+    overrides:
+      show_root_heading: false
+      type:
+        required: true
+
+#### `resources.gpu` { #resources-gpu data-toc-label="gpu" }
 
 #SCHEMA# dstack._internal.core.models.resources.GPUSpec
     overrides:
       show_root_heading: false
       type:
         required: true
 
-#### `resouces.disk` { #resources-disk data-toc-label="disk" }
+#### `resources.disk` { #resources-disk data-toc-label="disk" }
 
 #SCHEMA# dstack._internal.core.models.resources.DiskSpec
     overrides:

docs/docs/reference/environment-variables.md

@@ -117,8 +117,11 @@ For more details on the options below, refer to the [server deployment](../guide
      * `DSTACK_SERVER_MAX_OFFERS_TRIED` - Sets how many instance offers to try when starting a job.
        Setting a high value can degrade server performance.
      * `DSTACK_RUNNER_VERSION` – Sets exact runner version for debug. Defaults to `latest`. Ignored if `DSTACK_RUNNER_DOWNLOAD_URL` is set.
-     * `DSTACK_RUNNER_DOWNLOAD_URL` – Overrides `dstack-runner` binary download URL.
-     * `DSTACK_SHIM_DOWNLOAD_URL` – Overrides `dstack-shim` binary download URL.
+     * `DSTACK_RUNNER_DOWNLOAD_URL` – Overrides `dstack-runner` binary download URL. The URL can contain `{version}` and/or `{arch}` placeholders,
+      where `{version}` is `dstack` version in the `X.Y.Z` format or `latest`, and `{arch}` is either `amd64` or `arm64`, for example,
+      `https://dstack.example.com/{arch}/{version}/dstack-runner`.
+     * `DSTACK_SHIM_DOWNLOAD_URL` – Overrides `dstack-shim` binary download URL. The URL can contain `{version}` and/or `{arch}` placeholders,
+      see `DSTACK_RUNNER_DOWNLOAD_URL` for the details.
      * `DSTACK_DEFAULT_CREDS_DISABLED` – Disables default credentials detection if set. Defaults to `None`.
      * `DSTACK_LOCAL_BACKEND_ENABLED` – Enables local backend for debug if set. Defaults to `None`.
 

src/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]: