Lightning-AI / Lightning-AI/pytorch-lightning
Share/collect non-tensor objects between processes.
@dalek-who is already working on this.
Since Apr 17, 2023.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
## 🚀 Feature
### Motivation
Sometimes we need to share/collect some non-tensor objects between process in DDP mode, mainly includes two kinds of scenes:
- **share status**: for example, a random port to connect database, a timestamp when the experiment starts, an auto-created output directory, etc. These status are created during running and cannot be assigned manually.
- **collect outputs**: for example, in test/predict, collect all outputs to `gloabl_rank_zero` process in `on_{test,preidct}_{batch,epoch}_end` to write predict results and compute metrics over the full dataset.
Current functions cannot satisfy these requests:
- `LightningModule.all_gather` can only deal with tensors. Other objects (like string, id list, Path ...) cannot be gathered or shared.
- `DP` or `DDPSpawn`: they are not as "convenient" as DDP. DP is not encouraged by torch, and DDPSpawn sometimes cannot be started.
- `torch.multiprocessing.SimpleQueue`: there are no tutorials on how to use it with pytorch-lightning DDP mode.
### Pitch
### Alternatives
Consider something like this:
```python
def collect_all_objects(obj: object, global_zero_only: bool=False, default=None):
...
objects_from_process = {
0: obj_0,
1: obj_1,
global_rank_n: obj_n,
}
if not trainer.is_global_zero and global_zero_only:
return default
else:
return objects_from_process
```
This function can also be used to share status between processes: each process collect status and only use the status from rank_0.
**Important:**
When using DDP, the sampler may padding some examples in the last batch and cause the following problem:
- #11102
In this case, manually deduplication is necessary after collecting objects. However, deduplication is much more easier than dealing with inter-process synchronization.
### Additional context
There are some questions related to these requests:
1.**share status**:
- [Share state between DDP processes](https://forums.pytorchlightning.ai/t/share-state-between-ddp-processes/963)
- [Sync output dir between DDP processes](https://forums.pytorchlightning.ai/t/sync-output-dir-between-ddp-processes/779)
2.**collect outputs**:
- #5788
- #10618
- #11086
- #9259
cc @borda @awaelchli @rohitgr7 @akihironitta
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.