Make FlyteFile and FlyteDirectory pickleable #3030

cosmicBboy · 2025-01-03T22:22:33Z

Tracking issue

Why are the changes needed?

Because doing distributed training with the Elastic plugin fails when trying to run a task with FlyteDirectory or FlyteFile inputs.

What changes were proposed in this pull request?

pulls out the _downloader as a module-level function that's made into a partial when converting FlyteFile/Directory literals into python values.

How was this patch tested?

Unit tests were updated to test for round-trip pickling.

Check all the applicable boxes

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

Related PRs

Docs link

Summary by Bito

Implementation of pickleable downloaders for FlyteFile and FlyteDirectory by replacing lambda functions with functools.partial. This critical fix enables distributed training with the Elastic plugin by ensuring proper serialization of file handling objects. Changes include restructuring downloader implementation, moving FileAccessProvider import, and streamlining file access through context object.

Unit tests added: True

Estimated effort to review (1-5, lower is better): 2

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

flyte-bot · 2025-01-03T22:22:46Z

Code Review Agent Run #639caa

Actionable Suggestions - 3

tests/flytekit/unit/core/test_flyte_directory.py - 1
- Consider adding more pickle test assertions · Line 445-448
flytekit/types/directory/types.py - 2
- Consider adding error handling for downloader · Line 687-688
- Consider adding return type annotation · Line 687-687

Review Details

Files reviewed - 4 · Commit Range: 8a0804e..2060962
- flytekit/types/directory/types.py
- flytekit/types/file/file.py
- tests/flytekit/unit/core/test_flyte_directory.py
- tests/flytekit/unit/core/test_flyte_file.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

flyte-bot · 2025-01-03T22:29:04Z

Changelist by Bito

This pull request implements the following key changes.

Key Change	Files Impacted
Bug Fix - Make File Handling Objects Pickleable	- `types.py` - Replace lambda downloader with functools.partial for serialization - `file.py` - Refactor downloader implementation to use functools.partial - `test_flyte_directory.py` - Add pickle serialization tests for FlyteDirectory - `test_flyte_file.py` - Add pickle serialization tests for FlyteFile

flyte-bot · 2025-01-03T22:29:05Z

tests/flytekit/unit/core/test_flyte_directory.py

+    # test round trip pickling
+    pickled_input = pickle.dumps(downstream_input)
+    unpickled_input = pickle.loads(pickled_input)
+    assert downstream_input == unpickled_input


Consider adding more pickle test assertions

Consider adding more assertions to verify the pickled/unpickled FlyteDirectory object's properties like uri and other attributes are preserved correctly after deserialization.

Code suggestion

Check the AI-generated fix before applying

Suggested change

# test round trip pickling

pickled_input = pickle.dumps(downstream_input)

unpickled_input = pickle.loads(pickled_input)

assert downstream_input == unpickled_input

# test round trip pickling

pickled_input = pickle.dumps(downstream_input)

unpickled_input = pickle.loads(pickled_input)

assert downstream_input == unpickled_input

assert unpickled_input.uri == "s3://sample-path/directory"

Code Review Run #639caa

Is this a valid issue, or was it incorrectly flagged by the Agent?

it was incorrectly flagged

flyte-bot · 2025-01-03T22:29:06Z

flytekit/types/directory/types.py

+def _flyte_directory_downloader(file_access_provider: FileAccessProvider, uri: str, local_folder: str, batch_size: int):
+    return file_access_provider.get_data(uri, local_folder, is_multipart=True, batch_size=batch_size)


Consider adding error handling for downloader

Consider adding error handling for the case where file_access_provider.get_data() fails. The function should handle potential exceptions that may occur during data retrieval. A similar issue was also found in flytekit/types/file/file.py (line 717-718).

Code suggestion

Check the AI-generated fix before applying

Suggested change

def _flyte_directory_downloader(file_access_provider: FileAccessProvider, uri: str, local_folder: str, batch_size: int):

return file_access_provider.get_data(uri, local_folder, is_multipart=True, batch_size=batch_size)

def _flyte_directory_downloader(file_access_provider: FileAccessProvider, uri: str, local_folder: str, batch_size: int):

try:

