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

dstack/_internal/utils/event_loop.py

@@ -0,0 +1,30 @@
+import asyncio
+import threading
+from collections.abc import Awaitable
+from typing import TypeVar
+
+T = TypeVar("T")
+
+
+class DaemonEventLoop:
+    """
+    A wrapper around asyncio.EventLoop that runs the loop in a daemon thread.
+    The thread is started with the first `await_` call.
+    """
+
+    def __init__(self) -> None:
+        self._loop = asyncio.new_event_loop()
+        self._start_lock = threading.Lock()
+        self._started = False
+
+    def await_(self, awaitable: Awaitable[T]) -> T:
+        with self._start_lock:
+            if not self._started:
+                threading.Thread(target=self._loop.run_forever, daemon=True).start()
+                self._started = True
+        future = asyncio.run_coroutine_threadsafe(_coroutine(awaitable), self._loop)
+        return future.result()
+
+
+async def _coroutine(awaitable: Awaitable[T]) -> T:
+    return await awaitable

dstack/_internal/utils/files.py

@@ -0,0 +1,69 @@
+import tarfile
+from pathlib import Path
+from typing import BinaryIO
+
+import ignore
+import ignore.overrides
+
+from dstack._internal.utils.hash import get_sha256
+from dstack._internal.utils.path import PathLike, normalize_path
+
+
+def create_file_archive(root: PathLike, fp: BinaryIO) -> str:
+    """
+    Packs the directory or file to a tar archive and writes it to the file-like object.
+
+    Archives can be used to transfer file(s) (e.g., over the network) preserving
+    file properties such as permissions, timestamps, etc.
+
+    NOTE: `.gitignore` and `.dstackignore` are respected.
+
+    Args:
+        root: The absolute path to the directory or file.
+        fp: The binary file-like object.
+
+    Returns:
+        The SHA-256 hash of the archive as a hex string.
+
+    Raises:
+        ValueError: If the path is not absolute.
+        OSError: Underlying errors from the tarfile module
+    """
+    root = Path(root)
+    if not root.is_absolute():
+        raise ValueError(f"path must be absolute: {root}")
+    walk = (
+        ignore.WalkBuilder(root)
+        .overrides(ignore.overrides.OverrideBuilder(root).add("!/.git/").build())
+        .hidden(False)  # do not ignore files that start with a dot
+        .require_git(False)  # respect git ignore rules even if not a git repo
+        .add_custom_ignore_filename(".dstackignore")
+        .build()
+    )
+    # sort paths to ensure archive reproducibility
+    paths = sorted(entry.path() for entry in walk)
+    with tarfile.TarFile(mode="w", fileobj=fp) as t:
+        for path in paths:
+            arcname = str(path.relative_to(root.parent))
+            info = t.gettarinfo(path, arcname)
+            if info.issym():
+                # Symlinks are handled as follows: each symlink in the chain is checked, and
+                # * if the target is inside the root: keep relative links as is, replace absolute
+                #   links with relative ones;
+                # * if the target is outside the root: replace the link with the actual file.
+                target = Path(info.linkname)
+                if not target.is_absolute():
+                    target = path.parent / target
+                target = normalize_path(target)
+                try:
+                    target.relative_to(root)
+                except ValueError:
+                    # Adding as a file
+                    t.add(path.resolve(), arcname, recursive=False)
+                else:
+                    # Adding as a relative symlink
+                    info.linkname = str(target.relative_to(path.parent, walk_up=True))
+                    t.addfile(info)
+            else:
+                t.add(path, arcname, recursive=False)
+    return get_sha256(fp)

dstack/_internal/utils/gpu.py

