skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/utils/common_utils.py

@@ -5,7 +5,7 @@ import functools
 import getpass
 import hashlib
 import inspect
-import json
+import io
 import os
 import platform
 import random
@@ -16,20 +16,21 @@ import time
 from typing import Any, Callable, Dict, List, Optional, Union
 import uuid
 
-import colorama
 import jinja2
 import jsonschema
+import psutil
 import yaml
 
 from sky import exceptions
 from sky import sky_logging
 from sky.skylet import constants
+from sky.usage import constants as usage_constants
+from sky.utils import annotations
 from sky.utils import ux_utils
 from sky.utils import validator
 
 _USER_HASH_FILE = os.path.expanduser('~/.sky/user_hash')
 USER_HASH_LENGTH = 8
-USER_HASH_LENGTH_IN_CLUSTER_NAME = 4
 
 # We are using base36 to reduce the length of the hash. 2 chars -> 36^2 = 1296
 # possibilities. considering the final cluster name contains the prefix as well,
@@ -38,16 +39,12 @@ CLUSTER_NAME_HASH_LENGTH = 2
 
 _COLOR_PATTERN = re.compile(r'\x1b[^m]*m')
 
-_PAYLOAD_PATTERN = re.compile(r'<sky-payload>(.*)</sky-payload>')
-_PAYLOAD_STR = '<sky-payload>{}</sky-payload>'
-
 _VALID_ENV_VAR_REGEX = '[a-zA-Z_][a-zA-Z0-9_]*'
 
 logger = sky_logging.init_logger(__name__)
 
-_usage_run_id = None
-
 
+@annotations.lru_cache(scope='request')
 def get_usage_run_id() -> str:
     """Returns a unique run id for each 'run'.
 
@@ -55,42 +52,44 @@ def get_usage_run_id() -> str:
     and has called its CLI or programmatic APIs. For example, two successive
     `sky launch` are two runs.
     """
-    global _usage_run_id
-    if _usage_run_id is None:
-        _usage_run_id = str(uuid.uuid4())
-    return _usage_run_id
+    usage_run_id = os.getenv(usage_constants.USAGE_RUN_ID_ENV_VAR)
+    if usage_run_id is not None:
+        return usage_run_id
+    return str(uuid.uuid4())
+
+
+def _is_valid_user_hash(user_hash: Optional[str]) -> bool:
+    if user_hash is None:
+        return False
+    try:
+        int(user_hash, 16)
+    except (TypeError, ValueError):
+        return False
+    return len(user_hash) == USER_HASH_LENGTH
+
+
+def generate_user_hash() -> str:
+    """Generates a unique user-machine specific hash."""
+    hash_str = user_and_hostname_hash()
+    user_hash = hashlib.md5(hash_str.encode()).hexdigest()[:USER_HASH_LENGTH]
+    if not _is_valid_user_hash(user_hash):
+        # A fallback in case the hash is invalid.
+        user_hash = uuid.uuid4().hex[:USER_HASH_LENGTH]
+    return user_hash
 
 
-def get_user_hash(force_fresh_hash: bool = False) -> str:
+def get_user_hash() -> str:
     """Returns a unique user-machine specific hash as a user id.
 
     We cache the user hash in a file to avoid potential user_name or
     hostname changes causing a new user hash to be generated.
-
-    Args:
-        force_fresh_hash: Bypasses the cached hash in USER_HASH_FILE and the
-            hash in the USER_ID_ENV_VAR and forces a fresh user-machine hash
-            to be generated. Used by `kubernetes.ssh_key_secret_field_name` to
-            avoid controllers sharing the same ssh key field name as the
-            local client.
     """
+    user_hash = os.getenv(constants.USER_ID_ENV_VAR)
+    if _is_valid_user_hash(user_hash):
+        assert user_hash is not None
+        return user_hash
 
-    def _is_valid_user_hash(user_hash: Optional[str]) -> bool:
-        if user_hash is None:
-            return False
-        try:
-            int(user_hash, 16)
-        except (TypeError, ValueError):
-            return False
-        return len(user_hash) == USER_HASH_LENGTH
-
-    if not force_fresh_hash:
-        user_hash = os.getenv(constants.USER_ID_ENV_VAR)
-        if _is_valid_user_hash(user_hash):
-            assert user_hash is not None
-            return user_hash
-
-    if not force_fresh_hash and os.path.exists(_USER_HASH_FILE):
+    if os.path.exists(_USER_HASH_FILE):
         # Read from cached user hash file.
         with open(_USER_HASH_FILE, 'r', encoding='utf-8') as f:
             # Remove invalid characters.
