Support Elastic PyTorch in TrainJob

[andreyvelich]

What you would like to be added?

As part of Kubeflow Trainer API, we designed the ElasticPolicy API which should allow users to run PyTorch in elastic mode: https://docs.pytorch.org/docs/stable/elastic/run.html

It would be nice if someone could drive KEP for that, and propose API changes to the next release: Trainer v2.2

elasticPolicy:
  minNodes: 2
  maxNodes: 5
  metrics:
    - type: Resource
      resource:
        name: nvidia.com/gpu
        target:
          type: Utilization
          averageUtilization: 75

Potentially, we should support elastic JobSet for that: kubernetes-sigs/jobset#463

cc @kubeflow/kubeflow-trainer-team

/area api
/area controller
/help

Why is this needed?

We should support Elastic TrainJobs like in Training Operator V1 PyTorchJob.

Love this feature?

Give it a 👍 We prioritize the features with most 👍

[google-oss-prow]

@andreyvelich:
This request has been marked as needing help from a contributor.

Please ensure the request meets the requirements listed here.

If this request no longer meets these requirements, the label can be removed
by commenting with the /remove-help command.

Details

In response to this:

What you would like to be added?

As part of Kubeflow Trainer API, we designed the ElasticPolicy API which should allow users to run PyTorch in elastic mode: https://docs.pytorch.org/docs/stable/elastic/run.html

It would be nice if someone can drive this KEP, and propose API changes to the next release: Trainer v2.2

Potentially, we should support elastic JobSet for that: kubernetes-sigs/jobset#463

cc @kubeflow/kubeflow-trainer-team

/area api
/area controller
/help

Why is this needed?

We should support Elastic TrainJobs like in Training Operator V1 PyTorchJob.

Love this feature?

Give it a 👍 We prioritize the features with most 👍

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.

[andreyvelich]

cc @kannon92 @danielvegamyhre if we should learn something new from the torchrun Elastic capabilities.

[adity1raut]

/assign