skypilot-nightly 1.0.0.dev20250215__py3-none-any.whl → 1.0.0.dev20250217__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
@@ -23,6 +23,8 @@ 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
 
@@ -36,16 +38,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'.
 
@@ -53,42 +51,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.
@@ -96,19 +96,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
 
 
@@ -253,7 +244,46 @@ class Backoff:
         return self._backoff
 
 
-def get_pretty_entry_point() -> str:
+_current_command: Optional[str] = None
+_current_client_entrypoint: Optional[str] = None
+
+
+def set_client_entrypoint_and_command(client_entrypoint: Optional[str],
+                                      client_command: Optional[str]):
+    """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, _current_client_entrypoint
+    _current_command = client_command
+    _current_client_entrypoint = client_entrypoint
+
+
+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_pretty_entrypoint_cmd() -> str:
     """Returns the prettified entry point of this process (sys.argv).
 
     Example return values:
@@ -298,29 +328,51 @@ def user_and_hostname_hash() -> str:
     return f'{getpass.getuser()}-{hostname_hash}'
 
 
-def read_yaml(path: str) -> 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.
+
+    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: 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):
 
@@ -408,43 +460,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.
 
@@ -492,12 +507,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:
@@ -600,7 +617,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:
@@ -715,3 +732,26 @@ def hash_file(path: str, hash_alg: str) -> 'hashlib._Hash':
                 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

sky/utils/config_utils.py

