snowflake-sandbox-python 0.2.1a1__py3-none-any.whl

snowflake/sandbox/_runtime/_shims.py

@@ -0,0 +1,267 @@
+"""The generated python entrypoints the SDK ships into a container.
+
+Three programs, one per deploy shape, each written into the bundle and named as
+the container's entry:
+
+* `_session_runner_shim` — ``__session_runner__.py``, the ``@app.session`` daemon.
+* `_function_runner_source` — ``__app_runner__.py``, the ``@app.function`` runner.
+* `_detached_wrapper_source` — the ``python3 -c`` wrapper a detached
+  ``create(command=…)`` runs so ``wait()`` has a sentinel to read.
+
+All three are **source text**, not code this process runs, and the container's
+base image ships only the *published* SDK — possibly older than the one
+deploying. So none of them may import from ``snowflake.sandbox``: the session
+daemon inlines the whole mailbox loop rather than calling `session_loop`, and
+every shared value arrives baked in as a literal from `_protocol`.
+"""
+
+from __future__ import annotations
+
+import textwrap
+
+from snowflake.sandbox._runtime._protocol import (
+    _ENTER_FLAG,
+    _MAILBOX_PATH,
+    _REPLY_SENTINEL,
+    _RESULT_SENTINEL,
+)
+
+__all__ = [
+    "_detached_wrapper_source",
+    "_function_runner_source",
+    "_session_runner_shim",
+]
+
+
+def _session_runner_shim(module: str, name: str) -> str:
+    """Generate the daemon shim for a ``@app.session`` class.
+
+    The shim runs as the container's main process (the manifest entry is
+    ``python __session_runner__.py``). It **inlines the mailbox loop** — it does
+    NOT import ``snowflake.sandbox.session_loop``, because the container's base
+    image ships only the *published* SDK, which may not have it. The literal
+    mailbox path / reply-sentinel / enter-flag values are baked in from
+    ``_runtime._protocol`` so they stay defined in exactly one place.
+
+    Behaviour (mirrors ``session_loop`` but with the user's class as warm state):
+
+    1. Insert the runner's own dir on ``sys.path`` and import the user's class
+       by ``module:name`` qualname.
+    2. Instantiate it once — **the instance IS the warm state across turns**.
+    3. Run its cold-start hook once: the method flagged by ``@enter``
+       (``_ENTER_FLAG``), else a method named ``enter`` or ``setup`` if present.
+    4. Tail the mailbox for ``{"id", "msg"}`` lines; for each call
+       ``instance.on_message(msg)`` and print
+       ``__SANDBOX_REPLY__<id>__<base64(json(reply))>`` to stdout (flushed).
+       Malformed lines and handler exceptions never kill the daemon.
+    """
+    return textwrap.dedent(f"""\
+        import base64
+        import importlib
+        import json
+        import os
+        import signal
+        import sys
+        import time
+
+        _MAILBOX_PATH = {_MAILBOX_PATH!r}
+        _REPLY_SENTINEL = {_REPLY_SENTINEL!r}
+        _ENTER_FLAG = {_ENTER_FLAG!r}
+
+        # Per-session reply nonce: read it and remove it from the environment
+        # BEFORE importing the user's module, so nothing the user's code does at
+        # import time (or in on_message) can read it and forge an authenticated
+        # reply. Empty when the deploying SDK predates the nonce.
+        _REPLY_NONCE = os.environ.pop("SANDBOX_SESSION_NONCE", "")
+
+        sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+        _mod = importlib.import_module({module!r})
+        _cls = getattr(_mod, {name!r})
+        _instance = _cls()
+
+        def _find_hook(_flag, *_names):
+            for _attr in dir(_instance):
+                _m = getattr(_instance, _attr, None)
+                if callable(_m) and getattr(_m, _flag, False):
+                    return _m
+            for _attr in _names:
+                _m = getattr(_instance, _attr, None)
+                if callable(_m):
+                    return _m
+            return None
+
+        # --- cold-start hook (once): @enter-flagged method, else enter/setup ---
+        _enter = _find_hook(_ENTER_FLAG, "enter", "setup")
+        if _enter is not None:
+            _enter()
+
+        # --- survive the platform's own idle checkpoint -------------------
+        # Snowflake's idle auto-suspend sends SIGUSR1 immediately before the
+        # gVisor memory checkpoint and SIGUSR2 after restore. SIGUSR1's default
+        # action is *terminate*, so a daemon that installs nothing is killed on
+        # suspend; ignoring both keeps it alive. Nothing else is wired to these
+        # signals -- every reply is printed with flush=True, so there is no
+        # buffered output to drain, and warm state on _instance survives the
+        # checkpoint by itself.
+        try:
+            signal.signal(signal.SIGUSR1, signal.SIG_IGN)
+            signal.signal(signal.SIGUSR2, signal.SIG_IGN)
+        except (ValueError, OSError, AttributeError):
+            pass
+
+        # --- inline mailbox loop (no session_loop import) ---
+        _offset = 0
+        try:
+            open(_MAILBOX_PATH, "a").close()
+        except OSError:
+            pass
+
+        while True:
+            try:
+                _size = os.path.getsize(_MAILBOX_PATH)
+            except OSError:
+                time.sleep(0.5)
+                continue
+            if _size < _offset:
+                _offset = 0
+            if _size == _offset:
+                time.sleep(0.5)
+                continue
+            with open(_MAILBOX_PATH, "r") as _fh:
+                _fh.seek(_offset)
+                _new_lines = _fh.readlines()
+                _offset = _fh.tell()
+            for _line in _new_lines:
+                _line = _line.strip()
+                if not _line:
+                    continue
+                try:
+                    _rec = json.loads(_line)
+                    _msg_id = _rec["id"]
+                    _msg = _rec["msg"]
+                except (ValueError, KeyError, TypeError):
+                    continue
+                try:
+                    _reply = _instance.on_message(_msg)
+                except Exception as _exc:
+                    _reply = "handler error: " + type(_exc).__name__ + ": " + str(_exc)
+                if not isinstance(_reply, str):
+                    _reply = str(_reply)
+                _b64 = base64.b64encode(json.dumps(_reply).encode()).decode()
+                print(_REPLY_SENTINEL + _msg_id + "__" + _REPLY_NONCE + "__" + _b64, flush=True)
+    """)
+
+
+def _function_runner_source(module: str, name: str, mod_basename: str) -> str:
+    """Generate the ``__app_runner__.py`` text for function-mode.
+
+    Assembled by `Function._runner_shim`: *module* / *name* locate the user's
+    callable and *mod_basename* is the bundled source file to fall back to. See
+    `Function._runner_shim` for what the generated program does and why it writes
+    the result artifact.
+    """
+    return textwrap.dedent(f"""\
+        import importlib
+        import importlib.util
+        import json
+        import os
+        import sys
+
+        _here = os.path.dirname(os.path.abspath(__file__))
+        sys.path.insert(0, _here)
+        kwargs = json.loads(os.environ.get("SANDBOX_CALL_JSON", "{{}}"))
+
+        _MODULE = {module!r}
+        _MODFILE = {mod_basename!r}
+
+        def _load_user_module():
+            # Prefer import-by-name (resolves packages / sibling imports, since
+            # the bundle root is on sys.path). Fall back to loading the bundled
+            # source file directly — this is the ``__module__ == "__main__"``
+            # case (a script/REPL/notebook): importing "__main__" would resolve
+            # to THIS shim, not the user's module.
+            if _MODULE and _MODULE != "__main__":
+                try:
+                    return importlib.import_module(_MODULE)
+                except Exception:
+                    pass
+            _path = os.path.join(_here, _MODFILE)
+            if _MODFILE and os.path.isfile(_path):
+                _spec = importlib.util.spec_from_file_location(
+                    "__sandbox_user_module__", _path
+                )
+                _m = importlib.util.module_from_spec(_spec)
+                _spec.loader.exec_module(_m)
+                return _m
+            # No source file bundled: import by name and let it fail clearly.
+            return importlib.import_module(_MODULE)
+
+        _mod = _load_user_module()
+        _result = getattr(_mod, {name!r})(**kwargs)
+        # An ``async def`` function returns a coroutine; run it to completion so
+        # the app-function layer works for both sync and async callables (a bare
+        # coroutine would fail the json.dumps below and warn "never awaited").
+        import inspect as _inspect
+        if _inspect.iscoroutine(_result):
+            import asyncio as _asyncio
+            _result = _asyncio.run(_result)
+        _payload = {{"result": _result}}
+
+        # Typed-result artifact, read by .spawn()/Job.get(); the stdout line
+        # below is what .remote() reads.
+        _artifact = os.environ.get("SANDBOX_RESULT", "/sandbox/result.json")
+        try:
+            os.makedirs(os.path.dirname(_artifact) or ".", exist_ok=True)
+            with open(_artifact, "w") as _fh:
+                json.dump(_payload, _fh)
+        except OSError:
+            pass
+
+        print(json.dumps(_payload))
+    """)
+
+
+def _detached_wrapper_source(payload: str, nonce: str) -> str:
+    """Generate the ``python3 -c`` wrapper text for a detached ``create(command=…)``.
+
+    *payload* is the JSON-encoded inner argv and *nonce* the per-run token the
+    sentinel is stamped with. See `_args._wrap_detached`, which mints both, for
+    why the nonce never reaches the inner command's environment.
+    """
+    return f"""\
+import base64, json, os, signal, subprocess, sys
+argv = json.loads({payload!r})
+nonce = {nonce!r}
+# Survive the platform's idle checkpoint. Snowflake's auto-suspend sends SIGUSR1
+# immediately before the gVisor memory checkpoint and SIGUSR2 after restore, and
+# SIGUSR1's default action is *terminate* -- so without this the wrapper (and with it
+# the managed process) is killed the first time a detached sandbox goes idle, instead
+# of being suspended and resumed. The session shim and the warm-session loop already
+# do exactly this; this path was missed.
+try:
+    signal.signal(signal.SIGUSR1, signal.SIG_IGN)
+    signal.signal(signal.SIGUSR2, signal.SIG_IGN)
+except (ValueError, OSError, AttributeError):
+    pass
+rp = os.environ.get("SANDBOX_RESULT") or "/sandbox/result.json"
+try:
+    os.makedirs(os.path.dirname(rp), exist_ok=True)
+except OSError:
+    rp = "/tmp/_sandbox_result.json"
+os.environ["SANDBOX_RESULT"] = rp
+rc = subprocess.run(argv).returncode
+try:
+    with open(rp) as fh:
+        artifact = json.load(fh)
+except Exception:
+    artifact = None
+core = {{
+    "status": "succeeded" if rc == 0 else "failed",
+    "exit_code": rc,
+    "result": artifact if isinstance(artifact, dict) else {{"exit_code": rc}},
+    "nonce": nonce,
+}}
+sys.stdout.write({_RESULT_SENTINEL!r} + base64.b64encode(json.dumps(core).encode()).decode() + "\\n")
+sys.stdout.flush()
+"""

