kubetorch 0.2.0__py3-none-any.whl

kubetorch/resources/compute/websocket.py

@@ -0,0 +1,137 @@
+import random
+import socket
+import threading
+import time
+
+import websocket
+
+
+class WebSocketRsyncTunnel:
+    def __init__(self, local_port: int, ws_url: str):
+        self.requested_port = local_port
+        self.local_port = None  # Will be set in __enter__
+        self.ws_url = ws_url
+        self.running = False
+        self.server_socket = None
+
+    def __enter__(self):
+        self.running = True
+
+        # Add randomization to reduce concurrent collision probability
+        # Try multiple times with different socket instances
+        max_attempts = 20
+        port_range = 100  # Much wider range to avoid collisions
+
+        for attempt in range(max_attempts):
+            # Add random offset to spread out concurrent requests
+            random_offset = random.randint(0, 50)
+            start_port = self.requested_port + random_offset
+
+            # Create a new socket for each attempt
+            server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+
+            # Try ports in this range
+            for i in range(port_range):
+                port = start_port + i
+                try:
+                    server_socket.bind(("127.0.0.1", port))
+                    server_socket.listen(5)
+                    # Success! Save the socket and port
+                    self.server_socket = server_socket
+                    self.local_port = port
+                    break
+                except OSError as e:
+                    if e.errno == 98:  # Address already in use
+                        continue
+                    else:
+                        server_socket.close()
+                        raise
+            else:
+                # No ports available in this range, close socket and try again
+                server_socket.close()
+                # Add small random delay to reduce simultaneous collision probability
+                time.sleep(random.uniform(0.01, 0.05))
+                continue
+
+            # Websocket server successfully bound and is listening
+            break
+        else:
+            raise RuntimeError(
+                f"Could not find available port after {max_attempts} attempts starting from {self.requested_port}"
+            )
+
+        threading.Thread(target=self._accept_loop, daemon=True).start()
+
+        # Wait for ready
+        for _ in range(50):
+            try:
+                with socket.socket() as s:
+                    s.settimeout(0.1)
+                    if s.connect_ex(("127.0.0.1", self.local_port)) == 0:
+                        return self
+            except:
+                pass
+            time.sleep(0.1)
+        raise RuntimeError("Tunnel failed to start")
+
+    def __exit__(self, *args):
+        self.running = False
+        if self.server_socket:
+            try:
+                self.server_socket.close()
+            except:
+                pass  # Already closed
+
+    def _accept_loop(self):
+        while self.running:
+            try:
+                client_sock, _ = self.server_socket.accept()
+                threading.Thread(
+                    target=self._handle_client, args=(client_sock,), daemon=True
+                ).start()
+            except:
+                break
+
+    def _handle_client(self, client_sock):
+        ws = None
+        try:
+            ws = websocket.create_connection(self.ws_url)
+
+            def tcp_to_ws():
+                while self.running:
+                    try:
+                        data = client_sock.recv(65536)
+                        if not data:
+                            break
+                        ws.send_binary(data)
+                    except:
+                        break
+
+            def ws_to_tcp():
+                while self.running:
+                    try:
+                        data = ws.recv()
+                        if isinstance(data, bytes):
+                            client_sock.send(data)
+                    except:
+                        break
+
+            t1 = threading.Thread(target=tcp_to_ws, daemon=True)
+            t2 = threading.Thread(target=ws_to_tcp, daemon=True)
+            t1.start()
+            t2.start()
+            t1.join()
+            t2.join()
+
+        except Exception as e:
+            print(f"WebSocket connection error: {e}")
+
+        finally:
+            # close both connections
+            for conn in [client_sock, ws]:
+                if conn:
+                    try:
+                        conn.close()
+                    except:
+                        pass

kubetorch/resources/images/__init__.py

@@ -0,0 +1 @@
+from .images import *  # noqa: F403

kubetorch/resources/images/image.py