@@ -0,0 +1,59 @@
+import re
+
+
+def convert_nvidia_gpu_name(name: str) -> str:
+    """Convert gpu_name from nvidia-smi to short version"""
+    # https://github.com/NVIDIA/open-gpu-kernel-modules/
+    name = name.replace("NVIDIA ", "")
+    name = name.replace("Tesla ", "")
+    name = name.replace("Quadro ", "")
+    name = name.replace("GeForce ", "")
+
+    if "GH200" in name:
+        return "GH200"
+
+    if "RTX A" in name:
+        name = name.replace("RTX A", "A")
+        m = re.search(r"(A\d+)", name)
+        if m is not None:
+            return m.group(0)
+        return name.replace(" ", "")
+
+    name = name.replace(" Ti", "Ti")
+    name = name.replace(" NVL", "NVL")
+    name = name.replace(" Ada Generation", "Ada")
+    name = name.replace("RTX ", "RTX")
+    m = re.search(r"([AHLPTV]\d+\w*)", name)
+    if m is not None:
+        return m.group(0)
+    return name.replace(" ", "")
+
+
+def convert_amd_gpu_name(name: str) -> str:
+    """Convert asic.market_name from amd-smi to short version"""
+    if match := _AMD_INSTINCT_MARKET_NAME_REGEX.search(name):
+        name = match.group("name")
+    # https://github.com/ROCm/amdsmi/blob/52b3947/src/amd_smi/amd_smi_utils.cc#L558-L593
+    if name == "MI300X-O":
+        return "MI300X"
+    return name
+
+
+def convert_intel_accelerator_name(name: str) -> str:
+    """Convert name from hl-smi to market name"""
+    for model_name, market_name in _INTEL_GAUDI_MODELS.items():
+        if name.startswith(model_name):
+            return market_name
+    return name
+
+
+_AMD_INSTINCT_MARKET_NAME_REGEX = re.compile(
+    r"^(?:AMD )?(?:Instinct )?(?P<name>MI\d{1,3}[A-Z]?(?:-\w+)?)(?:\s|$)", flags=re.ASCII | re.I
+)
+
+_INTEL_GAUDI_MODELS = {
+    "HL-205": "Gaudi",
+    "HL-225": "Gaudi2",
+    "HL-325": "Gaudi3",  # OAM
+    "HL-338": "Gaudi3",  # PCIe
+}

dstack/_internal/utils/hash.py

@@ -0,0 +1,31 @@
+import hashlib
+import math
+import string
+from typing import BinaryIO
+
+base36chars = string.digits + string.ascii_lowercase
+
+
+def get_sha256(fp: BinaryIO, chunk_size: int = 65536) -> str:
+    sha256 = hashlib.sha256()
+    fp.seek(0)
+    chunk = fp.read(chunk_size)
+    while len(chunk) > 0:
+        sha256.update(chunk)
+        chunk = fp.read(chunk_size)
+    return sha256.hexdigest()
+
+
+def base36encode(value: bytes) -> str:
+    output = []
+    n = len(value) * math.ceil(math.log(8) / math.log(len(base36chars)))
+    value = int.from_bytes(value, "little", signed=False)
+    for _ in range(n):
+        value, rem = divmod(value, len(base36chars))
+        output.append(base36chars[rem])
+    return "".join(reversed(output))
+
+
+def slugify(prefix: str, unique_key: str, hash_size: int = 8) -> str:
+    full_hash = base36encode(hashlib.sha256(unique_key.encode()).digest())
+    return f"{prefix}-{full_hash[:hash_size]}"

dstack/_internal/utils/interpolator.py

