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

dstack/_internal/utils/random_names.py

@@ -0,0 +1,258 @@
+import random
+from typing import List
+
+__adjectives = [
+    "afraid",
+    "ancient",
+    "angry",
+    "average",
+    "bad",
+    "big",
+    "bitter",
+    "black",
+    "blue",
+    "brave",
+    "breezy",
+    "bright",
+    "brown",
+    "calm",
+    "chatty",
+    "chilly",
+    "clever",
+    "cold",
+    "cowardly",
+    "cuddly",
+    "curly",
+    "curvy",
+    "dangerous",
+    "dry",
+    "dull",
+    "empty",
+    "evil",
+    "fast",
+    "fat",
+    "fluffy",
+    "foolish",
+    "fresh",
+    "friendly",
+    "funny",
+    "fuzzy",
+    "gentle",
+    "giant",
+    "good",
+    "great",
+    "green",
+    "grumpy",
+    "happy",
+    "hard",
+    "heavy",
+    "helpless",
+    "honest",
+    "horrible",
+    "hot",
+    "hungry",
+    "itchy",
+    "jolly",
+    "kind",
+    "lazy",
+    "light",
+    "little",
+    "loud",
+    "lovely",
+    "lucky",
+    "massive",
+    "mean",
+    "mighty",
+    "modern",
+    "moody",
+    "nasty",
+    "neat",
+    "nervous",
+    "new",
+    "nice",
+    "odd",
+    "old",
+    "orange",
+    "ordinary",
+    "perfect",
+    "pink",
+    "plastic",
+    "polite",
+    "popular",
+    "pretty",
+    "proud",
+    "purple",
+    "quick",
+    "quiet",
+    "rare",
+    "red",
+    "rotten",
+    "rude",
+    "selfish",
+    "serious",
+    "shaggy",
+    "sharp",
+    "short",
+    "shy",
+    "silent",
+    "silly",
+    "slimy",
+    "slippery",
+    "smart",
+    "smooth",
+    "soft",
+    "sour",
+    "spicy",
+    "splendid",
+    "spotty",
+    "stale",
+    "strange",
+    "strong",
+    "stupid",
+    "sweet",
+    "swift",
+    "tall",
+    "tame",
+    "tasty",
+    "tender",
+    "terrible",
+    "thin",
+    "tidy",
+    "tiny",
+    "tough",
+    "tricky",
+    "ugly",
+    "unlucky",
+    "warm",
+    "weak",
+    "wet",
+    "white",
+    "wicked",
+    "wise",
+    "witty",
+    "wonderful",
+    "yellow",
+    "young",
+]
+
+
+__animals = [
+    "ape",
+    "baboon",
+    "badger",
+    "bat",
+    "bear",
+    "bird",
+    "bobcat",
+    "bulldog",
+    "bullfrog",
+    "cat",
+    "catfish",
+    "cheetah",
+    "chicken",
+    "chipmunk",
+    "cobra",
+    "cougar",
+    "cow",
+    "crab",
+    "deer",
+    "dingo",
+    "dodo",
+    "dog",
+    "dolphin",
+    "donkey",
+    "dragon",
+    "dragonfly",
+    "duck",
+    "eagle",
+    "earwig",
+    "eel",
+    "elephant",
+    "emu",
+    "falcon",
+    "fireant",
+    "firefox",
+    "fish",
+    "fly",
+    "fox",
+    "frog",
+    "gecko",
+    "goat",
+    "goose",
+    "grasshopper",
+    "horse",
+    "hound",
+    "husky",
+    "impala",
+    "insect",
+    "jellyfish",
+    "kangaroo",
+    "ladybug",
+    "liger",
+    "lion",
+    "lionfish",
+    "lizard",
+    "mayfly",
+    "mole",
+    "monkey",
+    "moose",
+    "moth",
+    "mouse",
+    "mule",
+    "newt",
+    "octopus",
+    "otter",
+    "owl",
+    "panda",
+    "panther",
+    "parrot",
+    "penguin",
+    "pig",
+    "puma",
+    "pug",
+    "quail",
+    "rabbit",
+    "rat",
+    "rattlesnake",
+    "robin",
+    "seahorse",
+    "sheep",
+    "shrimp",
+    "skunk",
+    "sloth",
+    "snail",
+    "snake",
+    "squid",
+    "starfish",
+    "stingray",
+    "swan",
+    "termite",
+    "tiger",
+    "treefrog",
+    "turkey",
+    "turtle",
+    "vampirebat",
+    "walrus",
+    "warthog",
+    "wasp",
+    "wolverine",
+    "wombat",
+    "yak",
+    "zebra",
+]
+
+
+def get_adjectives() -> List[str]:
+    return __adjectives
+
+
+def get_animals() -> List[str]:
+    return __animals
+
+
+def generate_name() -> str:
+    return generate_name_from_sets(__adjectives, __animals)
+
+
+def generate_name_from_sets(set1: List[str], set2: List[str]) -> str:
+    return random.choice(set1) + "-" + random.choice(set2)

