klio-compute 0.5.4__py3-none-any.whl

klio_compute/__init__.py

@@ -0,0 +1,57 @@
+"""Klio Compute SDK.
+
+Everything the underlying SDK exports publicly is available here::
+
+    from klio_compute import function, Image
+"""
+
+from importlib import import_module
+
+import beta9 as _sdk
+from beta9.config import set_settings as _set_settings
+
+from ._branding import install_terminal_wording as _install_terminal_wording
+from ._connect import install_connect_timeout as _install_connect_timeout
+from ._distribution import installed_version as _installed_version
+from ._errors import install_error_reports as _install_error_reports
+from ._gpu_access import install_gpu_access_check as _install_gpu_access_check
+from ._images import Image as _Image
+from ._settings import in_platform_container as _in_platform_container
+from ._settings import klio_settings as _klio_settings
+
+# Customer code imports this package inside containers too, where the platform's runner already
+# has what it needs from the environment and the SDK is mounted with no package metadata.
+if not _in_platform_container():
+    _install_connect_timeout()
+    _install_terminal_wording()
+    _install_error_reports()
+    _install_gpu_access_check()
+    # Point the SDK at Klio Compute before any of its code reads its settings.
+    _set_settings(_klio_settings())
+
+__version__ = _installed_version()
+
+__all__ = list(_sdk.__all__)
+
+# Names that Klio replaces with its own version of the SDK's.
+_OVERRIDES = {"Image": _Image}
+
+
+def __getattr__(name):
+    if name in _OVERRIDES:
+        globals()[name] = _OVERRIDES[name]
+        return _OVERRIDES[name]
+    if name in _sdk.__all__:
+        value = getattr(_sdk, name)
+        globals()[name] = value
+        return value
+    try:
+        return import_module(f"{__name__}.{name}")
+    except ModuleNotFoundError as error:
+        if error.name != f"{__name__}.{name}":
+            raise
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__():
+    return sorted(set(globals()) | set(__all__))

klio_compute/__main__.py

@@ -0,0 +1,3 @@
+from .cli import start
+
+start()

klio_compute/_branding.py

@@ -0,0 +1,63 @@
+"""Customer-facing wording: the SDK prints a few lines that name itself.
+
+The SDK writes everything the customer reads through ``beta9.terminal``. Its functions are
+looked up on that module at call time, so wrapping them here rewrites the handful of upstream
+names (the token variable in the endpoint's curl example, a ``beta9 <command>`` hint) without
+touching any other behaviour. Paths, URLs and module names are left alone.
+"""
+
+import functools
+import re
+from typing import Any, Callable, Tuple
+
+import beta9.terminal as sdk_terminal
+
+TOKEN_ENV_VAR = "KLIO_TOKEN"
+
+_NOT_A_NAME_PART = r"[\w./-]"
+
+# Most specific first: a later rule must not see text an earlier one already replaced.
+_RULES: Tuple[Tuple["re.Pattern[str]", str], ...] = (
+    (re.compile(r"\$\{BETA9_TOKEN:\?[^}]*\}"), "${KLIO_TOKEN:?Set KLIO_TOKEN to your Klio API key}"),
+    (
+        re.compile(r"Set the token from your selected profile:"),
+        "Set your Klio API key (create one in the console under Keys):",
+    ),
+    (re.compile(r"<your-profile-token>"), "<your-klio-key>"),
+    (re.compile(r"\bBETA9_TOKEN\b"), TOKEN_ENV_VAR),
+    (re.compile(rf"(?<!{_NOT_A_NAME_PART})beta9(?!{_NOT_A_NAME_PART})"), "klio"),
+    (re.compile(rf"(?<!{_NOT_A_NAME_PART})Beta9(?!{_NOT_A_NAME_PART})"), "Klio"),
+)
+
+# Lines about the platform's own capacity: nothing a customer can act on, so they are not shown.
+_NOT_SHOWN = re.compile(r"^GPU capacity for .* is currently")
+
+# Every output function of the SDK's terminal that takes customer-visible text.
+_WRAPPED = ("print", "detail", "debug", "warn", "success", "error", "header")
+
+
+def rewrite(text: Any) -> Any:
+    """The text as Klio's customer should read it; anything that is not a string passes through."""
+    if not isinstance(text, str):
+        return text
+    for pattern, replacement in _RULES:
+        text = pattern.sub(replacement, text)
+    return text
+
+
+def _wrap(function: Callable) -> Callable:
+    @functools.wraps(function)
+    def wrapper(*args, **kwargs):
+        if args and isinstance(args[0], str) and _NOT_SHOWN.match(args[0]):
+            return None
+        return function(*(rewrite(arg) for arg in args), **kwargs)
+
+    wrapper._klio_wrapped = True  # type: ignore[attr-defined]
+    return wrapper
+
+
+def install_terminal_wording() -> None:
+    for name in _WRAPPED:
+        function = getattr(sdk_terminal, name)
+        if not getattr(function, "_klio_wrapped", False):
+            setattr(sdk_terminal, name, _wrap(function))

