NotebookTask

Package: flyteplugins.papermill

A Flyte task that executes a Jupyter notebook via Papermill.

The notebook receives task inputs as parameters (injected into the cell tagged parameters) and produces outputs via record_outputs() called inside the notebook.

Inside notebooks/analyze.ipynb::

from flyteplugins.papermill import record_outputs

result = x + y  # x, y injected by papermill
record_outputs(result=int(result))

You can also call other Flyte tasks from within the notebook — just import and call them as usual::

from my_tasks import expensive_task

intermediate = await expensive_task(data=x)  # submitted to Flyte when running remotely
record_outputs(result=intermediate)

Spark example::

from flyteplugins.papermill import NotebookTask
from flyteplugins.spark import Spark

spark_nb = NotebookTask(
    name="spark_analyze",
    notebook_path="notebooks/spark_analysis.ipynb",
    task_environment=env,
    plugin_config=Spark(spark_conf={...}),
    inputs={"path": str},
    outputs={"count": int},
)

Parameters

class NotebookTask(
    name: str,
    notebook_path: str,
    task_environment: TaskEnvironment,
    plugin_config: Optional[Any],
    inputs: Optional[dict[str, Type]],
    outputs: Optional[dict[str, Type]],
    kernel_name: Optional[str],
    engine_name: Optional[str],
    log_output: bool,
    start_timeout: int,
    execution_timeout: Optional[int],
    report_mode: bool,
    request_save_on_cell_execute: bool,
    progress_bar: bool,
    language: Optional[str],
    engine_kwargs: Optional[dict[str, Any]],
    output_notebooks: bool,
    kwargs: **kwargs,
)

Properties

Methods

aio()

def aio(
    args: *args,
    kwargs: **kwargs,
) -> Coroutine[Any, Any, R] | R

The aio function allows executing “sync” tasks, in an async context. This helps with migrating v1 defined sync tasks to be used within an asyncio parent task. This function will also re-raise exceptions from the underlying task.

@env.task
def my_legacy_task(x: int) -> int:
    return x

@env.task
async def my_new_parent_task(n: int) -> List[int]:
    collect = []
    for x in range(n):
        collect.append(my_legacy_task.aio(x))
    return asyncio.gather(*collect)

config()

def config(
    sctx: SerializationContext,
) -> Dict[str, str]

Returns additional configuration for the task. This is a set of key-value pairs that can be used to configure the task execution environment at runtime. This is usually used by plugins.

container_args()

def container_args(
    serialize_context: SerializationContext,
) -> list[str]

Returns the container args for the task. This is a set of key-value pairs that can be used to configure the task execution environment at runtime. This is usually used by plugins.

custom_config()

def custom_config(
    sctx: SerializationContext,
) -> dict[str, Any]

Returns additional configuration for the task. This is a set of key-value pairs that can be used to configure the task execution environment at runtime. This is usually used by plugins.

data_loading_config()

def data_loading_config(
    sctx: SerializationContext,
) -> DataLoadingConfig

This configuration allows executing raw containers in Flyte using the Flyte CoPilot system Flyte CoPilot, eliminates the needs of sdk inside the container. Any inputs required by the users container are side-loaded in the input_path Any outputs generated by the user container - within output_path are automatically uploaded

execute()

def execute(
    args: *args,
    kwargs: **kwargs,
) -> Any

Execute the notebook within a Flyte task context.

forward()

def forward(
    args: *args,
    kwargs: **kwargs,
) -> Any

Execute the notebook locally (outside of a Flyte run context).

Called when the task is invoked directly as a Python function (e.g. in a test or script) rather than through the Flyte runner. Runs the notebook in-process and returns Python values.

override()

def override(
    short_name: Optional[str],
    resources: Optional[Resources],
    cache: Optional[CacheRequest],
    retries: Union[int, RetryStrategy],
    timeout: Optional[TimeoutType],
    reusable: Union[ReusePolicy, Literal['off'], None],
    env_vars: Optional[Dict[str, str]],
    secrets: Optional[SecretRequest],
    max_inline_io_bytes: int | None,
    pod_template: Optional[Union[str, PodTemplate]],
    queue: Optional[str],
    interruptible: Optional[bool],
    entrypoint: Optional[bool],
    links: Tuple[Link, ...],
    kwargs: **kwargs,
) -> TaskTemplate

Override various parameters of the task template. This allows for dynamic configuration of the task when it is called, such as changing the image, resources, cache policy, etc.

Returns: A new TaskTemplate instance with the overridden parameters.

post()

def post(
    return_vals: Any,
) -> Any

This is the postexecute function that will be called after the task is executed

pre()

def pre(
    args,
    kwargs,
) -> Dict[str, Any]

This is the preexecute function that will be called before the task is executed

sql()

def sql(
    sctx: SerializationContext,
) -> Optional[str]

Returns the SQL for the task. This is a set of key-value pairs that can be used to configure the task execution environment at runtime. This is usually used by plugins.