@@ -98,19 +97,10 @@ def get_user_hash(force_fresh_hash: bool = False) -> str:
         if _is_valid_user_hash(user_hash):
             return user_hash
 
-    hash_str = user_and_hostname_hash()
-    user_hash = hashlib.md5(hash_str.encode()).hexdigest()[:USER_HASH_LENGTH]
-    if not _is_valid_user_hash(user_hash):
-        # A fallback in case the hash is invalid.
-        user_hash = uuid.uuid4().hex[:USER_HASH_LENGTH]
+    user_hash = generate_user_hash()
     os.makedirs(os.path.dirname(_USER_HASH_FILE), exist_ok=True)
-    if not force_fresh_hash:
-        # Do not cache to file if force_fresh_hash is True since the file may
-        # be intentionally using a different hash, e.g. we want to keep the
-        # user_hash for usage collection the same on the jobs/serve controller
-        # as users' local client.
-        with open(_USER_HASH_FILE, 'w', encoding='utf-8') as f:
-            f.write(user_hash)
+    with open(_USER_HASH_FILE, 'w', encoding='utf-8') as f:
+        f.write(user_hash)
     return user_hash
 
 
@@ -183,7 +173,7 @@ def make_cluster_name_on_cloud(display_name: str,
             f'on the cloud, we convert it to {cluster_name_on_cloud}.')
     user_hash = ''
     if add_user_hash:
-        user_hash = get_user_hash()[:USER_HASH_LENGTH_IN_CLUSTER_NAME]
+        user_hash = get_user_hash()
         user_hash = f'-{user_hash}'
     user_hash_length = len(user_hash)
 
@@ -233,7 +223,7 @@ class Backoff:
     MULTIPLIER = 1.6
     JITTER = 0.4
 
-    def __init__(self, initial_backoff: int = 5, max_backoff_factor: int = 5):
+    def __init__(self, initial_backoff: float = 5, max_backoff_factor: int = 5):
         self._initial = True
         self._backoff = 0.0
         self._initial_backoff = initial_backoff
@@ -255,7 +245,62 @@ class Backoff:
         return self._backoff
 
 
-def get_pretty_entry_point() -> str:
+_current_command: Optional[str] = None
+_current_client_entrypoint: Optional[str] = None
+_using_remote_api_server: Optional[bool] = None
+
+
+def set_client_status(client_entrypoint: Optional[str],
+                      client_command: Optional[str],
+                      using_remote_api_server: bool):
+    """Override the current client entrypoint and command.
+
+    This is useful when we are on the SkyPilot API server side and we have a
+    client entrypoint and command from the client.
+    """
+    global _current_command
+    global _current_client_entrypoint
+    global _using_remote_api_server
+    _current_command = client_command
+    _current_client_entrypoint = client_entrypoint
+    _using_remote_api_server = using_remote_api_server
+
+
+def get_current_command() -> str:
+    """Returns the command related to this operation.
+
+    Normally uses get_pretty_entry_point(), but will use the client command on
+    the server side.
+    """
+    if _current_command is not None:
+        return _current_command
+
+    return get_pretty_entrypoint_cmd()
+
+
+def get_current_client_entrypoint(server_entrypoint: str) -> str:
+    """Returns the current client entrypoint.
+
+    Gets the client entrypoint from the context, if it is not set, returns the
+    server entrypoint.
+    """
+    if _current_client_entrypoint is not None:
+        return _current_client_entrypoint
+    return server_entrypoint
+
+
+def get_using_remote_api_server() -> bool:
+    """Returns whether the API server is remote."""
+    if _using_remote_api_server is not None:
+        return _using_remote_api_server
+    # This gets the right status for the local client.
+    # TODO(zhwu): This is to prevent circular import. We should refactor this.
+    # pylint: disable=import-outside-toplevel
+    from sky.server import common as server_common
+    return not server_common.is_api_server_local()
+
+
+def get_pretty_entrypoint_cmd() -> str:
     """Returns the prettified entry point of this process (sys.argv).
 
     Example return values:
@@ -300,28 +345,51 @@ def user_and_hostname_hash() -> str:
     return f'{getpass.getuser()}-{hostname_hash}'
 
 
-def read_yaml(path) -> Dict[str, Any]:
+def read_yaml(path: Optional[str]) -> Dict[str, Any]:
+    if path is None:
+        raise ValueError('Attempted to read a None YAML.')
     with open(path, 'r', encoding='utf-8') as f:
         config = yaml.safe_load(f)
     return config
 
 
+def read_yaml_all_str(yaml_str: str) -> List[Dict[str, Any]]:
+    stream = io.StringIO(yaml_str)
+    config = yaml.safe_load_all(stream)
+    configs = list(config)
+    if not configs:
+        # Empty YAML file.
+        return [{}]
+    return configs
+
+
 def read_yaml_all(path: str) -> List[Dict[str, Any]]:
     with open(path, 'r', encoding='utf-8') as f:
-        config = yaml.safe_load_all(f)
-        configs = list(config)
-        if not configs:
-            # Empty YAML file.
-            return [{}]
-        return configs
+        return read_yaml_all_str(f.read())
+
 
+def dump_yaml(path: str, config: Union[List[Dict[str, Any]],
+                                       Dict[str, Any]]) -> None:
+    """Dumps a YAML file.
 
