Introduce PerProcessContext to adjust DDP per process behavior (#913)

[breakds]

Motivation

• To make DDP as transparent to the user as possible, we should not expect the user to have to manually update a configuration to achieve parity between DDP mode and non DDP mode. More specifically, if single process has 64 environments in parallel, we would expect in dual process the per-process number of parallel environments to be reduced accordingly (to 32).
• Similar to only having process 0 (master process) writing checkpoint, we want to have only process 0 writing the summary. This is done by not enabling summary unless the process rank is 0 (or -1) when setting up PolicyTrainer.

Therefore, this PR targets the combined problem of #953 and #954

Solution

Introduce a global singleton called PerProcessContext. As its name suggests, this context maintains properties that are private to each process. In the current implementation, PerProcessContext allows accessing the rank (process ID) and the total number of processes from anywhere.

Design Choices:

1. Decided to use a global singleton to avoid having to pass context down to the very bottom of the calling stack, as discussed with @emailweixu in Adjust number of environments in DDP mode based on number of processes #953
2. Decided to disable summary by hijacking run_under_record_context as a whole instead of just hijacking _cond. This is because run_under_record_context will initialize the writer and the write the global steps before the first call to _cond, and we would like to avoid that as well.
3. Decided to transparently update the number of environments per process in the config, instead of providing a new configuration item, unlike what has been suggested by @le-horizon in Adjust number of environments in DDP mode based on number of processes #953. Both approaches have their benefits but I think having to update the config to run distributed training is a bit more work (for the user), and adding new configuration item may further damage the readability (of the config file). As long as the training/optimization is done the same, DDP and non-DDP shouldn't make a difference in terms of the training result.

Testing

1. Run training without DDP to make sure summary and number of environments are the same as before.
2. Run training with DDP with dual process to make sure

• Number of processes are halved per process
• There is only one events file being generated for summary

[emailweixu]

Should we add some unittest to ensure DDP can run correctly? (perhaps in a different PR). A possible place for this is alf/bin/train_play_test.py.

[breakds]

Should we add some unittest to ensure DDP can run correctly? (perhaps in a different PR). A possible place for this is alf/bin/train_play_test.py.

Agree. Created #956 for this.

[le-horizon]

LG