klio_compute/_compat_cli/__init__.py

@@ -0,0 +1,31 @@
+"""CLI verbs for Modal-style apps. Existing `klio` commands and their forms keep working as they did."""
+
+from typing import Iterator, Tuple
+
+import click
+
+from . import _app_cmd, _curl, _dashboard, _deploy, _migrate, _run, _secret_cmd, _shell, _volume_cmd
+
+
+def _command_holders(cli) -> Iterator[Tuple[click.Group, str]]:
+    for source in cli.common_group.sources:
+        for name in list(getattr(source, "commands", {})):
+            yield source, name
+
+
+def _replace(cli, name: str, wrap) -> None:
+    for group, command_name in _command_holders(cli):
+        if command_name == name:
+            group.commands[name] = wrap(group.commands[name])
+            return
+    raise RuntimeError(f"expected the SDK to have a command named {name}")
+
+
+def install(cli) -> None:
+    _replace(cli, "run", _run.wrap)
+    _replace(cli, "deploy", _deploy.wrap)
+    _replace(cli, "shell", _shell.wrap)
+    _replace(cli, "secret", _secret_cmd.extend)
+    _replace(cli, "volume", _volume_cmd.extend)
+    for command in (_app_cmd.app, _curl.curl, _dashboard.dashboard, _migrate.migrate):
+        cli.management_group.add_command(command)

klio_compute/_compat_cli/_app_cmd.py

@@ -0,0 +1,111 @@
+"""`klio app`: deployments grouped into apps. The platform's own list is one line per function."""
+
+import json
+from typing import Any, Callable, Dict, List, Optional
+
+import click
+from beta9 import terminal
+from rich.table import Column, Table, box
+
+from ..compat import _platform
+
+
+def _app_name_of(deployment: Any) -> str:
+    """Stubs from `klio deploy app.py` are named <type>/<app>.<function>. Others are their own app."""
+    last = deployment.stub_name.rsplit("/", 1)[-1]
+    if "." in last and ":" not in last:
+        return last.split(".", 1)[0]
+    return deployment.name
+
+
+def _apps() -> Dict[str, List[Any]]:
+    grouped: Dict[str, List[Any]] = {}
+    for deployment in _platform.current().deployment_rows():
+        if deployment.active:
+            grouped.setdefault(_app_name_of(deployment), []).append(deployment)
+    return grouped
+
+
+def _latest_versions(deployments: List[Any]) -> List[Any]:
+    """Earlier versions stay active while they drain. What an app is, is the newest of each."""
+    latest: Dict[str, Any] = {}
+    for deployment in deployments:
+        if deployment.name not in latest or deployment.version > latest[deployment.name].version:
+            latest[deployment.name] = deployment
+    return list(latest.values())
+
+
+def _app_or_exit(name: str) -> List[Any]:
+    deployments = _apps().get(name)
+    if not deployments:
+        raise click.ClickException(f"No app named {name} is running. See `klio app list`.")
+    return deployments
+
+
+def _run_logs_command(arguments: List[str]) -> None:
+    from beta9.cli import logs as sdk_logs
+
+    command = sdk_logs.common.commands["logs"]
+    parent = click.get_current_context()
+    delegated = command.make_context("logs", arguments, parent=parent)
+    with delegated:
+        command.invoke(delegated)
+
+
+def _web_url(deployment: Any) -> Optional[str]:
+    if "asgi" not in deployment.stub_type and "endpoint" not in deployment.stub_type:
+        return None
+    return _platform.current().web_url(deployment)
+
+
+def build_group(run_logs: Callable[[List[str]], None] = _run_logs_command) -> click.Group:
+    @click.group("app", help="List, inspect, stop and read the logs of apps.")
+    def app() -> None:
+        pass
+
+    @app.command("list", help="List running apps.")
+    @click.option("--json", "as_json", is_flag=True)
+    def list_apps(as_json: bool) -> None:
+        apps = _apps()
+        if as_json:
+            rows = [{"name": name, "app_id": items[0].app_id, "deployments": len(_latest_versions(items))} for name, items in sorted(apps.items())]
+            click.echo(json.dumps(rows))
+            return
+        table = Table(Column("App"), Column("App id"), Column("Deployments"), box=box.SIMPLE)
+        for name, items in sorted(apps.items()):
+            table.add_row(name, items[0].app_id, str(len(_latest_versions(items))))
+        terminal.print(table)
+
+    @app.command("info", help="Show an app's deployments and addresses.")
+    @click.argument("name")
+    def info(name: str) -> None:
+        for deployment in _latest_versions(_app_or_exit(name)):
+            url = _web_url(deployment)
+            click.echo(f"{deployment.name}  v{deployment.version}  {deployment.id}" + (f"  {url}" if url else ""))
+
+    @app.command("stop", help="Stop every deployment of an app.")
+    @click.argument("name")
+    @click.option("--yes", "-y", is_flag=True, help="Don't ask first.")
+    def stop(name: str, yes: bool) -> None:
+        deployments = _app_or_exit(name)
+        if not yes:
+            click.confirm(f"Stop {name} ({len(deployments)} deployments)?", abort=True)
+        platform = _platform.current()
+        for deployment in deployments:
+            platform.stop_deployment(deployment.id)
+        click.echo(f"Stopped {name}.")
+
+    @app.command(
+        "logs",
+        help="Read an app's logs. Takes the options of `klio logs`.",
+        context_settings={"ignore_unknown_options": True, "allow_extra_args": True},
+    )
+    @click.argument("name")
+    @click.argument("options", nargs=-1, type=click.UNPROCESSED)
+    def logs(name: str, options: List[str]) -> None:
+        run_logs(["--app-id", _app_or_exit(name)[0].app_id, *options])
+
+    return app
+
+
+app = build_group()

