Add mlflow example for PyTorch DDP

Author: mvsjober

mnist_ddp_mlflow.py

@@ -47,9 +47,10 @@ def train(num_epochs):
     verbose = dist.get_rank() == 0  # print only on global_rank==0
     if verbose:
         mlflow.set_tracking_uri("/scratch/project_2001659/mvsjober/mlruns")
-        slurm_id = os.getenv("SLURM_JOB_ID")
-        if slurm_id:
-            mlflow.start_run(run_name=slurm_id)
+        #mlflow.set_tracking_uri("sqlite:////scratch/project_2001659/mvsjober/mlruns.db")
+        #mlflow.set_tracking_uri('https://mats-mlflow2.rahtiapp.fi/')
+
+        mlflow.start_run(run_name=os.getenv("SLURM_JOB_ID"))
 
     model = ConvNet().cuda()
     batch_size = 100

run-ddp-gpu1-mlflow.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+#SBATCH --account=project_2001659
+#SBATCH --partition=gputest
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=10
+#SBATCH --mem=64G
+#SBATCH --time=15
+#SBATCH --gres=gpu:v100:1
+
+module purge
+module load pytorch
+
+# Old way with torch.distributed.run
+# srun python3 -m torch.distributed.run --standalone --nnodes=1 --nproc_per_node=4 mnist_ddp.py --epochs=100
+
+# New way with torchrun
+srun torchrun --standalone --nnodes=1 --nproc_per_node=1 mnist_ddp_mlflow.py --epochs=100