Pickle remote task for Jupyter Notebook Environment

[Mecoli1219]

Tracking issue

flyteorg/flyte#5907
#2579
Original PR: #2618

Why are the changes needed?

Some people may love to develop tasks and workflows in notebooks and run them from the notebook.
Feel free to check out previous Pull Requests. This PR's goal is just to split two functions into separate PRs.

What changes were proposed in this pull request?

Jupyter Notebooks store code in a different format, as a JSON structure, rather than a standard Python script. It is also difficult to locate the code during runtime because Jupyter does not provide direct access to the file path or the exact code location in a standard __file__ variable. This makes it challenging for Flytekit to extract the task code or module when registering a task, as it relies on standard Python modules or scripts for task registration.

Main Idea: Pickle the task function entity during runtime and upload it to a remote location (e.g., S3) when running in a Jupyter Notebook environment.

To implement this, the following modules were modified:

1. New interactive_mode_enabled field inSerializationSettings:
   • When serializing the task in flytekit/tools/translator.py::get_serializable_task, the Python task entity is pickled and uploaded to an S3 bucket. The FastSerializationSettings in SerializationSettings is updated to ensure that the remote task can download the entity from S3 before execution.
2. Update DefaultTaskResolver class in flyte/core/python_auto_container.py:
   • When loading the command arguments, we now check if interactive_mode_enabled is set in SerializationSettings. If so, we add an extra pkl argument to signal that the resolver should download the task entity from S3 in load_task().
3. Update download_distribution function in flytekit/tools/fast_registration.py:
   • The original function extracted tarred files from S3, but for pickled entities, I updated the logic to handle pkl.gz file as a special case that can bypass the distribution format check.
4. Update Promise class:
   • Pickling Promise objects in partial functions (commonly used in map tasks) can lead to an infinite loop. To address this, I implemented custom __getstate__ and __setstate__ methods to handle the pickling process properly as suggested here.
5. New interactive_mode_enabled property in FlyteRemote:
   • Adding this property to FlyteRemote allows it to be specified once during the construction of the FlyteRemote object. Afterward, FlyteRemote will handle pickling and uploading tasks during registration.
   • FlyteRemote also provides an upload_file function to manage how the Python task entity is uploaded to S3.

Map task:

For map_task, we will only pickle the actual_task of the map_task, and upload the pickled file to the remote environment. However, when initializing ArrayNodeMapTask in the remote, it will try to get the name of the array node and call extract_task_module function to determine some information about the task. This will cause errors as the pickled entity is not a Python code, so we couldn't determine the module with inspect.getmodule(f) in the _task_module_from_callable function. As the naming process in initializing ArrayNodeMapTask doesn't really need the module name, I'll simply make the function successful if the task is pickled.

Nevertheless, I believe we need to have a more determined behavior for extracting information about pickled tasks or tasks from the Jupyter Notebook cells.

TODO

• a new integration unit-test testing the Jupyter Notebook environment.
• Add translator unit-test if interactive_mode_enabled => check pickle & no dest-dir
• [0917] Fix Map task

How was this patch tested?

Setup process

• Adding the following codes into the Jupyter cells:

  from flytekit import task, workflow
  from flytekit.configuration import Config
  from flytekit.remote import FlyteRemote
  import nest_asyncio
  nest_asyncio.apply()
  
  remote = FlyteRemote(Config.auto(), default_project="flytesnacks", default_domain="development", interactive_mode_enabled=True)

Simple Task & Workflow & map_task:

@task(container_image=img)
def hello(name: str) -> str:
    return f"Hello {name}"

@task(container_image=img)
def world(pre: str) -> str:
    return f"{pre} Welcome to the world!"

@workflow
def wf(name: str) -> str:
    return world(pre=hello(name=name))

task:

out = remote.execute(hello, inputs={"name": "world"}, version="1.0.1", wait=True)
print(remote.generate_console_url(out))
print(out.outputs)

workflow:

