dstack 0.0.9__py3-none-any.whl → 0.20.7__py3-none-any.whl

dstack/_internal/server/testing/conf.py

@@ -0,0 +1,53 @@
+import pytest
+import pytest_asyncio
+from sqlalchemy import StaticPool
+from sqlalchemy.ext.asyncio import create_async_engine
+from testcontainers.postgres import PostgresContainer
+
+from dstack._internal.server import settings
+from dstack._internal.server.db import Database, override_db
+from dstack._internal.server.models import BaseModel
+
+
+@pytest.fixture(scope="session")
+def postgres_container():
+    with PostgresContainer("postgres:16-alpine", driver="asyncpg") as postgres:
+        yield postgres.get_connection_url()
+
+
+@pytest_asyncio.fixture
+async def test_db(request):
+    db_type = getattr(request, "param", "sqlite")
+    engine = None
+    if db_type == "sqlite":
+        db_url = "sqlite+aiosqlite://"
+        # For SQLite, allow accessing the in-memory DB from multiple threads:
+        # https://docs.sqlalchemy.org/en/13/dialects/sqlite.html#using-a-memory-database-in-multiple-threads
+        engine = create_async_engine(
+            db_url,
+            echo=settings.SQL_ECHO_ENABLED,
+            connect_args={"check_same_thread": False},
+            poolclass=StaticPool,
+        )
+    elif db_type == "postgres":
+        if not request.config.getoption("--runpostgres"):
+            pytest.skip("Skipping Postgres tests as --runpostgres was not provided")
+        db_url = request.getfixturevalue("postgres_container")
+    else:
+        raise ValueError(f"Unknown db_type {db_type}")
+    db = Database(db_url, engine=engine)
+    override_db(db)
+    async with db.engine.begin() as conn:
+        await conn.run_sync(BaseModel.metadata.drop_all)
+        await conn.run_sync(BaseModel.metadata.create_all)
+    yield db
+    async with db.engine.begin() as conn:
+        await conn.run_sync(BaseModel.metadata.drop_all)
+    await db.engine.dispose()
+
+
+@pytest_asyncio.fixture
+async def session(test_db):
+    db = test_db
+    async with db.get_session() as session:
+        yield session

dstack/_internal/server/testing/matchers.py

@@ -0,0 +1,31 @@
+import re
+import uuid
+
+
+class SomeUUID4Str:
+    """
+    A matcher that compares equal to any valid UUID4 string
+    """
+
+    # Simplified UUID regex: just checks the 8-4-4-4-12 hex structure
+    _uuid_regex = re.compile(
+        r"^[0-9a-f]{8}-"
+        r"[0-9a-f]{4}-"
+        r"[0-9a-f]{4}-"
+        r"[0-9a-f]{4}-"
+        r"[0-9a-f]{12}$"
+    )
+
+    def __eq__(self, other):
+        if isinstance(other, str):
+            if not self._uuid_regex.match(other):
+                return False
+            try:
+                return uuid.UUID(other).version == 4
+            except ValueError:
+                return False
+
+        return False
+
+    def __repr__(self):
+        return "SomeUUID4Str()"

dstack/_internal/server/utils/common.py

@@ -0,0 +1,55 @@
+import asyncio
+from typing import (
+    Awaitable,
+    Callable,
+    Iterable,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    TypeVar,
+    Union,
+)
+
+ItemT = TypeVar("ItemT")
+ResultT = TypeVar("ResultT")
+
+
+async def gather_map_async(
+    items: Sequence[ItemT],
+    func: Callable[[ItemT], Awaitable[ResultT]],
+    *,
+    return_exceptions: bool = False,
+) -> List[Tuple[ItemT, Union[ResultT, BaseException]]]:
+    """
+    A parallel wrapper around asyncio.gather that returns a list of tuples (item, result).
+    Args:
+        items: list of items to be processed
+        func: function to be applied to each item, return awaitable coroutine
+        return_exceptions: passed to asyncio.gather
+
+    Returns:
+        list of tuples (item, result) or (item, exception) if return_exceptions is True
+    """
+    return [
+        (item, result)
+        for item, result in zip(
+            items,
+            await asyncio.gather(
+                *(func(item) for item in items), return_exceptions=return_exceptions
+            ),
+        )
+    ]
+
+
+def join_byte_stream_checked(stream: Iterable[bytes], max_size: int) -> Optional[bytes]:
+    """
+    Join an iterable of `bytes` values into one `bytes` value,
+    unless its size exceeds `max_size`.
+    """
+    result = b""
+    for chunk in stream:
+        if len(result) + len(chunk) > max_size:
+            return None
+        result += chunk
+    return result

