konduktor-nightly 0.1.0.dev20251128104812__py3-none-any.whl

konduktor/utils/base64_utils.py

@@ -0,0 +1,95 @@
+"""
+Utility for (un)zip and encode/decoding k8s secrets in base64
+"""
+
+import base64
+import os
+import shutil
+import tempfile
+import zipfile
+from typing import List
+
+
+def zip_base64encode(files: List[str]) -> str:
+    """Zips files and encodes them in base64.
+
+    Args:
+        files: List of file paths to zip. Can include files and directories.
+
+    Returns:
+        Base64 encoded string of the zipped files.
+    """
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Copy all files/directories to temp dir preserving structure
+        for file_path in files:
+            src_path = os.path.expanduser(file_path)
+            if not os.path.exists(src_path):
+                continue
+            dst_path = os.path.join(temp_dir, os.path.basename(file_path))
+
+            if os.path.isdir(src_path):
+                shutil.copytree(src_path, dst_path)
+            else:
+                shutil.copy2(src_path, dst_path)
+
+        # Create zip file
+        zip_path = os.path.join(temp_dir, 'archive.zip')
+        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            for item in os.listdir(temp_dir):
+                if item == 'archive.zip':
+                    continue
+                item_path = os.path.join(temp_dir, item)
+                if os.path.isfile(item_path):
+                    zipf.write(item_path, item)
+                else:
+                    for root, _, files in os.walk(item_path):
+                        for file in files:
+                            if file == '.DS_Store':
+                                continue
+                            file_path = os.path.join(root, file)
+                            arcname = os.path.relpath(file_path, temp_dir)
+                            zipf.write(file_path, arcname)
+
+        # Read and encode zip file
+        with open(zip_path, 'rb') as f:
+            zip_str = f.read()
+            secret_value = base64.b64encode(zip_str).decode('utf-8')
+            # print("encoding")
+            # print(type(secret_value))
+            # print(len(secret_value))
+            # print(secret_value[-20:])
+            return secret_value
+
+
+def base64decode_unzip(secret_value: str, output_path: str) -> str:
+    """Decodes a base64 encoded string and unzips the files.
+
+    Args:
+        secret_value: Base64 encoded string of the zipped files.
+        output_path: Path where to extract the unzipped files.
+
+    Returns:
+        Path to the unzipped files.
+    """
+    # TODO(asaiacai): this is messy I know...
+    # Decode base64 string
+    # print("decoding")
+    # print(type(secret_value))
+    # print(len(secret_value))
+    # print(secret_value[-20:])
+    decoded_data = base64.b64decode(secret_value)
+
+    # Write decoded data to temporary zip file
+    with tempfile.TemporaryDirectory() as temp_dir:
+        zip_path = os.path.join(temp_dir, 'archive.zip')
+
+        with zipfile.ZipFile(zip_path, 'w') as zipf:
+            zipf.writestr('data.zip', decoded_data)
+
+        with zipfile.ZipFile(zip_path, 'r') as zipf:
+            zipf.extractall(path=output_path)
+
+        with zipfile.ZipFile(os.path.join(output_path, 'data.zip'), 'r') as zipf:
+            zipf.extractall(path=output_path)
+
+    return output_path

konduktor/utils/common_utils.py

@@ -0,0 +1,426 @@
+# 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 random
+import re
+import socket
+import sys
+import uuid
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import jinja2
+import jsonschema
+import yaml  # type: ignore
+
+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.
+    """
+
+    override = os.environ.get('KONDUKTOR_TEST_USER_HASH')
+    if override:
+        return override
+
+    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}'
+
+
+class Backoff:
+    """Exponential backoff with jittering."""
+
+    MULTIPLIER = 1.6
+    JITTER = 0.4
+
+    def __init__(self, initial_backoff: float = 5, max_backoff_factor: int = 5):
+        self._initial = True
+        self._backoff = 0.0
+        self._initial_backoff = initial_backoff
+        self._max_backoff = max_backoff_factor * self._initial_backoff
+
+    # https://github.com/grpc/grpc/blob/2d4f3c56001cd1e1f85734b2f7c5ce5f2797c38a/doc/connection-backoff.md
+    # https://github.com/grpc/grpc/blob/5fc3ff82032d0ebc4bf252a170ebe66aacf9ed9d/src/core/lib/backoff/backoff.cc
+
+    def current_backoff(self) -> float:
+        """Backs off once and returns the current backoff in seconds."""
+        if self._initial:
+            self._initial = False
+            self._backoff = min(self._initial_backoff, self._max_backoff)
+        else:
+            self._backoff = min(self._backoff * self.MULTIPLIER, self._max_backoff)
+        self._backoff += random.uniform(
+            -self.JITTER * self._backoff, self.JITTER * self._backoff
+        )
+        return self._backoff

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)}