@@ -0,0 +1,91 @@
+import string
+from typing import Dict, Iterable, List, Literal, Optional, Tuple, Union, overload
+
+
+class Pattern:
+    opening = "${{"
+    closing = "}}"
+
+
+class Name:
+    first_char = set(string.ascii_letters + "_")
+    char = first_char | set(string.digits + ".")
+
+
+class InterpolatorError(ValueError):
+    """
+    Raised when interpolation fails.
+
+    May be shown to the users, should not contain sensitive information,
+    such as variable values.
+    """
+
+    pass
+
+
+class VariablesInterpolator:
+    def __init__(
+        self, namespaces: Dict[str, Dict[str, str]], *, skip: Optional[Iterable[str]] = None
+    ):
+        self.skip = set(skip) if skip is not None else set()
+        self.variables = {f"{ns}.{k}": v for ns in namespaces for k, v in namespaces[ns].items()}
+
+    @overload
+    def interpolate(self, s: str, return_missing: Literal[False] = False) -> str: ...
+
+    @overload
+    def interpolate(self, s: str, return_missing: Literal[True]) -> Tuple[str, List[str]]: ...
+
+    def interpolate(
+        self, s: str, return_missing: bool = False
+    ) -> Union[str, Tuple[str, List[str]]]:
+        tokens = []
+        missing = []
+        start = 0
+        while start < len(s):
+            dollar = s.find("$", start)
+            if dollar == -1 or dollar == len(s) - 1:
+                tokens.append(s[start:])
+                break
+            if s[dollar + 1] == "$":  # escaped $$
+                tokens.append(s[start : dollar + 1])
+                start = dollar + 2
+                continue
+
+            opening = s.find(Pattern.opening, start)
+            if opening == -1:
+                tokens.append(s[start:])
+                break
+            tokens.append(s[start:opening])
+            closing = s.find(Pattern.closing, opening)
+            if closing == -1:
+                raise InterpolatorError(f"No pattern closing: {s[opening:]}")
+
+            name = s[opening + len(Pattern.opening) : closing].strip()
+            if not self.validate_name(name):
+                raise InterpolatorError(f"Illegal reference name: {name}")
+            if name.split(".")[0] in self.skip:
+                tokens.append(s[opening : closing + len(Pattern.closing)])
+            elif name in self.variables:
+                tokens.append(self.variables[name])
+            else:
+                missing.append(name)
+            start = closing + len(Pattern.closing)
+        s = "".join(tokens)
+        return (s, missing) if return_missing else s
+
+    def interpolate_or_error(self, s: str) -> str:
+        res, missing = self.interpolate(s, return_missing=True)
+        if len(missing) == 0:
+            return res
+        raise InterpolatorError(f"Failed to interpolate due to missing vars: {missing}")
+
+    @staticmethod
+    def validate_name(s: str) -> bool:
+        if s.count(".") != 1 or not (0 < s.index(".") < len(s) - 1):
+            return False
+        if s[0] not in Name.first_char or s[s.index(".") + 1] not in Name.first_char:
+            return False
+        if any((c not in Name.char) for c in s):
+            return False
+        return True

dstack/_internal/utils/json_schema.py

@@ -0,0 +1,11 @@
+def add_extra_schema_types(schema_property: dict, extra_types: list[dict]):
+    if "allOf" in schema_property:
+        refs = [schema_property.pop("allOf")[0]]
+    elif "anyOf" in schema_property:
+        refs = schema_property.pop("anyOf")
+    elif "type" in schema_property:
+        refs = [{"type": schema_property.pop("type")}]
+    else:
+        refs = [{"$ref": schema_property.pop("$ref")}]
+    refs.extend(extra_types)
+    schema_property["anyOf"] = refs

dstack/_internal/utils/json_utils.py

@@ -0,0 +1,54 @@
+from typing import Any
+
+import orjson
+from pydantic import BaseModel
+
+FREEZEGUN = True
+try:
+    from freezegun.api import FakeDatetime
+except ImportError:
+    FREEZEGUN = False
+
+
+ASYNCPG = True
+try:
+    import asyncpg.pgproto.pgproto
+except ImportError:
+    ASYNCPG = False
+
+
+def pydantic_orjson_dumps(v: Any, *, default: Any) -> str:
+    return orjson.dumps(
+        v,
+        option=get_orjson_default_options(),
+        default=orjson_default,
+    ).decode()
+
+
+def pydantic_orjson_dumps_with_indent(v: Any, *, default: Any) -> str:
+    return orjson.dumps(
+        v,
+        option=get_orjson_default_options() | orjson.OPT_INDENT_2,
+        default=orjson_default,
+    ).decode()
+
+
+def orjson_default(obj):
+    if isinstance(obj, float):
+        # orjson does not convert float subclasses be default
+        return float(obj)
+    if isinstance(obj, BaseModel):
+        # Allows calling orjson.dumps() on pydantic models
+        # (e.g. to return from the API)
+        return obj.dict()
+    if ASYNCPG:
+        if isinstance(obj, asyncpg.pgproto.pgproto.UUID):
+            return str(obj)
+    if FREEZEGUN:
+        if isinstance(obj, FakeDatetime):
+            return obj.isoformat()
+    raise TypeError
+
+
+def get_orjson_default_options() -> int:
+    return orjson.OPT_NON_STR_KEYS

dstack/_internal/utils/logging.py

@@ -0,0 +1,5 @@
+import logging
+
+
+def get_logger(name: str) -> logging.Logger:
+    return logging.getLogger(name)

dstack/_internal/utils/nested_list.py

