[Fix] Issue when using FlyteDirectory with Elastic

[machichima]

Tracking issue

Follow up #3313

Why are the changes needed?

As described in #3313

What changes were proposed in this pull request?

Ensure that everytime we call download, we will get the current_context of the current process, to ensure calling the get_data synchronously.

How was this patch tested?

Test with script

from flytekit import task, workflow, FlyteDirectory, ImageSpec, Resources
from flytekitplugins.kfpytorch.task import Elastic, PyTorch


DEFAULT_REMOTE_PATH = "s3://my-s3-bucket/test_dir/"

image_spec = ImageSpec(
    registry="localhost:30000",
    packages=[
        "flytekitplugins-kfpytorch[elastic]",
        "numpy",
        "torch",
    ],
)


@task(
    task_config=Elastic(nnodes=1, nproc_per_node=4),
    container_image=image_spec,
    requests=Resources(cpu="4", mem="4Gi"),
    limits=Resources(cpu="4", mem="4Gi"),
)
def my_task(dataset: FlyteDirectory):
    path = dataset.download()
    print(f"path: {path}", flush=True)

    assert Path(path).is_file()  # will pass now

@workflow
def wf():
    file = FlyteDirectory(path=DEFAULT_REMOTE_PATH)

    outputs = my_task(dataset=file)

Setup process

Screenshots

Result

Check all the applicable boxes

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

Related PRs

Follow up #3313

Docs link

Summary by Bito

This pull request introduces a new downloader class for the FlyteDirectory, enhancing the synchronous data retrieval mechanism. It replaces the previous downloader function with a class instance, improving reliability and ensuring correct context utilization during downloads, as per issue #3313.

[codecov]

Codecov Report

❌ Patch coverage is 27.27273% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 45.75%. Comparing base (bcdfaae) to head (7b3553c).
⚠️ Report is 3 commits behind head on master.

Files with missing lines	Patch %	Lines
flytekit/types/directory/types.py	27.27%	8 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #3320      +/-   ##
==========================================
- Coverage   52.64%   45.75%   -6.89%     
==========================================
  Files         215      328     +113     
  Lines       22543    27565    +5022     
  Branches     2950     2950              
==========================================
+ Hits        11868    12613     +745     
- Misses       9968    14855    +4887     
+ Partials      707       97     -610     

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

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[flyte-bot]

Bito Automatic Review Failed - Technical Failure

Bito encountered technical difficulties while generating code feedback . To retry, type /review in a comment and save. If the issue persists, contact [email protected] and provide the following details:

Agent Run ID: a025f168-79e2-4a5d-82c4-17d3dd60fd0c