kubetorch 0.2.5__py3-none-any.whl

kubetorch/resources/callables/utils.py

@@ -0,0 +1,203 @@
+import importlib.metadata as metadata
+import inspect
+import json
+import os
+from pathlib import Path
+from typing import Callable, Optional, Type, Union
+
+from kubetorch.logger import get_logger
+
+logger = get_logger(__name__)
+
+SHELL_COMMANDS = {"ssh", "run_bash", "rsync"}
+
+
+class NotebookError(Exception):
+    """Raised when a function defined in a notebook environment cannot be properly handled."""
+
+    pass
+
+
+def prepare_notebook_fn(fn_pointers, name):
+    """Handle a function defined in a notebook by writing it out to a dedicated .py file to be imported
+    on the cluster."""
+    module_path = Path.cwd() / (f"{name}_fn.py" if name else "sent_fn.py")
+    logger.info(
+        f"Function is defined in a notebook, writing it out to {str(module_path)} "
+        f"to make it importable. Please make sure the function does not rely on any local variables, "
+        f"including imports (which should be moved inside the function body). "
+        f"This restriction does not apply to functions defined in normal Python files."
+    )
+    try:
+        # Try to pull the frame variable for the function by name
+        user_fn_name = fn_pointers[2]
+        frame = inspect.stack()[2].frame
+        user_fn = frame.f_globals.get(user_fn_name) or frame.f_locals.get(user_fn_name)
+        source = inspect.getsource(user_fn).strip() if user_fn else None
+        if source is None:
+            raise NotebookError(
+                f"Failed to load source code for function {user_fn_name}. "
+                f"Please ensure the function is defined in the notebook and not relying on local variables."
+            )
+    except Exception as e:
+        raise NotebookError(str(e))
+
+    with module_path.open("w") as f:
+        f.write(source)
+
+    return fn_pointers[0], module_path.stem, fn_pointers[2]
+
+
+def extract_pointers(raw_cls_or_fn: Union[Type, Callable]):
+    """Get the path to the module, module name, and function name to be able to import it on the server"""
+    if not (isinstance(raw_cls_or_fn, Type) or isinstance(raw_cls_or_fn, Callable)):
+        raise TypeError(f"Expected Type or Callable but received {type(raw_cls_or_fn)}")
+
+    # (root_path, module_name, cls_or_fn_name)
+    return _get_module_import_info(raw_cls_or_fn)
+
+
+def _get_module_import_info(raw_cls_or_fn: Union[Type, Callable]):
+    """
+    Given a class or function in Python, get all the information needed to import it in another Python process.
+    """
+
+    # Background on all these dunders: https://docs.python.org/3/reference/import.html
+    py_module = inspect.getmodule(raw_cls_or_fn)
+
+    # Need to resolve in case just filename is given
+    module_path = _extract_module_path(raw_cls_or_fn)
+
+    # TODO better way of detecting if in a notebook or interactive Python env
+    if not module_path or module_path.endswith("ipynb"):
+        # The only time __file__ wouldn't be present is if the function is defined in an interactive
+        # interpreter or a notebook. We can't import on the server in that case, so we need to cloudpickle
+        # the fn to send it over. The __call__ function will serialize the function if we return it this way.
+        # This is a short-term hack.
+        # return None, "notebook", raw_fn.__name__
+        root_path = os.getcwd()
+        module_name = "notebook"
+        cls_or_fn_name = raw_cls_or_fn.__name__
+    else:
+        root_path = os.path.dirname(module_path)
+        module_name = inspect.getmodulename(module_path)
+        # TODO __qualname__ doesn't work when fn is aliased funnily, like torch.sum
+        cls_or_fn_name = getattr(raw_cls_or_fn, "__qualname__", raw_cls_or_fn.__name__)
+
+        # Adapted from https://github.com/modal-labs/modal-client/blob/main/modal/_function_utils.py#L94
+        if getattr(py_module, "__package__", None):
+            module_path = os.path.abspath(py_module.__file__)
+            package_paths = [os.path.abspath(p) for p in __import__(py_module.__package__).__path__]
+            base_dirs = [
+                base_dir for base_dir in package_paths if os.path.commonpath((base_dir, module_path)) == base_dir
+            ]
+
+            if len(base_dirs) != 1:
+                raise Exception("Wasn't able to find the package directory!")
+            root_path = os.path.dirname(base_dirs[0])
+            module_name = py_module.__spec__.name
+
+    return root_path, module_name, cls_or_fn_name
+
+
+def _extract_module_path(raw_cls_or_fn: Union[Type, Callable]):
+    py_module = inspect.getmodule(raw_cls_or_fn)
+
+    # Need to resolve in case just filename is given
+    module_path = str(Path(inspect.getfile(py_module)).resolve()) if hasattr(py_module, "__file__") else None
+
+    return module_path
+
+
+def locate_working_dir(start_dir=None):
+    """
+    Locate the working directory of the project.
+
+    Args:
+        start_dir (str, optional): The directory to start searching from. Defaults to the current working directory.
+
+    Returns:
+        tuple: A tuple containing the working directory and a boolean indicating if a project directory was found.
+    """
+    if start_dir is None:
+        start_dir = os.getcwd()
+
+    # Search first for anything that represents a Python package
+    target_files = [
+        ".git",
+        "setup.py",
+        "setup.cfg",
+        "pyproject.toml",
+        "requirements.txt",
+    ]
+
+    dir_with_target = _find_directory_containing_any_file(start_dir, target_files, searched_dirs=set())
+    found_project_dir = dir_with_target is not None
+    return (dir_with_target if found_project_dir else start_dir), found_project_dir
+
+
+def _find_directory_containing_any_file(dir_path, files, searched_dirs=None):
+    if Path(dir_path) == Path.home() or dir_path == Path("/"):
+        return None
+
+    if any(Path(dir_path, file).exists() for file in files):
+        return str(dir_path)
+
+    searched_dirs.add(dir_path)
+    parent_path = Path(dir_path).parent
+    if parent_path in searched_dirs:
+        return None
+    return _find_directory_containing_any_file(parent_path, files, searched_dirs=searched_dirs)
+
+
+def get_local_install_path(package_name: str) -> Optional[str]:
+    from importlib.metadata import distributions
+
+    for dist in distributions():
+        direct_url_json = dist.read_text("direct_url.json")
+        if direct_url_json and dist.metadata["Name"].lower() == package_name.lower():
+            try:
+                url = json.loads(direct_url_json).get("url", None)
+                if url:
+                    if url.startswith("file://"):
+                        return url[len("file://") :]
+            except json.JSONDecodeError:
+                pass
+    return None
+
+
+def find_locally_installed_version(package_name: str) -> Optional[str]:
+    try:
+        return metadata.version(package_name)
+    except metadata.PackageNotFoundError:
+        return None
+
+
+def get_names_for_reload_fallbacks(name: str, prefixes: list[str] = []):
+    from kubetorch.globals import config
+    from kubetorch.servers.http.utils import clean_and_validate_k8s_name
+    from kubetorch.utils import current_git_branch, validate_username
+
+    current_prefix = config.username
+
+    if prefixes:
+        fallback_prefixes = prefixes
+    else:
+        # try reloading based on current username or current git branch (in that order)
+        branch = current_git_branch()
+        if branch:
+            # Ensure that we use the truncated branch name that was used to create the service initially
+            valid_branch = validate_username(branch)
+            # Note: username/prefix takes precedence over branch (in the event they differ)
+            fallback_prefixes = [v for v in (current_prefix, valid_branch) if v is not None]
+        else:
+            fallback_prefixes = [current_prefix] if current_prefix else []
+
+    potential_names = [
+        clean_and_validate_k8s_name(f"{prefix}-{name}", allow_full_length=True) for prefix in fallback_prefixes
+    ]
+    if not prefixes and name not in potential_names:
+        # try loading the bare name (i.e. prod mode) last, but only if we're not looking for specific prefixes
+        potential_names.append(name)
+
+    return potential_names