snowflake/sandbox/_sandbox_state.py

@@ -0,0 +1,303 @@
+"""The sandbox state layer shared by both clients.
+
+`AsyncSandbox` and `Sandbox` differ in how they talk to the server, not in what
+they remember about it. This holds the part that is identical: the fields a
+sandbox carries, the read-only view onto them, and the absorption of a server
+response into them.
+
+It was duplicated verbatim — 12 methods, 169 lines, none of them async — and two
+of them (`_absorb_read_fields`, `_absorb_resource_echo`) parse the wire, which is
+where a divergence produces a *wrong value* rather than a crash. `_SERVER_STATUS`
+already drifted that way: a status added on one side only made a resumed sandbox
+read as ``"unknown"`` on the other, with a warning blaming the caller's SDK
+version. Sharing the code removes that failure mode instead of testing for it.
+
+Not a public base class: it is an implementation detail of the two clients, and
+neither exposes it. It deliberately holds no transport and performs no I/O, so it
+stays mode-agnostic — anything that needs to await belongs in the subclass.
+"""
+
+from __future__ import annotations
+
+import warnings
+from collections.abc import Mapping
+from typing import Any
+
+from snowflake.sandbox._env import is_platform_env_key
+from snowflake.sandbox._wire import _MEMORY_TIERS
+from snowflake.sandbox.exceptions import SandboxContractWarning, SandboxError
+from snowflake.sandbox.types import TERMINAL_STATUSES, MemoryTier, SandboxStatus
+
+__all__ = ["_SandboxState"]
+
+
+class _SandboxState:
+    """Fields both clients carry, and the read-only view onto them.
+
+    Subclasses set these in their own ``__init__`` — they are declared here so the
+    shared methods below type-check against them, and so the state a sandbox holds
+    is stated in one place.
+    """
+
+    _id: str | None
+    _status: SandboxStatus
+    _image: str
+    _memory: MemoryTier
+    _cpu: float | None
+    _gpu: dict[str, object] | None
+    _name: str | None
+    _role: str | None
+    _tags: dict[str, str]
+    _exit_code: int | None
+    _generation: int | None
+    _created_at: int | None
+    _stopped_at: int | None
+    _sandbox_env: list[str]
+    _error_message: str | None
+    _async_pending: bool
+
+    def _add_platform_env(self, key: str, value: str) -> None:
+        """Register one SDK-minted platform env var on the create body's `sandbox_env`.
+
+        This LIST is the single chokepoint every SDK-injected platform key goes
+        through.
+
+        Platform keys carry a server-reserved prefix (SNOWFLAKE_/SANDBOX_/
+        CNG_SECRET_), which the server rejects (400) on the user `env` map but accepts
+        on `sandbox_env`. Asserting the prefix here — and routing the session reply
+        nonce, the code-delivery URL and the SANDBOX_JOB_* runner vars through this one
+        method — is what keeps a reserved-env key from silently being dropped: an
+        SDK-minted platform key can reach the wire only via the channel the server
+        accepts, never the `env` map. User env
+        is untouched by this path (it stays on `self._env`; the server owns rejecting
+        a user-supplied reserved key there).
+        """
+        if not is_platform_env_key(key):
+            raise SandboxError(
+                f"internal error: {key!r} is not a platform-reserved env key "
+                "(prefix SNOWFLAKE_/SANDBOX_/CNG_SECRET_); only SDK-minted platform "
+                "keys belong on the sandbox_env channel — user env goes on the `env` map"
+            )
+        self._sandbox_env.append(f"{key}={value}")
+
+    def _absorb_read_fields(self, payload: Mapping[str, Any]) -> None:
+        """Adopt the read-only fields the server echoes (SDK-CONTRACT).
+
+        Those are image, gpu, name, exit_code, generation, created_at, labels,
+        and role. Each is taken only when the server actually sends it, so this
+        is forward-compatible: harmless against a server that omits them (the
+        field stays at its prior value / None).
+        """
+        img = payload.get("image")
+        # Only overwrite when the server sends a real image AND we don't already
+        # hold a concrete one — a hydrated handle seeds "unknown" it should replace,
+        # but a live handle's requested image must not be clobbered by an omission.
+        if isinstance(img, str) and img and self._image in ("", "unknown"):
+            self._image = img
+        gpu = payload.get("gpu")
+        if isinstance(gpu, Mapping):
+            self._gpu = dict(gpu)
+        name = payload.get("name")
+        if isinstance(name, str) and name:
+            self._name = name
+        ec = payload.get("exit_code")
+        if isinstance(ec, int) and not isinstance(ec, bool):
+            self._exit_code = ec
+        gen = payload.get("generation")
+        if isinstance(gen, int) and not isinstance(gen, bool):
+            self._generation = gen
+        created_at = payload.get("created_at")
+        if isinstance(created_at, int) and not isinstance(created_at, bool):
+            self._created_at = created_at
+        stopped_at = payload.get("stopped_at")
+        if isinstance(stopped_at, int) and not isinstance(stopped_at, bool):
+            self._stopped_at = stopped_at
+        # `labels` -> tags. Adopt only a non-empty map the server actually sent:
+        # `omitempty` means an untagged (or controller-reconstructed) record carries
+        # no `labels` key at all, so an omission must not wipe tags this
+        # handle already holds — same forward-compat rule as `image` above.
+        labels = payload.get("labels")
+        if isinstance(labels, Mapping) and labels:
+            self._tags = {str(k): str(v) for k, v in labels.items()}
+        # role is the first-class server field added in sandbox-api to echo the
+        # creating role. Adopt only when present and non-empty — omission must
+        # not wipe a live handle's known role (same forward-compat rule as image).
+        server_role = payload.get("role")
+        if isinstance(server_role, str) and server_role:
+            self._role = server_role
+        error_msg = payload.get("error_message")
+        if isinstance(error_msg, str) and error_msg:
+            self._error_message = error_msg
+
+    def _absorb_resource_echo(
+        self, payload: Mapping[str, Any], *, compare_to_request: bool = True
+    ) -> None:
+        """Adopt the resource shape the server says it actually served.
+
+        ``memory`` / ``cpu`` in a create or get response are the *effective*
+        limits, not the request echoed back (sandbox-api
+        ``containers/models.go`` ``ContainerFromInfo``). Reading them is the only
+        way the SDK can see the unknown-key-ignored drift: ask a server that
+        predates the ``memory`` rename for ``64g`` and it serves the default tier
+        with a ``200``, while ``sb.memory`` would keep reporting the 64g that was
+        requested and never delivered. A divergence is warned, never hidden.
+
+        ``compare_to_request=False`` for a handle being hydrated from a list/get
+        response, where there is no local request to diverge from.
+        """
+        raw_memory = payload.get("memory") or payload.get("memory_limit")
+        if isinstance(raw_memory, str) and raw_memory:
+            if raw_memory in _MEMORY_TIERS:
+                if compare_to_request and raw_memory != self._memory:
+                    warnings.warn(
+                        f"requested memory={self._memory!r} but the server served "
+                        f"{raw_memory!r} for container {self._id}",
+                        SandboxContractWarning,
+                        stacklevel=3,
+                    )
+                self._memory = raw_memory  # type: ignore[assignment]
+            else:
+                warnings.warn(
+                    f"server reported an unrecognized memory tier {raw_memory!r} for "
+                    f"container {self._id}; keeping {self._memory!r} as the local view",
+                    SandboxContractWarning,
+                    stacklevel=3,
+                )
+        raw_cpu = payload.get("cpu")
+        if isinstance(raw_cpu, (int, float)) and not isinstance(raw_cpu, bool) and raw_cpu > 0:
+            served_cpu = float(raw_cpu)
+            if compare_to_request and self._cpu is not None and abs(served_cpu - self._cpu) > 1e-9:
+                warnings.warn(
+                    f"requested cpu={self._cpu} but the server served {served_cpu} "
+                    f"for container {self._id}",
+                    SandboxContractWarning,
+                    stacklevel=3,
+                )
+            self._cpu = served_cpu
+
+    @property
+    def status(self) -> SandboxStatus:
+        """The last status the server reported for this container.
+
+        This is a cached value, not a live read: it is set by ``create()`` /
+        ``connect()`` and by each `refresh()`. Call `refresh()` when you need
+        the current state -- a container that has since crashed or been stopped
+        still reads as whatever it was when this handle last heard from the
+        server.
+        """
+        return self._status
+
+    @property
+    def is_running(self) -> bool:
+        """True if the sandbox has not yet terminated (cached status is not terminal).
+
+        This checks the cached `status`, not a live read. Call `refresh()` first
+        if you need the current state.
+        """
+        return self._status not in TERMINAL_STATUSES
+
+    @property
+    def is_terminated(self) -> bool:
+        """True if the sandbox has terminated (cached status is terminal).
+
+        This checks the cached `status`, not a live read. Call `refresh()` first
+        if you need the current state.
+        """
+        return self._status in TERMINAL_STATUSES
+
+    @property
+    def name(self) -> str | None:
+        """The sandbox's name, once the server echoes it.
+
+        Either the caller's own ``name=``, or -- if that was omitted -- a name
+        the server generated (e.g. ``"frosty-tricorder"``). ``None`` only on a
+        server that does not yet store/echo it.
+        """
+        return self._name
+
+    @property
+    def role(self) -> str | None:
+        """The Snowflake role this sandbox was created under, or ``None`` if not set.
+
+        Populated from the server's response field. Available immediately after
+        ``Sandbox.create()`` and after ``Sandbox.connect()`` or any list call that
+        returns the container record.
+        """
+        return self._role
+
+    @property
+    def tags(self) -> Mapping[str, str]:
+        """The tags (server ``labels``) attached to this sandbox, as last known here.
+
+        A cached view: set from ``tags=`` / `set_tags()` before create, then
+        refreshed from the server's echo on `refresh()` / `get_tags()`. May be
+        empty even for a tagged sandbox if the server rebuilt the record from the
+        controller. Returns a copy; mutating it does not relabel.
+        """
+        return dict(self._tags)
+
+    @property
+    def cpu(self) -> float | None:
+        """The effective CPU allocation in cores.
+
+        Set from the server's echo once the container exists (the server reports
+        the *effective* limit, which may differ from an override that was asked
+        for); ``None`` before create when no override was requested.
+        """
+        return self._cpu
+
+    @property
+    def gpu(self) -> Mapping[str, object] | None:
+        """The GPU spec passed at construction, or `None` if no GPU was requested.
+
+        Example:
+            if sb.gpu:
+                print("GPU:", sb.gpu)
+        """
+        return self._gpu
+
+    @property
+    def exit_code(self) -> int | None:
+        """The managed process's exit code, or ``None`` if it has not exited.
+
+        Also ``None`` when the server has not reported one yet. Populated by
+        `refresh()` / `wait()` from the cheap status route where the server
+        exposes one.
+        """
+        return self._exit_code
+
+    @property
+    def generation(self) -> int | None:
+        """A monotonic restart/resume counter, or ``None`` if the server reports none.
+
+        A bump means Snowflake cold-restarted the sandbox (a suspended container
+        resumed, or an unhealthy one recreated): in-container ``/tmp`` and
+        ``/snowflake/stages`` were wiped, so re-stage anything you cannot
+        recompute. Compare across `refresh()` calls to detect it.
+        """
+        return self._generation
+
+    @property
+    def created_at(self) -> int | None:
+        """The server-reported container creation timestamp, if available.
+
+        This is copied verbatim from the API payload when the backend includes
+        it. The SDK does not reinterpret units here; CLI/rendering code decides
+        how to present it.
+        """
+        return self._created_at
+
+    @property
+    def stopped_at(self) -> int | None:
+        """The server-reported stop timestamp, if available.
+
+        Only known for a stopped sandbox (the backend omits it while running or
+        suspended), copied verbatim from the API payload. ``None`` for a live
+        sandbox or a server that does not report it.
+        """
+        return self._stopped_at
+
+    def __repr__(self) -> str:
+        role_part = f", role={self._role!r}" if self._role else ""
+        name_part = f", name={self._name!r}" if self._name else ""
+        return f"Sandbox(id={self._id!r}, status={self._status!r}{name_part}{role_part})"