Adding per rank mlflow tracking location to fix race condition and up…

…dating fcn_afno config for easier benchmarking (#121)

Signed-off-by: Akshay Subramaniam <[email protected]>

examples/weather/fcn_afno/conf/config.yaml

@@ -19,7 +19,9 @@ hydra:
     dir: ./outputs/
 
 wb_artifacts: False
+use_mlflow: True
 
 start_epoch: 0
+max_epoch: 80
 
 num_samples_per_year_train: 1456

examples/weather/fcn_afno/train_era5.py

@@ -109,7 +109,7 @@ def main(cfg: DictConfig) -> None:
         user_name="Modulus User",
         mode="offline",
     )
-    LaunchLogger.initialize(use_mlflow=True)  # Modulus launch logger
+    LaunchLogger.initialize(use_mlflow=cfg.use_mlflow)  # Modulus launch logger
     logger = PythonLogger("main")  # General python logger
 
     datapipe = ERA5HDF5Datapipe(
@@ -199,7 +199,7 @@ def train_step_forward(my_model, invar, outvar):
         return loss
 
     # Main training loop
-    max_epoch = 80
+    max_epoch = cfg.max_epoch
     for epoch in range(max(1, loaded_epoch + 1), max_epoch + 1):
         # Wrap epoch in launch logger for console / WandB logs
         with LaunchLogger(

modulus/launch/logging/mlflow.py

@@ -90,7 +90,7 @@ def initialize_mlflow(
         group_name = f"{run_name}_{time_string}"
     # Set default value here for Hydra
     if tracking_location is None:
-        tracking_location = str(Path("./mlruns").absolute())
+        tracking_location = str(Path(f"./mlruns_{dist.rank}").absolute())
 
     # Set up URI (remote or local)
     if mode == "online":