I have the following issue with FlyteFile + single node/multi-GPU PyTorch plugin: 1. When task is decorated only with

@task()

, the following code works just fine:

@task()
def my_task(dataset: FlyteFile):
    path = dataset.download()
    # works as intended
    assert Path(path).is_file()

@workflow
def my_workflow():
    file = FlyteFile(
        path="<s3://path/to/my/file.csv>"
    )

    outputs = my_task(dataset=file)

2. When task is decorated with

@task(task_config=Elastic(nnodes=1, nproc_per_node=4)

, it breaks

@task(task_config=Elastic(nnodes=1, nproc_per_node=4)
def my_task(dataset: FlyteFile):
    path = dataset.download()
    # .download() immediately returns and file is not there
    assert Path(path).is_file() # this will raise

@workflow
def my_workflow():
    file = FlyteFile(
        path="<s3://path/to/my/file.csv>"
    )

    outputs = my_task(dataset=file)

There's also a warning raised:

.venv/lib/python3.12/site-packages/flytekit/types/file/file.py:356: RuntimeWarning: coroutine 'FileAccessProvider.async_get_data' was never awaited

It seems like the FlyteFile does not play well with underlying multiprocessing spawn This happens on

Flyte 1.15.3

Good catch. It should not stay open in the spawn

It should not stay open in the spawn

WDYM?

For the multi you train we spawn torchrun spawns many processes

I was asking about "It should not stay open in the spawn"

But this is just a warning

No, the warning is "additional" - the main reason I wrote this because it actually breaks

Ok I need to reproduce

Sure! I'll look into this

Cc @thankful-minister-83577 fyi for v2? Shouldn’t need it

Awesome @echoing-account-76888, thanks for fixing it! 🎖️ @freezing-airport-6809 will v2 still support PyTorch Elastic? (via decorator, same as in v1)?

Absolutely

I has quietly hoping for a decorator for PyTorch that wouldn't require installing 3rd party software in K8s besides Flyte itself 🤞

Wdym

Yes - right now, if you want to do distributed training on > 1 node, you have to either install KF training operator or Ray on the cluster, which adds complexity. PyTorch is THE FRAMEWORK for neural networks right now, so having 1st party support for it in Flyte itself would be 🎖️

@flat-waiter-82487 we have started working on honest support

Aw yeah! 🦜

will take some time

Functionally I would like the following things: 1. Ability to decorate a task to make it distributed (so 2 params in it at least: num_nodes, num_proc_per_node - same as now?) 2. Support for TorchElastic and fault tolerance during training with auto-recovery 3. No requirement of installing 3rd party extensions to the cluster besides Flyte and its CRDs (CRDs coming from Flyte itself would be fine) 4. Ability to specify num_nodes/num_proc_per_node dynamically (e.g. via task param or something) - right now as far as I know, it has to be "baked in" into the task definition = somehow hardcoded - this is a painpoint during development, because in development we usually use small machines with 2-4 GPUs and then we go full throttle on larger machines with 8 GPUs and we have to re-configure the code manually to do it instead of just swapping

--num_proc_per_node=8

or sth.