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

snowflake/sandbox/lifecycle.py

@@ -0,0 +1,67 @@
+"""Process-wide transport lifecycle — the public home for ``shutdown()``.
+
+The SDK pools sockets on a process-wide HTTP client so that repeated calls do not
+pay a fresh TLS handshake each time. That pool is created lazily on first use and
+outlives any individual `Sandbox`, so something has to be able to release it:
+``shutdown()`` is that something. Call it at the end of a program, a notebook
+kernel, or a request handler.
+
+`shutdown` was previously reachable only as ``snowflake.sandbox._transport.shutdown``
+— the package's one publicly-documented lifecycle call, defined in a private module.
+The transport *internals* are genuinely private (connection policy, retry budgets,
+SSE framing); the lifecycle verb over them is not, so it lives here instead. The
+transports still own their process-wide singletons, and this module only detaches and
+closes them.
+
+Two shapes, one per client:
+
+* `shutdown` — ``await`` it; closes the async transport's client. This is the one
+  exported as ``snowflake.sandbox.shutdown``.
+* `shutdown_sync` — for programs that only ever used the sync `Sandbox` API and so
+  have no event loop to await on. It is documented in the API reference alongside
+  `shutdown`, but is deliberately NOT a top-level ``snowflake.sandbox`` export the
+  way `shutdown` is: the sync pool is also closed on interpreter exit, so explicit
+  release is a rare opt-in rather than the default lifecycle call, and promoting it
+  to the package root would imply a symmetry with the four sync-by-default data
+  helpers that it does not share. It lives here (not buried in ``_sync_transport``)
+  so both halves of the sync/async axis have a public home a caller can reach.
+"""
+
+from __future__ import annotations
+
+__all__ = ["shutdown", "shutdown_sync"]
+
+
+async def shutdown() -> None:
+    """Close the process-wide async transport's HTTP client(s) and drop it.
+
+    The SDK pools sockets on a process-wide client; this is the public way to
+    release them (the alternative was a private ``get_transport().aclose()``).
+    Call ``await snowflake.sandbox.shutdown()`` at the end of a program, a
+    notebook kernel, or a request handler.
+
+    Safe to call more than once and safe when nothing was ever created.
+    """
+    # Imported inside the function so this module keeps no module-scope package
+    # imports: ``_transport`` imports ``shutdown`` from here to re-export it, and a
+    # module-scope import in this direction would close that into a cycle.
+    from snowflake.sandbox._transport import _take_all_transports
+
+    # Every pooled transport, not just the default slot: a role-scoped transport
+    # (``create(role=...)``) or one bound by a ``using()`` block is equally a
+    # process-wide pooled client, and leaving it open is the leak this call exists
+    # to close.
+    for transport in _take_all_transports():
+        await transport.aclose()
+
+
+def shutdown_sync() -> None:
+    """Close the process-wide *sync* transport's HTTP client and drop it.
+
+    The blocking counterpart to `shutdown`, for callers that only used the sync
+    `Sandbox` API and have no running event loop. Safe to call more than once and
+    safe when nothing was ever created.
+    """
+    from snowflake.sandbox._sync_transport import reset_sync_transport
+
+    reset_sync_transport()

snowflake/sandbox/log_stream.py