out = remote.execute(wf, inputs={"name": "world"}, version="1.1.1", wait=True)
print(remote.generate_console_url(out))
print(out.outputs)

map_task:

from flytekit import task, workflow, map_task
from functools import partial

@task(container_image=img)
def fn(x: int, y: int) -> int:
    return x + y

@workflow
def workflow_with_maptask(data: list[int], y: int) -> list[int]:
    partial_fn = partial(fn, y=y)
    return map_task(partial_fn)(x=data)

d: list[int] = [1, 2, 3]
out = remote.execute(workflow_with_maptask, inputs={"data": d, "y": 3}, version="1.2.1", wait=True)
print(out.outputs)

Run k8s plugin (pytorch)

Moving the code from flytesnacks/kfpytorch_plugin/pytorch_mnist.py into a Jupyter notebook cell.

Run File Sensor Agent

Follow the link of the doc, and copy the code to the Jupyter Notebook.

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

#2579
#2618

Docs link

[Mecoli1219]

@pingsutw I noticed the map_task wasn’t using the DefaultNotebookTaskResolver, so I’ve updated it. Could you please approve the CI test?

[Mecoli1219]

When I run #2733 (comment), I am getting this error in a Jupyter notebook:

we have to use nest_asyncio for now if we still want to use asyncio.run during registration.

[Mecoli1219]

Update comment:

Map task:

For map_task, we will only pickle the actual_task of the map_task, and upload the pickled file to the remote environment. However, when initializing ArrayNodeMapTask in the remote, it will try to get the name of the array node and call extract_task_module function to determine some information about the task. This will cause errors as the pickled entity is not a Python code, so we couldn't determine the module with inspect.getmodule(f) in the _task_module_from_callable function. As the naming process in initializing ArrayNodeMapTask doesn't really need the module name, I'll simply make the function successful if the task is pickled.

Nevertheless, I believe we need to have a more determined behavior for extracting information about pickled tasks or tasks from the Jupyter Notebook cells.

[thomasjpfan]

we have to use nest_asyncio for now if we still want to use asyncio.run during registration.

I do not want to use nest_asyncio because https://github.com/erdewit/nest_asyncio is not maintained anymore.

I implemented a way to run in a jupyter notebook without nest_asyncio in another project. I'll introduce the technique into FlyteRemote in another PR and test that it works with your PR.

[Mecoli1219]

@thomasjpfan Thank you for your support. I completely understand the concerns around using nest_asyncio. Let me know if there’s anything I need to update or adjust to ensure everything works smoothly.

[kumare3]

@thomasjpfan / @Mecoli1219 i would love to get this merged :). how can we move to get this done?

[Mecoli1219]

@kumare3 The code for interactive mode in the Jupyter notebook is ready. However, the current approach requires using nest_asyncio to run Flyte’s event loop inside Jupyter’s event loop, which is not ideal since nest_asyncio is no longer actively maintained. @thomasjpfan is planning to implement a different approach that avoids using nest_asyncio in a separate PR. Once that’s in place, I can update my PR to ensure compatibility and add the necessary tests.

Let me know if there’s anything else I can assist with to get this merged!

[thomasjpfan]

I opened #2784 to update FlyteRemote to use a event loop on another thread.

I tested this PR with my updates to FlyteRemote and it works! (No longer requires nest_asyncio)

[wild-endeavor]

hey @Mecoli1219 few comments on the pr. first thank you though for taking it across the finish line. had a few questions and wanted to revisit some things though...
shorter comments:

• I see that you added code and a test for pickling promise objects, but could you explain when this is needed?
• In the translator.py... the _update_serialization_settings_for_ipython function - can we make it not mutating? Ideally the serialization settings object is frozen - it's too late to do that, but we can keep new code copy-on-write i think.
• In tracker.py in the debug statements, can we add the values of the variables that will be returned? (like add basename to logger.debug). This file is confusing enough and has caused us enough headaches that more printing is okay.

