Lightning-AI / Lightning-AI/pytorch-lightning

Refactor `_SubprocessScriptLauncher` process launching strategy

Open
#17,248 11 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

design distributed fabric pl refactor
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Outline & Motivation

Hi all, firstly thanks for your work on the library, it's great!

# Existing code
The `_SubprocessScriptLauncher` results in a process hierarchy as follows, for `N` GPUs:
```
LocalRank=0
├── LocalRank=1
├── LocalRank=2
├── LocalRank=3
├── ...
└── LocalRank=N-1
```
where `LocalRank=0` launches each of `LocalRank=1` to `LocakRank=N-1` via `subprocess.Popen`.

# Issues
This can make a few things challenging, namely:

- **Defunct/zombie processes**: if any of `LocalRank=1` to `N-1` fail, then the process enters a zombie state, as `LocalRank=0`(its parent), never reads the exit code. This means that a container running multi-GPU training will not exit on the failure of a non-zero process until DDP times out. For users where termination leads to freeing of expensive compute resources this is problematic.
- **Exit code checking**: in Unix you can only read exit codes of a process from its parent. As a result, the user cannot read the exit code of `LocalRank=1..N-1`.
- **Local rank** environment variable: a user can only identify `LocalRank=0` from the lack of an environment variable. This makes any code relying on this less robust, especially, if it also runs in CPU-only environments where the environment variable is correctly missing.
- **Process re-parenting**: if `LocalRank=0` fails, then `LocalRank=1,..,N-1` are assigned to be children of `pid=1`. This makes cleaning up of processes a bit more challenging.

# Proposal
There are two main approaches to rectify this:
1. Add a thread in `LocalRank=0` that checks the status of the remaining processes and tidies up all the processes on a failure. Not sure I like this, however, since the thread isn't strictly speaking guaranteed to execute due to the GIL.
2. Refactor process structure to be something like
```
Watcher
├── LocalRank=0
├── LocalRank=1
├── LocalRank=2
├── LocalRank=3
├── ...
└── LocalRank=N-1
```
where the role of the `Watcher` process is to
a) launch the processes for each of the local ranks
b) monitor the status of the processes
c) terminate all the local rank processes, upon termination of a single one of them (and exit with the same exit code)

This would address all of the issues above and we could also explicitly set `LOCAL_RANK=0` in the zeroth process.

### Pitch

I feel that proposal 2. would add some robustness to multi-GPU training, especially in the case described above.

I would be happy to submit a PR that does this either in a new `_Launcher` implementation or as a refactor of `_SubprocessScriptLauncher`, if this sounds reasonable. Thanks!

### Additional context

_No response_

cc @awaelchli @justusschock @tchaton @borda @carmocca

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Locate `_SubprocessScriptLauncher` and inspect how it launches local-rank processes and handles their exit status. Compare the current hierarchy with the proposed watcher structure, then define tests for failure, cleanup, exit-code propagation, and `LOCAL_RANK`; done means one failed local process terminates the others and the watcher returns its exit code.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.