kubetorch/resources/compute/app.py

@@ -0,0 +1,253 @@
+import os
+import re
+import signal
+import sys
+import threading
+import time
+from datetime import datetime, timezone
+from typing import Dict
+
+from kubetorch.globals import config
+from kubetorch.logger import get_logger
+
+from kubetorch.resources.callables.module import Module
+from kubetorch.resources.compute.compute import Compute
+from kubetorch.resources.compute.utils import ServiceTimeoutError
+from kubetorch.servers.http.utils import is_running_in_kubernetes
+from kubetorch.utils import get_kt_install_url
+
+logger = get_logger(__name__)
+
+
+class App(Module):
+    MODULE_TYPE = "app"
+
+    def __init__(
+        self,
+        compute: Compute,
+        cli_command: str,
+        pointers: tuple,
+        name: str = None,
+        run_async: bool = False,
+    ):
+        """
+        Initialize an App object for remote execution.
+
+        .. note::
+
+            To create an App, please use the factory method :func:`app` in conjunction with the `kt run` CLI command.
+
+        Args:
+            compute (Compute): Compute
+            cli_command (str): CLI command to run on the compute.
+            pointers (tuple): A tuple containing references needed to locate the app file, of the format
+                (current working directory, path of file relative to cwd, None)
+            name (str, optional): Name to assign the app. If not provided, will be based on the name of the file in
+                which the app was defined.
+            run_async (bool, optional): Whether to run the app async. (Default: ``False``)
+        """
+        super().__init__(name=name, pointers=pointers)
+        self.cli_command = cli_command
+        self.pointers = pointers
+        self.name = name or self.module_name
+        self._compute = compute
+        self._run_async = run_async
+        self._remote_pointers = None
+
+        self._http_client = None
+
+    @property
+    def module_name(self):
+        return os.path.splitext(self.pointers[1])[0]
+
+    def from_name(self):
+        raise ValueError("Reloading app is not supported.")
+
+    def setup_signal_handlers(self):
+        signal.signal(signal.SIGINT, self.handle_termination_signal)
+        signal.signal(signal.SIGTERM, self.handle_termination_signal)
+
+    def handle_termination_signal(self, signum, frame):
+        red = "\u001b[31m"
+        reset = "\u001b[0m"
+
+        logger.info(f"{red}Received {signal.Signals(signum).name}. Exiting parent process.{reset}")
+        self._print_kt_cmds()
+        sys.exit(0)
+
+    def deploy(self):
+        """
+        Deploy the app to the compute specified by the app arguments.
+        """
+        self.compute.service_name = self.service_name
+
+        install_url, use_editable = get_kt_install_url(self.compute.freeze)
+        if not self.compute.freeze:
+            deployment_timestamp = datetime.now(timezone.utc).isoformat()
+            self._rsync_repo_and_image_patches(install_url, use_editable, init_args={})
+        else:
+            deployment_timestamp = None
+
+        self.setup_signal_handlers()
+
+        stream_logs = not self._run_async
+        self._launch_service(install_url, use_editable, deployment_timestamp, stream_logs)
+
+    def _get_service_dockerfile(self, metadata_env_vars):
+        image_instructions = super()._get_service_dockerfile(metadata_env_vars)
+
+        remote_script = os.path.join(self.remote_pointers[0], self.remote_pointers[1])
+        local_script = r"\b" + re.escape(self.remote_pointers[1]) + r"\b"
+        remote_cmd = re.sub(local_script, remote_script, self.cli_command)
+
+        image_instructions += f"CMD {remote_cmd}\n"
+        return image_instructions
+
+    def _launch_service(
+        self,
+        install_url,
+        use_editable,
+        deployment_timestamp,
+        stream_logs,
+    ):
+        trigger_reload = self.compute.is_up()
+        if self._run_async:
+            thread = threading.Thread(
+                target=super()._launch_service,
+                args=(
+                    install_url,
+                    use_editable,
+                    {},
+                    deployment_timestamp,
+                    stream_logs,
+                    config.log_verbosity,
+                    False,
+                ),
+            )
+            thread.start()
+
+            if trigger_reload:
+                self._update_service(stream_logs, deployment_timestamp)
+                time.sleep(1)
+            else:
+                # wait for pods to be ready before exiting out
+                start_time = time.time()
+                while not self.compute.is_up() and time.time() - start_time < 60:
+                    time.sleep(5)
+
+                if not self.compute.is_up():
+                    raise ServiceTimeoutError(f"Service {self.service_name} is not up after 60 seconds.")
+        else:
+            super()._launch_service(
+                install_url,
+                use_editable,
+                init_args={},
+                deployment_timestamp=deployment_timestamp,
+                stream_logs=stream_logs,
+                verbosity=config.log_verbosity,
+                dryrun=False,
+            )
+
+            if trigger_reload:
+                self._update_service(stream_logs, deployment_timestamp)
+
+    def _update_service(self, stream_logs, deployment_timestamp):
+        client = self._client()
+
+        if self._run_async:
+            thread = threading.Thread(
+                target=client.call_method,
+                args=(
+                    self.endpoint(),
+                    stream_logs,
+                ),
+                kwargs={"headers": {"X-Deployed-As-Of": deployment_timestamp}},
+            )
+            thread.start()
+            time.sleep(1)
+            sys.exit()
+        else:
+            client.call_method(
+                self.endpoint(),
+                stream_logs=stream_logs,
+                headers={"X-Deployed-As-Of": deployment_timestamp},
+            )
+
+    def _print_kt_cmds(self):
+        logger.info(f"To see logs, run: kt logs {self.service_name}.")
+        logger.info(f"To teardown service, run: kt teardown {self.service_name}")
+
+    def endpoint(self):
+        return f"{self.base_endpoint}/_reload_image"
+
+
+def app(
+    name: str = None,
+    port: int = None,
+    health_check: str = None,
+    **kwargs: Dict,
+):
+    """
+    Builds and deploys an instance of :class:`App`.
+
+    Args:
+        name (str, optional): Name to give the remote app. If not provided, will be based off the name of the file in
+            which the app was defined.
+        port (int, optional): Server port to expose, if the app starts an HTTP server.
+        health_check (str, optional): Health check endpoint, if running a server, to check when server is up and ready.
+        **kwargs: Compute kwargs, to define the compute on which to run the app on.
+
+    Examples:
+
+    Define the ``kt.app`` object and compute in your Python file:
+
+    .. code-block:: python
+
+        import kubetorch as kt
+
+        # Define the app at the top of the Python file to deploy
+        # train.py
+        kt.app(name="my-app", image=kt.Image("docker-latest"), cpus="0.01")
+
+        if __name__ == "__main__":
+            ...
+
+    Deploy and run the app remotely using the ``kt run`` CLI command:
+
+    .. code-block:: bash
+
+        kt run python train.py --epochs 5
+        kt run fastapi run my_app.py --name fastapi-app
+    """
+    if not os.getenv("KT_RUN") == "1" or is_running_in_kubernetes():
+        return None
+
+    if name and os.getenv("KT_RUN_NAME") and not (name == os.getenv("KT_RUN_NAME")):
+        raise ValueError(
+            f"Name mismatch between kt.App definition ({name}) and kt run command ({os.getenv('KT_RUN_NAME')})."
+        )
+    name = name or os.getenv("KT_RUN_NAME")
+    cli_command = os.getenv("KT_RUN_CMD")  # set in kt run
+    run_async = os.getenv("KT_RUN_ASYNC") == 1
+
+    env_vars = kwargs.get("env_vars", {})
+    if port:
+        env_vars["KT_APP_PORT"] = port
+    if health_check:
+        env_vars["KT_APP_HEALTHCHECK"] = health_check
+    kwargs["env_vars"] = env_vars
+    compute = Compute(**kwargs)
+
+    main_file = os.getenv("KT_RUN_FILE") or os.path.abspath(sys.modules["__main__"].__file__)
+    relative_path = os.path.relpath(main_file, os.getcwd())
+    pointers = [os.getcwd(), relative_path, None]
+    relative_cli_command = re.sub(main_file, relative_path, cli_command)
+
+    kt_app = App(
+        compute=compute,
+        cli_command=relative_cli_command,
+        pointers=pointers,
+        name=name,
+        run_async=run_async,
+    )
+    return kt_app