dstack/_internal/utils/ssh.py

@@ -0,0 +1,346 @@
+import io
+import os
+import re
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+from typing import Dict, Optional, Union
+
+import paramiko
+from filelock import FileLock
+from paramiko.config import SSHConfig
+from paramiko.pkey import PKey, PublicBlob
+from paramiko.ssh_exception import SSHException
+
+from dstack._internal.compat import IS_WINDOWS
+from dstack._internal.utils.logging import get_logger
+from dstack._internal.utils.path import FilePath, PathLike
+
+logger = get_logger(__name__)
+
+
+default_ssh_config_path = "~/.ssh/config"
+
+SUPPORTED_KEY_TYPES = (paramiko.RSAKey, paramiko.ECDSAKey, paramiko.Ed25519Key)
+
+
+def get_public_key_fingerprint(text: str) -> str:
+    pb = PublicBlob.from_string(text)
+    pk = PKey.from_type_string(pb.key_type, pb.key_blob)
+    return pk.fingerprint
+
+
+def get_host_config(hostname: str, ssh_config_path: PathLike = default_ssh_config_path) -> dict:
+    ssh_config_path = os.path.expanduser(ssh_config_path)
+    if os.path.exists(ssh_config_path):
+        config = SSHConfig.from_path(ssh_config_path)
+        return config.lookup(hostname)
+    return {}
+
+
+def make_ssh_command_for_git(identity_file: PathLike) -> str:
+    # No need to use :func:`find_ssh_client()` even on Windows even if `ssh` not in
+    # Windows `PATH` -- MSYS2 git (from Git for Windows) always has access to it,
+    # see https://www.msys2.org/docs/environments/ ("MSYS environment [...] is always active")
+    return (
+        f'ssh -F none -i "{normalize_path(identity_file)}"'
+        " -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o IdentitiesOnly=yes"
+    )
+
+
+def make_git_env(
+    *,
+    disable_prompt: bool = True,
+    disable_askpass: bool = False,
+    disable_config: bool = False,
+    identity_file: Optional[PathLike] = None,
+) -> dict[str, str]:
+    env: dict[str, str] = {}
+    if disable_prompt:
+        # Fail with error instead of prompting on the terminal (e.g., when asking for
+        # HTTP authentication)
+        env["GIT_TERMINAL_PROMPT"] = "0"
+    if disable_askpass:
+        env["GIT_ASKPASS"] = ""
+        env["SSH_ASKPASS"] = ""
+    if disable_config:
+        # Disable system-wide config (usually /etc/gitconfig)
+        env["GIT_CONFIG_SYSTEM"] = os.devnull
+        # Disable user (aka "global") config ($XDG_CONFIG_HOME/git/config or ~/.git/config)
+        env["GIT_CONFIG_GLOBAL"] = os.devnull
+        # Disable repo (aka "local") config (./.git/config)
+        env["GIT_DIR"] = os.devnull
+    if identity_file is not None:
+        env["GIT_SSH_COMMAND"] = make_ssh_command_for_git(identity_file)
+    return env
+
+
+def try_ssh_key_passphrase(identity_file: PathLike, passphrase: str = "") -> bool:
+    ssh_keygen = find_ssh_util("ssh-keygen")
+    if ssh_keygen is None:
+        logger.warning("ssh-keygen not found")
+        return False
+    r = subprocess.run(
+        [ssh_keygen, "-y", "-P", passphrase, "-f", identity_file],
+        stdout=subprocess.DEVNULL,
+        stderr=sys.stdout.buffer,
+    )
+    return r.returncode == 0
+
+
+def normalize_path(path: PathLike, *, collapse_user: bool = False) -> str:
+    """
+    Converts a path to the most compatible format.
+    On Windows, replaces backslashes with slashes.
+    Additionally, if `collapse_user` is `True`, tries to replace the user home part of the path
+    with `~`.
+
+    :param path: Path object or string
+    :param collapse_user: try to replace user home prefix with `~`. `False` by default.
+    :return: Normalized path as string
+    """
+    if collapse_user:
+        # The following "reverse" expanduser operation not only makes paths shorter and "nicer",
+        # but also fixes one specific issue with OpenSSH bundled with Git for Windows (MSYS2),
+        # see :func:`include_ssh_config` for details.
+        try:
+            path = Path(path).relative_to(Path.home())
+            path = f"~/{path}"
+        except ValueError:
+            pass
+    if IS_WINDOWS:
+        # Git for Windows ssh (based on MSYS2, but there may be subtle differences between
+        # vanilla MSYS2 ssh and Git for Windows ssh) supports:
+        #   * C:\\Users\\User
+        #   * C:/Users/User
+        #   * /c/Users/User
+        # does not support:
+        #   * C:\Users\User (as pathllib.WindowsPath is rendered)
+        # OpenSSH_for_Windows supports:
+        #   * C:\Users\User
+        #   * C:\\Users\\User
+        #   * C:/Users/User
+        # does not support:
+        #   * /c/User/User
+        # We use C:/Users/User format as the safest (supported by both ssh builds;
+        # no backslash-escaping pitfalls)
+        return str(path).replace("\\", "/")
+    return str(path)
+
+
+def include_ssh_config(path: PathLike, ssh_config_path: PathLike = default_ssh_config_path):
+    """
+    Adds Include entry on top of the default ssh config file
+    :param path: Absolute path to config
+    :param ssh_config_path: ~/.ssh/config
+    """
+    ssh_config_path = os.path.expanduser(ssh_config_path)
+    Path(ssh_config_path).parent.mkdir(0o600, parents=True, exist_ok=True)
+    # MSYS2 OpenSSH accepts only /c/Users/User/... format in the Include directive (although
+    # it accepts C:/Users/User/... in other directives). We try to work around this issue
+    # converting the path to ~/.dstack/... format.
+    path = normalize_path(path, collapse_user=True)
+    include = f"Include {path}\n"
+    content = ""
+    try:
+        with FileLock(str(ssh_config_path) + ".lock"):
+            if os.path.exists(ssh_config_path):
+                with open(ssh_config_path, "r") as f:
+                    content = f.read()
+            if include not in content:
+                with open(ssh_config_path, "w") as f:
+                    f.write(include + content)
+    except PermissionError:
+        logger.warning(
+            f"Couldn't update `{ssh_config_path}` due to a permissions problem.\n"
+            f"The `vscode://vscode-remote/ssh-remote+<run name>/<working_dir>` and "
+            f"`cursor://vscode-remote/ssh-remote+<run name>/<working_dir>` links and "
+            f"the `ssh <run name>` command won't work.\n"
+            f"To fix this, make sure `{ssh_config_path}` is writable, or add "
+            f"`Include {path}` to the top of `{ssh_config_path}` manually.",
+            extra={"markup": True},
+        )
+
+
+def get_ssh_config(path: PathLike, host: str) -> Optional[Dict[str, str]]:
+    if os.path.exists(path):
+        config = {}
+        current_host = None
+
+        with open(path, "r") as f:
+            for line in f:
+                line = line.strip()
+
+                if not line or line.startswith("#"):
+                    continue
+
+                if line.startswith("Host "):
+                    current_host = line.split(" ")[1]
+                    config[current_host] = {}
+                else:
+                    key, value = line.split(maxsplit=1)
+                    config[current_host][key] = value
+        return config.get(host)
+    else:
+        return None
+
+
+def update_ssh_config(path: PathLike, host: str, options: Dict[str, Union[str, int, FilePath]]):
+    Path(path).parent.mkdir(parents=True, exist_ok=True)
+    with FileLock(str(path) + ".lock"):
+        copy_mode = True
+        content = ""
+        if os.path.exists(path):
+            with open(path, "r") as f:
+                for line in f:
+                    m = re.match(r"^Host\s+(\S+)$", line.strip())
+                    if m is not None:
+                        copy_mode = m.group(1) != host
+                    if copy_mode:
+                        content += line
+        with open(path, "w") as f:
+            f.write(content)
+            if options:
+                f.write(f"Host {host}\n")
+                for k, v in options.items():
+                    if isinstance(v, FilePath):
+                        v = normalize_path(v.path, collapse_user=True)
+                    f.write(f"    {k} {v}\n")
+            f.flush()
+
+
+def convert_ssh_key_to_pem(private_string: str) -> str:
+    if not private_string.startswith("-----BEGIN PRIVATE KEY-----"):
+        return private_string
+
+    with tempfile.NamedTemporaryFile(mode="w+") as key_file:
+        key_file.write(private_string)
+        key_file.flush()
+        if ssh_keygen := find_ssh_util("ssh-keygen"):
+            cmd = [ssh_keygen, "-p", "-m", "PEM", "-f", key_file.name, "-y", "-q", "-N", ""]
+            try:
+                subprocess.run(
+                    cmd,
+                    check=True,
+                    capture_output=True,
+                    text=True,
+                )
+            except subprocess.CalledProcessError as e:
+                logger.error("Fail to convert ssh key: stdout=%s, stderr=%s", e.stdout, e.stderr)
+        else:
+            logger.error("Use a PEM key or install ssh-keygen to convert it automatically")
+        key_file.seek(0)
+        private_string = key_file.read()
+    return private_string
+
+
+def pkey_from_str(private_string: str) -> PKey:
+    for key_type in SUPPORTED_KEY_TYPES:
+        try:
+            key_file = io.StringIO(private_string.strip())
+            pkey = key_type.from_private_key(key_file)
+            key_file.close()
+            return pkey
+        except (SSHException, ValueError):
+            pass
+
+    raise ValueError("Unsupported key type")
+
+
+def generate_public_key(private_key: PKey) -> str:
+    public_key = f"{private_key.get_name()} {private_key.get_base64()}"
+    return public_key
+
+
+def check_required_ssh_version() -> bool:
+    try:
+        result = subprocess.run(["ssh", "-V"], capture_output=True, text=True)
+    except subprocess.CalledProcessError:
+        logger.error("Failed to get ssh version information")
+    else:
+        if result.returncode == 0:
+            # Extract the version number from stderr in unix-like and window systems
+            version_output = result.stderr if result.stderr.strip() else result.stdout.strip()
+
+            match = re.search(r"_(\d+\.\d+)", version_output)
+            if match:
+                ssh_version = float(match.group(1))
+                if ssh_version >= 8.4:
+                    return True
+                else:
+                    return False
+
+    return False
+
+
+def find_ssh_client() -> Optional[Path]:
+    """
+    Finds and returns an absolute path of `ssh` executable or `None` if not found.
+
+    If the `DSTACK_SSH_CLIENT` environment variable is set, return its value, otherwise:
+        * on POSIX, look for `ssh` executable in `PATH` and return it (if any).
+        * on Windows, first look for OpenSSH bundled with Git for Windows checking
+        a known directory structure, then check `PATH`, and finally check a well-known location of
+        OpenSSH for Windows.
+    """
+    path_str = os.getenv("DSTACK_SSH_CLIENT")
+    if path_str:
+        path = Path(path_str)
+        if path.exists():
+            return path.resolve()
+        logger.warning("DSTACK_SSH_CLIENT=%s does not exist", path_str)
+        return None
+    if not IS_WINDOWS:
+        path_str = shutil.which("ssh")
+        if path_str:
+            return Path(path_str)
+        return None
+    # First, we check for ssh bundled with Git for Windows (MSYS2/MinGW-w64-built OpenSSH Portable)
+    # as a preferred client. It supports ForkAfterAuthentication; ControlMaster is only partially
+    # supported, we don't use it.
+    git_path_str = shutil.which("git")
+    if git_path_str:
+        # C:\Program Files\Git\cmd\git.exe -> C:\Program Files\Git\usr\bin\ssh.exe
+        path = Path(git_path_str).parent.parent / "usr" / "bin" / "ssh.exe"
+        if path.exists():
+            return path
+    # Then we check for any ssh client in PATH. It can be anything, but most likely it will be
+    # OpenSSH for Windows (see below). Nonetheless, it's worth trying since it's also may be
+    # MSYS2/Cygwin OpenSSH Portable.
+    path_str = shutil.which("ssh")
+    if path_str:
+        return Path(path_str)
+    # Finally we check for OpenSSH for Windows (Microsoft's fork of OpenSSH Portable).
+    # It does not support some features, namely ControlMaster and ForkAfterAuthentication.
+    windir_str = os.getenv("WINDIR")
+    if windir_str:
+        path = Path(windir_str) / "System32" / "OpenSSH" / "ssh.exe"
+        if path.exists():
+            return path
+    return None
+
+
+_ssh_util_dir: Optional[Path] = None
+
+
+def find_ssh_util(name: str) -> Optional[Path]:
+    """
+    Returns an absolute path of a given `ssh*` utility or `None` if not found.
+
+    :param name: a utility binary name without `.exe` suffix, e.g., `ssh-keygen`, `ssh-copy-id`.
+    :return: a Path object.
+    """
+    global _ssh_util_dir
+    if _ssh_util_dir is None:
+        ssh_client_path = find_ssh_client()
+        if ssh_client_path is None:
+            return None
+        _ssh_util_dir = ssh_client_path.parent
+    if IS_WINDOWS:
+        name = f"{name}.exe"
+    path = _ssh_util_dir / name
+    if path.exists():
+        return path
+    return None

