train.py

Namely:

custom_pod_template = PodTemplate(
    primary_container_name="flytesnacks-pytorch-lightning",
    pod_spec=V1PodSpec(
        containers=[
            V1Container(
                name="flytesnacks-pytorch-lightning",
                volume_mounts=[V1VolumeMount(mount_path="/dev/shm", name="dshm")]
            )
        ],
        volumes=[
            V1Volume(
                name="dshm",
                empty_dir=V1EmptyDirVolumeSource(medium="", size_limit="200Gi")
            )
        ]
    ),
)

And then the task def is:

@task(
    container_image=custom_image,
    task_config=Elastic(
        nnodes=NUM_NODES,
        nproc_per_node=NUM_DEVICES,
        rdzv_configs={"timeout": 36000, "join_timeout": 36000},
        max_restarts=3,
    ),
    accelerator=T4,
    requests=Resources(mem="32Gi", cpu="48", gpu="8", ephemeral_storage="100Gi"),
    pod_template=custom_pod_template,
)
def train_model() -> FlyteDirectory:

I’d like to get rid of the boilerplate and automatically add this to the

Elastic

task config, perhaps exposing some options like

dshm_mount_path

,

dshm_size_limit

, etc.

This would require adding

kubernetes

as a dependency to the

flytekitplugins-kfpytorch

plugin Also need to do this in a way that nicely handles

primary_container_name

in the

PodTemplate

definition and specifying

pod_template

in the

@task

definition would override the one provided by the

Elastic

task config

Absolutely agree with the goal.

Let’s talk about details 🙂

2. primary container name > Also need to do this in a way that nicely handles

primary_container_name

in the

PodTemplate

definition In our pod template in the flyte namespace we just name the container

default

and flyteplugins/pytorch just renames the container to pytorch. What’s the use case for naming the container something different like

flytesnacks-pytorch-lightning

?

3. kubernetes dependency Flytekit was able to remove this dependency, right? I wouldn’t mind adding it to the elastic plugin though …

4. timeouts

rdzv_configs={"timeout": 36000, "join_timeout": 36000},

I agree with setting default timeouts that we expect should work in a default case for most users. I personally feel an hour is quite a lot though - gpus might end up idling for a long time. For the join timeout I feel we should consider the scenario that some workers have a hot start (node is up and image is cached) while other workers have a cold start, i.e. node needs to be scaled up and image has to be pulled. How long would this take, maybe 15 minutes tops? I personally feel that if the join timeout needs to be longer than this, e.g. due to resource constraints/nodes not being able to scale up immediately, one should maybe consider gang scheduling so that all of them start at the same time (see “coscheduling” here). Do you happen to know whether the

timeout

is the same timeout as one specifies to

torch.distributed.init_process_group

? If yes, I personally find an hour also quite a lot because if the training deadlocked, it would be nice to bring it down more quickly than 1h. (For this use case, we set

NCCL_ASYNC_ERROR_HANDLING=1

as env var in the pod template. We could consider doing this in the plugin as well!)