return file_access_provider.get_data(uri, local_folder, is_multipart=True, batch_size=batch_size)

except Exception as e:

raise TypeTransformerFailedError(

f"Failed to download directory from {uri} to {local_folder}: {str(e)}"

)

Code Review Run #639caa

Is this a valid issue, or was it incorrectly flagged by the Agent?

it was incorrectly flagged

flyte-bot · 2025-01-03T22:29:07Z

flytekit/types/directory/types.py

@@ -683,4 +684,8 @@ def guess_python_type(self, literal_type: LiteralType) -> typing.Type[FlyteDirec
        raise ValueError(f"Transformer {self} cannot reverse {literal_type}")


+def _flyte_directory_downloader(file_access_provider: FileAccessProvider, uri: str, local_folder: str, batch_size: int):


Consider adding return type annotation

Consider adding type hints for the return value of _flyte_directory_downloader. The function appears to return a result from get_data() but its return type is not specified.

Code suggestion

Check the AI-generated fix before applying

Suggested change

def _flyte_directory_downloader(file_access_provider: FileAccessProvider, uri: str, local_folder: str, batch_size: int):

def _flyte_directory_downloader(file_access_provider: FileAccessProvider, uri: str, local_folder: str, batch_size: int) -> Any:

Code Review Run #639caa

Is this a valid issue, or was it incorrectly flagged by the Agent?

it was incorrectly flagged

codecov · 2025-01-04T00:29:42Z

Codecov Report

Attention: Patch coverage is 60.00000% with 4 lines in your changes missing coverage. Please review.

Project coverage is 51.81%. Comparing base (f3996f6) to head (6d17e5c).
Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
flytekit/types/directory/types.py	60.00%	2 Missing ⚠️
flytekit/types/file/file.py	60.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #3030      +/-   ##
==========================================
+ Coverage   51.28%   51.81%   +0.52%     
==========================================
  Files         201      201              
  Lines       21281    21289       +8     
  Branches     2734     2734              
==========================================
+ Hits        10915    11031     +116     
+ Misses       9767     9739      -28     
+ Partials      599      519      -80

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

flyte-bot · 2025-01-04T01:28:58Z

Code Review Agent Run #f65ef4

Actionable Suggestions - 1

flytekit/types/file/file.py - 1
- Consider naming the downloader lambda function · Line 741-743

Review Details

Files reviewed - 10 · Commit Range: 2060962..6d17e5c
- flytekit/clis/sdk_in_container/run.py
- flytekit/types/file/file.py
- plugins/flytekit-inference/flytekitplugins/inference/__init__.py
- plugins/flytekit-inference/flytekitplugins/inference/vllm/serve.py
- plugins/flytekit-inference/setup.py
- plugins/flytekit-inference/tests/test_vllm.py
- tests/flytekit/integration/remote/test_remote.py
- tests/flytekit/integration/remote/utils.py
- tests/flytekit/integration/remote/workflows/basic/flytefile.py
- tests/flytekit/integration/remote/workflows/basic/pydantic_wf.py
Files skipped - 1
- plugins/flytekit-inference/README.md - Reason: Filter setting
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

flyte-bot · 2025-01-04T01:38:45Z

flytekit/types/file/file.py

        ff = FlyteFile.__class_getitem__(expected_format)(
            path=local_path, downloader=lambda: self.downloader(ctx=ctx, remote_path=uri, local_path=local_path)
        )


Consider naming the downloader lambda function

Consider using a more descriptive lambda function name instead of an anonymous lambda. A named function would better document the purpose of the downloader callback.

Code suggestion

Check the AI-generated fix before applying

Suggested change

ff = FlyteFile.__class_getitem__(expected_format)(

path=local_path, downloader=lambda: self.downloader(ctx=ctx, remote_path=uri, local_path=local_path)

)

download_callback = partial(self.downloader, ctx=ctx, remote_path=uri, local_path=local_path)

ff = FlyteFile.__class_getitem__(expected_format)(

path=local_path, downloader=download_callback

)

Code Review Run #f65ef4

Is this a valid issue, or was it incorrectly flagged by the Agent?

it was incorrectly flagged

thomasjpfan · 2025-01-05T15:16:23Z

flytekit/types/directory/types.py

@@ -665,8 +667,7 @@ async def async_to_python_value(

        batch_size = get_batch_size(expected_python_type)

-        def _downloader():
-            return ctx.file_access.get_data(uri, local_folder, is_multipart=True, batch_size=batch_size)
+        _downloader = partial(_flyte_directory_downloader, ctx.file_access, uri, local_folder, batch_size)


Looking at this again, I think we can partial ctx.file_access itself and not need to define _flyte_directory_downloader.

_downloader = partial(ctx.file_access.get_data, url, local_folder, is_multipart=True, batch_size=batch_size)

thomasjpfan · 2025-01-05T15:17:03Z

flytekit/types/file/file.py

@@ -732,12 +734,14 @@ async def async_to_python_value(

        # For the remote case, return an FlyteFile object that can download
        local_path = ctx.file_access.get_random_local_path(uri)
+
+        _downloader = partial(_flyte_file_downloader, ctx.file_access, uri, local_path)


Same here regarding partialing ctx.file_access.get_data directly.

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

flyte-bot · 2025-01-06T21:39:19Z

Code Review Agent Run #be429e

Actionable Suggestions - 2

flytekit/types/directory/types.py - 2
- Consider adding error handling for downloads · Line 687-688
- Missing return type hint in function · Line 687-687

Review Details

Files reviewed - 4 · Commit Range: 8a0804e..6adf14b
- flytekit/types/directory/types.py
- flytekit/types/file/file.py
- tests/flytekit/unit/core/test_flyte_directory.py
- tests/flytekit/unit/core/test_flyte_file.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

flyte-bot · 2025-01-06T21:44:04Z

flytekit/types/directory/types.py

+
 TypeEngine.register(FlyteDirToMultipartBlobTransformer())


Consider adding error handling for downloads

Consider adding error handling for the case where file_access_provider.get_data() fails. The function should handle potential exceptions and provide appropriate error messages.

Code suggestion

Check the AI-generated fix before applying

Suggested change

TypeEngine.register(FlyteDirToMultipartBlobTransformer())

def _flyte_directory_downloader(file_access_provider: FileAccessProvider, uri: str, local_folder: str, batch_size: int):

try:

return file_access_provider.get_data(uri, local_folder, is_multipart=True, batch_size=batch_size)

except Exception as e:

raise TypeTransformerFailedError(

f"Failed to download directory from {uri} to {local_folder}: {str(e)}"

)

Code Review Run #be429e

Is this a valid issue, or was it incorrectly flagged by the Agent?

it was incorrectly flagged

flyte-bot · 2025-01-06T21:44:05Z

flytekit/types/directory/types.py

+def _flyte_directory_downloader(file_access_provider: FileAccessProvider, uri: str, local_folder: str, batch_size: int):
+    return file_access_provider.get_data(uri, local_folder, is_multipart=True, batch_size=batch_size)
+
+


Missing return type hint in function

Consider adding type hints for the return value of _flyte_directory_downloader function to improve code clarity and maintainability.

Code suggestion

Check the AI-generated fix before applying

Suggested change

def _flyte_directory_downloader(file_access_provider: FileAccessProvider, uri: str, local_folder: str, batch_size: int) -> Any:

Code Review Run #be429e

Is this a valid issue, or was it incorrectly flagged by the Agent?

it was incorrectly flagged

flyte-bot · 2025-01-06T22:49:18Z

Code Review Agent Run #5c2427

Actionable Suggestions - 0

Review Details

Files reviewed - 2 · Commit Range: 6adf14b..37f61e2
- flytekit/types/directory/types.py
- flytekit/types/file/file.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

flyte-bot · 2025-01-07T16:06:19Z

Code Review Agent Run #d95336

Actionable Suggestions - 1

flytekit/types/file/file.py - 1
- Consider adding type hints to downloader · Line 742-742

Review Details

Files reviewed - 2 · Commit Range: 37f61e2..a60450b
- flytekit/types/directory/types.py
- flytekit/types/file/file.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

flyte-bot · 2025-01-07T16:09:12Z

flytekit/types/file/file.py

        expected_format = FlyteFilePathTransformer.get_format(expected_python_type)
        ff = FlyteFile.__class_getitem__(expected_format)(
-            path=local_path, downloader=lambda: self.downloader(ctx=ctx, remote_path=uri, local_path=local_path)
+            path=local_path,
+            downloader=partial(self.downloader, ctx.file_access, remote_path=uri, local_path=local_path),


Consider adding type hints to downloader

Consider updating the downloader partial function call to include type hints for better code maintainability and IDE support. The ctx.file_access parameter type could be explicitly specified as FileAccessProvider.

Code suggestion

Check the AI-generated fix before applying

Suggested change

downloader=partial(self.downloader, ctx.file_access, remote_path=uri, local_path=local_path),

downloader=partial(self.downloader, typing.cast(FileAccessProvider, ctx.file_access), remote_path=uri, local_path=local_path),

Code Review Run #d95336

Is this a valid issue, or was it incorrectly flagged by the Agent?

it was incorrectly flagged

thomasjpfan · 2025-01-07T16:58:36Z

flytekit/types/file/file.py

        return ff

    @staticmethod
    def downloader(
-        ctx: FlyteContext, remote_path: typing.Union[str, os.PathLike], local_path: typing.Union[str, os.PathLike]
+        file_access_provider: FileAccessProvider,


This breaks backward compatibility. Is it safe to change the signature here?

As far as I can tell this staticmethod was added in a PR 4 days ago: https://github.com/flyteorg/flytekit/pull/2991/files

And the only call site was in this module itself

I rather not have this be part of the public API. Can this be renamed to _remote_downloader?

we can also revert this to the older implementation with the pure partial

That works for me.

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

flyte-bot · 2025-01-08T22:04:35Z

Code Review Agent Run #0644db

Actionable Suggestions - 0

Review Details

Files reviewed - 1 · Commit Range: a60450b..d0b1d3c
- flytekit/types/file/file.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

* make _downloader function in FlyteFile/Directory pickleable Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com> * make FlyteFile and Directory pickleable Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com> * remove unnecessary helper functions Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com> * fix lint Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com> * use partials instead of lambda Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com> * fix lint Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com> * remove unneeded helper function Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com> * update FlyteFilePathTransformer.downloader method Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com> * remove downloader staticmethod Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com> * fix lint Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com> --------- Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com> Signed-off-by: Shuying Liang <shuying.liang@gmail.com>

cosmicBboy added 2 commits January 2, 2025 16:34

make _downloader function in FlyteFile/Directory pickleable

8a0804e

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

make FlyteFile and Directory pickleable

2060962

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

cosmicBboy requested review from wild-endeavor, kumare3, eapolinario, pingsutw, samhita-alla, thomasjpfan and Future-Outlier as code owners January 3, 2025 22:22

flyte-bot reviewed Jan 3, 2025

View reviewed changes

flyte-bot reviewed Jan 4, 2025

View reviewed changes

thomasjpfan reviewed Jan 5, 2025

View reviewed changes

remove unnecessary helper functions

dddde88

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

cosmicBboy force-pushed the bugfix/pickleable-downloader branch from 6d17e5c to dddde88 Compare January 6, 2025 20:36

cosmicBboy added 2 commits January 6, 2025 15:38

Merge branch 'master' into bugfix/pickleable-downloader

cc49802

fix lint

6adf14b

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

use partials instead of lambda

37f61e2

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

flyte-bot reviewed Jan 6, 2025

View reviewed changes

cosmicBboy added 3 commits January 7, 2025 10:03

fix lint

6152ce4

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

remove unneeded helper function

3aaf9ac

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

update FlyteFilePathTransformer.downloader method

a60450b

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

flyte-bot reviewed Jan 7, 2025

View reviewed changes

thomasjpfan reviewed Jan 7, 2025

View reviewed changes

cosmicBboy added 2 commits January 8, 2025 16:03

remove downloader staticmethod

2e2e5d9

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

fix lint

d0b1d3c

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

thomasjpfan approved these changes Jan 8, 2025

View reviewed changes

cosmicBboy changed the title ~~Bugfix/pickleable downloader~~ Make FlyteFile and FlyteDirectory pickleable Jan 8, 2025

cosmicBboy merged commit 556595a into master Jan 8, 2025
104 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Make FlyteFile and FlyteDirectory pickleable #3030

Make FlyteFile and FlyteDirectory pickleable #3030

cosmicBboy commented Jan 3, 2025 •

edited by flyte-bot

Loading

flyte-bot commented Jan 3, 2025 •

edited

Loading

Code Review Agent Run #639caa

flyte-bot commented Jan 3, 2025 •

edited

Loading

Changelist by Bito

flyte-bot Jan 3, 2025

flyte-bot Jan 3, 2025

flyte-bot Jan 3, 2025

codecov bot commented Jan 4, 2025 •

edited

Loading

flyte-bot commented Jan 4, 2025 •

edited

Loading

Code Review Agent Run #f65ef4

flyte-bot Jan 4, 2025

thomasjpfan Jan 5, 2025

thomasjpfan Jan 5, 2025

flyte-bot commented Jan 6, 2025 •

edited

Loading

Code Review Agent Run #be429e

flyte-bot Jan 6, 2025

flyte-bot Jan 6, 2025

flyte-bot commented Jan 6, 2025 •

edited

Loading

Code Review Agent Run #5c2427

flyte-bot commented Jan 7, 2025 •

edited

Loading

Code Review Agent Run #d95336

flyte-bot Jan 7, 2025

thomasjpfan Jan 7, 2025

cosmicBboy Jan 7, 2025

thomasjpfan Jan 7, 2025

cosmicBboy Jan 7, 2025

thomasjpfan Jan 7, 2025

flyte-bot commented Jan 8, 2025 •

edited

Loading

Code Review Agent Run #0644db

		def _flyte_directory_downloader(file_access_provider: FileAccessProvider, uri: str, local_folder: str, batch_size: int):
		return file_access_provider.get_data(uri, local_folder, is_multipart=True, batch_size=batch_size)

		@@ -683,4 +684,8 @@ def guess_python_type(self, literal_type: LiteralType) -> typing.Type[FlyteDirec
		raise ValueError(f"Transformer {self} cannot reverse {literal_type}")


		def _flyte_directory_downloader(file_access_provider: FileAccessProvider, uri: str, local_folder: str, batch_size: int):

	downloader=partial(self.downloader, ctx.file_access, remote_path=uri, local_path=local_path),
	downloader=partial(self.downloader, typing.cast(FileAccessProvider, ctx.file_access), remote_path=uri, local_path=local_path),

Make FlyteFile and FlyteDirectory pickleable #3030

Make FlyteFile and FlyteDirectory pickleable #3030

Conversation

cosmicBboy commented Jan 3, 2025 • edited by flyte-bot Loading

Tracking issue

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Check all the applicable boxes

Related PRs

Docs link

Summary by Bito

flyte-bot commented Jan 3, 2025 • edited Loading

Code Review Agent Run #639caa

flyte-bot commented Jan 3, 2025 • edited Loading

Changelist by Bito

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Jan 4, 2025 • edited Loading

Codecov Report

flyte-bot commented Jan 4, 2025 • edited Loading

Code Review Agent Run #f65ef4

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

flyte-bot commented Jan 6, 2025 • edited Loading

Code Review Agent Run #be429e

Choose a reason for hiding this comment

Choose a reason for hiding this comment

flyte-bot commented Jan 6, 2025 • edited Loading

Code Review Agent Run #5c2427

flyte-bot commented Jan 7, 2025 • edited Loading

Code Review Agent Run #d95336

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

flyte-bot commented Jan 8, 2025 • edited Loading

Code Review Agent Run #0644db

cosmicBboy commented Jan 3, 2025 •

edited by flyte-bot

Loading

flyte-bot commented Jan 3, 2025 •

edited

Loading

flyte-bot commented Jan 3, 2025 •

edited

Loading

codecov bot commented Jan 4, 2025 •

edited

Loading

flyte-bot commented Jan 4, 2025 •

edited

Loading

flyte-bot commented Jan 6, 2025 •

edited

Loading

flyte-bot commented Jan 6, 2025 •

edited

Loading

flyte-bot commented Jan 7, 2025 •

edited

Loading

flyte-bot commented Jan 8, 2025 •

edited

Loading