-def dump_yaml(path, config) -> None:
+    Args:
+        path: the path to the YAML file.
+        config: the configuration to dump.
+    """
     with open(path, 'w', encoding='utf-8') as f:
         f.write(dump_yaml_str(config))
 
 
-def dump_yaml_str(config):
+def dump_yaml_str(config: Union[List[Dict[str, Any]], Dict[str, Any]]) -> str:
+    """Dumps a YAML string.
+
+    Args:
+        config: the configuration to dump.
+
+    Returns:
+        The YAML string.
+    """
+
     # https://github.com/yaml/pyyaml/issues/127
     class LineBreakDumper(yaml.SafeDumper):
 
@@ -331,9 +399,9 @@ def dump_yaml_str(config):
                 super().write_line_break()
 
     if isinstance(config, list):
-        dump_func = yaml.dump_all
+        dump_func = yaml.dump_all  # type: ignore
     else:
-        dump_func = yaml.dump
+        dump_func = yaml.dump  # type: ignore
     return dump_func(config,
                      Dumper=LineBreakDumper,
                      sort_keys=False,
@@ -362,7 +430,6 @@ def make_decorator(cls, name_or_fn: Union[str, Callable],
 
             @functools.wraps(f)
             def _record(*args, **kwargs):
-                nonlocal name_or_fn
                 with cls(name_or_fn, **ctx_kwargs):
                     return f(*args, **kwargs)
 
@@ -376,7 +443,6 @@ def make_decorator(cls, name_or_fn: Union[str, Callable],
 
         @functools.wraps(name_or_fn)
         def _record(*args, **kwargs):
-            nonlocal name_or_fn
             f = name_or_fn
             func_name = getattr(f, '__qualname__', f.__name__)
             module_name = getattr(f, '__module__', '')
@@ -411,43 +477,6 @@ def retry(method, max_retries=3, initial_backoff=1):
     return method_with_retries
 
 
-def encode_payload(payload: Any) -> str:
-    """Encode a payload to make it more robust for parsing.
-
-    This makes message transfer more robust to any additional strings added to
-    the message during transfer.
-
-    An example message that is polluted by the system warning:
-    "LC_ALL: cannot change locale (en_US.UTF-8)\n<sky-payload>hello, world</sky-payload>" # pylint: disable=line-too-long
-
-    Args:
-        payload: A str, dict or list to be encoded.
-
-    Returns:
-        A string that is encoded from the payload.
-    """
-    payload_str = json.dumps(payload)
-    payload_str = _PAYLOAD_STR.format(payload_str)
-    return payload_str
-
-
-def decode_payload(payload_str: str) -> Any:
-    """Decode a payload string.
-
-    Args:
-        payload_str: A string that is encoded from a payload.
-
-    Returns:
-        A str, dict or list that is decoded from the payload string.
-    """
-    matched = _PAYLOAD_PATTERN.findall(payload_str)
-    if not matched:
-        raise ValueError(f'Invalid payload string: \n{payload_str}')
-    payload_str = matched[0]
-    payload = json.loads(payload_str)
-    return payload
-
-
 def class_fullname(cls, skip_builtins: bool = True):
     """Get the full name of a class.
 
