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

snowflake/sandbox/function.py

@@ -0,0 +1,567 @@
+"""``Function`` — one deployable unit of work.
+
+Created by the ``@app.function`` decorator; ``.remote()`` materialises a
+`DeploySpec` from the decorator metadata and calls ``deploy_spec()`` on it.
+
+Two execution modes
+-------------------
+**Process-mode** (``entry=`` specified): kwargs from ``.remote(**kwargs)``
+become env-var overrides. The function body is the *local* implementation; it
+runs only via ``fn()`` or ``fn.local()``.
+
+**Function-mode** (no ``entry=``): the function body runs *in the container*.
+``.remote(**kwargs)`` JSON-encodes kwargs into ``SANDBOX_CALL_JSON`` and injects
+a generated ``__app_runner__.py`` shim as the entry.
+"""
+
+from __future__ import annotations
+
+import functools
+import json
+import tempfile
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from snowflake.sandbox._assemble import Bundle
+from snowflake.sandbox._deploy_spec import DeploySpec
+from snowflake.sandbox._runtime._shims import _function_runner_source
+from snowflake.sandbox.egress import Egress
+from snowflake.sandbox.image import Image
+from snowflake.sandbox.mount import StageMount
+from snowflake.sandbox.secret import Secret
+from snowflake.sandbox.types import MemoryTier
+
+if TYPE_CHECKING:
+    from snowflake.sandbox.app import App
+    from snowflake.sandbox.client import AsyncSandbox
+    from snowflake.sandbox.deploy import DeployResult
+    from snowflake.sandbox.sync_client import Sandbox
+
+__all__ = ["Function", "FunctionSpec"]
+
+
+@dataclass(frozen=True)
+class FunctionSpec:
+    """All decorator metadata for a registered function.
+
+    ``image`` is normalised to ``Image`` at decoration time (bare strings
+    are coerced).  ``entry=None`` selects function-mode; an explicit ``entry``
+    selects process-mode.
+    """
+
+    image: Image
+    memory: MemoryTier = "4g"
+    cpu: float | None = None
+    bundle: Bundle | None = None
+    secrets: tuple[Secret, ...] = field(default_factory=tuple)
+    egress: Egress | None = None
+    env: Mapping[str, str] = field(default_factory=dict)
+    timeout_s: float = 3600.0
+    entry: tuple[str, ...] | None = None
+    code_stage: str | None = None
+    stage_mounts: tuple[StageMount, ...] = field(default_factory=tuple)
+
+    def to_deploy_spec(
+        self,
+        *,
+        project_name: str,
+        entry: Sequence[str],
+        extra_env: Mapping[str, str] | None = None,
+        bundle: Bundle | None = None,
+    ) -> DeploySpec:
+        """Compile this decorator metadata into the deploy IR.
+
+        *entry* is explicit because the two execution modes differ: process-mode
+        uses ``spec.entry``, function-mode the generated runner shim.
+        """
+        env = {**self.env, **(extra_env or {})}
+        return DeploySpec(
+            image=self.image.name,
+            entry=tuple(entry),
+            memory=self.memory,
+            cpu=self.cpu,
+            code_stage=self.code_stage,
+            env={k: str(v) for k, v in env.items() if str(v) != ""},
+            egress=self.egress,
+            secrets=self.secrets,
+            bundle=bundle if bundle is not None else self.bundle,
+            stage_mounts=self.stage_mounts,
+            project_name=project_name,
+            timeout_s=self.timeout_s,
+        )
+
+
+class Function:
+    """A deployable function with ``.remote()`` and ``.local()`` calling conventions.
+
+    Constructed by the ``@app.function`` decorator; do not instantiate directly.
+    ``functools.update_wrapper`` preserves the wrapped function's ``__name__``,
+    ``__doc__``, ``__module__``, ``__qualname__``, ``__annotations__``, and
+    ``__wrapped__`` so type-checkers and ``help()`` see the original signature.
+    """
+
+    def __init__(self, fn: Callable[..., Any], spec: FunctionSpec, app: App) -> None:
+        self._fn = fn
+        self._spec = spec
+        self._app = app
+        functools.update_wrapper(self, fn)
+        self.__name__ = fn.__name__
+
+    # ------------------------------------------------------------------
+    # Properties
+    # ------------------------------------------------------------------
+
+    @property
+    def spec(self) -> FunctionSpec:
+        return self._spec
+
+    @property
+    def app(self) -> App:
+        return self._app
+
+    # ------------------------------------------------------------------
+    # Local calling conventions
+    # ------------------------------------------------------------------
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Call the function locally (no sandbox, no deployment)."""
+        return self._fn(*args, **kwargs)
+
+    def local(self, *args: Any, **kwargs: Any) -> Any:
+        """Alias for ``__call__`` — explicit counterpart to ``.remote()``."""
+        return self._fn(*args, **kwargs)
+
+    # ------------------------------------------------------------------
+    # Remote calling convention
+    # ------------------------------------------------------------------
+
+    async def remote(
+        self,
+        /,
+        *args: Any,
+        on_line: Callable[[str, str], None] | None = None,
+        **kwargs: Any,
+    ) -> DeployResult:
+        """Deploy and run in a Snowflake sandbox; return the `DeployResult`.
+
+        Parameters
+        ----------
+        *args:
+            Positional args are not supported; use keyword arguments only.
+            Raises ``TypeError`` if any positional args are passed (avoids
+            silent kwargs-vs-args confusion when the caller uses
+            ``fn.remote("platform", "JAD-50")`` instead of the expected
+            ``fn.remote(TRIAGE_TEAM="platform", TARGET_JIRA="JAD-50")``).
+        on_line:
+            Optional line callback ``(stream: str, data: str) -> None``
+            for progressive streaming output.  Passed through to
+            ``deploy_spec()``'s ``on_line`` parameter.
+        **kwargs:
+            **Process-mode**: kwargs become env-var overrides passed directly
+            to ``deploy_spec()`` (``TRIAGE_TEAM="platform"`` →
+            ``env_overrides={"TRIAGE_TEAM": "platform"}``).
+
+            **Function-mode**: kwargs are JSON-encoded into
+            ``SANDBOX_CALL_JSON=<json>``; the auto-generated shim calls
+            ``fn(**kwargs)`` inside the container.
+        """
+        self._reject_positional(args, "remote")
+        from snowflake.sandbox.deploy import deploy_spec
+
+        with self._remote_deploy_target(kwargs) as (dspec, source_dir, platform_env):
+            return await deploy_spec(
+                dspec,
+                source_dir=source_dir,
+                on_line=on_line,
+                timeout=self._spec.timeout_s,
+                run_preflight=False,
+                platform_env=platform_env or None,
+            )
+
+    def remote_sync(
+        self,
+        /,
+        *args: Any,
+        on_line: Callable[[str, str], None] | None = None,
+        **kwargs: Any,
+    ) -> DeployResult:
+        """Deploy and run in a Snowflake sandbox; return the `DeployResult` —
+        synchronous counterpart of `remote`.
+
+        Same two execution modes and the same bundle as `remote` — both build
+        their payload through ``_remote_deploy_target`` — but blocks instead of
+        awaiting. Keyword-only, exactly like `remote` (positional args raise
+        ``TypeError``). Safe with no running event loop and never starts one, so
+        it does not deadlock inside a notebook.
+
+        Example:
+            result = fn.remote_sync(TRIAGE_TEAM="platform")
+            print(result.stdout)
+        """
+        self._reject_positional(args, "remote_sync")
+        from snowflake.sandbox.deploy import deploy_spec_sync
+
+        with self._remote_deploy_target(kwargs) as (dspec, source_dir, platform_env):
+            return deploy_spec_sync(
+                dspec,
+                source_dir=source_dir,
+                on_line=on_line,
+                timeout=self._spec.timeout_s,
+                run_preflight=False,
+                platform_env=platform_env or None,
+            )
+
+    async def spawn(
+        self,
+        /,
+        *args: Any,
+        webhook: str | None = None,
+        **kwargs: Any,
+    ) -> AsyncSandbox:
+        """Spawn this function detached and return a live `AsyncSandbox`.
+
+        The fire-and-forget sibling of `remote()`. It assembles the *exact*
+        same bundle/manifest/env that `remote()` does for this function, but
+        routes it through `deploy_async()` instead of ``deploy_spec`` — so
+        you get a `AsyncSandbox` handle back without blocking on completion,
+        and pull the result later with ``await sb.wait()`` (Modal
+        ``FunctionCall.spawn`` parity). The sandbox *is* the handle: persist
+        ``sb.id`` and reattach with `AsyncSandbox.connect()`.
+
+        Parameters
+        ----------
+        *args:
+            Positional args are not supported (same as `remote()`).
+        webhook:
+            Optional URL the in-sandbox job runner best-effort POSTs the result
+            payload to on completion (passed through to `deploy_async()`).
+        **kwargs:
+            **Process-mode**: env-var overrides (``TRIAGE_TEAM="platform"`` →
+            ``env_overrides={"TRIAGE_TEAM": "platform"}``), identical to
+            `remote()`.
+
+            **Function-mode**: JSON-encoded into ``SANDBOX_CALL_JSON``; the same
+            auto-generated ``__app_runner__.py`` shim runs in the container. The
+            shim writes its return value to ``$SANDBOX_RESULT`` as
+            ``{"result": <ret>}``, which the job-runner's ``__SANDBOX_RESULT__``
+            sentinel carries back — so ``(await sb.wait()).result == {"result":
+            <ret>}`` (Modal ``FunctionCall.get`` parity), in addition to the
+            ``{"result": ...}`` line in the logs.
+        """
+        self._reject_positional(args, "spawn")
+        from snowflake.sandbox.jobs import deploy_async
+
+        with self._spawn_deploy_target(kwargs) as (dspec, platform_env):
+            job = await deploy_async(
+                dspec,
+                run_preflight=False,
+                webhook=webhook,
+                extra_platform_env=platform_env or None,
+            )
+
+        # The sandbox *is* the handle — wrap the detached run's id in a live
+        # AsyncSandbox (no extra round-trip); poll with ``await sb.wait()``. Carry the
+        # job's result nonce so wait()'s sentinel parsing authenticates the result
+        # rather than trusting any __SANDBOX_RESULT__ line the workload prints.
+        from snowflake.sandbox.client import AsyncSandbox
+
+        sb = AsyncSandbox(image=self._spec.image.name, _internal=True)
+        sb._id = job.id
+        sb._created = True
+        sb._result_nonce = job.nonce
+        return sb
+
+    def spawn_sync(
+        self,
+        /,
+        *args: Any,
+        webhook: str | None = None,
+        **kwargs: Any,
+    ) -> Sandbox:
+        """Spawn this function detached and return a live `Sandbox` —
+        synchronous counterpart of `spawn`.
+
+        The fire-and-return-a-handle sibling of `remote_sync`: assembles the exact
+        same bundle/manifest/env as `spawn` (both go through
+        ``_spawn_deploy_target``), but routes it through the synchronous
+        detached-deploy core, so you get a synchronous `Sandbox` back without
+        blocking on completion. Pull the result later with ``sb.wait()``; the
+        sandbox *is* the handle — persist ``sb.id`` and reattach with
+        `Sandbox.connect()`. Keyword-only, like `spawn`. Safe with no running
+        event loop and never starts one.
+
+        Example:
+            sb = fn.spawn_sync(TRIAGE_TEAM="platform")
+            result = sb.wait()
+        """
+        self._reject_positional(args, "spawn_sync")
+        from snowflake.sandbox.jobs import deploy_async_sync
+
+        with self._spawn_deploy_target(kwargs) as (dspec, platform_env):
+            job = deploy_async_sync(
+                dspec,
+                run_preflight=False,
+                webhook=webhook,
+                extra_platform_env=platform_env or None,
+            )
+
+        # The sandbox *is* the handle — mirror `spawn`, but wrap the detached run's
+        # id in a synchronous `Sandbox` and carry the job's result nonce so
+        # ``sb.wait()``'s sentinel parsing authenticates the result.
+        from snowflake.sandbox.sync_client import Sandbox
+
+        sb = Sandbox(image=self._spec.image.name, _internal=True)
+        sb._id = job.id
+        sb._created = True
+        sb._result_nonce = job.nonce
+        return sb
+
+    # ------------------------------------------------------------------
+    # Manifest generation
+    # ------------------------------------------------------------------
+
+    def deploy_spec(self, **env_overrides: str) -> DeploySpec:
+        """The `DeploySpec` this function would deploy.
+
+        The reviewable artifact that ``to_manifest()`` used to produce, minus the
+        file format — print it, diff it, or hand it to `deploy.deploy_spec()`.
+        """
+        spec = self._spec
+        entry = list(spec.entry) if spec.entry else ["python", _RUNNER_FILENAME]
+        return spec.to_deploy_spec(
+            project_name=self._app.name,
+            entry=entry,
+            extra_env=env_overrides or None,
+            bundle=spec.bundle or Bundle.from_dir(self._bundle_default_root()),
+        )
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _reject_positional(self, args: tuple[Any, ...], method: str) -> None:
+        """Reject positional call args on a remote invocation.
+
+        ``remote``/``spawn`` and their ``_sync`` twins are keyword-only so
+        ``fn.remote(TEAM="platform")`` can never be silently mistaken for
+        ``fn.remote("platform")``. Raises ``TypeError`` naming the *method* the caller
+        actually used.
+        """
+        if args:
+            raise TypeError(
+                f"{self.__name__}.{method}() does not accept positional arguments; "
+                "use keyword arguments only."
+            )
+
+    @contextmanager
+    def _remote_deploy_target(
+        self, kwargs: dict[str, Any]
+    ) -> Iterator[tuple[DeploySpec, str | Path, dict[str, str]]]:
+        """Yield the ``(DeploySpec, source_dir, platform_env)`` a remote *run* deploys.
+
+        The single source of the two execution modes' payload, shared by `remote`
+        and `remote_sync` so they cannot build different bundles. **Process-mode**
+        turns *kwargs* into env overrides over the user's tree; **function-mode**
+        stages the module + declared bundle, drops the generated runner shim in,
+        and JSON-encodes *kwargs* into ``SANDBOX_CALL_JSON``.
+
+        ``SANDBOX_CALL_JSON`` is a server-reserved ``SANDBOX_`` key, so it is returned
+        in ``platform_env`` (the ``sandbox_env`` channel), NOT baked into the user
+        ``env`` map: the create path rejects a ``SANDBOX_`` key on ``env`` with
+        ``400 ... "SANDBOX_CALL_JSON" is reserved``, which otherwise breaks every
+        function-mode ``remote()``. This mirrors the job runner's ``SANDBOX_JOB_*``
+        vars (see jobs.py). Process-mode returns an empty ``platform_env``. The staged
+        temp dir is removed when the context exits.
+        """
+        spec = self._spec
+        if spec.entry is not None:
+            # Process-mode: kwargs are env overrides; the bundle is the user's tree.
+            yield (
+                spec.to_deploy_spec(
+                    project_name=self._app.name,
+                    entry=list(spec.entry),
+                    extra_env={k: str(v) for k, v in kwargs.items()},
+                ),
+                self._bundle_default_root(),
+                {},
+            )
+        else:
+            # Function-mode: the body runs in the container via a generated shim,
+            # layered over a copy of the user's module + declared bundle so
+            # ``import <module>`` resolves in the container. SANDBOX_CALL_JSON rides
+            # platform_env (reserved SANDBOX_ key), never the user env map.
+            import shutil
+
+            staged = self._stage_function_bundle()
+            try:
+                (staged / _RUNNER_FILENAME).write_text(self._runner_shim())
+                yield (
+                    spec.to_deploy_spec(
+                        project_name=self._app.name,
+                        entry=["python", _RUNNER_FILENAME],
+                        bundle=Bundle.from_dir(str(staged)),
+                    ),
+                    staged,
+                    {"SANDBOX_CALL_JSON": json.dumps(kwargs)},
+                )
+            finally:
+                shutil.rmtree(staged, ignore_errors=True)
+
+    @contextmanager
+    def _spawn_deploy_target(
+        self, kwargs: dict[str, Any]
+    ) -> Iterator[tuple[DeploySpec, dict[str, str]]]:
+        """Yield the ``(DeploySpec, platform_env)`` a detached *spawn* deploys.
+
+        The single source of `spawn` / `spawn_sync`'s payload — the same bundling
+        as `_remote_deploy_target`, except the bundle is baked into the spec (a
+        detached job has no streamed ``source_dir``): process-mode falls back to
+        the declared bundle or the function's source dir; function-mode stages the
+        shim bundle.
+
+        As in `_remote_deploy_target`, function-mode's ``SANDBOX_CALL_JSON`` is a
+        reserved ``SANDBOX_`` key and is returned in ``platform_env`` (passed to
+        ``deploy_async`` as ``extra_platform_env``, merged with the ``SANDBOX_JOB_*``
+        runner vars on the sandbox_env channel), never the user ``env`` map — the
+        create path rejects a ``SANDBOX_`` key there. Process-mode returns an empty
+        ``platform_env``. The staged temp dir is removed when the context exits.
+        """
+        spec = self._spec
+        if spec.entry is not None:
+            # Process-mode (mirrors remote(); kwargs -> env overrides).
+            yield (
+                spec.to_deploy_spec(
+                    project_name=self._app.name,
+                    entry=list(spec.entry),
+                    extra_env={k: str(v) for k, v in kwargs.items()},
+                    bundle=spec.bundle or Bundle.from_dir(self._bundle_default_root()),
+                ),
+                {},
+            )
+        else:
+            # Function-mode (mirrors remote(); same shim + bundling — A1/A2).
+            # SANDBOX_CALL_JSON rides platform_env (reserved SANDBOX_ key).
+            import shutil
+
+            staged = self._stage_function_bundle()
+            try:
+                (staged / _RUNNER_FILENAME).write_text(self._runner_shim())
+                yield (
+                    spec.to_deploy_spec(
+                        project_name=self._app.name,
+                        entry=["python", _RUNNER_FILENAME],
+                        bundle=Bundle.from_dir(str(staged)),
+                    ),
+                    {"SANDBOX_CALL_JSON": json.dumps(kwargs)},
+                )
+            finally:
+                shutil.rmtree(staged, ignore_errors=True)
+
+    def _bundle_default_root(self) -> str:
+        """Resolve the bundle root: explicit ``bundle.root`` or the function's
+        source-file directory (so the user's module is importable)."""
+        if self._spec.bundle is not None:
+            return self._spec.bundle.root
+        try:
+            import inspect
+
+            src = inspect.getsourcefile(self._fn) or inspect.getfile(self._fn)
+            return str(Path(src).resolve().parent)
+        except (TypeError, OSError):
+            return "."
+
+    def _fn_source_file(self) -> Path | None:
+        """The function's own source file, if it has one on disk."""
+        try:
+            import inspect
+
+            src = inspect.getsourcefile(self._fn) or inspect.getfile(self._fn)
+        except (TypeError, OSError):
+            return None
+        if not src:
+            return None
+        p = Path(src).resolve()
+        return p if p.is_file() else None
+
+    def _stage_function_bundle(self) -> Path:
+        """Stage the function-mode bundle in a fresh temp dir and return it.
+
+        Function-mode needs more than the generated ``__app_runner__.py`` in the
+        bundle: without the user's module and the declared ``bundle`` beside it,
+        ``import <module>`` in the container raises ``ModuleNotFoundError``. Mirroring
+        ``warm_session._layer_shim``: copy the function's own module directory (so the
+        module is importable) plus the declared bundle (root tree, or ``include``
+        slice) into a throwaway dir, credential-filtered — the caller then drops
+        the shim in. Returns a ``mkdtemp`` path the caller must ``rmtree``.
+        """
+        import shutil
+
+        from snowflake.sandbox._assemble import (
+            _DEFAULT_EXCLUDES,
+            _assemble_bundle,
+            _copy_tree_into,
+        )
+
+        bundle = self._spec.bundle
+        excludes = list(_DEFAULT_EXCLUDES) + (list(bundle.exclude) if bundle else [])
+        staged = Path(tempfile.mkdtemp(prefix="sandbox_app_"))
+
+        # (1) The function's own module directory — guarantees the module ships and
+        #     is importable inside the container.
+        fn_file = self._fn_source_file()
+        fn_dir = fn_file.parent if fn_file is not None else None
+        if fn_dir is not None and fn_dir.is_dir():
+            _copy_tree_into(fn_dir, staged, excludes)
+
+        # (2) The declared bundle, layered on top. Include globs assemble a monorepo
+        #     slice; otherwise the declared root ships wholesale. Both are
+        #     credential-filtered (via _copy_tree_into / the fixed _assemble_bundle).
+        if bundle is not None:
+            root = Path(bundle.root).expanduser().resolve()
+            if bundle.include:
+                assembled = _assemble_bundle(
+                    root,
+                    {
+                        "root": str(root),
+                        "include": list(bundle.include),
+                        "exclude": list(bundle.exclude),
+                    },
+                )
+                try:
+                    _copy_tree_into(assembled, staged, excludes)
+                finally:
+                    shutil.rmtree(assembled, ignore_errors=True)
+            elif root.is_dir() and root != fn_dir:
+                _copy_tree_into(root, staged, excludes)
+
+        return staged
+
+    def _runner_shim(self) -> str:
+        """Generate the __app_runner__.py shim for function-mode.
+
+        The shim runs the function and emits its return value two ways:
+
+        * **stdout** ``{"result": <ret>}`` — the ``.remote()`` path reads it
+          from ``DeployResult.stdout`` (unchanged, always present).
+        * **``$SANDBOX_RESULT``** ``{"result": <ret>}`` (default
+          ``/sandbox/result.json``) — the typed-result artifact the async
+          job-runner picks up. When this function is ``.spawn()``'d, the
+          job-runner's ``__SANDBOX_RESULT__`` sentinel carries this dict, so
+          ``Job.get().result == {"result": <ret>}`` (Modal ``FunctionCall.get``
+          parity). Writing it is additive: it ``mkdir -p``'s the parent first so
+          it works whether or not the job-runner wrapped the entry, and never
+          affects the ``.remote()`` stdout path. Errors writing the artifact are
+          swallowed so a missing/read-only ``/sandbox`` can't fail the call.
+        """
+        module = getattr(self._fn, "__module__", None) or "__main__"
+        name = getattr(self._fn, "__name__", None) or "fn"
+        fn_file = self._fn_source_file()
+        mod_basename = fn_file.name if fn_file is not None else ""
+        return _function_runner_source(module, name, mod_basename)
+
+
+# Name of the generated daemon shim inside the bundle.
+_RUNNER_FILENAME = "__app_runner__.py"

snowflake/sandbox/image.py

@@ -0,0 +1,46 @@
+"""The ``Image`` reference — a named catalog base image.
+
+Import-cheap (no httpx, no pydantic): a frozen dataclass naming the catalog
+base a sandbox launches on.
+
+    from snowflake.sandbox import Image
+
+    image = Image.from_catalog("sandbox-base")
+
+The layered image *builder* (`pip_install`, `run_commands`, `copy_local_*`,
+`env`, `workdir`, `build`) has been removed while the backing build service is
+not yet generally available; it will return in a later release. Today an
+``Image`` is a reference to a catalog base only.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from snowflake.sandbox.exceptions import SandboxError
+
+__all__ = ["Image"]
+
+
+@dataclass(frozen=True)
+class Image:
+    """A reference to a catalog base image a sandbox launches on.
+
+    ``Image`` is an immutable value — a thin handle to a named base image in the
+    Snowflake image catalog. Pass one to ``Sandbox.create(image=...)`` (a plain
+    catalog-name string works too).
+
+    Example:
+        image = Image.from_catalog("sandbox-base")
+        with Sandbox.create(image=image) as sb:
+            result = sb.exec(["python", "-c", "print(40 + 2)"])
+    """
+
+    name: str
+
+    @staticmethod
+    def from_catalog(name: str) -> Image:
+        """Named constructor — identical to ``Image(name=name)``."""
+        if not name or not name.strip():
+            raise SandboxError("Image.from_catalog requires a non-empty base name")
+        return Image(name=name)