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

snowflake/sandbox/deploy.py

@@ -0,0 +1,499 @@
+"""Deploy an agent to a sandbox from a `DeploySpec`.
+
+This is the deterministic core behind ``snow sandbox deploy`` (and, later, a
+``snow`` plugin). It is intentionally click-independent so tests and other
+front-ends call it directly. Deploy is deterministic from the spec forward: what
+this module does is decided by the spec, never by a model. (The diagnostics that
+explain a failure — including the ones that do call a model — live in
+``_diagnostics``; the local bundle packaging lives in ``_assemble``.)
+
+Two run modes:
+  * run-to-completion (default): create the container with a keepalive command,
+    ``exec`` the spec's ``entry``, stream output, return its exit code. Right
+    for one-shot agents (e.g. jira triage) and for tests.
+  * detach: make ``entry`` the container's main process and return the id. Right
+    for long-running daemons (tail with ``logs()``).
+
+Preflight runs first (cheap, client-side) and fails loudly with a fix instead of
+letting a misconfig surface as a deep, cryptic error. ``diagnose_failure``
+interprets the one trap preflight can't see from the client: the platform's
+real-credentials requirement for presigned code delivery.
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+
+if TYPE_CHECKING:
+    from snowflake.sandbox.config import ConnectionLike
+    from snowflake.sandbox.types import MemoryTier
+
+from snowflake.sandbox._assemble import _assemble_bundle, _bundle_files, _clone_git_ref
+from snowflake.sandbox._deploy_spec import DeploySpec
+from snowflake.sandbox._diagnostics import Problem, diagnose_failure, preflight_checks
+from snowflake.sandbox.client import AsyncSandbox
+from snowflake.sandbox.exceptions import (
+    SandboxError,
+    SandboxExecError,
+    SandboxExecTimeoutError,
+)
+from snowflake.sandbox.sync_client import Sandbox
+
+# The names *defined here*, plus the three imported from ``_diagnostics`` above,
+# which are re-exported because ``deploy_spec`` itself calls them.
+#
+# The detached-run / warm-session names this module also serves (``Job``,
+# ``RunResult``, ``deploy_async``, ``Session``, ``agent_session``,
+# ``session_loop``) resolve through the module ``__getattr__`` at the bottom of
+# this file and are deliberately left out: they are not bound at module scope, so
+# listing them would be an undefined-name export (ruff F822). They stay reachable
+# the way they are documented — an explicit
+# ``from snowflake.sandbox.deploy import deploy_async`` — which ``__all__`` does
+# not gate.
+__all__ = [
+    "DeployPlan",
+    "DeployResult",
+    "Problem",
+    "deploy_spec",
+    "deploy_spec_sync",
+    "diagnose_failure",
+    "preflight_checks",
+]
+
+
+@dataclass
+class DeployPlan:
+    """What a deploy *would* do (used by --dry-run and surfaced on deploy)."""
+
+    image: str
+    memory: str
+    cpu: float | None
+    entry: Sequence[str]
+    code_stage: str | None
+    egress: Mapping[str, object] | None
+    env_keys: Sequence[str]
+    bundle_files: Sequence[str]
+
+
+@dataclass
+class DeployResult:
+    plan: DeployPlan
+    problems: list[Problem] = field(default_factory=list)
+    sandbox_id: str = ""
+    # The sandbox's vanity name — the user-facing identity surfaced by the CLI
+    # instead of sandbox_id (which stays the internal wire key).
+    name: str = ""
+    exit_code: int | None = None
+    stdout: str = ""
+    stderr: str = ""
+    diagnosis: str = ""
+    deployed: bool = False
+    resolved_ref: str = ""
+
+
+@dataclass
+class _DeployPrep:
+    """The output of `_prepare_deploy` — everything the spec decides, resolved
+    before any container I/O, plus the temp dirs the caller must clean up.
+
+    ``stop`` is set when the deploy should return ``result`` without launching a
+    container (a preflight error or ``dry_run``). ``clone_root`` / ``assembled``
+    are temp trees the async and sync deploy tails hand to `_cleanup_prep` in a
+    ``finally``. This is the "spec -> plan" half CONTRIBUTING calls the
+    deterministic core: it is pure of container I/O and single-sourced so the
+    async and sync deploys cannot decide a deploy differently.
+    """
+
+    result: DeployResult
+    bundle_root: Path
+    env: dict[str, str]
+    egress_body: Mapping[str, object] | None
+    effective_timeout: float | None
+    clone_root: Path | None
+    assembled: Path | None
+    stop: bool
+
+
+def _prepare_deploy(
+    spec: DeploySpec,
+    *,
+    source_dir: str | Path | None,
+    git_repo: str | None,
+    git_ref: str | None,
+    git_subdir: str | None,
+    run_preflight: bool,
+    dry_run: bool,
+    timeout: float | None,
+) -> _DeployPrep:
+    """Resolve a `DeploySpec` to a launch-ready `_DeployPrep` (no container I/O).
+
+    Shared by `deploy_spec` and `deploy_spec_sync`: source resolution, bundle
+    assembly, the `DeployPlan`, preflight, and the dry-run short-circuit all live
+    here so the two calling styles cannot drift on what a spec deploys. Only the
+    container create/exec differs between them, and that stays in each tail.
+
+    On any failure after a clone/assembly the temp trees are removed before the
+    error propagates, since the caller never receives the prep to clean up.
+    """
+    clone_root: Path | None = None
+    resolved_ref: str = ""
+    effective_timeout = spec.timeout_s if timeout is None else timeout
+
+    # Two levels, as the manifest had: *base* is the project dir holding the entry
+    # script (it ships wholesale), and bundle.root is the root that include globs
+    # resolve against — often a sibling or repo root. They coincide when no
+    # source_dir is given.
+    if git_repo is not None:
+        clone_root, effective_src, resolved_ref = _clone_git_ref(git_repo, git_ref, git_subdir)
+        base = Path(effective_src)
+    elif source_dir is not None:
+        base = Path(source_dir).expanduser().resolve()
+    elif spec.bundle is not None:
+        base = Path(spec.bundle.root).expanduser().resolve()
+    else:
+        raise SandboxError(
+            "deploy_spec needs a bundle (Bundle.from_dir(...)), a source_dir, or a git_repo"
+        )
+
+    assembled: Path | None = None
+    try:
+        if not base.is_dir():
+            raise SandboxError(f"bundle base not found: {base}")
+
+        # Include globs mean a slice to assemble on top of *base*; a bare exclude
+        # also forces assembly, since without it the raw dir ships as-is and the
+        # excludes are a silent no-op. Otherwise the directory itself is the bundle.
+        if spec.bundle is not None and (spec.bundle.include or spec.bundle.exclude):
+            include_root = Path(spec.bundle.root).expanduser()
+            if not include_root.is_absolute():
+                include_root = (base / include_root).resolve()
+            assembled = _assemble_bundle(
+                base,
+                {
+                    "root": str(include_root),
+                    "include": list(spec.bundle.include),
+                    "exclude": list(spec.bundle.exclude),
+                },
+            )
+        bundle_root = assembled or base
+
+        env = {k: str(v) for k, v in spec.env.items() if str(v) != ""}
+        egress_body = spec.egress_body()
+        plan = DeployPlan(
+            image=spec.image,
+            memory=spec.memory,
+            cpu=spec.cpu,
+            entry=list(spec.entry),
+            code_stage=spec.code_stage,
+            egress=egress_body,
+            env_keys=sorted(env.keys()),
+            bundle_files=_bundle_files(bundle_root),
+        )
+        result = DeployResult(plan=plan, resolved_ref=resolved_ref)
+
+        stop = False
+        if run_preflight:
+            result.problems = preflight_checks(spec, bundle_root)
+            if any(p.severity == "error" for p in result.problems):
+                stop = True  # abort before packaging anything
+        if dry_run:
+            stop = True
+
+        return _DeployPrep(
+            result=result,
+            bundle_root=bundle_root,
+            env=env,
+            egress_body=egress_body,
+            effective_timeout=effective_timeout,
+            clone_root=clone_root,
+            assembled=assembled,
+            stop=stop,
+        )
+    except BaseException:
+        # The caller never gets the prep, so it cannot run _cleanup_prep — clean
+        # the temp trees here before the failure propagates.
+        if assembled is not None:
+            shutil.rmtree(assembled, ignore_errors=True)
+        if clone_root is not None:
+            shutil.rmtree(clone_root, ignore_errors=True)
+        raise
+
+
+def _cleanup_prep(prep: _DeployPrep) -> None:
+    """Remove the temp trees `_prepare_deploy` created (assembled bundle, clone)."""
+    if prep.assembled is not None:
+        shutil.rmtree(prep.assembled, ignore_errors=True)
+    if prep.clone_root is not None:
+        shutil.rmtree(prep.clone_root, ignore_errors=True)
+
+
+def _deploy_command(
+    spec: DeploySpec, *, detach: bool, force_fresh: bool, keepalive_ttl: int
+) -> list[str]:
+    """The container's main command: the entry itself when detached, otherwise a
+    keepalive the entry is ``exec``-ed into.
+
+    Pure; shared so the async and sync deploys launch byte-identical containers.
+    ``force_fresh`` appends a random token to the keepalive so an otherwise
+    identical create does not warm-reuse a previous sandbox.
+    """
+    if detach:
+        return list(spec.entry)
+    command = ["sleep", str(keepalive_ttl)]
+    if force_fresh:
+        command.append(os.urandom(4).hex())  # unique id; dodge warm reuse
+    return command
+
+
+async def deploy_spec(
+    spec: DeploySpec,
+    *,
+    source_dir: str | Path | None = None,
+    git_repo: str | None = None,
+    git_ref: str | None = None,
+    git_subdir: str | None = None,
+    on_line: Callable[[str, str], None] | None = None,
+    dry_run: bool = False,
+    run_preflight: bool = True,
+    detach: bool = False,
+    force_fresh: bool = False,
+    keepalive_ttl: int = 86400,
+    timeout: float | None = None,
+    platform_env: Mapping[str, str] | None = None,
+    connection: ConnectionLike | None = None,
+) -> DeployResult:
+    """Deploy a `DeploySpec` and (unless detached) run its entry to
+    completion, streaming output via ``on_line(stream, data)``.
+
+    This is the deploy core — the decorator API and the CLI both land here.
+    There is no manifest file: the spec *is* the contract.
+
+    The bundle comes from ``spec.bundle`` (assembled from its include globs when
+    it has any, otherwise used as the root directly); ``source_dir`` is the
+    fallback root when the spec has no bundle.  With ``git_repo`` the repo is
+    shallow-cloned at ``git_ref``, the deploy runs against
+    ``<clone>/<git_subdir or .>``, and the clone is removed in a ``finally``.
+    The resolved SHA lands on ``DeployResult.resolved_ref``.
+
+    For a blocking caller with no event loop, `deploy_spec_sync` is the exact
+    synchronous counterpart — it shares this function's spec->plan core.
+
+    ``connection`` picks the Snowflake connection this deploy runs against — a name
+    from ``~/.snowflake/connections.toml``, a live connector connection, or a `Config`.
+    Omit it to use the enclosing ``using()`` block, else the default connection. See
+    `config.using` for the resolution rules.
+    """
+    prep = _prepare_deploy(
+        spec,
+        source_dir=source_dir,
+        git_repo=git_repo,
+        git_ref=git_ref,
+        git_subdir=git_subdir,
+        run_preflight=run_preflight,
+        dry_run=dry_run,
+        timeout=timeout,
+    )
+    try:
+        if prep.stop:
+            return prep.result
+        result = prep.result
+        command = _deploy_command(
+            spec, detach=detach, force_fresh=force_fresh, keepalive_ttl=keepalive_ttl
+        )
+        sb = await AsyncSandbox.from_local(
+            prep.bundle_root,
+            connection=connection,
+            command=command,
+            image=spec.image,
+            memory=cast("MemoryTier", spec.memory),
+            cpu=spec.cpu,
+            code_stage=spec.code_stage,
+            env=prep.env,
+            platform_env=platform_env,
+            egress=prep.egress_body,
+            stage_mounts=spec.stage_mounts,
+        )
+
+        if detach:
+            await sb._ensure_created()
+            result.sandbox_id = sb.id
+            result.name = getattr(sb, "name", "") or ""
+            result.deployed = True
+            result.exit_code = 0
+            return result
+
+        async with sb:
+            result.sandbox_id = sb.id
+            result.name = getattr(sb, "name", "") or ""
+            result.deployed = True
+            try:
+                r = await sb.exec(list(spec.entry), timeout=prep.effective_timeout)
+                result.stdout = r.stdout or ""
+                result.stderr = r.stderr or ""
+                result.exit_code = r.exit_code if r.exit_code is not None else 0
+            except (SandboxExecError, SandboxExecTimeoutError) as exc:
+                # A server-side deadline now raises SandboxExecTimeoutError, which is a
+                # SandboxError but NOT a SandboxExecError; catch both so a timed-out
+                # entrypoint still yields a graceful DeployResult rather than escaping.
+                result.stdout = exc.stdout or ""
+                result.stderr = exc.stderr or ""
+                result.exit_code = exc.exit_code or 1
+            if on_line:
+                if result.stdout:
+                    on_line("stdout", result.stdout)
+                if result.stderr:
+                    on_line("stderr", result.stderr)
+
+        if result.exit_code != 0:
+            result.diagnosis = diagnose_failure(result.exit_code, result.stdout, result.stderr)
+        return result
+    finally:
+        _cleanup_prep(prep)
+
+
+def deploy_spec_sync(
+    spec: DeploySpec,
+    *,
+    source_dir: str | Path | None = None,
+    git_repo: str | None = None,
+    git_ref: str | None = None,
+    git_subdir: str | None = None,
+    on_line: Callable[[str, str], None] | None = None,
+    dry_run: bool = False,
+    run_preflight: bool = True,
+    detach: bool = False,
+    force_fresh: bool = False,
+    keepalive_ttl: int = 86400,
+    timeout: float | None = None,
+    platform_env: Mapping[str, str] | None = None,
+    connection: ConnectionLike | None = None,
+) -> DeployResult:
+    """Deploy a `DeploySpec` synchronously — the blocking counterpart of `deploy_spec`.
+
+    Same contract, same spec->plan core (via `_prepare_deploy`), and the same
+    requests; it blocks on the synchronous ``Sandbox`` instead of awaiting
+    ``AsyncSandbox``, so it works from an ordinary function with no running event
+    loop (a notebook cell, a CLI, a plain script) without spinning one up.
+
+    ``connection`` picks the Snowflake connection this deploy runs against — a name
+    from ``~/.snowflake/connections.toml``, a live connector connection, or a `Config`.
+    Omit it to use the enclosing ``using()`` block, else the default connection. See
+    `config.using` for the resolution rules.
+
+    Example:
+        result = deploy_spec_sync(spec, source_dir="./agent", detach=True)
+        print(result.sandbox_id, result.deployed)
+    """
+    prep = _prepare_deploy(
+        spec,
+        source_dir=source_dir,
+        git_repo=git_repo,
+        git_ref=git_ref,
+        git_subdir=git_subdir,
+        run_preflight=run_preflight,
+        dry_run=dry_run,
+        timeout=timeout,
+    )
+    try:
+        if prep.stop:
+            return prep.result
+        result = prep.result
+        command = _deploy_command(
+            spec, detach=detach, force_fresh=force_fresh, keepalive_ttl=keepalive_ttl
+        )
+        sb = Sandbox.from_local(
+            prep.bundle_root,
+            connection=connection,
+            command=command,
+            image=spec.image,
+            memory=cast("MemoryTier", spec.memory),
+            cpu=spec.cpu,
+            code_stage=spec.code_stage,
+            env=prep.env,
+            platform_env=platform_env,
+            egress=prep.egress_body,
+            stage_mounts=spec.stage_mounts,
+        )
+
+        if detach:
+            sb._ensure_created()
+            result.sandbox_id = sb.id
+            result.name = getattr(sb, "name", "") or ""
+            result.deployed = True
+            result.exit_code = 0
+            return result
+
+        with sb:
+            result.sandbox_id = sb.id
+            result.name = getattr(sb, "name", "") or ""
+            result.deployed = True
+            try:
+                r = sb.exec(list(spec.entry), timeout=prep.effective_timeout)
+                result.stdout = r.stdout or ""
+                result.stderr = r.stderr or ""
+                result.exit_code = r.exit_code if r.exit_code is not None else 0
+            except (SandboxExecError, SandboxExecTimeoutError) as exc:
+                # A server-side deadline now raises SandboxExecTimeoutError, which is a
+                # SandboxError but NOT a SandboxExecError; catch both so a timed-out
+                # entrypoint still yields a graceful DeployResult rather than escaping.
+                result.stdout = exc.stdout or ""
+                result.stderr = exc.stderr or ""
+                result.exit_code = exc.exit_code or 1
+            if on_line:
+                if result.stdout:
+                    on_line("stdout", result.stdout)
+                if result.stderr:
+                    on_line("stderr", result.stderr)
+
+        if result.exit_code != 0:
+            result.diagnosis = diagnose_failure(result.exit_code, result.stdout, result.stderr)
+        return result
+    finally:
+        _cleanup_prep(prep)
+
+
+# ── Detached run surface ────────────────────────────────────────────────────
+# The fire-and-forget run surface lives in jobs.py and the warm-session surface
+# in warm_session.py, both re-exported here so callers can keep doing
+# ``from snowflake.sandbox.deploy import deploy_async``.
+#
+# The re-export is lazy (module __getattr__) because both surfaces are built ON
+# deploy: jobs.py imports DeployPlan from here at its module top and
+# warm_session.py calls deploy_spec. An alias pointing back at them is therefore
+# an inverted edge, and making it eager would be a real module-scope import
+# cycle. Resolving it on first attribute access, after the target module is fully
+# loaded, is what keeps the graph acyclic in the direction that reflects the
+# layering.
+_JOBS_REEXPORTS = (
+    "Job",
+    "SyncJob",
+    "RunResult",
+    "deploy_async",
+    "deploy_async_sync",
+    "_parse_result_sentinel",
+    "_SANDBOX_JOB_RUNNER",
+)
+
+# Warm session surface lives in warm_session.py; it reuses the detached
+# ``deploy_spec(detach=True)`` path above. Lazy for the same reason, plus one of
+# its own: warm_session.py defers its ``deploy_spec`` import so that reaching
+# ``snowflake.sandbox.session_loop`` (a _LAZY entry) does not drag deploy.py +
+# httpx into a bare ``import snowflake.sandbox``.
+_SESSION_REEXPORTS = ("agent_session", "Session", "session_loop", "_parse_reply_sentinel")
+
+
+def __getattr__(name: str) -> Any:  # noqa: D401 - module-level lazy attribute hook
+    if name in _JOBS_REEXPORTS:
+        import importlib
+
+        return getattr(importlib.import_module("snowflake.sandbox.jobs"), name)
+    if name in _SESSION_REEXPORTS:
+        import importlib
+
+        return getattr(importlib.import_module("snowflake.sandbox.warm_session"), name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")