@@ -478,11 +507,9 @@ def format_exception(e: Union[Exception, SystemExit, KeyboardInterrupt],
     Returns:
         A string that represents the exception.
     """
-    bright = colorama.Style.BRIGHT
-    reset = colorama.Style.RESET_ALL
     if use_bracket:
-        return f'{bright}[{class_fullname(e.__class__)}]{reset} {e}'
-    return f'{bright}{class_fullname(e.__class__)}:{reset} {e}'
+        return f'[{class_fullname(e.__class__)}] {e}'
+    return f'{class_fullname(e.__class__)}: {e}'
 
 
 def remove_color(s: str):
@@ -497,12 +524,14 @@ def remove_color(s: str):
     return _COLOR_PATTERN.sub('', s)
 
 
-def remove_file_if_exists(path: str):
+def remove_file_if_exists(path: Optional[str]):
     """Delete a file if it exists.
 
     Args:
         path: The path to the file.
     """
+    if path is None:
+        return
     try:
         os.remove(path)
     except FileNotFoundError:
@@ -581,7 +610,10 @@ def validate_schema(obj, schema, err_msg_prefix='', skip_none=True):
                            e.message)
             else:
                 err_msg = err_msg_prefix
+                assert isinstance(e.schema, dict), 'Schema must be a dictionary'
                 known_fields = set(e.schema.get('properties', {}).keys())
+                assert isinstance(e.instance,
+                                  dict), 'Instance must be a dictionary'
                 for field in e.instance:
                     if field not in known_fields:
                         most_similar_field = difflib.get_close_matches(
@@ -602,7 +634,7 @@ def validate_schema(obj, schema, err_msg_prefix='', skip_none=True):
 
     if err_msg:
         with ux_utils.print_exception_no_traceback():
-            raise ValueError(err_msg)
+            raise exceptions.InvalidSkyPilotConfigError(err_msg)
 
 
 def get_cleaned_username(username: str = '') -> str:
@@ -634,7 +666,7 @@ def get_cleaned_username(username: str = '') -> str:
     return username
 
 
-def fill_template(template_name: str, variables: Dict,
+def fill_template(template_name: str, variables: Dict[str, Any],
                   output_path: str) -> None:
     """Create a file from a Jinja template and return the filename."""
     assert template_name.endswith('.j2'), template_name
@@ -678,3 +710,182 @@ def deprecated_function(
         return func(*args, **kwargs)
 
     return new_func
+
+
+def truncate_long_string(s: str, max_length: int = 35) -> str:
+    """Truncate a string to a maximum length, preserving whole words."""
+    if len(s) <= max_length:
+        return s
+    splits = s.split(' ')
+    if len(splits[0]) > max_length:
+        return splits[0][:max_length] + '...'  # Use '…'?
+    # Truncate on word boundary.
+    i = 0
+    total = 0
+    for i, part in enumerate(splits):
+        total += len(part)
+        if total >= max_length:
+            break
+    prefix = ' '.join(splits[:i])
+    if len(prefix) < max_length:
+        prefix += s[len(prefix):max_length]
+    return prefix + '...'
+
+
+def hash_file(path: str, hash_alg: str) -> 'hashlib._Hash':
+    # In python 3.11, hashlib.file_digest is available, but for <3.11 we have to
+    # do it manually.
+    # This implementation is simplified from the implementation in CPython.
+    # TODO(cooperc): Use hashlib.file_digest once we move to 3.11+.
+    # Beware of f.read() as some files may be larger than memory.
+    with open(path, 'rb') as f:
+        file_hash = hashlib.new(hash_alg)
+        buf = bytearray(2**18)
+        view = memoryview(buf)
+        while True:
+            size = f.readinto(buf)
+            if size == 0:
+                # EOF
+                break
+            file_hash.update(view[:size])
+        return file_hash
+
+
+def is_port_available(port: int, reuse_addr: bool = True) -> bool:
+    """Check if a TCP port is available for binding on localhost.
+
+    Args:
+        port: The port number to check.
+        reuse_addr: If True, sets SO_REUSEADDR socket option to allow reusing
+            ports in TIME_WAIT state. Servers like multiprocessing.Manager set
+            SO_REUSEADDR by default to accelerate restart. The option should be
+            coordinated in check.
+
+    Returns:
+        bool: True if the port is available for binding, False otherwise.
+    """
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        if reuse_addr:
+            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        try:
+            s.bind(('localhost', port))
+            return True
+        except OSError:
+            return False
+
+
+def get_cpu_count() -> int:
+    """Get the number of CPUs, with cgroup awareness."""
+    # This env-var is kept since it is still useful for limiting the resource
+    # of SkyPilot in non-containerized environments.
+    cpu_count = os.getenv('SKYPILOT_POD_CPU_CORE_LIMIT')
+    if cpu_count is not None:
+        try:
+            return int(float(cpu_count))
+        except ValueError as e:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    f'Failed to parse the number of CPUs from {cpu_count}'
+                ) from e
+    return _cpu_count()
+
+
+def get_mem_size_gb() -> float:
+    """Get the memory size in GB, with cgroup awareness."""
+    mem_size = os.getenv('SKYPILOT_POD_MEMORY_GB_LIMIT')
+    if mem_size is not None:
+        try:
+            return float(mem_size)
+        except ValueError as e:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    f'Failed to parse the memory size from {mem_size}') from e
+    return _mem_size_gb()
+
+
+def _cpu_count() -> int:
+    # host cpu cores (logical)
+    cpu = psutil.cpu_count()
+    # cpu affinity on Linux
+    if hasattr(os, 'sched_getaffinity'):
+        # just for safe, length of CPU set should always <= logical cpu cores
+        cpu = min(cpu, len(os.sched_getaffinity(0)))
+    cgroup_cpu = _get_cgroup_cpu_limit()
+    if cgroup_cpu is not None:
+        cpu = min(cpu, int(cgroup_cpu))
+    return cpu
+
+
+def _mem_size_gb() -> float:
+    # host memory limit
+    mem = psutil.virtual_memory().total
+    cgroup_mem = _get_cgroup_memory_limit()
+    if cgroup_mem is not None:
+        mem = min(mem, cgroup_mem)
+    return mem / (1024**3)
+
+
+# Refer to:
+# - https://docs.kernel.org/admin-guide/cgroup-v1/index.html
+# - https://docs.kernel.org/admin-guide/cgroup-v2.html
+# for the standards of handler files in cgroupv1 and v2.
+# Since all those paths are well-known standards that are unlikely to change,
+# we use string literals instead of defining extra constants.
+def _get_cgroup_cpu_limit() -> Optional[float]:
+    """Return cpu limit from cgroups in cores.
+
+    Returns:
+        The cpu limit in cores as a float (can be fractional), or None if there
+        is no limit in cgroups.
+    """
+    try:
+        if _is_cgroup_v2():
+            with open('/sys/fs/cgroup/cpu.max', 'r', encoding='utf-8') as f:
+                quota_str, period_str = f.read().strip().split()
+                if quota_str == 'max':
+                    return None
+                quota = float(quota_str)
+                period = float(period_str)
+                return quota / period if quota > 0 else None
+        else:
+            # cgroup v1
+            with open('/sys/fs/cgroup/cpu/cpu.cfs_quota_us',
+                      'r',
+                      encoding='utf-8') as f:
+                quota = float(f.read().strip())
+            with open('/sys/fs/cgroup/cpu/cpu.cfs_period_us',
+                      'r',
+                      encoding='utf-8') as f:
+                period = float(f.read().strip())
+            # Return unlimited if cpu quota is not set.
+            # Note that we do not use cpu.shares since it is a relative weight
+            # instead of a hard limit. It is okay to get CPU throttling under
+            # high contention. And unlimited enables the server to use as much
+            # CPU as available if there is no contention.
+            return quota / period if (quota > 0 and period > 0) else None
+    except (OSError, ValueError):
+        return None
+
+
+def _get_cgroup_memory_limit() -> Optional[int]:
+    """Return memory limit from cgroups in bytes.
+
+    Returns:
+        The memory limit in bytes, or None if there is no limit in cgroups.
+    """
+    try:
+        path = ('/sys/fs/cgroup/memory.max' if _is_cgroup_v2() else
+                '/sys/fs/cgroup/memory/memory.limit_in_bytes')
+        with open(path, 'r', encoding='utf-8') as f:
+            value = f.read().strip()
+            if value == 'max' or not value:
+                return None
+            limit = int(value)
+            return limit if limit > 0 else None
+    except (OSError, ValueError):
+        return None
+
+
+def _is_cgroup_v2() -> bool:
+    """Return True if the environment is running cgroup v2."""
+    return os.path.isfile('/sys/fs/cgroup/cgroup.controllers')