@@ -0,0 +1,204 @@
+"""Utilities for nested config."""
+import copy
+from typing import Any, Dict, List, Optional, Tuple
+
+from sky import sky_logging
+
+logger = sky_logging.init_logger(__name__)
+
+
+class Config(Dict[str, Any]):
+    """SkyPilot config that supports setting/getting values with nested keys."""
+
+    def get_nested(
+        self,
+        keys: Tuple[str, ...],
+        default_value: Any,
+        override_configs: Optional[Dict[str, Any]] = None,
+        allowed_override_keys: Optional[List[Tuple[str, ...]]] = None,
+        disallowed_override_keys: Optional[List[Tuple[str,
+                                                      ...]]] = None) -> Any:
+        """Gets a nested key.
+
+        If any key is not found, or any intermediate key does not point to a
+        dict value, returns 'default_value'.
+
+        Args:
+            keys: A tuple of strings representing the nested keys.
+            default_value: The default value to return if the key is not found.
+            override_configs: A dict of override configs with the same schema as
+                the config file, but only containing the keys to override.
+            allowed_override_keys: A list of keys that are allowed to be
+                overridden.
+            disallowed_override_keys: A list of keys that are disallowed to be
+                overridden.
+
+        Returns:
+            The value of the nested key, or 'default_value' if not found.
+        """
+        config = copy.deepcopy(self)
+        if override_configs is not None:
+            config = _recursive_update(config, override_configs,
+                                       allowed_override_keys,
+                                       disallowed_override_keys)
+        return _get_nested(config, keys, default_value, pop=False)
+
+    def set_nested(self, keys: Tuple[str, ...], value: Any) -> None:
+        """In-place sets a nested key to value.
+
+        Like get_nested(), if any key is not found, this will not raise an
+        error.
+        """
+        override = {}
+        for i, key in enumerate(reversed(keys)):
+            if i == 0:
+                override = {key: value}
+            else:
+                override = {key: override}
+        _recursive_update(self, override)
+
+    def pop_nested(self, keys: Tuple[str, ...], default_value: Any) -> Any:
+        """Pops a nested key."""
+        return _get_nested(self, keys, default_value, pop=True)
+
+    @classmethod
+    def from_dict(cls, config: Optional[Dict[str, Any]]) -> 'Config':
+        if config is None:
+            return cls()
+        return cls(**config)
+
+
+def _check_allowed_and_disallowed_override_keys(
+    key: str,
+    allowed_override_keys: Optional[List[Tuple[str, ...]]] = None,
+    disallowed_override_keys: Optional[List[Tuple[str, ...]]] = None
+) -> Tuple[Optional[List[Tuple[str, ...]]], Optional[List[Tuple[str, ...]]]]:
+    allowed_keys_with_matched_prefix: Optional[List[Tuple[str, ...]]] = []
+    disallowed_keys_with_matched_prefix: Optional[List[Tuple[str, ...]]] = []
+    if allowed_override_keys is not None:
+        for nested_key in allowed_override_keys:
+            if key == nested_key[0]:
+                if len(nested_key) == 1:
+                    # Allowed key is fully matched, no need to check further.
+                    allowed_keys_with_matched_prefix = None
+                    break
+                assert allowed_keys_with_matched_prefix is not None
+                allowed_keys_with_matched_prefix.append(nested_key[1:])
+        if (allowed_keys_with_matched_prefix is not None and
+                not allowed_keys_with_matched_prefix):
+            raise ValueError(f'Key {key} is not in allowed override keys: '
+                             f'{allowed_override_keys}')
+    else:
+        allowed_keys_with_matched_prefix = None
+
+    if disallowed_override_keys is not None:
+        for nested_key in disallowed_override_keys:
+            if key == nested_key[0]:
+                if len(nested_key) == 1:
+                    raise ValueError(
+                        f'Key {key} is in disallowed override keys: '
+                        f'{disallowed_override_keys}')
+                assert disallowed_keys_with_matched_prefix is not None
+                disallowed_keys_with_matched_prefix.append(nested_key[1:])
+    else:
+        disallowed_keys_with_matched_prefix = None
+    return allowed_keys_with_matched_prefix, disallowed_keys_with_matched_prefix
+
+
+def _recursive_update(
+        base_config: Config,
+        override_config: Dict[str, Any],
+        allowed_override_keys: Optional[List[Tuple[str, ...]]] = None,
+        disallowed_override_keys: Optional[List[Tuple[str,
+                                                      ...]]] = None) -> Config:
+    """Recursively updates base configuration with override configuration"""
+    for key, value in override_config.items():
+        (next_allowed_override_keys, next_disallowed_override_keys
+        ) = _check_allowed_and_disallowed_override_keys(
+            key, allowed_override_keys, disallowed_override_keys)
+        if key == 'kubernetes' and key in base_config:
+            merge_k8s_configs(base_config[key], value,
+                              next_allowed_override_keys,
+                              next_disallowed_override_keys)
+        elif (isinstance(value, dict) and key in base_config and
+              isinstance(base_config[key], dict)):
+            _recursive_update(base_config[key], value,
+                              next_allowed_override_keys,
+                              next_disallowed_override_keys)
+        else:
+            base_config[key] = value
+    return base_config
+
+
+def _get_nested(configs: Optional[Dict[str, Any]],
+                keys: Tuple[str, ...],
+                default_value: Any,
+                pop: bool = False) -> Any:
+    if configs is None:
+        return default_value
+    curr = configs
+    for i, key in enumerate(keys):
+        if isinstance(curr, dict) and key in curr:
+            value = curr[key]
+            if i == len(keys) - 1:
+                if pop:
+                    curr.pop(key, default_value)
+            curr = value
+        else:
+            return default_value
+    logger.debug(f'User config: {".".join(keys)} -> {curr}')
+    return curr
+
+
+def merge_k8s_configs(
+        base_config: Dict[Any, Any],
+        override_config: Dict[Any, Any],
+        allowed_override_keys: Optional[List[Tuple[str, ...]]] = None,
+        disallowed_override_keys: Optional[List[Tuple[str,
+                                                      ...]]] = None) -> None:
+    """Merge two configs into the base_config.
+
+    Updates nested dictionaries instead of replacing them.
+    If a list is encountered, it will be appended to the base_config list.
+
+    An exception is when the key is 'containers', in which case the
+    first container in the list will be fetched and merge_dict will be
+    called on it with the first container in the base_config list.
+    """
+    for key, value in override_config.items():
+        (next_allowed_override_keys, next_disallowed_override_keys
+        ) = _check_allowed_and_disallowed_override_keys(
+            key, allowed_override_keys, disallowed_override_keys)
+        if isinstance(value, dict) and key in base_config:
+            merge_k8s_configs(base_config[key], value,
+                              next_allowed_override_keys,
+                              next_disallowed_override_keys)
+        elif isinstance(value, list) and key in base_config:
+            assert isinstance(base_config[key], list), \
+                f'Expected {key} to be a list, found {base_config[key]}'
+            if key in ['containers', 'imagePullSecrets']:
+                # If the key is 'containers' or 'imagePullSecrets, we take the
+                # first and only container/secret in the list and merge it, as
+                # we only support one container per pod.
+                assert len(value) == 1, \
+                    f'Expected only one container, found {value}'
+                merge_k8s_configs(base_config[key][0], value[0],
+                                  next_allowed_override_keys,
+                                  next_disallowed_override_keys)
+            elif key in ['volumes', 'volumeMounts']:
+                # If the key is 'volumes' or 'volumeMounts', we search for
+                # item with the same name and merge it.
+                for new_volume in value:
+                    new_volume_name = new_volume.get('name')
+                    if new_volume_name is not None:
+                        destination_volume = next(
+                            (v for v in base_config[key]
+                             if v.get('name') == new_volume_name), None)
+                        if destination_volume is not None:
+                            merge_k8s_configs(destination_volume, new_volume)
+                        else:
+                            base_config[key].append(new_volume)
+            else:
+                base_config[key].extend(value)
+        else:
+            base_config[key] = value

sky/utils/control_master_utils.py

@@ -1,8 +1,7 @@
 """Utils to check if the ssh control master should be disabled."""
 
-import functools
-
 from sky import sky_logging
+from sky.utils import annotations
 from sky.utils import subprocess_utils
 
 logger = sky_logging.init_logger(__name__)
@@ -34,7 +33,7 @@ def is_tmp_9p_filesystem() -> bool:
     return filesystem_types[1].lower() == '9p'
 
 
-@functools.lru_cache
+@annotations.lru_cache(scope='global')
 def should_disable_control_master() -> bool:
     """Whether disable ssh control master based on file system.