@@ -0,0 +1,47 @@
+from dataclasses import dataclass, field
+from typing import Optional
+
+
+@dataclass
+class NestedListItem:
+    label: str
+    children: list["NestedListItem"] = field(default_factory=list)
+
+    def render(self, indent: int = 0, visited: Optional[set[int]] = None) -> str:
+        if visited is None:
+            visited = set()
+
+        item_id = id(self)
+        if item_id in visited:
+            raise ValueError(f"Cycle detected at item: {self.label}")
+
+        visited.add(item_id)
+        prefix = "  " * indent + "- "
+        output = f"{prefix}{self.label}\n"
+        for child in self.children:
+            # `visited.copy()` so that we only detect cycles within each path,
+            # rather than duplicate items in unrelated paths
+            output += child.render(indent + 1, visited.copy())
+        return output
+
+
+@dataclass
+class NestedList:
+    """
+    A nested list that can be rendered in Markdown-like format:
+
+    - Item 1
+    - Item 2
+      - Item 2.1
+      - Item 2.2
+        - Item 2.2.1
+    - Item 3
+    """
+
+    children: list[NestedListItem] = field(default_factory=list)
+
+    def render(self) -> str:
+        output = ""
+        for child in self.children:
+            output += child.render()
+        return output

dstack/_internal/utils/network.py

@@ -0,0 +1,50 @@
+import ipaddress
+from typing import List, Optional, Sequence
+
+
+def get_ip_from_network(network: Optional[str], addresses: Sequence[str]) -> Optional[str]:
+    ip_addresses = []
+    for address in addresses:
+        try:
+            interface = ipaddress.IPv4Interface(address)
+            ip_addresses.append(interface.ip)
+        except ipaddress.AddressValueError:
+            continue
+
+    internal_network = None
+    if network is not None:
+        if not ip_addresses:
+            return None
+        try:
+            internal_network = ipaddress.ip_network(network)
+        except ValueError:
+            return None
+
+        if not internal_network.is_private:
+            return None
+
+        for internal_ip in ip_addresses:
+            if internal_ip in internal_network:
+                return str(internal_ip)
+        else:
+            return None
+
+    # return any ipv4
+    internal_ip = str(ip_addresses[0]) if ip_addresses else None
+    return internal_ip
+
+
+def is_ip_among_addresses(ip_address: str, addresses: Sequence[str]) -> bool:
+    ip_addresses = get_ips_from_addresses(addresses)
+    return ip_address in ip_addresses
+
+
+def get_ips_from_addresses(addresses: Sequence[str]) -> List[str]:
+    ip_addresses = []
+    for address in addresses:
+        try:
+            interface = ipaddress.IPv4Interface(address)
+            ip_addresses.append(interface.ip)
+        except ipaddress.AddressValueError:
+            continue
+    return [str(ip) for ip in ip_addresses]

dstack/_internal/utils/path.py

@@ -0,0 +1,57 @@
+import os
+from dataclasses import dataclass
+from pathlib import Path, PurePath, PurePosixPath
+from typing import Union
+
+PathLike = Union[str, os.PathLike]
+
+
+@dataclass
+class FilePath:
+    path: PathLike
+
+
+@dataclass
+class FileContent:
+    content: str
+
+
+FilePathOrContent = Union[FilePath, FileContent]
+
+
+def path_in_dir(path: PathLike, directory: PathLike) -> bool:
+    try:
+        Path(path).resolve().relative_to(Path(directory).resolve())
+        return True
+    except ValueError:
+        return False
+
+
+def normalize_path(path: PathLike) -> PurePath:
+    path = PurePath(path)
+    stack = []
+    for part in path.parts:
+        if part == "..":
+            if not stack:
+                raise ValueError("Path is outside of the top directory")
+            stack.pop()
+        else:
+            stack.append(part)
+    return PurePath(*stack)
+
+
+def resolve_relative_path(path: PathLike) -> PurePath:
+    path = PurePath(path)
+    if path.is_absolute():
+        raise ValueError("Path should be relative")
+    try:
+        return normalize_path(path)
+    except ValueError:
+        raise ValueError("Path is outside of the repo")
+
+
+def is_absolute_posix_path(path: PathLike) -> bool:
+    # Passing Windows path leads to undefined behavior
+    if str(path).startswith("~"):
+        return True
+    return PurePosixPath(path).is_absolute()