Hey everyone- I have a question about running multi-node pytorch workflows and error/exception handling. We're currently defining our training task as something like this:

@task(
    task_config=task_config,
    cache=False,
    container_image=container_image,
    pod_template=pod_template,
    timeout=timeout,
    retries=max_retries,
)
def flyte_training_main_task():
  ...

with the task_config being (note that we don't really need the elastic part of things - we just need to launch a multi-node pytorch task):

task_config = Elastic(
    nnodes=num_nodes,
    nproc_per_node=8,
)

Now imagine that a rank in the distributed training has an error of some sort - is there a way for us to configure our task so that the whole task/workflow is terminated (including all the pods corresponding to it) as soon as a single rank errors? Currently it seems like it requires all the ranks to exit/error until the task/workflow is terminated, which we often don't want (because other ranks might be stuck until NCCL timeout or might be stuck for other reasons). I've tried raising special exception types like

SignalException

or

ChildFailedError

, but it seems like it always waits until all the ranks exit. One hacky workaround I could think of is to manually terminate the workflow, but that also does not seem ideal. Thanks!!

Hmm interesting- let me understand more - are you using elastic or not?

hi @fancy-twilight-30247, could you try setting

max_restarts=0

in the

Elastic

config?

task_config = Elastic(
    nnodes=num_nodes,
    nproc_per_node=8,
    max_restarts=0,
)

this ensures that the job won’t restart if any rank fails. also you can try lowering the NCCL timeout so that the failure is detected more quickly and the job reaches a failed state faster.

Let me try that! I know that decreasing the NCCL timeout can be one potential option, but that has a whole set of confounding variables across other workloads so I'd like to avoid that 🙂 Ideally what I'd be looking for is a user-controller way to terminate the whole workflow from any particular rank, while (optionally) still allowing for restarts to happen

yeah, this plugin uses the kubeflow training operator v1 under the hood, and i don’t think there’s a way to configure that behavior. for now, it’d have to be workarounds. once we migrate to v2, you should be able to set custom failure policies.

Got it! Excited about v2 then haha 🙂

if you’re interested in contributing (we’ve got a couple other folks interested in this), would love to loop you in!

@fancy-twilight-30247 let me add you to our contributor channel. Just posting all suggestions of improvements would be fantabulous

Thank you! Yeah would be super interested to check out what the current plans for v2 are - happy to chime in for sure!

Also let me know if you want to try v2 on union