[SPARK-41592][PYTHON][ML] Pytorch file Distributed Training

[rithwik-db]

What changes were proposed in this pull request?

This is an addition to #39188 to add support for multi node training using PyTorch files. The users would follow the second workflow in the design document to run training on the executors. I added some new utility functions as well as built on top of current functions. This is largely WIP so testing will be added very soon.

Why are the changes needed?

Look at the main ticket for more details.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

Tested with a pseudo-integration test. Integration tests will be added in a future PR.

[AmplabJenkins]

Can one of the admins verify this patch?

[rithwik-db]

This is actually no longer needed since if num_processes > len(gpus_owned), then we set num_processes = len(gpus_owned)

[lu-wang-dl]

What does this mean?

[rithwik-db]

Something like the following seems to error:

import socket
sock = socket.socket()
sock.bind((master_address, 0))
port = sock.getsockname()[1]

So I just find a port using randomness.

[lu-wang-dl]

What happened if two processes choose the same port?

[rithwik-db]

I believe it will raise a RuntimeError: Address already in use

[lu-wang-dl]

Do we need to do this?

[rithwik-db]

I think it should be added because in the case the user runs training with TorchDistributor(use_gpu=False, **kwargs).run(train_fn) but accidentally has some PyTorch Lightning code like pl.Trainer(accelerator="gpu") in their train_fn, an error should be raised saying no cuda devices available even though you specified a gpu accelerator.

We already have a check in get_num_tasks that checks when use_gpu=True but no GPUs are available, and I think this code addresses the case when use_gpu=False but the internal code has usage of GPUs.

[rithwik-db]

@WeichenXu123 @lu-wang-dl is my logic reasonable here or did I misunderstand anything?

[lu-wang-dl]

I still don't understand. If the user run something like pl.Trainer(accelerator="gpu") on a CPU cluster, what is the behavior from PyTorch lighting?

[rithwik-db]

PyTorch Lightning will raise a MisconfigurationException: No supported gpu backend found!. This is what we expect to see if the user sets use_gpu=False and calls pl.Trainer(accelerator="gpu") My understanding is that if a user runs this code on a local cluster with GPUs on each node without os.environ[CUDA_VISIBLE_DEVICES] = "", then the task will be assigned a GPU even when use_gpu=True.

[lu-wang-dl]

Why not just define the function here?

[rithwik-db]

I guess just for the sake of modularity. We could just define the function here.

[lu-wang-dl]

Overall LGTM. Just some minor comments/questions.

[WeichenXu123]

Shall we add a sleep(0.1) in the loop body ?

[WeichenXu123]

and I recommend to set maximum retry number (e.g. 100) for get_free_port loop, to avoid dead loop in some unexpected cases.

[WeichenXu123]

We can simplify set_gpus function:

if CUDA_VISIBLE_DEVICES env var exists, do nothing (spark already set CUDA_VISIBLE_DEVICES properly
otherwise generates CUDA_VISIBLE_DEVICES from taskcontext.resources["gpu"].addresses

[rithwik-db]

https://github.com/apache/spark/pull/39267/files#diff-76c395a6b98138662faaec37460ccda966f5cc0df0bccd224dfefcd81b2a7a79R459 <- Is this what you were suggesting?

[WeichenXu123]

LGTM, thanks!

[HyukjinKwon]

Merged to master.

[mattoh91]

@rithwik-db Can I clarify that the num_processors attribute of the TorchDistributor class refers to the number of spark.executor.cores used, and not the number of spark.executor.instances?
Trying to use >1 num_processors seems to take up more cores / slots on a single executor during training (using spark operator on k8s).

[rithwik-db]

If we are using CPUs for training, num_processors attribute refers to the number of spark tasks that will be created for training, and each spark task can use >1 CPU depending on what spark.task.cpus says. (This function is where this logic is defined).