flyteplugins-papermill 2.1.8__py3-none-any.whl

flyteplugins/papermill/__init__.py

@@ -0,0 +1,15 @@
+__all__ = [
+    "NotebookTask",
+    "load_dataframe",
+    "load_dir",
+    "load_file",
+    "record_outputs",
+]
+
+from flyteplugins.papermill.notebook import (
+    load_dataframe,
+    load_dir,
+    load_file,
+    record_outputs,
+)
+from flyteplugins.papermill.task import NotebookTask

flyteplugins/papermill/_setup.py

@@ -0,0 +1,97 @@
+"""Notebook kernel setup — called from the injected setup cell.
+
+This module is imported inside the notebook kernel subprocess. Keep imports
+at function scope so the module itself has no side-effects on import.
+"""
+
+from __future__ import annotations
+
+
+def initialize_context() -> None:
+    """Initialize the Flyte runtime context inside a notebook kernel.
+
+    Reads the serialized task context from the ``_FLYTE_NB_CTX`` environment
+    variable (set by the parent task runner before launching papermill) and
+    reconstructs a ``TaskContext`` + ContextVar so that Flyte APIs work
+    normally inside the notebook.
+
+    Does nothing if the environment variable is not set.
+    """
+    import json
+    import os
+
+    raw = os.environ.get("_FLYTE_NB_CTX")
+    if not raw:
+        return
+
+    try:
+        data = json.loads(raw)
+
+        import flyte.report as _report
+        from flyte._context import Context, ContextData, root_context_var
+        from flyte._internal.imagebuild.image_builder import ImageCache
+        from flyte.models import ActionID, CodeBundle, RawDataPath, TaskContext
+
+        if data.get("mode") == "local":
+            # Local mode: no controller or remote connection needed.
+            # Set _init_config directly to avoid @syncify running in a
+            # background thread, which can cause visibility issues with
+            # module-level globals in the kernel's main thread.
+            import flyte._initialize as _init_mod
+
+            with _init_mod._init_lock:
+                if _init_mod._init_config is None:
+                    from pathlib import Path
+
+                    from flyte._initialize import _InitConfig
+
+                    _init_mod._init_config = _InitConfig(root_dir=Path.cwd())
+        else:
+            # Remote mode: use the same init + controller pattern as runtime.py.
+            from flyte._initialize import init_in_cluster
+            from flyte._internal.controllers import create_controller
+
+            controller_kwargs = init_in_cluster()
+            create_controller(ct="remote", **controller_kwargs)
+
+        action = ActionID(
+            name=data["action_name"],
+            run_name=data["run_name"],
+            project=data["project"],
+            domain=data["domain"],
+            org=data["org"],
+        )
+
+        cb = None
+        if "code_bundle" in data:
+            cbd = data["code_bundle"]
+            cb = CodeBundle(
+                tgz=cbd.get("tgz"),
+                pkl=cbd.get("pkl"),
+                destination=cbd.get("destination", "."),
+                computed_version=cbd.get("computed_version", ""),
+            )
+
+        ic = None
+        if "image_cache" in data:
+            ic = ImageCache.from_transport(data["image_cache"])
+
+        tctx = TaskContext(
+            action=action,
+            version=data["version"],
+            output_path=data["output_path"],
+            run_base_dir=data["run_base_dir"],
+            raw_data_path=RawDataPath(path=data["raw_data_path"]),
+            mode=data["mode"],
+            interactive_mode=data.get("interactive_mode", False),
+            report=_report.Report(name=action.name),
+            code_bundle=cb,
+            compiled_image_cache=ic,
+        )
+        root_context_var.set(Context(data=ContextData(task_context=tctx)))
+        print("[flyte-notebook] Context initialized successfully")
+    except Exception as err:
+        import traceback
+
+        print(f"[flyte-notebook] WARNING: Failed to initialize context: {err}")
+        traceback.print_exc()

flyteplugins/papermill/notebook.py