klio_compute/_compat_cli/_curl.py

@@ -0,0 +1,85 @@
+"""`klio curl`: calls a private endpoint with your key, and only ever sends it to an endpoint's own address."""
+
+import os
+import shutil
+import subprocess
+from typing import List, Optional
+from urllib.parse import urlsplit
+
+import click
+
+ENDPOINT_DOMAIN = ".run.klio.tech"
+# Only these are passed on. Anything else could send the request, and with it the key, somewhere else.
+FLAGS = set("sSfiIvgG")
+LONG_FLAGS = {
+    "--silent", "--show-error", "--fail", "--include", "--head", "--verbose", "--globoff", "--get", "--compressed",
+    "--no-buffer",
+}  # fmt: skip
+OPTIONS_WITH_VALUE = {
+    "-X", "--request", "-H", "--header", "-d", "--data", "--data-raw", "--data-binary", "--data-urlencode", "--json",
+    "-F", "--form", "-o", "--output", "-w", "--write-out", "-m", "--max-time", "--connect-timeout", "-A", "--user-agent",
+}  # fmt: skip
+# Before the user's arguments and after our checks: https only, no redirects, no proxy from the environment.
+PINNED = ["--proto", "=https", "--max-redirs", "0", "--noproxy", "*"]
+
+
+def _token() -> str:
+    from beta9.abstractions.base import get_channel
+
+    try:
+        return get_channel().config.token or ""
+    except Exception:  # no saved login: the message below says what to do
+        return ""
+
+
+def _is_endpoint_url(url: str) -> bool:
+    parts = urlsplit(url)
+    host = parts.hostname or ""
+    return parts.scheme == "https" and host.endswith(ENDPOINT_DOMAIN) and not parts.username and not parts.password
+
+
+def _refuse_option(option: str) -> None:
+    raise click.ClickException(f"klio curl doesn't take {option}, so that your key only goes to the endpoint you name.")
+
+
+def _check(arguments: List[str]) -> str:
+    """The URL these arguments call. Everything else must be an option known to stay on that URL."""
+    url: Optional[str] = None
+    remaining = iter(arguments)
+    for argument in remaining:
+        option = argument.split("=", 1)[0]
+        if option in OPTIONS_WITH_VALUE:
+            if "=" not in argument or not option.startswith("--"):
+                next(remaining, None)
+        elif option in LONG_FLAGS or (option.startswith("-") and not option.startswith("--") and set(option[1:]) <= FLAGS and len(option) > 1):
+            continue
+        elif option.startswith("-"):
+            _refuse_option(option)
+        elif url is None and _is_endpoint_url(argument):
+            url = argument
+        else:
+            raise click.ClickException(f"klio curl takes one address, and only sends your key to https://*{ENDPOINT_DOMAIN}.")
+    if url is None:
+        raise click.ClickException(f"Give the endpoint's address: klio curl https://NAME{ENDPOINT_DOMAIN}/path")
+    return url
+
+
+@click.command(
+    "curl",
+    help="Call a private endpoint with your key. Takes curl's options.",
+    context_settings={"ignore_unknown_options": True, "allow_extra_args": True, "help_option_names": []},
+)
+@click.argument("arguments", nargs=-1, type=click.UNPROCESSED)
+def curl(arguments: List[str]) -> None:
+    arguments = list(arguments)
+    _check(arguments)
+    token = os.environ.get("KLIO_TOKEN") or _token()
+    if not token:
+        raise click.ClickException("You aren't signed in. Run `klio login`, or set KLIO_TOKEN.")
+    binary = shutil.which("curl")
+    if not binary:
+        raise click.ClickException("curl isn't installed on this machine.")
+    # Through a config on stdin, so the key is never in the process list.
+    config = f'header = "Authorization: Bearer {token}"\n'
+    done = subprocess.run([binary, "-K", "-", *PINNED, *arguments], input=config, text=True)
+    click.get_current_context().exit(done.returncode)

