feat: Added the option to run inference in parallel

docs/parallel.rst

@@ -1,54 +1,64 @@
-###################
+####################
  Parallel Inference
-###################
+####################
 
-If the memory requirements of your model are too large to fit within a single GPU, you run Anemoi-Inference in parallel across multiple GPUs.
+If the memory requirements of your model are too large to fit within a
+single GPU, you run Anemoi-Inference in parallel across multiple GPUs.
 
-Parallel inference requires SLURM to launch the parallel processes and to determine information about your network environment. If SLURM is not available to you, please create an issue on the Anemoi-Inference github page.
+Parallel inference requires SLURM to launch the parallel processes and
+to determine information about your network environment. If SLURM is not
+available to you, please create an issue on the Anemoi-Inference github
+page.
 
-**************
+***************
  Configuration
-**************
+***************
 
-To run in parallel, you must add '`runner:parallel`' to your inference config file.
+To run in parallel, you must add '`runner:parallel`' to your inference
+config file.
 
 .. code:: yaml
 
-  checkpoint: /path/to/inference-last.ckpt
-  lead_time: 60
-  runner: parallel
-  input:
-    grib: /path/to/input.grib
-  output:
-    grib: /path/to/output.grib
+   checkpoint: /path/to/inference-last.ckpt
+   lead_time: 60
+   runner: parallel
+   input:
+     grib: /path/to/input.grib
+   output:
+     grib: /path/to/output.grib
 
-
-
-
-******************************
+*******************************
  Running inference in parallel
-******************************
+*******************************
 
-Below is an example SLURM batch script to launch a parallel inference job across 4 GPUs. 
+Below is an example SLURM batch script to launch a parallel inference
+job across 4 GPUs.
 
 .. code:: bash
 
-        #!/bin/bash
-        #SBATCH --nodes=1
-        #SBATCH --ntasks-per-node=4
-        #SBATCH --gpus-per-node=4
-        #SBATCH --cpus-per-task=8
-        #SBATCH --time=0:05:00
-        #SBATCH --output=outputs/paralell_inf.%j.out
+   #!/bin/bash
+   #SBATCH --nodes=1
+   #SBATCH --ntasks-per-node=4
+   #SBATCH --gpus-per-node=4
+   #SBATCH --cpus-per-task=8
+   #SBATCH --time=0:05:00
+   #SBATCH --output=outputs/paralell_inf.%j.out
 
-        source /path/to/venv/bin/activate
-        srun anemoi-inference run parallel.yaml
+   source /path/to/venv/bin/activate
+   srun anemoi-inference run parallel.yaml
 
 .. warning::
-   If you specify '`runner:parallel`' but you don't launch with '`srun`', your anemoi-inference job may hang as only 1 process will be launched.
+
+   If you specify '`runner:parallel`' but you don't launch with
+   '`srun`', your anemoi-inference job may hang as only 1 process will
+   be launched.
 
 .. note::
-   By default, anemoi-inference will determine your systems master address and port itself. If this fails (i.e. when running Anemoi-Inference inside a container), you can instead set these values yourself via environment variables in your SLURM batch script:
+
+   By default, anemoi-inference will determine your systems master
+   address and port itself. If this fails (i.e. when running
+   Anemoi-Inference inside a container), you can instead set these
+   values yourself via environment variables in your SLURM batch script:
 
    .. code:: bash
 
@@ -57,4 +67,3 @@ Below is an example SLURM batch script to launch a parallel inference job across
       export MASTER_PORT=$((10000 + RANDOM % 10000))
 
       srun anemoi-inference run parallel.yaml
-