Map over notebook task

[pingsutw]

TL;DR

This PR allows map over notebook task

Type

• Bug Fix
• Feature
• Plugin

Are all requirements met?

• Code completed
• Smoke tested
• Unit tests added
• Code documentation added
• Any pending items have an associated Issue

Complete description

import os
import pathlib
import typing

from flytekit import kwtypes, task, workflow, ImageSpec, map_task
from flytekitplugins.papermill import NotebookTask

new_flytekit = "git+https://github.com/flyteorg/flytekit@022ade936d2fe1eb6a2d705c319cce7d2c66ade6"
new_papermill = "git+https://github.com/flyteorg/flytekit.git@022ade936d2fe1eb6a2d705c319cce7d2c66ade6#subdirectory=plugins/flytekit-papermill"
image = ImageSpec(registry="pingsutw", packages=[new_papermill, new_flytekit], apt_packages=["git"])

nb = NotebookTask(
    name="simple-nb",
    notebook_path=os.path.join(
        pathlib.Path(__file__).parent.absolute(), "nb-simple.ipynb"
    ),
    render_deck=True,
    inputs=kwtypes(a=float),
    outputs=kwtypes(square=float),
    container_image=image,
)


@workflow
def wf(a: float) -> typing.List[float]:
    return map_task(nb)(a=[a])


if __name__ == "__main__":
    print(wf(a=3.14))

Tracking Issue

https://flyte-org.slack.com/archives/CP2HDHKE1/p1684471512753359

Follow-up issue

NA

[kumare3]

You don't need the second check instance is subclass of pythontask

[kumare3]

This seems to special for notebook- we should make it for instance tasks in general

[kumare3]

Are you just removing these?

[pingsutw]

cc @rahul-theorem I know you're using papermill plugin, do you use these paths in the output?

[codecov]

Codecov Report

Merging #1650 (ab0375a) into master (ba70f46) will decrease coverage by 0.01%.
The diff coverage is 25.00%.

@@            Coverage Diff             @@
##           master    #1650      +/-   ##
==========================================
- Coverage   71.02%   71.01%   -0.01%     
==========================================
  Files         336      336              
  Lines       30724    30743      +19     
  Branches     5567     5568       +1     
==========================================
+ Hits        21821    21832      +11     
- Misses       8360     8365       +5     
- Partials      543      546       +3     

Impacted Files	Coverage Δ
flytekit/core/map_task.py	48.07% <25.00%> (-1.60%)	⬇️

... and 10 files with indirect coverage changes

[wild-endeavor]

the only thing is that this will change the interface for the user right? we will need to add a disclaimer to the release notes.

[wild-endeavor]

Suggested change

	if outputs and output_notebooks:
	if output_notebooks:

[wild-endeavor]

what if you just want a notebook? don't want to have to make a fake output.

[pingsutw]

updated it

[wild-endeavor]

what happens if we don't do this?

[pingsutw]

There will be a mismatch between the output type and the downstream task's input type.

[wild-endeavor]

is there a difference between this and isinstance(actual_task, PythonInstanceTask)

[pingsutw]

updated it to isinstance(actual_task, PythonInstanceTask)

[wild-endeavor]

should this be actual_task.instantiated_in? what happens if there are two of the same notebook tasks, named the same, with the same interface, but in two different .py files? will there be confusion?

[pingsutw]

used lhs instead

[wild-endeavor]

can do outputs = outputs or {}