dstack/_internal/utils/tags.py

@@ -0,0 +1,42 @@
+import re
+from typing import Dict, Optional
+
+# dstack resource tags allow alphanumeric tags with some special symbols in values.
+# Should be valid across most backends (e.g. AWS).
+# Does not guarantee that they are valid across all backends (e.g. GCP).
+# So backends need to filter bad tags out.
+TAG_KEY_PATTERN = re.compile(r"^[_\-a-zA-Z0-9]{1,60}$")
+TAG_VALUE_PATTERN = re.compile(r"^[a-zA-Z0-9 .:/=_\-+@]{0,256}$")
+
+
+def tags_validator(tags: Optional[Dict[str, str]]) -> Optional[Dict[str, str]]:
+    if tags is None:
+        return
+    validate_tags(tags)
+    return tags
+
+
+def validate_tags(tags: Dict[str, str]):
+    for k, v in tags.items():
+        _validate_tag(k, v)
+
+
+def _validate_tag(key: str, value: str):
+    if not is_valid_tag_key(key):
+        raise ValueError(
+            f"Invalid tag key {key}. The key must match regex '{TAG_KEY_PATTERN.pattern}'"
+        )
+    if not is_valid_tag_value(value):
+        raise ValueError(
+            f"Invalid tag value {value}. The value must match regex '{TAG_VALUE_PATTERN.pattern}'"
+        )
+
+
+def is_valid_tag_key(name: str) -> bool:
+    match = re.match(TAG_KEY_PATTERN, name)
+    return match is not None
+
+
+def is_valid_tag_value(value: str) -> bool:
+    match = re.match(TAG_VALUE_PATTERN, value)
+    return match is not None

