konduktor-nightly 0.1.0.dev20250209104336__py3-none-any.whl → 0.1.0.dev20250313070642__py3-none-any.whl

konduktor/utils/common_utils.py

@@ -0,0 +1,393 @@
+# Proprietary Changes made for Trainy under the Trainy Software License
+# Original source: skypilot: https://github.com/skypilot-org/skypilot
+# which is Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import datetime
+import difflib
+import functools
+import getpass
+import hashlib
+import inspect
+import os
+import re
+import socket
+import sys
+import uuid
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import jinja2
+import jsonschema
+import yaml
+
+from konduktor.utils import annotations, constants, ux_utils, validator
+
+_USER_HASH_FILE = os.path.expanduser('~/.konduktor/user_hash')
+_usage_run_id = None
+_VALID_ENV_VAR_REGEX = '[a-zA-Z_][a-zA-Z0-9_]*'
+USER_HASH_LENGTH = 8
+USER_HASH_LENGTH_IN_CLUSTER_NAME = 4
+
+
+def get_timestamp() -> str:
+    return datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
+
+
+def user_and_hostname_hash() -> str:
+    """Returns a string containing <user>-<hostname hash last 4 chars>.
+
+    For uniquefying user workloads on a shared-k8s cluster.
+
+    Using uuid.getnode() instead of gethostname() is incorrect; observed to
+    collide on Macs.
+    """
+    hostname_hash = hashlib.md5(socket.gethostname().encode()).hexdigest()[-4:]
+    return f'{getpass.getuser()}-{hostname_hash}'
+
+
+def base36_encode(hex_str: str) -> str:
+    """Converts a hex string to a base36 string."""
+    int_value = int(hex_str, 16)
+
+    def _base36_encode(num: int) -> str:
+        if num == 0:
+            return '0'
+        alphabet = '0123456789abcdefghijklmnopqrstuvwxyz'
+        base36 = ''
+        while num != 0:
+            num, i = divmod(num, 36)
+            base36 = alphabet[i] + base36
+        return base36
+
+    return _base36_encode(int_value)
+
+
+def get_cleaned_username(username: str = '') -> str:
+    """Cleans the username. Underscores are allowed, as we will
+     handle it when mapping to the cluster_name_on_cloud in
+     common_utils.make_cluster_name_on_cloud.
+
+    Clean up includes:
+     1. Making all characters lowercase
+     2. Removing any non-alphanumeric characters (excluding hyphens and
+        underscores)
+     3. Removing any numbers and/or hyphens at the start of the username.
+     4. Removing any hyphens at the end of the username
+     5. Truncate the username to 63 characters, as requested by GCP labels
+
+    Dots are removed due to: https://cloud.google.com/compute/docs/labeling-resources#requirements
+
+    e.g. 1SkY-PiLot2- becomes sky-pilot2
+
+    Returns:
+      A cleaned username.
+    """  # noqa: E501
+    username = username or getpass.getuser()
+    username = username.lower()
+    username = re.sub(r'[^a-z0-9-_]', '', username)
+    username = re.sub(r'^[0-9-]+', '', username)
+    username = re.sub(r'-$', '', username)
+    username = username[:63]
+    return username
+
+
+def is_valid_env_var(name: str) -> bool:
+    """Checks if the task environment variable name is valid."""
+    return bool(re.fullmatch(_VALID_ENV_VAR_REGEX, name))
+
+
+def get_pretty_entry_point() -> str:
+    """Returns the prettified entry point of this process (sys.argv).
+
+    Example return values:
+        $ konduktor launch app.yaml  # 'konduktor launch app.yaml'
+        $ python examples/app.py  # 'app.py'
+    """
+    argv = sys.argv
+    basename = os.path.basename(argv[0])
+    if basename == 'konduktor':
+        # Turn '/.../anaconda/envs/py36/bin/sky' into 'konduktor', but keep other
+        # things like 'examples/app.py'.
+        argv[0] = basename
+    return ' '.join(argv)
+
+
+@annotations.lru_cache(scope='request')
+def get_usage_run_id() -> str:
+    """Returns a unique run id for each 'run'.
+
+    A run is defined as the lifetime of a process that has imported `sky`
+    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
+
+
+def make_decorator(cls, name_or_fn: Union[str, Callable], **ctx_kwargs) -> Callable:
+    """Make the cls a decorator.
+
+    class cls:
+        def __init__(self, name, **kwargs):
+            pass
+        def __enter__(self):
+            pass
+        def __exit__(self, exc_type, exc_value, traceback):
+            pass
+
+    Args:
+        name_or_fn: The name of the event or the function to be wrapped.
+        message: The message attached to the event.
+    """
+    if isinstance(name_or_fn, str):
+
+        def _wrapper(f):
+            @functools.wraps(f)
+            def _record(*args, **kwargs):
+                with cls(name_or_fn, **ctx_kwargs):
+                    return f(*args, **kwargs)
+
+            return _record
+
+        return _wrapper
+    else:
+        if not inspect.isfunction(name_or_fn):
+            raise ValueError('Should directly apply the decorator to a function.')
+
+        @functools.wraps(name_or_fn)
+        def _record(*args, **kwargs):
+            f = name_or_fn
+            func_name = getattr(f, '__qualname__', f.__name__)
+            module_name = getattr(f, '__module__', '')
+            if module_name:
+                full_name = f'{module_name}.{func_name}'
+            else:
+                full_name = func_name
+            with cls(full_name, **ctx_kwargs):
+                return f(*args, **kwargs)
+
+        return _record
+
+
+def get_user_hash(force_fresh_hash: bool = False) -> 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.
+    """
+
+    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):
+        # Read from cached user hash file.
+        with open(_USER_HASH_FILE, 'r', encoding='utf-8') as f:
+            # Remove invalid characters.
+            user_hash = f.read().strip()
+        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]
+    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)
+    return user_hash
+
+
+def read_yaml(path: str) -> Dict[str, Any]:
+    with open(path, 'r', encoding='utf-8') as f:
+        config = yaml.safe_load(f)
+    return config
+
+
+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
+
+
+def validate_schema(obj, schema, err_msg_prefix='', skip_none=True):
+    """Validates an object against a given JSON schema.
+
+    Args:
+        obj: The object to validate.
+        schema: The JSON schema against which to validate the object.
+        err_msg_prefix: The string to prepend to the error message if
+          validation fails.
+        skip_none: If True, removes fields with value None from the object
+          before validation. This is useful for objects that will never contain
+          None because yaml.safe_load() loads empty fields as None.
+
+    Raises:
+        ValueError: if the object does not match the schema.
+    """
+    if skip_none:
+        obj = {k: v for k, v in obj.items() if v is not None}
+    err_msg = None
+    try:
+        validator.SchemaValidator(schema).validate(obj)
+    except jsonschema.ValidationError as e:
+        if e.validator == 'additionalProperties':
+            if tuple(e.schema_path) == ('properties', 'envs', 'additionalProperties'):
+                # Hack. Here the error is Task.envs having some invalid keys. So
+                # we should not print "unsupported field".
+                #
+                # This will print something like:
+                # 'hello world' does not match any of the regexes: <regex>
+                err_msg = (
+                    err_msg_prefix
+                    + 'The `envs` field contains invalid keys:\n'
+                    + 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(
+                            field, known_fields, 1
+                        )
+                        if most_similar_field:
+                            err_msg += (
+                                f'Instead of {field!r}, did you mean '
+                                f'{most_similar_field[0]!r}?'
+                            )
+                        else:
+                            err_msg += f'Found unsupported field {field!r}.'
+        else:
+            message = e.message
+            # Object in jsonschema is represented as dict in Python. Replace
+            # 'object' with 'dict' for better readability.
+            message = message.replace("type 'object'", "type 'dict'")
+            # Example e.json_path value: '$.resources'
+            err_msg = (
+                err_msg_prefix
+                + message
+                + f'. Check problematic field(s): {e.json_path}'
+            )
+
+    if err_msg:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(err_msg)
+
+
+def dump_yaml(path: str, config: Union[List[Dict[str, Any]], Dict[str, Any]]) -> None:
+    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:
+    # https://github.com/yaml/pyyaml/issues/127
+    class LineBreakDumper(yaml.SafeDumper):
+        def write_line_break(self, data=None):
+            super().write_line_break(data)
+            if len(self.indents) == 1:
+                super().write_line_break()
+
+    if isinstance(config, list):
+        dump_func = yaml.dump_all  # type: ignore
+    else:
+        dump_func = yaml.dump  # type: ignore
+    return dump_func(
+        config, Dumper=LineBreakDumper, sort_keys=False, default_flow_style=False
+    )
+
+
+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
+    root_dir = os.path.dirname(os.path.dirname(__file__))
+    template_path = os.path.join(root_dir, 'templates', template_name)
+    if not os.path.exists(template_path):
+        raise FileNotFoundError(f'Template "{template_name}" does not exist.')
+    with open(template_path, 'r', encoding='utf-8') as fin:
+        template = fin.read()
+    output_path = os.path.abspath(os.path.expanduser(output_path))
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+
+    # Write out yaml config.
+    j2_template = jinja2.Template(template)
+    content = j2_template.render(**variables)
+    with open(output_path, 'w', encoding='utf-8') as fout:
+        fout.write(content)
+
+
+def class_fullname(cls, skip_builtins: bool = True):
+    """Get the full name of a class.
+
+    Example:
+        >>> e = konduktor.exceptions.FetchClusterInfoError()
+        >>> class_fullname(e.__class__)
+        'konduktor.exceptions.FetchClusterInfoError'
+
+    Args:
+        cls: The class to get the full name.
+
+    Returns:
+        The full name of the class.
+    """
+    module_name = getattr(cls, '__module__', '')
+    if not module_name or (module_name == 'builtins' and skip_builtins):
+        return cls.__name__
+    return f'{cls.__module__}.{cls.__name__}'
+
+
+def format_exception(
+    e: Union[Exception, SystemExit, KeyboardInterrupt], use_bracket: bool = False
+) -> str:
+    """Format an exception to a string.
+
+    Args:
+        e: The exception to format.
+
+    Returns:
+        A string that represents the exception.
+    """
+    if use_bracket:
+        return f'[{class_fullname(e.__class__)}] {e}'
+    return f'{class_fullname(e.__class__)}: {e}'

konduktor/utils/constants.py

@@ -0,0 +1,5 @@
+# The name for the environment variable that stores KONDUKTOR user hash
+USER_ID_ENV_VAR = 'KONDUKTOR_USER_ID'
+
+# The name for the environment variable that stores KONDUKTOR user name.
+USER_ENV_VAR = 'KONDUKTOR_USER'

konduktor/utils/env_options.py

@@ -0,0 +1,55 @@
+# Proprietary Changes made for Trainy under the Trainy Software License
+# Original source: skypilot: https://github.com/skypilot-org/skypilot
+# which is Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Global environment options for konduktor."""
+
+import enum
+import os
+from typing import Dict
+
+
+class Options(enum.Enum):
+    """Environment variables for SkyPilot."""
+
+    # (env var name, default value)
+    IS_DEVELOPER = ('KONDUKTOR_DEV', False)
+    SHOW_DEBUG_INFO = ('KONDUKTOR_DEBUG', True)
+    DISABLE_LOGGING = ('KONDUKTOR_DISABLE_USAGE_COLLECTION', False)
+    MINIMIZE_LOGGING = ('KONDUKTOR_MINIMIZE_LOGGING', False)
+    SUPPRESS_SENSITIVE_LOG = ('KONDUKTOR_SUPPRESS_SENSITIVE_LOG', False)
+    # Internal: this is used to skip the cloud user identity check, which is
+    # used to protect cluster operations in a multi-identity scenario.
+    # Currently, this is only used in the job and serve controller, as there
+    # will not be multiple identities, and skipping the check can increase
+    # robustness.
+    SKIP_CLOUD_IDENTITY_CHECK = ('KONDUKTOR_SKIP_CLOUD_IDENTITY_CHECK', False)
+
+    def __init__(self, env_var: str, default: bool) -> None:
+        self.env_var = env_var
+        self.default = default
+
+    def __repr__(self) -> str:
+        return self.env_var
+
+    def get(self) -> bool:
+        """Check if an environment variable is set to True."""
+        return os.getenv(self.env_var, str(self.default)).lower() in ('true', '1')
+
+    @property
+    def env_key(self) -> str:
+        """The environment variable key name."""
+        return self.value[0]
+
+    @classmethod
+    def all_options(cls) -> Dict[str, bool]:
+        """Returns all options as a dictionary."""
+        return {option.env_key: option.get() for option in list(Options)}

konduktor/utils/exceptions.py

@@ -0,0 +1,226 @@
+# Proprietary Changes made for Trainy under the Trainy Software License
+# Original source: skypilot: https://github.com/skypilot-org/skypilot
+# which is Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Exceptions."""
+
+import builtins
+import traceback
+import types
+from typing import Any, Dict
+
+# Return code for keyboard interruption and SIGTSTP
+KEYBOARD_INTERRUPT_CODE = 130
+SIGTSTP_CODE = 146
+RSYNC_FILE_NOT_FOUND_CODE = 23
+# Arbitrarily chosen value. Used in SkyPilot's storage mounting scripts
+MOUNT_PATH_NON_EMPTY_CODE = 42
+# Arbitrarily chosen value. Used to provision Kubernetes instance in Skypilot
+INSUFFICIENT_PRIVILEGES_CODE = 52
+# Return code when git command is ran in a dir that is not git repo
+GIT_FATAL_EXIT_CODE = 128
+
+
+def is_safe_exception(exc: Exception) -> bool:
+    """Returns True if the exception is safe to send to clients.
+
+    Safe exceptions are:
+    1. Built-in exceptions
+    2. Konduktor's own exceptions
+    """
+    module = type(exc).__module__
+
+    # Builtin exceptions (e.g., ValueError, RuntimeError)
+    if module == 'builtins':
+        return True
+
+    # Konduktor's own exceptions
+    if module.startswith('sky.'):
+        return True
+
+    return False
+
+
+def wrap_exception(exc: Exception) -> Exception:
+    """Wraps non-safe exceptions into Konduktor exceptions
+
+    This is used to wrap exceptions that are not safe to deserialize at clients.
+
+    Examples include exceptions from cloud providers whose packages are not
+    available at clients.
+    """
+    if is_safe_exception(exc):
+        return exc
+
+    return CloudError(
+        message=str(exc),
+        cloud_provider=type(exc).__module__.split('.')[0],
+        error_type=type(exc).__name__,
+    )
+
+
+def serialize_exception(e: Exception) -> Dict[str, Any]:
+    """Serialize the exception.
+
+    This function also wraps any unsafe exceptions (e.g., cloud exceptions)
+    into Konduktor's CloudError before serialization to ensure clients can
+    deserialize them without needing cloud provider packages installed.
+    """
+    # Wrap unsafe exceptions before serialization
+    e = wrap_exception(e)
+
+    stacktrace = getattr(e, 'stacktrace', None)
+    attributes = e.__dict__.copy()
+    if 'stacktrace' in attributes:
+        del attributes['stacktrace']
+    for attr_k in list(attributes.keys()):
+        attr_v = attributes[attr_k]
+        if isinstance(attr_v, types.TracebackType):
+            attributes[attr_k] = traceback.format_tb(attr_v)
+
+    data = {
+        'type': e.__class__.__name__,
+        'message': str(e),
+        'args': e.args,
+        'attributes': attributes,
+        'stacktrace': stacktrace,
+    }
+    return data
+
+
+def deserialize_exception(serialized: Dict[str, Any]) -> Exception:
+    """Deserialize the exception."""
+    exception_type = serialized['type']
+    if hasattr(builtins, exception_type):
+        exception_class = getattr(builtins, exception_type)
+    else:
+        exception_class = globals().get(exception_type, None)
+    if exception_class is None:
+        # Unknown exception type.
+        return Exception(f'{exception_type}: {serialized["message"]}')
+    e = exception_class(*serialized['args'], **serialized['attributes'])
+    if serialized['stacktrace'] is not None:
+        setattr(e, 'stacktrace', serialized['stacktrace'])
+    return e
+
+
+class CloudError(Exception):
+    """Wraps cloud-specific errors into a SkyPilot exception."""
+
+    def __init__(self, message: str, cloud_provider: str, error_type: str):
+        super().__init__(message)
+        self.cloud_provider = cloud_provider
+        self.error_type = error_type
+
+    def __str__(self):
+        return (
+            f'{self.cloud_provider} error ({self.error_type}): ' f'{super().__str__()}'
+        )
+
+
+class CommandError(Exception):
+    pass
+
+
+class NotSupportedError(Exception):
+    """Raised when a feature is not supported."""
+
+    pass
+
+
+class StorageError(Exception):
+    pass
+
+
+class StorageSpecError(ValueError):
+    # Errors raised due to invalid specification of the Storage object
+    pass
+
+
+class StorageInitError(StorageError):
+    # Error raised when Initialization fails - either due to permissions,
+    # unavailable name, or other reasons.
+    pass
+
+
+class StorageBucketCreateError(StorageInitError):
+    # Error raised when bucket creation fails.
+    pass
+
+
+class StorageBucketGetError(StorageInitError):
+    # Error raised if attempt to fetch an existing bucket fails.
+    pass
+
+
+class StorageBucketDeleteError(StorageError):
+    # Error raised if attempt to delete an existing bucket fails.
+    pass
+
+
+class StorageUploadError(StorageError):
+    # Error raised when bucket is successfully initialized, but upload fails,
+    # either due to permissions, ctrl-c, or other reasons.
+    pass
+
+
+class StorageSourceError(StorageSpecError):
+    # Error raised when the source of the storage is invalid. E.g., does not
+    # exist, malformed path, or other reasons.
+    pass
+
+
+class StorageNameError(StorageSpecError):
+    # Error raised when the source of the storage is invalid. E.g., does not
+    # exist, malformed path, or other reasons.
+    pass
+
+
+class StorageModeError(StorageSpecError):
+    # Error raised when the storage mode is invalid or does not support the
+    # requested operation (e.g., passing a file as source to MOUNT mode)
+    pass
+
+
+class StorageExternalDeletionError(StorageBucketGetError):
+    # Error raised when the bucket is attempted to be fetched while it has been
+    # deleted externally.
+    pass
+
+
+class NonExistentStorageAccountError(StorageExternalDeletionError):
+    # Error raise when storage account provided through config.yaml or read
+    # from store handle(local db) does not exist.
+    pass
+
+
+class NetworkError(Exception):
+    """Raised when network fails."""
+
+    pass
+
+
+class CloudUserIdentityError(Exception):
+    """Raised when the cloud identity is invalid."""
+
+    pass
+
+
+class ClusterOwnerIdentityMismatchError(Exception):
+    """The cluster's owner identity does not match the current user identity."""
+
+    pass
+
+
+class NoCloudAccessError(Exception):
+    """Raised when all clouds are disabled."""
+
+    pass

konduktor/utils/kubernetes_enums.py

@@ -0,0 +1,8 @@
+import enum
+
+
+class KubernetesAutoscalerType(enum.Enum):
+    """Enum for the different types of cluster autoscalers for Kubernetes."""
+
+    GKE = 'gke'
+    GENERIC = 'generic'