Cap nproc_per_node based on the CPU resources of the node for PyTorch TrainJob

[astefanutti]

PyTorch relies on the number of CPUs on the physical host when determining the "local world size" if nproc_per_node is set to auto and the node is a CPU-only device.

In that configuration, which is used by the preset torch-distributed training runtime, the number of processes equals the number of CPUs on the host, and leads to the following problems:

• Out of memory issues for worker Pods scheduled on nodes with large number of CPUs
• Deadlocks when the CPU limit set for the container is less than the actual number of CPUs

To mitigate these issues, nproc_per_node should default to the CPU limit when set, or fallback to 1 when the PyTorch ML policy defines numProcPerNode: auto.

This has been discussed in details in #2387 (comment)

[andreyvelich]

As we discussed, we should transition the numProcPerNode assignment to the ML Plugins from the client SDK: #2470 (comment).

If user doesn't explicitly set .trainer.numProcPerNode value in TrainJob, we should automatically calculate this value based on container resources and device type: cpu, gpu, tpu.

/good-first-issue
/area controller
/kind enhancement

[google-oss-prow]

@andreyvelich:
This request has been marked as suitable for new contributors.

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-good-first-issue command.

In response to this:

As we discussed, we should transition the numProcPerNode assignment to the ML Plugins from the client SDK: #2470 (comment).

If user doesn't explicitly set .trainer.numProcPerNode value in TrainJob, we should automatically calculate this value based on container resources and device type: cpu, gpu, tpu.

/good-first-issue
/area controller
/kind enhancement

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.