dstack/_internal/server/utils/logging.py

@@ -0,0 +1,51 @@
+import asyncio
+import logging
+import sys
+
+from pythonjsonlogger.json import JsonFormatter
+
+from dstack._internal.cli.utils.common import console
+from dstack._internal.cli.utils.rich import DstackRichHandler
+from dstack._internal.server import settings
+
+
+class AsyncioCancelledErrorFilter(logging.Filter):
+    def filter(self, record: logging.LogRecord) -> bool:
+        if not record.exc_info:
+            return True
+        if isinstance(record.exc_info[1], asyncio.CancelledError):
+            return False
+        return True
+
+
+def configure_logging():
+    formatters = {
+        "rich": logging.Formatter(fmt="%(message)s", datefmt="[%X]"),
+        "standard": logging.Formatter(
+            fmt="%(levelname)s %(asctime)s.%(msecs)03d %(name)s %(message)s",
+            datefmt="%Y-%m-%dT%H:%M:%S",
+        ),
+        "json": JsonFormatter(
+            "%(asctime)s %(name)s %(levelname)s %(message)s",
+            json_ensure_ascii=False,
+            rename_fields={"name": "logger", "asctime": "timestamp", "levelname": "level"},
+        ),
+    }
+    handlers: dict[str, logging.Handler] = {
+        "rich": DstackRichHandler(console=console),
+        "standard": logging.StreamHandler(stream=sys.stdout),
+        "json": logging.StreamHandler(stream=sys.stdout),
+    }
+    if settings.LOG_FORMAT not in formatters:
+        raise ValueError(f"Invalid settings.LOG_FORMAT: {settings.LOG_FORMAT}")
+    formatter = formatters[settings.LOG_FORMAT]
+    handler = handlers[settings.LOG_FORMAT]
+    handler.setFormatter(formatter)
+    handler.addFilter(AsyncioCancelledErrorFilter())
+    root_logger = logging.getLogger(None)
+    root_logger.addHandler(handler)
+    root_logger.setLevel(settings.ROOT_LOG_LEVEL)
+    dstack_logger = logging.getLogger("dstack")
+    dstack_logger.setLevel(settings.LOG_LEVEL)
+    # paramiko emits error logs in cases handled by dstack
+    logging.getLogger("paramiko").setLevel(logging.CRITICAL)

dstack/_internal/server/utils/provisioning.py