@@ -0,0 +1,414 @@
+from enum import Enum
+from typing import Any, Dict, List, Union
+
+# Internal class to represent the image construction process
+class ImageSetupStepType(Enum):
+    """Enum for valid Image setup step types"""
+
+    CMD_RUN = "cmd_run"
+    RSYNC = "rsync"
+    PIP_INSTALL = "pip_install"
+    SYNC_PACKAGE = "sync_package"
+    SET_ENV_VARS = "set_env_vars"
+
+
+class ImageSetupStep:
+    def __init__(
+        self,
+        step_type: ImageSetupStepType,
+        **kwargs: Dict[str, Any],
+    ):
+        """
+        A component of the Kubetorch Image, consisting of the step type (e.g. packages, set_env_vars),
+        along with arguments to provide to the function corresponding to the step type.
+
+        Args:
+            step_type (ImageSetupStepType): Type of setup step used to provide the Image.
+            kwargs (Dict[str, Any]): Please refer to the corresponding functions in ``Image`` to determine
+                the correct keyword arguments to provide.
+        """
+        self.step_type = step_type
+        self.kwargs = kwargs
+
+
+class Image:
+    def __init__(
+        self,
+        name: str = None,
+        image_id: str = None,
+        python_path: str = None,
+        install_cmd: str = None,
+    ):
+        """
+        Kubetorch Image object, specifying cluster setup properties and steps.
+
+        Args:
+            name (str, optional): Name to assign the Kubetorch image.
+            image_id (str, optional): Machine image to use, if any. (Default: ``None``)
+            python_path (str, optional): Absolute path to the Python executable to use for remote server and installs.
+                (Default: ``None``)
+            install_cmd (str, optional): Custom pip/uv install command to use for package installations.
+                If not provided, will be inferred based on python_path and available tools (preferring to use uv).
+                Examples: "uv pip install", "python -m pip install", "/path/to/.venv/bin/python -m uv pip install"
+                (Default: ``None``)
+
+        Note:
+            For convenience, Kubetorch provides ready-to-use base images under ``kt.images``.
+            These cover common environments like Python, CUDA, and Ray:
+
+              * ``kt.images.Python310()``, ``kt.images.Python311()``, ``kt.images.Python312()``
+              * ``kt.images.Debian()``
+              * ``kt.images.Ubuntu()``
+              * ``kt.images.Ray()`` (defaults to the latest Ray release)
+
+            You can also use flexible factories:
+
+              * ``kt.images.python("3.12")``
+              * ``kt.images.pytorch()`` (defaults to "nvcr.io/nvidia/pytorch:23.12-py3")
+              * ``kt.images.ray("2.32.0-py311")``
+
+            These base images can be further customized with methods like
+            ``.pip_install()``, ``.set_env_vars()``, ``.sync_package()``, etc.
+
+        Example:
+
+            .. code-block:: python
+
+                import kubetorch as kt
+
+                custom_image = (
+                    kt.Image(name="base_image")
+                    .pip_install(["numpy", "pandas"])
+                    .set_env_vars({"OMP_NUM_THREADS": 1})
+                )
+                debian_image = (
+                    kt.images.Debian()
+                    .pip_install(["numpy", "pandas"])
+                    .set_env_vars({"OMP_NUM_THREADS": 1})
+                )
+        """
+
+        self.name = name
+        self.image_id = image_id
+        self.python_path = python_path
+        self.install_cmd = install_cmd
+
+        self.setup_steps = []
+        self.docker_secret = None
+
+    @staticmethod
+    def _setup_step_config(step: ImageSetupStep):
+        """Get ImageSetupStep config"""
+        config = {
+            "step_type": step.step_type.value,
+            "kwargs": step.kwargs,
+        }
+        return config
+
+    @staticmethod
+    def _setup_step_from_config(step: Dict):
+        """Convert setup step config (dict) to ImageSetupStep object"""
+        step_type = step["step_type"]
+        kwargs = step["kwargs"]
+        return ImageSetupStep(
+            step_type=ImageSetupStepType(step_type),
+            **kwargs,
+        )
+
+    def from_docker(self, image_id: str):
+        """Set up and use an existing Docker image.
+
+        Args:
+            image_id (str): Docker image in the following format ``"<registry>/<image>:<tag>"``
+        """
+        if self.image_id:
+            raise ValueError(
+                "Setting both a machine image and docker image is not yet supported."
+            )
+        self.image_id = image_id
+        return self
+
+    ########################################################
+    # Steps to build the image
+    ########################################################
+
+    def pip_install(
+        self,
+        reqs: List[Union["Package", str]],
+        force: bool = False,
+    ):
+        """Pip install the given packages.
+
+        Args:
+            reqs (List[Package or str]): List of packages to pip install on cluster and env.
+                Each string is passed directly to the pip/uv command, allowing full control
+                over pip arguments. Examples:
+                - Simple package: ``"numpy"``
+                - Version constraint: ``"pandas>=1.2.0"``
+                - With pip flags: ``"--pre torch==2.0.0rc1"``
+                - Multiple flags: ``"--index-url https://pypi.org/simple torch"``
+            force (bool, optional): Whether to force re-install a package, if it already exists on the compute. (Default: ``False``)
+
+        Example:
+            .. code-block:: python
+
+                import kubetorch as kt
+
+                image = (
+                    kt.images.Debian()
+                    .pip_install([
+                        "numpy>=1.20",
+                        "pandas",
+                        "--pre torchmonarch==0.1.0rc7",  # Install pre-release
+                        "--index-url https://test.pypi.org/simple/ mypackage"
+                    ])
+                )
+        """
+
+        self.setup_steps.append(
+            ImageSetupStep(
+                step_type=ImageSetupStepType.PIP_INSTALL,
+                reqs=reqs,
+                force=force,
+            )
+        )
+        return self
+
+    def set_env_vars(self, env_vars: Dict):
+        """Set environment variables with support for variable expansion.
+
+        Environment variables can reference other variables using shell-style expansion:
+        - ``$VAR`` or ``${VAR}`` syntax to reference existing variables
+        - Variables are expanded when the container starts
+        - Variables are expanded in the order they are defined
+
+        Args:
+            env_vars (Dict): Dict of environment variables and values to set.
+                Values can include references to other environment variables.
+
+        Example:
+            .. code-block:: python
+
+                import kubetorch as kt
+
+                image = (
+                    kt.images.Debian()
+                    .set_env_vars({
+                        "BASE_PATH": "/usr/local",
+                        "BIN_PATH": "$BASE_PATH/bin",  # Expands to /usr/local/bin
+                        "PATH": "$BIN_PATH:$PATH",      # Prepends to existing PATH
+                        "LD_LIBRARY_PATH": "/opt/lib:${LD_LIBRARY_PATH}",  # Appends to existing
+                        "CUSTOM": "${HOME}/data",       # Uses HOME from container
+                    })
+                )
+
+        Note:
+            - Variables are expanded using Python's ``os.path.expandvars()``
+            - Undefined variables remain as literal strings (e.g., ``$UNDEFINED`` stays as ``$UNDEFINED``)
+            - To include a literal ``$``, escape it with backslash: ``\\$``
+        """
+        # TODO - support .env files
+        self.setup_steps.append(
+            ImageSetupStep(
+                step_type=ImageSetupStepType.SET_ENV_VARS,
+                env_vars=env_vars,
+            )
+        )
+        return self
+
+    def sync_package(
+        self,
+        package: str,
+        force: bool = False,
+    ):
+        """Sync local package over and add to path.
+
+        Args:
+            package (Package or str): Package to sync. Either the name of a local editably installed package, or
+                the path to the folder to sync over.
+            force (bool, optional): Whether to re-sync the package over, if already previously synced over. (Default: ``False``)
+        """
+        self.setup_steps.append(
+            ImageSetupStep(
+                step_type=ImageSetupStepType.SYNC_PACKAGE,
+                package=package,
+                force=force,
+            )
+        )
+        return self
+
+    def run_bash(
+        self,
+        command: str,
+        force: bool = False,
+    ):
+        """Run bash commands during image setup.
+
+        Executes shell commands during container initialization. Commands run in the
+        order they are defined and can be used to install software, configure the
+        environment, or start background services.
+
+        Args:
+            command (str): Shell command(s) to run on the cluster. Supports:
+                - Single commands: ``"apt-get update"``
+                - Chained commands: ``"apt-get update && apt-get install -y curl"``
+                - Background processes: ``"jupyter notebook --no-browser &"``
+            force (bool): Whether to rerun the command on the cluster, if previously run in image setup already. (Default: ``False``)
+
+        Example:
+            .. code-block:: python
+
+                import kubetorch as kt
+
+                image = (
+                    kt.images.Debian()
+                    .run_bash("apt-get update && apt-get install -y vim")
+                    .run_bash("pip install jupyter")
+                    .run_bash("jupyter notebook --no-browser --port=8888 &")  # Runs in background
+                )
+
+        Note:
+            - Commands ending with ``&`` run in the background and won't block image setup
+            - Background processes continue running after setup completes
+            - The setup waits 0.5s to catch immediate failures in background processes
+            - Use ``&&`` to chain commands that depend on each other
+            - Use ``;`` to run commands sequentially regardless of success/failure
+        """
+        self.setup_steps.append(
+            ImageSetupStep(
+                step_type=ImageSetupStepType.CMD_RUN,
+                command=command,
+                force=force,
+            )
+        )
+        return self
+
+    def rsync(
+        self,
+        source: str,
+        dest: str = None,
+        contents: bool = False,
+        filter_options: str = None,
+        force: bool = False,
+    ):
+        """Sync files or directories from local machine to the remote container.
+
+        This method efficiently transfers files to remote containers using rsync,
+        which only copies changed files and supports compression. Files are first
+        uploaded to a jump pod and then distributed to worker pods on startup.
+
+        Args:
+            source (str): Path to the local file or directory to sync. Supports:
+                - Absolute paths: ``/path/to/file``
+                - Relative paths: ``./data/file.txt``
+                - Home directory paths: ``~/documents/data.csv``
+            dest (str, optional): Target path on the remote container. Supports:
+                - Absolute paths: ``/data/config.yaml`` (places file at exact location)
+                - Relative paths: ``configs/settings.json`` (relative to working directory)
+                - Tilde paths: ``~/results/output.txt`` (relative to working directory, ~ is stripped)
+                - None: Uses the basename of source in working directory
+            contents (bool, optional): For directories only - whether to copy the contents
+                or the directory itself.
+                If ``True`` the contents of the source directory are copied to the destination,
+                and the source directory itself is not created at the destination.
+                If ``False`` the source directory along with its contents are copied to the
+                destination, creating an additional directory layer at the destination.
+                (Default: ``False``)
+            filter_options (str, optional): Additional rsync filter options. These are added
+                to (not replacing) the default filters. By default, rsync excludes:
+
+                - Files from ``.gitignore`` (if present)
+                - Files from ``.ktignore`` (if present)
+                - Common Python artifacts: ``*.pyc``, ``__pycache__``
+                - Virtual environments: ``.venv``
+                - Git metadata: ``.git``
+
+                Your filter_options are appended after these defaults. Examples:
+
+                - Exclude more patterns: ``"--exclude='*.log' --exclude='temp/'"``
+                - Include specific files: ``"--include='important.log' --exclude='*.log'"``
+                - Override all defaults: Set ``KT_RSYNC_FILTERS`` environment variable
+
+                (Default: ``None``)
+            force (bool, optional): When ``True``, forces rsync to transfer all files
+                regardless of modification times by using ``--ignore-times`` flag. This ensures
+                all files are copied even if timestamps suggest they haven't changed.
+                Useful when timestamp-based change detection is unreliable.
+                Note: Files are always synced when deploying with ``.to()``, this flag just
+                affects how rsync determines which files need updating.
+                (Default: ``False``)
+
+        Returns:
+            Image: Returns self for method chaining.
+
+        Examples:
+            .. code-block:: python
+
+                import kubetorch as kt
+
+                # Basic file sync
+                image = (
+                    kt.images.Debian()
+                    .rsync("./config.yaml", "app/config.yaml")
+                )
+
+                # Sync to absolute path
+                image = (
+                    kt.images.Python312()
+                    .rsync("./model_weights.pth", "/models/weights.pth")
+                )
+
+                # No destination specified - uses basename
+                image = (
+                    kt.images.Ubuntu()
+                    .rsync("/local/data/dataset.csv")  # Goes to ./dataset.csv
+                )
+
+                # Directory sync - copy directory itself
+                image = (
+                    kt.images.Debian()
+                    .rsync("./src", "app")  # Creates app/src/
+                )
+
+                # Directory sync - copy contents only
+                image = (
+                    kt.images.Debian()
+                    .rsync("./src", "app", contents=True)  # Contents go directly into app/
+                )
+
+                # Multiple rsync operations with filtering
+                image = (
+                    kt.images.Python312()
+                    .rsync("./data", "/data", filter_options="--exclude='*.tmp'")
+                    .rsync("./configs", "~/configs")
+                    .rsync("./scripts")
+                    .pip_install(["numpy", "pandas"])
+                )
+
+                # Force re-sync for development
+                image = (
+                    kt.images.Debian()
+                    .rsync("./rapidly_changing_code", "app", force=True)
+                )
+
+        Note:
+            - Absolute destination paths (starting with ``/``) place files at exact locations
+            - Relative paths and ``~/`` paths are relative to the container's working directory
+            - The ``contents`` parameter only affects directory sources, not files
+            - Default exclusions include ``.gitignore`` patterns, ``__pycache__``, ``.venv``, and ``.git``
+            - User-provided ``filter_options`` are added to (not replacing) the default filters
+            - To completely override filters, set the ``KT_RSYNC_FILTERS`` environment variable
+            - Use ``force=True`` to bypass timestamp checks and transfer all files
+        """
+
+        self.setup_steps.append(
+            ImageSetupStep(
+                step_type=ImageSetupStepType.RSYNC,
+                source=source,
+                dest=dest,
+                contents=contents,
+                filter_options=filter_options,
+                force=force,
+            )
+        )
+        return self

