Distributed training

[frostedoyster]

This PR adds utilities and infrastructure changes to allow distributed, multi-GPU training. The SOAP-BPNN model now supports it.

---

📚 Documentation preview 📚: https://metatensor-models--179.org.readthedocs.build/en/179/

[Luthaf]

How does the error message (following the discussion in #194) looks like if you set the port to something already taken? You can test this by starting python -m http.server 39591 somewhere and set the port to 39591 here as well.

[frostedoyster]

Here is the error:

[2024-05-22 14:14:38][INFO] - This log is also available in 'outputs/2024-05-22/14-14-38/train.log'.
[2024-05-22 14:14:39][INFO] - random seed of this run is 42
[2024-05-22 14:14:39][INFO] - Setting up training set
[2024-05-22 14:14:39][INFO] - Forces found in section 'energy'. Forces are taken for training!
[2024-05-22 14:14:39][WARNING] - No Stress found in section 'energy'. Continue without stress!
[2024-05-22 14:14:39][INFO] - Setting up test set
[2024-05-22 14:14:39][INFO] - Setting up validation set
[2024-05-22 14:14:39][INFO] - Setting up model
[2024-05-22 14:14:39][INFO] - Calling architecture trainer
[W socket.cpp:464] [c10d] The server socket has failed to bind to [::]:39591 (errno: 98 - Address already in use).
[W socket.cpp:464] [c10d] The server socket has failed to bind to 0.0.0.0:39591 (errno: 98 - Address already in use).
[E socket.cpp:500] [c10d] The server socket has failed to listen on any local network address.
�[31mERROR: The error below most likely originates from an architecture. If you think this is a bug, please contact its maintainer (see the architecture's documentation).

The server socket has failed to listen on any local network address. The server socket has failed to bind to [::]:39591 (errno: 98 - Address already in use). The server socket has failed to bind to 0.0.0.0:39591 (errno: 98 - Address already in use).�[0m
srun: error: i70: task 0: Exited with exit code 1
srun: Terminating StepId=2083332.0
slurmstepd: error: *** STEP 2083332.0 ON i70 CANCELLED AT 2024-05-22T14:14:39 ***
10.91.27.70 - - [22/May/2024 14:14:39] code 400, message Bad HTTP/0.9 request type ('\x00Î÷')
10.91.27.70 - - [22/May/2024 14:14:39] "�Î÷�<���������init/��������" 400 -
srun: error: i70: task 1: Terminated
srun: Force Terminated StepId=2083332.0

[frostedoyster]

The final part is due to the http server