klio_compute/_compat_cli/_dashboard.py

@@ -0,0 +1,17 @@
+"""`klio dashboard`: the console."""
+
+import webbrowser
+
+import click
+
+CONSOLE_URL = "https://compute.klio.tech/dashboard"
+
+
+def _open(url: str) -> None:
+    webbrowser.open(url)
+
+
+@click.command("dashboard", help="Open the Klio Compute console.")
+def dashboard() -> None:
+    click.echo(CONSOLE_URL)
+    _open(CONSOLE_URL)

klio_compute/_compat_cli/_deploy.py

@@ -0,0 +1,125 @@
+"""`klio deploy app.py`: every function and class of an app. `klio deploy app.py:handler` stays as it was."""
+
+import dataclasses
+from typing import Any, Dict, Iterator, List, Optional, Tuple
+
+import click
+
+from ..compat import _names, _platform
+from ..compat._app import App
+from ..compat._options import StubSpec
+from ..compat.exception import Error
+from . import _loader
+
+REFUSED_FLAGS = {"--tag", "-e", "--env", "--stream-logs", "--strategy", "--timeout"}
+HELP = """Deploy an app from a file, or one handler.
+
+\b
+  klio deploy app.py              every function and class in the file's App
+  klio deploy app.py --name NAME  the same, under another app name
+  klio deploy app.py:handler      one handler, as before
+"""
+
+
+def _deployments_of(app: App, app_name: str) -> Iterator[Tuple[str, StubSpec]]:
+    for name, function in app.registered_functions.items():
+        spec = dataclasses.replace(function.spec, app_name=app_name)
+        label = spec.web.label if spec.web else None
+        yield (label or _names.deployment_name(app_name, name)), spec
+    for name, wrapped in app.registered_classes.items():
+        if wrapped.methods:
+            yield _names.deployment_name(app_name, name), dataclasses.replace(wrapped.spec, app_name=app_name)
+        for method_name, marked in wrapped.web_methods.items():
+            spec = dataclasses.replace(wrapped.web_spec(method_name), app_name=app_name)
+            yield (marked.web.label or _names.deployment_name(app_name, f"{name}.{method_name}")), spec
+
+
+def deploy_app(app: App, app_name: Optional[str] = None) -> List[Dict[str, Any]]:
+    platform = _platform.current()
+    name = app_name or app.name
+    deployed = []
+    for deployment_name, spec in _deployments_of(app, name):
+        reference = platform.deploy(spec, deployment_name)
+        deployed.append(
+            {
+                "name": deployment_name,
+                "deployment_id": reference.id,
+                "url": platform.web_url(reference) if spec.web else None,
+            }
+        )
+    return deployed
+
+
+def _print_summary(app_name: str, deployed: List[Dict[str, Any]]) -> None:
+    click.echo(f"Deployed {app_name}: {len(deployed)} deployment{'' if len(deployed) == 1 else 's'}")
+    for item in deployed:
+        click.echo(f"  {item['name']}" + (f"  {item['url']}" if item["url"] else ""))
+
+
+def _flag_value(arguments: List[str], flag: str) -> Optional[str]:
+    for index, argument in enumerate(arguments):
+        if argument == flag and index + 1 < len(arguments):
+            return arguments[index + 1]
+        if argument.startswith(f"{flag}="):
+            return argument.split("=", 1)[1]
+    return None
+
+
+def _file_target(arguments: List[str]) -> Optional[_loader.Target]:
+    for argument in arguments:
+        if not argument.startswith("-"):
+            target = _loader.parse_target(argument)
+            if target.is_python_file:
+                return target
+    return None
+
+
+def _wants_compat(target: Optional[_loader.Target]) -> bool:
+    if target is None or target.name:
+        return False
+    return bool(_loader.load(target.path).apps)
+
+
+def _refuse_flags(arguments: List[str]) -> None:
+    for argument in arguments:
+        flag = argument.split("=", 1)[0]
+        if flag in REFUSED_FLAGS:
+            _loader.refuse_flag(flag)
+
+
+def wrap(original: click.Command) -> click.Command:
+    def callback(arguments: Tuple[str, ...]) -> None:
+        context = click.get_current_context()
+        items = list(arguments)
+        target = _file_target(items)
+        if not _wants_compat(target):
+            delegated = original.make_context(original.name, items, parent=context)
+            with delegated:
+                original.invoke(delegated)
+            return
+        _refuse_flags(items)
+        name = _flag_value(items, "--name") or _flag_value(items, "-n")
+        app = _loader.only_app(_loader.load(target.path))
+        try:
+            deployed = deploy_app(app, name)
+        except Error as error:
+            raise click.ClickException(str(error)) from None
+        _print_summary(name or app.name, deployed)
+
+    def show_help(ctx: click.Context, _param: Any, value: bool) -> None:
+        if value and not ctx.resilient_parsing:
+            click.echo(HELP)
+            click.echo(original.get_help(click.Context(original, info_name="klio deploy", parent=ctx.parent)))
+            ctx.exit()
+
+    return click.Command(
+        original.name,
+        help=HELP,
+        callback=callback,
+        params=[
+            click.Option(["--help", "-h"], is_flag=True, is_eager=True, expose_value=False, callback=show_help),
+            click.Argument(["arguments"], nargs=-1, type=click.UNPROCESSED),
+        ],
+        add_help_option=False,
+        context_settings={"ignore_unknown_options": True, "allow_extra_args": True},
+    )