kubetorch/resources/images/images.py

@@ -0,0 +1,74 @@
+from .image import Image
+
+
+def debian() -> Image:
+    """Return a Kubetorch Debian slim image.
+
+    This uses the default Kubetorch server image, which is built on top of
+    a minimal Debian base.
+    """
+    import kubetorch.serving.constants as serving_constants
+    from kubetorch.globals import config
+
+    if config.tracing_enabled:
+        image_id = serving_constants.SERVER_IMAGE_WITH_OTEL
+    else:
+        image_id = serving_constants.SERVER_IMAGE_MINIMAL
+    return Image(name="debian", image_id=image_id)
+
+
+def ubuntu() -> Image:
+    """Return a Kubetorch ubuntu image."""
+    import kubetorch.serving.constants as serving_constants
+    from kubetorch.globals import config
+
+    if config.tracing_enabled:
+        image_id = serving_constants.UBUNTU_IMAGE_WITH_OTEL
+    else:
+        image_id = serving_constants.UBUNTU_IMAGE_MINIMAL
+    return Image(name="ubuntu", image_id=image_id)
+
+
+def python(version: str) -> Image:
+    """Return a Python slim base image, e.g. ``python('3.12')``."""
+    tag = version.replace(".", "")
+    return Image(name=f"python{tag}", image_id=f"python:{version}-slim")
+
+
+def ray(version: str = "latest") -> Image:
+    """Return a Ray base image, defaults to ``ray:latest``."""
+    return Image(
+        name=f"ray{version if version != 'latest' else ''}".strip(),
+        image_id=f"rayproject/ray:{version}",
+    )
+
+
+def pytorch(version: str = "23.12-py3") -> Image:
+    """Return an NVIDIA PyTorch base image. Defaults to ``nvcr.io/nvidia/pytorch:23.12-py3``."""
+    tag = version.replace(".", "").replace("-", "")
+    return Image(name=f"pytorch{tag}", image_id=f"nvcr.io/nvidia/pytorch:{version}")
+
+
+# Predefined convenience aliases for common versions
+Python310 = lambda: python("3.10")
+Python311 = lambda: python("3.11")
+Python312 = lambda: python("3.12")
+Ray = lambda: ray("latest")
+Pytorch2312 = lambda: pytorch("23.12-py3")
+Debian = lambda: debian()
+Ubuntu = lambda: ubuntu()
+
+__all__ = [
+    "python",
+    "ray",
+    "pytorch",
+    "debian",
+    "ubuntu",
+    "Python310",
+    "Python311",
+    "Python312",
+    "Ray",
+    "Pytorch2312",
+    "Debian",
+    "Ubuntu",
+]

kubetorch/resources/secrets/__init__.py

@@ -0,0 +1,2 @@
+from .secret import Secret  # noqa: F401
+from .secret_factory import secret  # noqa: F401