dstack/_internal/utils/typing.py

@@ -0,0 +1,14 @@
+from typing import Any, Protocol, TypeVar, Union
+
+_T_contra = TypeVar("_T_contra", contravariant=True)
+
+
+class SupportsDunderLT(Protocol[_T_contra]):
+    def __lt__(self, other: _T_contra, /) -> bool: ...
+
+
+class SupportsDunderGT(Protocol[_T_contra]):
+    def __gt__(self, other: _T_contra, /) -> bool: ...
+
+
+SupportsRichComparison = Union[SupportsDunderLT[Any], SupportsDunderGT[Any]]

dstack/_internal/utils/version.py

@@ -0,0 +1,22 @@
+from typing import Optional
+
+import packaging.version
+
+
+def parse_version(version_string: str) -> Optional[packaging.version.Version]:
+    """
+    Returns a `packaging.version.Version` instance or `None` if the version is dev/latest.
+
+    Values parsed as the dev/latest version:
+    * the "latest" literal
+    * any "0.0.0" release, e.g., "0.0.0", "0.0.0a1", "0.0.0.dev0"
+    """
+    if version_string == "latest":
+        return None
+    try:
+        version = packaging.version.parse(version_string)
+    except packaging.version.InvalidVersion as e:
+        raise ValueError(f"Invalid version: {version_string}") from e
+    if version.release == (0, 0, 0):
+        return None
+    return version