more structural stuff:

• I feel like the addition of the file_uploader function to the Options dataclass a bit out of place. If you take a look at that object, it's really meant to be used alongside launch plans. Having a function, that makes a network call, that serializes tasks, feels pretty odd.
  • Extending that thought though, the addition of a network call in general changes the pattern for serialization a bit. I don't remember if there are other network calls in translator but the original idea at least was that the translating step is only pure functions, and just converts the Python objects into Flyte protos. Adding the file_uploader now changes that and feels like a bit of a shortcut. Think about how fast registration works in the pyflyte package case. When you fast package workflows, flytekit will create a fast register tar.gz file that contains the code to be injected into the image. This is the same as pyflyte run. But when you do package, this file is not uploaded... it's the responsibility of flytectl register to do the upload (and then inject the upload location into the various task definitions [something I think it doesn't do 100% correctly today actually but that's beside the point]). Could you revisit this part a bit? Not necessarily saying it must change - but I think worth spending time thinking about, and seeing if you can find a better pattern.
  • Along the same line, could you explain what the issue is for dynamic tasks? I read the TODO you left. What backend work needs to happen? The comment suggests adding another uploader to the flytecontext - would definitely like to avoid this if at all possible. Having a callable in there also feels strange.

Thank you!

[Mecoli1219]

@wild-endeavor Thank you for your comments! I apologize for any inconvenience this PR may have caused. Here are some clarifications and further thoughts based on our discussion:

Quick Response

I see that you added code and a test for pickling promise objects, but could you explain when this is needed?

you’re right about the Promise part. We don’t need to pickle the Promise object anymore. The issue was caused by some previously unorganized and faulty code, which has since been cleaned up.
I just tested it with __setstate__ and __getstate__ removed, and it worked. We should remove these two functions.

In tracker.py in the debug statements, can we add the values of the variables that will be returned? (like add basename to logger.debug). This file is confusing enough and has caused us enough headaches that more printing is okay.

I agree that we should add the basename to logger.debug

Discussion

In the translator.py... the _update_serialization_settings_for_ipython function - can we make it not mutating?

I’m not entirely sure what you mean by “non-mutating” in this context. In this code snippet the serialization settings are being updated as well.

Nevertheless, after revisiting the code, I think there is another possible alternative is that we can pickle the file within FlyteRemote::execute_local_task in remote.py and set the fast_serialization during construction of the SerializationSettings if self.interactive_mode_enabled is True.

structural stuff & file_uploader issue

Currently, file_uploader uses the FlyteRemote::upload_file function. If we integrate the uploading logic into FlyteRemote itself as mentioned in the previous section, we wouldn’t need to pass file_uploader into the Options dataclass.

I believe this restructuring could simplify the code, but I’m not sure if it’s too late to implement such changes at this stage. What are your thoughts?

[Mecoli1219]

Update:

I found that pickling the file within FlyteRemote::execute_local_task isn’t sufficient because it only makes remote task execution succeed, while workflow execution still fails. The main issue with moving upload_file outside of get_serializable_task() is that, at the workflow level, the fast_serialization_setting will package the entire workflow (and its tasks) into a single tarred module directory, which is then used as an additional distribution in the sandbox.
In the Jupyter Notebook environment, the main challenge is that we don’t have an actual module or directory that we can compress into a tar.gz file. All we can do is read the entity from memory and dump it via cloudpickle, which raises two concerns:

1. We need to update or add the fast_serialization_setting in SerializationSettings for each task during workflow registration.
2. We can only modify the serialization setting within the get_serialization() function (specifically, get_serialization_task()) in translator.py during workflow registration (as shown in this PR via get_serializable_task()).

I understand the concern around using network calls in translator.py, but I couldn’t find a better solution for now. Nevertheless, I think that using file_uploader in the Options could be addressed by using the FlyteRemote’s upload_file function directly, as this is the only use case currently. However, how to send this function inside get_serialization_task() is also another problem.