Add Hot Aisle bare metal support

mkdocs/docs/concepts/backends.md

@@ -976,20 +976,18 @@ projects:
 </div>
 
 ??? info "Required permissions"
-    The API key must have the following roles assigned:
+    The API key requires the `owner` role for your user and the `operator` and `user` roles for the team specified in `team_handle`.
 
-    * **Owner role for the user** - Required for creating and managing SSH keys
-    * **Operator role for the team** - Required for managing virtual machines within the team
+??? info "Instance types"
+    `dstack` supports Hot Aisle VMs (`vm-mi300x-1`, `vm-mi300x-2`, `vm-mi300x-4`, `vm-mi300x-8`) and bare metal servers (`bm-mi300x-8`). To use a specific type, set [`instance_types`](../reference/dstack.yml/fleet.md#instance_types) in the fleet configuration.
 
-??? info "Pricing"
-    `dstack` shows the hourly price for Hot Aisle instances. Some instances also require an upfront payment for a minimum reservation period, which is usually a few hours. You will be charged for the full minimum period even if you stop the instance early.
-
-    See the Hot Aisle API for the minimum reservation period for each instance type:
+    Some instances are prepaid for a minimum period (8 hours for bare metal). To avoid releasing them early, use a fleet with a fixed number of [`nodes`](../concepts/fleets.md#nodes), or at least set [`idle_duration`](../reference/dstack.yml/fleet.md#idle_duration) to cover the period. To check it for instance types in stock:
 
     <div class="termy">
 
     ```shell
-    $ curl -H "Authorization: Token $API_KEY" https://admin.hotaisle.app/api/teams/$TEAM_HANDLE/virtual_machines/available/ | jq ".[] | {gpus: .Specs.gpus, MinimumReservationMinutes}"
+    $ curl -sS -H "Authorization: Token $API_KEY" https://admin.hotaisle.app/api/teams/$TEAM_HANDLE/virtual_machines/available/ | jq ".[]? | {gpus: .Specs.gpus, MinimumReservationMinutes}"
+    $ curl -sS -H "Authorization: Token $API_KEY" https://admin.hotaisle.app/api/teams/$TEAM_HANDLE/bare_metal/available/ | jq ".[]? | {gpus: .Specs.gpus, MinimumReservationMinutes}"
     ```
 
     </div>

pyproject.toml

@@ -33,7 +33,8 @@ dependencies = [
     "python-multipart>=0.0.16",
     "filelock",
     "psutil",
-    "gpuhunt==0.1.30",
+    # TODO: Release gpuhunt with Hot Aisle bare metal support and pin it.
+    "gpuhunt @ git+https://github.com/dstackai/gpuhunt.git@pr_hotaisle_bare_metal",
     "argcomplete>=3.5.0",
     "ignore-python>=0.2.0",
     "apscheduler<4",
@@ -57,6 +58,10 @@ dstack = "dstack._internal.cli.main:main"
 [tool.hatch.version]
 path = "src/dstack/version.py"
 
+[tool.hatch.metadata]
+# TODO: Remove after pinning the released gpuhunt package.
+allow-direct-references = true
+
 [tool.hatch.build.targets.sdist]
 artifacts = [
     "src/dstack/_internal/server/statics/**",

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

@@ -30,6 +30,7 @@
     "gcp-dws-calendar-mode",
     "runpod-cpu",
     "runpod-cluster",
+    "hotaisle-bm",
 ]
 
 

src/dstack/_internal/core/backends/hotaisle/api_client.py

@@ -3,6 +3,7 @@
 import requests
 
 from dstack._internal.core.backends.base.configurator import raise_invalid_credentials_error
+from dstack._internal.core.errors import BackendError, NoCapacityError
 from dstack._internal.utils.logging import get_logger
 
 API_URL = "https://admin.hotaisle.app/api"
@@ -88,6 +89,38 @@ def terminate_virtual_machine(self, vm_name: str) -> None:
             return
         response.raise_for_status()
 
+    def reserve_bare_metal_server(self, specs: Dict[str, Any], description: str) -> Dict[str, Any]:
+        url = f"{API_URL}/teams/{self.team_handle}/bare_metal/"
+        payload = {"specs": specs, "description": description}
+        response = self._make_request("POST", url, json=payload)
+        # 403: the team's bare metal server limit is reached or the API key lacks permissions.
+        # 404: no available server matches the specs, e.g. another team reserved it.
+        if response.status_code in [403, 404]:
+            raise NoCapacityError(response.text)
+        response.raise_for_status()
+        return response.json()
+
+    def get_bare_metal_server(self, server_id: str) -> Dict[str, Any]:
+        url = f"{API_URL}/teams/{self.team_handle}/bare_metal/{server_id}/"
+        response = self._make_request("GET", url)
+        response.raise_for_status()
+        return response.json()
+
+    def release_bare_metal_server(self, server_id: str, force: bool = True) -> None:
+        url = f"{API_URL}/teams/{self.team_handle}/bare_metal/{server_id}/"
+        # force releases even if min reservation time not met
+        params = {"force": "true"} if force else None
+        response = self._make_request("DELETE", url, params=params)
+        if response.status_code == 404:
+            logger.debug("Hot Aisle bare metal server %s not found", server_id)
+            return
+        if response.status_code == 400 and not force:
+            raise BackendError(
+                f"Hot Aisle refused to release bare metal server {server_id}"
+                f" before its minimum reservation period ends: {response.text}"
+            )
+        response.raise_for_status()
+
     def _make_request(
         self,
         method: str,

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

@@ -1,7 +1,6 @@
 import shlex
 import subprocess
 import tempfile
-from threading import Thread
 from typing import Any, List, Optional
 
 import gpuhunt
@@ -18,6 +17,7 @@
 from dstack._internal.core.backends.base.offers import get_catalog_offers
 from dstack._internal.core.backends.hotaisle.api_client import HotAisleAPIClient
 from dstack._internal.core.backends.hotaisle.models import HotAisleConfig
+from dstack._internal.core.errors import ProvisioningError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import (
     CoreModel,
@@ -32,12 +32,16 @@
 )
 from dstack._internal.core.models.placement import PlacementGroup
 from dstack._internal.core.models.runs import JobProvisioningData
+from dstack._internal.settings import FeatureFlags
+from dstack._internal.utils.common import get_or_error
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
 
 
 SUPPORTED_GPUS = ["MI300X"]
+SSH_CONNECT_TIMEOUT_SECONDS = 10
+SSH_LAUNCH_TIMEOUT_SECONDS = 60
 
 
 class HotAisleCompute(
@@ -81,11 +85,24 @@ def create_instance(
         offer_backend_data = validate_extra_ignore(
             HotAisleOfferBackendData, instance_offer.backend_data
         )
-        vm_data = self.api_client.create_virtual_machine(offer_backend_data.vm_specs)
+        if offer_backend_data.bare_metal_specs is not None:
+            server_data = self.api_client.reserve_bare_metal_server(
+                specs=offer_backend_data.bare_metal_specs,
+                description=instance_config.instance_name,
+            )
+            # The deployment ID identifies this reservation, the name identifies the server.
+            instance_id = server_data["deployment_id"]
+            ip_address = server_data["ip_address"]
+        else:
+            vm_data = self.api_client.create_virtual_machine(
+                get_or_error(offer_backend_data.vm_specs)
+            )
+            instance_id = vm_data["name"]
+            ip_address = vm_data["ip_address"]
         return JobProvisioningData(
             backend=instance_offer.backend,
             instance_type=instance_offer.instance,
-            instance_id=vm_data["name"],
+            instance_id=instance_id,
             hostname=None,
             internal_ip=None,
             region=instance_offer.region,
@@ -95,7 +112,8 @@ def create_instance(
             dockerized=True,
             ssh_proxy=None,
             backend_data=HotAisleInstanceBackendData(
-                ip_address=vm_data["ip_address"]
+                ip_address=ip_address,
+                bare_metal=offer_backend_data.bare_metal_specs is not None,
             ).model_dump_json(),
         )
 
@@ -105,75 +123,125 @@ def update_provisioning_data(
         project_ssh_public_key: str,
         project_ssh_private_key: str,
     ):
-        vm_state = self.api_client.get_vm_state(provisioning_data.instance_id)
-        if vm_state == "running":
-            if provisioning_data.hostname is None and provisioning_data.backend_data:
-                backend_data = HotAisleInstanceBackendData.load(provisioning_data.backend_data)
-                provisioning_data.hostname = backend_data.ip_address
-            commands = get_shim_commands(arch=provisioning_data.instance_type.resources.cpu_arch)
-            launch_command = "sudo sh -c " + shlex.quote(" && ".join(commands))
-            thread = Thread(
-                target=_start_runner,
-                kwargs={
-                    "hostname": provisioning_data.hostname,
-                    "project_ssh_private_key": project_ssh_private_key,
-                    "launch_command": launch_command,
-                },
-                daemon=True,
-            )
-            thread.start()
+        backend_data = HotAisleInstanceBackendData.load(provisioning_data.backend_data)
+        hostname = backend_data.ip_address
+        port = 22
+        if backend_data.bare_metal:
+            server_data = self.api_client.get_bare_metal_server(provisioning_data.instance_id)
+            os_install_status = (server_data.get("os_status") or {}).get("os_install_status")
+            if os_install_status == "failed":
+                raise ProvisioningError("Hot Aisle bare metal server OS installation failed")
+            if os_install_status != "installed":
+                return
+            # The bare metal server's ip_address is private, SSH is exposed via ssh_access.
+            ssh_access = server_data.get("ssh_access") or {}
+            hostname = ssh_access.get("ip_address") or hostname
+            port = ssh_access.get("port") or port
+        elif self.api_client.get_vm_state(provisioning_data.instance_id) != "running":
+            return
+        # Retried on the next check until the shim starts.
+        if not _start_runner(
+            hostname=hostname,
+            port=port,
+            project_ssh_private_key=project_ssh_private_key,
+            arch=provisioning_data.instance_type.resources.cpu_arch,
+        ):
+            return
+        provisioning_data.hostname = hostname
+        provisioning_data.ssh_port = port
 
     def terminate_instance(
         self, instance_id: str, region: str, backend_data: Optional[str] = None
     ):
+        if backend_data is not None and HotAisleInstanceBackendData.load(backend_data).bare_metal:
+            self.api_client.release_bare_metal_server(
+                instance_id, force=not FeatureFlags.HOTAISLE_BARE_METAL_NO_FORCE_RELEASE
+            )
+            return
         vm_name = instance_id
         self.api_client.terminate_virtual_machine(vm_name)
 
 
 def _start_runner(
     hostname: str,
+    port: int,
     project_ssh_private_key: str,
-    launch_command: str,
-):
-    _launch_runner(
+    arch: Optional[str],
+) -> bool:
+    commands = get_shim_commands(arch=arch)
+    launch_command = "sudo sh -c " + shlex.quote(" && ".join(commands))
+    return _launch_runner(
         hostname=hostname,
+        port=port,
         ssh_private_key=project_ssh_private_key,
         launch_command=launch_command,
     )
 
 
 def _launch_runner(
     hostname: str,
+    port: int,
     ssh_private_key: str,
     launch_command: str,
-):
-    daemonized_command = f"{launch_command.rstrip('&')} >/tmp/dstack-shim.log 2>&1 & disown"
-    _run_ssh_command(
+) -> bool:
+    # nohup instead of disown, which isn't available in all shells and would fail the exit code.
+    daemonized_command = (
+        f"nohup {launch_command.rstrip('&')} >/tmp/dstack-shim.log 2>&1 </dev/null &"
+    )
+    return _run_ssh_command(
         hostname=hostname,
+        port=port,
         ssh_private_key=ssh_private_key,
         command=daemonized_command,
     )
 
 
-def _run_ssh_command(hostname: str, ssh_private_key: str, command: str):
+def _run_ssh_command(hostname: str, port: int, ssh_private_key: str, command: str) -> bool:
     with tempfile.NamedTemporaryFile("w+", 0o600) as f:
         f.write(ssh_private_key)
         f.flush()
-        subprocess.run(
-            [
-                "ssh",
-                "-F",
-                "none",
-                "-o",
-                "StrictHostKeyChecking=no",
-                "-i",
-                f.name,
-                f"hotaisle@{hostname}",
-                command,
-            ],
-            stdout=subprocess.DEVNULL,
-            stderr=subprocess.DEVNULL,
+        try:
+            proc = subprocess.run(
+                [
+                    "ssh",
+                    "-F",
+                    "none",
+                    "-o",
+                    "BatchMode=yes",
+                    "-o",
+                    f"ConnectTimeout={SSH_CONNECT_TIMEOUT_SECONDS}",
+                    "-o",
+                    "ConnectionAttempts=1",
+                    "-o",
+                    "StrictHostKeyChecking=no",
+                    "-o",
+                    "UserKnownHostsFile=/dev/null",
+                    "-o",
+                    "LogLevel=ERROR",
+                    "-i",
+                    f.name,
+                    "-p",
+                    str(port),
+                    f"hotaisle@{hostname}",
+                    command,
+                ],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.PIPE,
+                text=True,
+                timeout=SSH_LAUNCH_TIMEOUT_SECONDS,
+            )
+        except subprocess.TimeoutExpired:
+            logger.debug("Timed out running SSH command on Hot Aisle instance %s", hostname)
+            return False
+    if proc.returncode != 0:
+        logger.debug(
+            "SSH command failed on Hot Aisle instance %s: exit_code=%s stderr=%r",
+            hostname,
+            proc.returncode,
+            proc.stderr[-1000:],
         )
+        return False
+    return True
 
 
 def _supported_instances(offer: InstanceOffer) -> bool:
@@ -184,6 +252,7 @@ def _supported_instances(offer: InstanceOffer) -> bool:
 
 class HotAisleInstanceBackendData(CoreModel):
     ip_address: str
+    bare_metal: bool = False
 
     @classmethod
     def load(cls, raw: Optional[str]) -> "HotAisleInstanceBackendData":
@@ -192,4 +261,5 @@ def load(cls, raw: Optional[str]) -> "HotAisleInstanceBackendData":
 
 
 class HotAisleOfferBackendData(CoreModel):
-    vm_specs: dict[str, Any]
+    vm_specs: Optional[dict[str, Any]] = None
+    bare_metal_specs: Optional[dict[str, Any]] = None

src/dstack/_internal/settings.py

@@ -50,3 +50,12 @@ class FeatureFlags:
     """If DSTACK_FF_CLI_PRINT_JOB_CONNECTION_INFO enabled, `dstack apply` command prints server-provided
     IDE URL(s) and SSH command(s) before job logs (for dev-environments only).
     """
+
+    # TODO: Change the default to "0" before merging https://github.com/dstackai/dstack/pull/4325.
+    HOTAISLE_BARE_METAL_NO_FORCE_RELEASE = (
+        os.getenv("DSTACK_FF_HOTAISLE_BARE_METAL_NO_FORCE_RELEASE", "1") != "0"
+    )
+    """Enabled unless set to `0`. If enabled, Hot Aisle bare metal servers are deleted without `force`,
+    so a server still within its minimum reservation period isn't deleted and must be released
+    manually. This prevents accidentally losing a prepaid server.
+    """