@@ -0,0 +1,368 @@
+import io
+import json
+import time
+from contextlib import contextmanager, nullcontext
+from textwrap import dedent
+from typing import Any, Dict, Generator, List, Optional
+
+import paramiko
+from gpuhunt import AcceleratorVendor, correct_gpu_memory_gib
+
+from dstack._internal.core.backends.base.compute import (
+    DSTACK_SHIM_RESTART_INTERVAL_SECONDS,
+    GoArchType,
+    normalize_arch,
+)
+from dstack._internal.core.consts import DSTACK_SHIM_HTTP_PORT
+
+# FIXME: ProvisioningError is a subclass of ComputeError and should not be used outside of Compute
+from dstack._internal.core.errors import ProvisioningError
+from dstack._internal.core.models.instances import (
+    Disk,
+    Gpu,
+    InstanceType,
+    Resources,
+    SSHConnectionParams,
+)
+from dstack._internal.utils.gpu import (
+    convert_amd_gpu_name,
+    convert_intel_accelerator_name,
+    convert_nvidia_gpu_name,
+)
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+SSH_CONNECT_TIMEOUT = 10
+
+DSTACK_SHIM_ENV_FILE = "shim.env"
+
+HOST_INFO_FILE = "host_info.json"
+
+
+def detect_cpu_arch(client: paramiko.SSHClient) -> GoArchType:
+    cmd = "uname -m"
+    try:
+        _, stdout, stderr = client.exec_command(cmd, timeout=20)
+    except (paramiko.SSHException, OSError) as e:
+        raise ProvisioningError(f"detect_cpu_arch: {e}") from e
+    out = stdout.read().strip().decode()
+    err = stderr.read().strip().decode()
+    if err:
+        raise ProvisioningError(f"detect_cpu_arch: {cmd} failed, stdout: {out}, stderr: {err}")
+    try:
+        return normalize_arch(out)
+    except ValueError as e:
+        raise ProvisioningError(f"detect_cpu_arch: failed to normalize arch: {e}") from e
+
+
+def sftp_upload(client: paramiko.SSHClient, path: str, body: str) -> None:
+    try:
+        sftp = client.open_sftp()
+        channel = sftp.get_channel()
+        if channel is not None:
+            channel.settimeout(10)
+        sftp.putfo(io.BytesIO(body.encode()), path)
+        sftp.close()
+    except (paramiko.SSHException, OSError) as e:
+        raise ProvisioningError(f"sft_upload failed: {e}") from e
+
+
+def upload_envs(client: paramiko.SSHClient, working_dir: str, envs: Dict[str, str]) -> None:
+    envs["DSTACK_SERVICE_MODE"] = "1"  # make host_info.json on start
+    dot_env = "\n".join(f'{key}="{value.strip()}"' for key, value in envs.items())
+    tmp_file_path = f"/tmp/{DSTACK_SHIM_ENV_FILE}"
+    sftp_upload(client, tmp_file_path, dot_env)
+    try:
+        cmd = f"sudo mkdir -p {working_dir} && sudo mv {tmp_file_path} {working_dir}/"
+        _, stdout, stderr = client.exec_command(cmd, timeout=20)
+        out = stdout.read().strip().decode()
+        err = stderr.read().strip().decode()
+        if out or err:
+            raise ProvisioningError(
+                f"The command 'upload_envs' didn't work. stdout: {out}, stderr: {err}"
+            )
+    except (paramiko.SSHException, OSError) as e:
+        raise ProvisioningError(f"upload_envs failed: {e}") from e
+
+
+def run_pre_start_commands(
+    client: paramiko.SSHClient, shim_pre_start_commands: List[str], authorized_keys: List[str]
+) -> None:
+    try:
+        authorized_keys_content = "\n".join(authorized_keys).strip()
+        _, stdout, stderr = client.exec_command(
+            f"echo '\n{authorized_keys_content}' >> ~/.ssh/authorized_keys", timeout=5
+        )
+        out = stdout.read().strip().decode()
+        err = stderr.read().strip().decode()
+        if out or err:
+            raise ProvisioningError(
+                f"The command 'authorized_keys' didn't work. stdout: {out}, stderr: {err}"
+            )
+    except (paramiko.SSHException, OSError) as e:
+        raise ProvisioningError(f"upload authorized_keys failed: {e}") from e
+
+    script = " && ".join(shim_pre_start_commands)
+    try:
+        _, stdout, stderr = client.exec_command(f"sudo sh -c '{script}'", timeout=120)
+        out = stdout.read().strip().decode()
+        err = stderr.read().strip().decode()
+        if out or err:
+            raise ProvisioningError(
+                f"The command 'run_pre_start_commands' didn't work. stdout: {out}, stderr: {err}"
+            )
+    except (paramiko.SSHException, OSError) as e:
+        raise ProvisioningError(f"run_pre-start_commands failed: {e}") from e
+
+
+def run_shim_as_systemd_service(
+    client: paramiko.SSHClient, binary_path: str, working_dir: str, dev: bool
+) -> None:
+    # Stop restart attempts after ≈ 1 hour
+    start_limit_interval_seconds = 3600
+    start_limit_burst = int(
+        start_limit_interval_seconds / DSTACK_SHIM_RESTART_INTERVAL_SECONDS * 0.9
+    )
+    shim_service = dedent(f"""\
+        [Unit]
+        Description=dstack-shim
+        After=network-online.target
+        StartLimitIntervalSec={start_limit_interval_seconds}
+        StartLimitBurst={start_limit_burst}
+
+        [Service]
+        Type=simple
+        User=root
+        Restart=always
+        RestartSec={DSTACK_SHIM_RESTART_INTERVAL_SECONDS}
+        WorkingDirectory={working_dir}
+        EnvironmentFile={working_dir}/{DSTACK_SHIM_ENV_FILE}
+        ExecStart={binary_path}
+
+        [Install]
+        WantedBy=multi-user.target
+    """)
+
+    sftp_upload(client, "/tmp/dstack-shim.service", shim_service)
+
+    try:
+        cmd = """\
+            sudo mv /tmp/dstack-shim.service /etc/systemd/system/dstack-shim.service && \
+            sudo systemctl daemon-reload && \
+            sudo systemctl --quiet enable dstack-shim && \
+            sudo systemctl restart dstack-shim
+        """
+        _, stdout, stderr = client.exec_command(cmd, timeout=100)
+        out = stdout.read().strip().decode()
+        err = stderr.read().strip().decode()
+        if out or err:
+            raise ProvisioningError(
+                f"The command 'run_shim_as_systemd_service' didn't work. stdout: {out}, stderr: {err}"
+            )
+    except (paramiko.SSHException, OSError) as e:
+        raise ProvisioningError(f"run_shim_as_systemd failed: {e}") from e
+
+
+def check_dstack_shim_service(client: paramiko.SSHClient):
+    try:
+        _, stdout, _ = client.exec_command("sudo systemctl status dstack-shim.service", timeout=10)
+        status = stdout.read()
+    except (paramiko.SSHException, OSError) as e:
+        raise ProvisioningError(f"Checking dstack-shim.service status failed: {e}") from e
+
+    for raw_line in status.splitlines():
+        line = raw_line.decode()
+        if line.strip().startswith("Active: failed"):
+            raise ProvisioningError(f"The dstack-shim service doesn't start: {line.strip()}")
+
+
+def remove_host_info_if_exists(client: paramiko.SSHClient, working_dir: str) -> None:
+    file_path = f"{working_dir}/{HOST_INFO_FILE}"
+    try:
+        _, _, stderr = client.exec_command(
+            f"sudo test -e {file_path} && sudo rm {file_path}", timeout=10
+        )
+        err = stderr.read().decode().strip()
+        if err:
+            logger.debug(f"{HOST_INFO_FILE} hasn't been removed: %s", err)
+    except (paramiko.SSHException, OSError) as e:
+        raise ProvisioningError(f"remove_host_info_if_exists failed: {e}")
+
+
+def remove_dstack_runner_if_exists(client: paramiko.SSHClient, path: str) -> None:
+    try:
+        _, _, stderr = client.exec_command(f"sudo test -e {path} && sudo rm {path}", timeout=10)
+        err = stderr.read().decode().strip()
+        if err:
+            logger.debug(f"{path} hasn't been removed: %s", err)
+    except (paramiko.SSHException, OSError) as e:
+        raise ProvisioningError(f"remove_dstack_runner_if_exists failed: {e}")
+
+
+def get_host_info(client: paramiko.SSHClient, working_dir: str) -> Dict[str, Any]:
+    # wait host_info
+    retries = 60
+    iter_delay = 3
+    for _ in range(retries):
+        try:
+            _, stdout, stderr = client.exec_command(
+                f"sudo cat {working_dir}/{HOST_INFO_FILE}", timeout=10
+            )
+            err = stderr.read().decode().strip()
+            if err:
+                logger.debug("Retry after error: %s", err)
+                time.sleep(iter_delay)
+                continue
+        except (paramiko.SSHException, OSError) as e:
+            logger.debug(f"Cannot run `cat {HOST_INFO_FILE}` in the remote instance: %s", e)
+        else:
+            try:
+                host_info_json = stdout.read()
+                host_info = json.loads(host_info_json)
+                return host_info
+            except ValueError:  # JSON parse error
+                check_dstack_shim_service(client)
+                raise ProvisioningError("Cannot parse host_info")
+        time.sleep(iter_delay)
+    else:
+        check_dstack_shim_service(client)
+        raise ProvisioningError("Cannot get host_info")
+
+
+def get_shim_healthcheck(client: paramiko.SSHClient) -> str:
+    retries = 20
+    iter_delay = 3
+    for _ in range(retries):
+        healthcheck = _get_shim_healthcheck(client)
+        if healthcheck is not None:
+            return healthcheck
+        logger.debug("healthcheck is empty. retry")
+        time.sleep(iter_delay)
+    raise ProvisioningError("Cannot get HealthcheckResponse")
+
+
+def _get_shim_healthcheck(client: paramiko.SSHClient) -> Optional[str]:
+    try:
+        _, stdout, stderr = client.exec_command(
+            f"curl -s http://localhost:{DSTACK_SHIM_HTTP_PORT}/api/healthcheck", timeout=15
+        )
+        out = stdout.read().strip().decode()
+        err = stderr.read().strip().decode()
+    except (paramiko.SSHException, OSError) as e:
+        raise ProvisioningError(f"get_shim_healthcheck failed: {e}") from e
+    if err:
+        raise ProvisioningError(f"get_shim_healthcheck didn't work. stdout: {out}, stderr: {err}")
+    if not out:
+        return None
+    return out
+
+
+def host_info_to_instance_type(host_info: Dict[str, Any], arch: GoArchType) -> InstanceType:
+    gpu_count = host_info.get("gpu_count", 0)
+    if gpu_count > 0:
+        gpu_vendor = AcceleratorVendor.cast(host_info.get("gpu_vendor", "nvidia"))
+        gpu_name = host_info["gpu_name"]
+        if gpu_vendor == AcceleratorVendor.NVIDIA:
+            gpu_name = convert_nvidia_gpu_name(gpu_name)
+        elif gpu_vendor == AcceleratorVendor.AMD:
+            gpu_name = convert_amd_gpu_name(gpu_name)
+        elif gpu_vendor == AcceleratorVendor.INTEL:
+            gpu_name = convert_intel_accelerator_name(gpu_name)
+        gpu_memory_mib = host_info["gpu_memory"]
+        if isinstance(gpu_memory_mib, str):
+            # older shim versions report gpu_memory as a string
+            gpu_memory_mib = float(gpu_memory_mib.lower().replace("mib", "").strip())
+        else:
+            # newer shim versions report gpu_memory as an integer
+            gpu_memory_mib = float(gpu_memory_mib)
+        gpu_memory_mib = correct_gpu_memory_gib(gpu_name, gpu_memory_mib) * 1024
+        gpus = [Gpu(vendor=gpu_vendor, name=gpu_name, memory_mib=gpu_memory_mib)] * gpu_count
+    else:
+        gpus = []
+    instance_type = InstanceType(
+        name="instance",
+        resources=Resources(
+            cpu_arch=arch.to_cpu_architecture(),
+            cpus=host_info["cpus"],
+            memory_mib=host_info["memory"] / 1024 / 1024,
+            spot=False,
+            gpus=gpus,
+            disk=Disk(size_mib=host_info["disk_size"] / 1024 / 1024),
+        ),
+    )
+    return instance_type
+
+
+@contextmanager
+def get_paramiko_connection(
+    ssh_user: str,
+    host: str,
+    port: int,
+    pkeys: List[paramiko.PKey],
+    proxy: Optional[SSHConnectionParams] = None,
+    proxy_pkeys: Optional[list[paramiko.PKey]] = None,
+) -> Generator[paramiko.SSHClient, None, None]:
+    if proxy is not None:
+        if proxy_pkeys is None:
+            raise ProvisioningError("Missing proxy private keys")
+        proxy_ctx = get_paramiko_connection(
+            proxy.username, proxy.hostname, proxy.port, proxy_pkeys
+        )
+    else:
+        proxy_ctx = nullcontext()
+    conn_url = f"{ssh_user}@{host}:{port}"
+    with proxy_ctx as proxy_client, paramiko.SSHClient() as client:
+        proxy_channel: Optional[paramiko.Channel] = None
+        if proxy_client is not None:
+            transport = proxy_client.get_transport()
+            assert transport is not None
+            try:
+                proxy_channel = transport.open_channel("direct-tcpip", (host, port), ("", 0))
+            except (paramiko.SSHException, OSError) as e:
+                raise ProvisioningError(f"Proxy channel failed: {e}") from e
+        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        for pkey in pkeys:
+            logger.debug("Try to connect to %s with key %s", conn_url, pkey.fingerprint)
+            connected = _paramiko_connect(client, ssh_user, host, port, pkey, proxy_channel)
+            if connected:
+                yield client
+                return
+            logger.debug(
+                f'Authentication failed to connect to "{conn_url}" and {pkey.fingerprint}'
+            )
+        keys_fp = ", ".join(f"{pk.fingerprint!r}" for pk in pkeys)
+        raise ProvisioningError(
+            f"SSH connection to the {conn_url} with keys [{keys_fp}] was unsuccessful"
+        )
+
+
+def _paramiko_connect(
+    client: paramiko.SSHClient,
+    user: str,
+    host: str,
+    port: int,
+    pkey: paramiko.PKey,
+    channel: Optional[paramiko.Channel] = None,
+) -> bool:
+    """
+    Returns `True` if connected, `False` if auth failed, and raises `ProvisioningError`
+    on other errors.
+    """
+    try:
+        client.connect(
+            username=user,
+            hostname=host,
+            port=port,
+            pkey=pkey,
+            look_for_keys=False,
+            allow_agent=False,
+            timeout=SSH_CONNECT_TIMEOUT,
+            sock=channel,
+        )
+        return True
+    except paramiko.AuthenticationException:
+        return False
+    except (paramiko.SSHException, OSError) as e:
+        raise ProvisioningError(f"Connect failed: {e}") from e