klio_compute/_compat_cli/_loader.py

@@ -0,0 +1,91 @@
+"""Finds what a file defines: the object a CLI target names, and which apps live there."""
+
+import importlib
+import os
+import sys
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import click
+
+from ..compat import _cls, _function, _names
+from ..compat._app import App, registered_apps
+from ..compat._refuse import SUFFIX
+from ..compat.exception import Error
+
+PYTHON_SUFFIX = ".py"
+
+
+@dataclass
+class Loaded:
+    path: str
+    module: Any
+    apps: List[App]
+
+
+@dataclass(frozen=True)
+class Target:
+    path: str
+    name: str
+
+    @property
+    def is_python_file(self) -> bool:
+        return self.path.endswith(PYTHON_SUFFIX)
+
+
+def parse_target(text: str) -> Target:
+    """`file.py`, `file.py::name` or `file.py:name`."""
+    if "::" in text:
+        path, _, name = text.partition("::")
+    else:
+        path, _, name = text.partition(":")
+    return Target(path, name)
+
+
+def load(path: str) -> Loaded:
+    if not os.path.exists(path):
+        raise click.ClickException(f"Can't find {path}.")
+    root = os.getcwd()
+    if root not in sys.path:
+        sys.path.insert(0, root)
+    before = set(id(app) for app in registered_apps())
+    name = _names.module_for_file(path)
+    try:
+        module = importlib.import_module(name)
+    except (Error, click.ClickException):
+        raise
+    except Exception as error:
+        raise click.ClickException(
+            f"Couldn't import {path}: {type(error).__name__}: {error}\n"
+            "Its imports must be installed on this machine. Wrap packages that exist only in the image in "
+            "`with image.imports():`."
+        ) from error
+    apps = [app for app in registered_apps() if id(app) not in before]
+    if not apps:
+        apps = [value for value in vars(module).values() if isinstance(value, App)]
+    return Loaded(path, module, apps)
+
+
+def is_compat_object(value: Any) -> bool:
+    return isinstance(value, (App, _function.Function, _cls.Cls)) or (callable(value) and _entrypoint_of(value))
+
+
+def _entrypoint_of(value: Any) -> Optional[App]:
+    for app in registered_apps():
+        if value in app.registered_entrypoints.values():
+            return app
+    return None
+
+
+def only_app(loaded: Loaded) -> App:
+    if not loaded.apps:
+        raise click.ClickException(f"{loaded.path} doesn't define an App.")
+    if len(loaded.apps) > 1:
+        names = ", ".join(app.name for app in loaded.apps)
+        raise click.ClickException(f"{loaded.path} defines several apps ({names}). Keep one per file.")
+    return loaded.apps[0]
+
+
+def refuse_flag(flag: str, hint: str = "") -> None:
+    message = f"{flag} {SUFFIX}"
+    raise click.ClickException(f"{message} {hint}".strip())