@@ -0,0 +1,219 @@
+"""Iterable readers over a command-container's captured output.
+
+`AsyncSandbox.logs()` returns a flattened text snapshot. That discards two things
+the server already sends per line — the ``stream`` channel and the ``ts`` — so
+callers cannot separate stdout from stderr, and cannot follow output as it is
+produced.
+
+`LogStream` restores both. It is what `AsyncSandbox.stdout` and
+`AsyncSandbox.stderr` return, giving a command container the same
+``.stdout`` / ``.stderr`` shape an exec'd process already has via
+`exec_stream.ExecStream`.
+
+There is no server-side follow endpoint for logs, so this polls
+``GET containers/{id}/logs?since=<ts>`` and yields only lines newer than the last
+one seen. That is an implementation detail: the iteration contract is the same as
+a push stream, and the polling can be swapped for a streaming endpoint later
+without changing callers.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+from typing import TYPE_CHECKING, Any, Self
+
+from snowflake.sandbox._runtime._protocol import _RESULT_SENTINEL
+from snowflake.sandbox.exceptions import SandboxConflictError, SandboxError
+from snowflake.sandbox.types import TERMINAL_STATUSES, StreamName
+
+if TYPE_CHECKING:  # pragma: no cover - typing only
+    from snowflake.sandbox.client import AsyncSandbox
+
+__all__ = ["LogStream"]
+
+# Poll cadence while following. Fast enough to feel live, slow enough not to
+# hammer the API for a quiet container.
+_POLL_INTERVAL_S = 1.0
+
+# How long to keep polling after the container reports a terminal status, so a
+# process that writes just before exiting does not lose its last lines.
+_DRAIN_AFTER_TERMINAL_S = 2.0
+
+
+class LogStream(AsyncIterator[str]):
+    """Async iterator over one channel of a command container's output.
+
+    Yields line text with the trailing newline stripped, oldest first. Iteration
+    ends once the container is in a terminal state and no further lines arrive,
+    so ``async for`` over a finished container terminates rather than hanging.
+
+    Only meaningful for command containers (created with ``command=``); an
+    ``exec``-only sandbox produces no managed-process output and the stream ends
+    immediately.
+    """
+
+    def __init__(self, sandbox: AsyncSandbox, channel: StreamName, *, since_ts_ms: int = 0) -> None:
+        self._sandbox = sandbox
+        self._channel = channel
+        self._since_ts_ms = since_ts_ms
+        self._pending: list[str] = []
+        self._done = False
+        # Set once the managed process's exit marker is seen in the log stream.
+        self._process_done = False
+        self._terminal_since: float | None = None
+        # How many lines carrying no usable `ts` we have already consumed. The
+        # server filters on `since`, so a line it never timestamped can never be
+        # filtered out and comes back on every poll. Counting them dedupes
+        # positionally, which -- unlike matching on text -- keeps a genuinely
+        # repeated log line.
+        self._untimestamped_seen = 0
+
+    def __aiter__(self) -> Self:
+        return self
+
+    async def __anext__(self) -> str:
+        while True:
+            if self._pending:
+                return self._pending.pop(0)
+            if self._done:
+                raise StopAsyncIteration
+            await self._fill()
+
+    async def _fill(self) -> None:
+        """Fetch any lines newer than the last one seen, or decide we are done."""
+        lines = await self._sandbox._fetch_log_lines(since_ts_ms=self._since_ts_ms)
+
+        # Snapshot the skip threshold once. Advancing it inside the loop drops
+        # later lines in the same batch that share a `ts` (same millisecond) with
+        # an earlier one -- they would compare `ts < self._since_ts_ms` and be
+        # skipped as "already seen". The watermark is advanced once, after the loop.
+        start_ts = self._since_ts_ms
+        next_ts = self._since_ts_ms
+        got = False
+        untimestamped = 0
+        for ln in lines:
+            text = str(ln.get("text") or "")
+            # The injected runner prints `__SANDBOX_RESULT__<b64>` (on stdout) when
+            # the managed process exits. That is the real end-of-stream signal: the
+            # container itself stays `ready` after the process ends, so refresh()
+            # never reaches a terminal status and the loop below would otherwise
+            # poll forever. Check it BEFORE the watermark/dedup logic so a sentinel
+            # sharing a timestamp with the final output line can never be skipped as
+            # "already seen"; check on every channel's poll so both the stdout and
+            # stderr streams terminate; and suppress the internal marker from output.
+            if _RESULT_SENTINEL in text:
+                self._process_done = True
+                continue
+            ts = int(ln.get("ts") or 0)
+            if ts:
+                # Compare against the batch-start watermark, not one mutated
+                # mid-loop; `since` is inclusive server-side, so this skips only
+                # what earlier polls already yielded.
+                if ts < start_ts:
+                    continue
+                next_ts = max(next_ts, ts + 1)
+            else:
+                # No usable timestamp: advancing the watermark is impossible, so
+                # this line is re-delivered by every poll for the life of the
+                # stream. Without positional dedupe it was re-yielded forever and
+                # `read()` grew without bound.
+                untimestamped += 1
+                if untimestamped <= self._untimestamped_seen:
+                    continue
+                self._untimestamped_seen = untimestamped
+            if ln.get("stream") != self._channel:
+                continue
+            self._pending.append(text)
+            got = True
+
+        # Advance the watermark once, after the whole batch (see start_ts above).
+        self._since_ts_ms = next_ts
+
+        if got:
+            self._terminal_since = None
+            return
+
+        # Managed process has exited (its result marker was seen) and everything
+        # newer than the watermark is drained -- stop now. Waiting for the container
+        # lifecycle to reach `dead` would hang: a command container stays `ready`
+        # after its process exits.
+        if self._process_done:
+            self._done = True
+            return
+
+        # Fallback for a stream with no result marker (older runtime, or a raw
+        # command): end once the container is finished AND has stayed quiet for a
+        # moment -- a process can write immediately before exiting.
+        #
+        # This must be a live read. `AsyncSandbox.status` is a cached value set at
+        # create time, and nothing on the container's own lifecycle ever mutates
+        # it, so iterating over a finished container polled forever: `async for
+        # line in sb.stdout` never returned and `read()` never completed.
+        if await self._sandbox.refresh() in TERMINAL_STATUSES:
+            loop = asyncio.get_running_loop()
+            now = loop.time()
+            if self._terminal_since is None:
+                self._terminal_since = now
+            elif now - self._terminal_since >= _DRAIN_AFTER_TERMINAL_S:
+                self._done = True
+                return
+
+        await asyncio.sleep(_POLL_INTERVAL_S)
+
+    async def read(self) -> str:
+        """Drain the stream and return it as one newline-joined string.
+
+        Convenience for the non-streaming case; equivalent to joining the
+        iteration. Blocks until the stream ends, so only use it on a container
+        that will finish.
+        """
+        return "\n".join([line async for line in self])
+
+
+def _log_lines_from_payload(payload: Any) -> list[dict[str, Any]]:
+    """Extract the ``lines`` array from a logs response, tolerating an empty body."""
+    if not isinstance(payload, dict):
+        return []
+    lines = payload.get("lines")
+    return list(lines) if isinstance(lines, list) else []
+
+
+# Marker in the server's 409 body when the logs route is asked about a sandbox that
+# has no managed process. Matched as a substring because it is the only signal the
+# response carries: the error `type` is the generic "invalid_request_error", and the
+# status is a bare 409, both of which the logs route also uses for other reasons.
+_NO_COMMAND_MARKER = "not created with a command"
+
+
+def _explain_logs_conflict(exc: SandboxConflictError) -> SandboxError:
+    """Translate the logs route's 409 into an error that says what to do instead.
+
+    Only command containers have captured output: the logs route reads the managed
+    process's stdout/stderr, and a sandbox created without ``command=`` has no such
+    process. Asking for its logs is a usage error, but the server reports it as a
+    bare 409 carrying raw JSON:
+
+        SandboxConflictError: resource conflict: {"error":{"message":"container was
+        not created with a command","type":"invalid_request_error"}}
+
+    which reads like a transient conflict worth retrying rather than a permanent
+    "you wanted a different method". Callers hit this because `logs()` is the
+    obvious-looking way to read output from any sandbox, and nothing in that message
+    points at `exec_stream()`.
+
+    Returns the original exception unchanged when the marker is absent, so a genuine
+    409 from this route is not mislabelled.
+    """
+    if _NO_COMMAND_MARKER not in str(exc):
+        return exc
+    return SandboxError(
+        "logs() is only available on a sandbox created with command=: the logs "
+        "route reads that managed process's captured stdout/stderr, and this "
+        "sandbox has none.\n\n"
+        "To stream output from a command you run yourself, use exec_stream():\n"
+        '    for line in sb.exec_stream(["python", "report.py"]):\n'
+        "        print(line.stream, line.data)\n\n"
+        "To get a sandbox whose logs() works, pass the command at create time:\n"
+        '    sb = Sandbox.create(image=..., command=["python", "worker.py"])'
+    )