@@ -0,0 +1,122 @@
+"""Helpers for use inside Jupyter notebooks executed by NotebookTask.
+
+These functions are meant to be called from within a notebook that is
+being run as a Flyte task via papermill.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def record_outputs(**kwargs: Any) -> str:
+    """Record output values from a notebook for use by downstream Flyte tasks.
+
+    Call this as the **last expression** in a cell tagged ``"outputs"``.
+    The returned protobuf text is captured by Jupyter as the cell output
+    and later extracted by ``NotebookTask``.
+
+    Values are serialized as Flyte Literals, so any type supported by
+    Flyte's type system works — primitives, ``File``, ``Dir``,
+    ``DataFrame``, dataclasses, etc.
+
+    Example (cell tagged ``"outputs"``)::
+
+        from flyteplugins.papermill import record_outputs
+
+        record_outputs(result=42, summary="done")
+
+    Args:
+        **kwargs: Output name/value pairs.
+
+    Returns:
+        Protobuf text representation of a ``LiteralMap``. Jupyter captures
+        this as the cell's text/plain output.
+    """
+    from flyte.syncify import syncify
+    from flyte.types import TypeEngine
+    from google.protobuf import text_format
+
+    @syncify
+    async def _to_literal_map(values: dict[str, Any]):
+        from flyteidl2.core.literals_pb2 import LiteralMap
+
+        literals = {}
+        for name, val in values.items():
+            py_type = type(val)
+            lit_type = TypeEngine.to_literal_type(py_type)
+            literals[name] = await TypeEngine.to_literal(val, py_type, lit_type)
+        return LiteralMap(literals=literals)
+
+    literal_map = _to_literal_map(kwargs)
+    return text_format.MessageToString(literal_map)
+
+
+def load_file(path: str):
+    """Load a ``flyte.io.File`` from a serialized path inside a notebook.
+
+    When a ``File`` is passed as an input to a ``NotebookTask``, it is
+    serialized to its remote path string for papermill injection.  Use
+    this helper to reconstruct the ``File`` object inside the notebook::
+
+        from flyteplugins.papermill import load_file
+
+        f = load_file(my_file_path)  # my_file_path injected by papermill
+        with f.open_sync() as fh:
+            data = fh.read()
+
+    Args:
+        path: The remote path string (injected as a papermill parameter).
+
+    Returns:
+        A ``flyte.io.File`` instance pointing at the remote path.
+    """
+    from flyte.io import File
+
+    return File(path=path)
+
+
+def load_dir(path: str):
+    """Load a ``flyte.io.Dir`` from a serialized path inside a notebook.
+
+    When a ``Dir`` is passed as an input to a ``NotebookTask``, it is
+    serialized to its remote path string.  Use this helper to
+    reconstruct it::
+
+        from flyteplugins.papermill import load_dir
+
+        d = load_dir(my_dir_path)
+
+    Args:
+        path: The remote path string (injected as a papermill parameter).
+
+    Returns:
+        A ``flyte.io.Dir`` instance pointing at the remote path.
+    """
+    from flyte.io import Dir
+
+    return Dir(path=path)
+
+
+def load_dataframe(uri: str, fmt: str = "parquet"):
+    """Load a ``flyte.io.DataFrame`` from a serialized URI inside a notebook.
+
+    When a ``DataFrame`` is passed as an input to a ``NotebookTask``, it is
+    serialized to its remote URI for papermill injection.  Use this helper
+    to reconstruct it::
+
+        from flyteplugins.papermill import load_dataframe
+
+        df = load_dataframe(my_df_uri)
+        pandas_df = df.all()  # materializes as pandas DataFrame
+
+    Args:
+        uri: The remote URI string (injected as a papermill parameter).
+        fmt: The storage format (default ``"parquet"``).
+
+    Returns:
+        A ``flyte.io.DataFrame`` instance pointing at the remote URI.
+    """
+    from flyte.io import DataFrame
+
+    return DataFrame(uri=uri, format=fmt)

flyteplugins/papermill/resolver.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import json
+import os
+import pathlib
+import sys
+
+from flyte._task import TaskTemplate
+
+
+class NotebookTaskResolver:
+    """Resolver for NotebookTask instances.
+
+    Serializes all task state (notebook path, type schemas, execution config)
+    into the loader args at serialization time so the task can be reconstructed
+    in the container without importing the user's module. This lets NotebookTask
+    be defined inline (e.g. inside a task) rather than only at module level.
+
+    Loader args format (key-value pairs):
+        notebook  <relative-or-absolute notebook path>
+        name      <task name>
+        input-schema   <JSON: {field: LiteralType dict}>
+        output-schema  <JSON: {field: LiteralType dict}>
+        config         <JSON: execution params>
+    """
+
+    @property
+    def import_path(self) -> str:
+        return "flyteplugins.papermill.resolver.NotebookTaskResolver"
+
+    def load_task(self, loader_args: list[str]) -> TaskTemplate:
+        # Ensure all IO type transformers are registered before attempting
+        # guess_python_type — they are registered as side-effects of import.
+        import flyte.io  # noqa: F401
+        from flyte.types import TypeEngine
+        from flyteidl2.core.types_pb2 import LiteralType
+        from google.protobuf import json_format
+
+        from flyteplugins.papermill.task import NotebookTask
+
+        # Parse flat key-value list
+        it = iter(loader_args)
+        args_dict: dict[str, str] = {}
+        for key in it:
+            try:
+                args_dict[key] = next(it)
+            except StopIteration:
+                raise ValueError(f"Odd number of loader args — missing value for key '{key}'")
+
+        notebook_path = args_dict["notebook"]
+        name = args_dict["name"]
+
+        # Relative paths are stored relative to the bundle root (root_dir at
+        # serialization time). At execution time, the bundle extraction dir is
+        # prepended to sys.path by download_code_bundle, so we search there.
+        if not os.path.isabs(notebook_path):
+            for p in sys.path:
+                candidate = os.path.abspath(os.path.join(p, notebook_path))
+                if os.path.exists(candidate):
+                    notebook_path = candidate
+                    break
+            else:
+                # CWD fallback: covers the common destination="." case where
+                # the bundle is extracted into the current working directory.
+                notebook_path = os.path.abspath(notebook_path)
+
+        def _schema_to_types(schema_json: str) -> dict | None:
+            schema = json.loads(schema_json)
+            if not schema:
+                return None
+            result: dict = {}
+            for field_name, lt_dict in schema.items():
+                lt = LiteralType()
+                json_format.ParseDict(lt_dict, lt)
+                result[field_name] = TypeEngine.guess_python_type(lt)
+            return result
+
+        inputs = _schema_to_types(args_dict.get("input-schema", "{}"))
+        outputs = _schema_to_types(args_dict.get("output-schema", "{}"))
+        config: dict = json.loads(args_dict.get("config", "{}"))
+
+        return NotebookTask(
+            name=name,
+            notebook_path=notebook_path,
+            task_environment=None,
+            inputs=inputs,
+            outputs=outputs,
+            **config,
+        )
+
+    def loader_args(self, task: TaskTemplate, root_dir: pathlib.Path | None) -> list[str]:
+        from flyte.types import TypeEngine
+        from google.protobuf import json_format
+
+        from flyteplugins.papermill.task import NotebookTask
+
+        if not isinstance(task, NotebookTask):
+            raise TypeError(f"NotebookTaskResolver only handles NotebookTask, got {type(task)}")
+
+        # If the notebook is inside the bundle root, store a relative path so
+        # it resolves correctly wherever the bundle is extracted in the container.
+        # If outside root_dir (or no root_dir), preserve the path exactly as the
+        # user wrote it: absolute paths are expected to exist in the container
+        # image; relative paths remain relative and must be resolvable from CWD
+        # or sys.path at execution time.
+        nb_path = pathlib.Path(task._resolved_notebook_path)
+        if root_dir is not None:
+            try:
+                notebook_arg = str(nb_path.relative_to(pathlib.Path(root_dir).resolve()))
+            except ValueError:
+                notebook_arg = str(nb_path)
+        else:
+            # No bundle root: use the original path as the user wrote it.
+            notebook_arg = task.notebook_path
+
+        def _types_to_schema(types_dict: dict) -> dict:
+            schema: dict = {}
+            for field_name, typ in types_dict.items():
+                lt = TypeEngine.to_literal_type(typ)
+                schema[field_name] = json_format.MessageToDict(lt)
+            return schema
+
+        # Inputs
+        input_types = {name: typ for name, (typ, _) in (task.interface.inputs or {}).items()}
+        input_schema = _types_to_schema(input_types)
+
+        # Outputs: strip the auto-added notebook File outputs so the reconstructed
+        # task can re-add them via output_notebooks=True in the config.
+        skip_outputs = {"output_notebook", "output_notebook_executed"} if task.output_notebooks else set()
+        output_types = {name: typ for name, typ in (task.interface.outputs or {}).items() if name not in skip_outputs}
+        output_schema = _types_to_schema(output_types)
+
+        # Execution config (plugin_config is serialization-only; not needed at execution time)
+        config: dict = {
+            "log_output": task.log_output,
+            "start_timeout": task.start_timeout,
+            "report_mode": task.report_mode,
+            "request_save_on_cell_execute": task.request_save_on_cell_execute,
+            "progress_bar": task.progress_bar,
+            "output_notebooks": task.output_notebooks,
+        }
+        if task.kernel_name is not None:
+            config["kernel_name"] = task.kernel_name
+        if task.engine_name is not None:
+            config["engine_name"] = task.engine_name
+        if task.execution_timeout is not None:
+            config["execution_timeout"] = task.execution_timeout
+        if task.language is not None:
+            config["language"] = task.language
+        if task.engine_kwargs:
+            config["engine_kwargs"] = task.engine_kwargs
+
+        return [
+            "notebook",
+            notebook_arg,
+            "name",
+            task.name,
+            "input-schema",
+            json.dumps(input_schema),
+            "output-schema",
+            json.dumps(output_schema),
+            "config",
+            json.dumps(config),
+        ]