[2810] Allow private config to override output directories and filename

packages/common/src/weathergen/common/config.py

@@ -208,15 +208,18 @@ def save(config: Config, mini_epoch: int | None):
         f.write(json_str)
 
 
-def load_run_config(run_id: str, mini_epoch: int | None, model_path: str | None) -> Config:
+def load_run_config(
+    run_id: str, mini_epoch: int | None, model_path: str | None, *, config: Config | None = None
+) -> Config:
     """
     Load a configuration file from a given run_id and mini_epoch.
     If run_id is a full path, loads it from the full path.
 
     Args:
         run_id: Run ID of the pretrained WeatherGenerator model
         mini_epoch: Mini_epoch of the checkpoint to load. -1 indicates last checkpoint available.
-        model_path: Path to the model directory. If None, uses the model_path from private config.
+        model_path: Parent model directory containing run-id subdirectories.
+        config: Active configuration for resolving path_model when model_path is None.
 
     Returns:
         Configuration object loaded from the specified run and mini_epoch.
@@ -228,7 +231,7 @@ def load_run_config(run_id: str, mini_epoch: int | None, model_path: str | None)
     else:
         # Load model config here. In case model_path is not provided, get it from private conf
         if model_path is None:
-            path = get_path_model(run_id=run_id)
+            path = get_path_model(config, run_id=run_id)
         else:
             path = Path(model_path) / run_id
 
@@ -255,9 +258,9 @@ def load_run_config(run_id: str, mini_epoch: int | None, model_path: str | None)
     with fname.open() as f:
         json_str = f.read()
 
-    config = OmegaConf.create(json.loads(json_str))
+    loaded_config = OmegaConf.create(json.loads(json_str))
 
-    return _apply_fixes(config)
+    return _apply_fixes(loaded_config)
 
 
 def _get_model_config_file_write_name(run_id: str, mini_epoch: int | None):
@@ -451,7 +454,7 @@ def load_merge_configs(
     if from_run_id is None:
         base_config = _load_base_conf(base)
     else:
-        base_config = load_run_config(from_run_id, mini_epoch, None)
+        base_config = load_run_config(from_run_id, mini_epoch, None, config=private_config)
         from_run_id = get_run_id_from_config(base_config)
     with open_dict(base_config):
         base_config.from_run_id = from_run_id
@@ -698,9 +701,23 @@ def load_streams(streams_directory: Path) -> Config:
     return OmegaConf.create(streams)
 
 
+def _get_output_directory(config: Config, key: str, folder: str, run_id: str) -> Path:
+    """Use an exact directory override, or the existing shared per-run location."""
+    if config.get(key) is not None:
+        return Path(config[key])
+    working_dir = config.get("path_shared_working_dir")
+    root = Path(working_dir) if working_dir is not None else _get_shared_wg_path()
+    return root / folder / run_id
+
+
+def get_path_logs(config: Config) -> Path:
+    """Get the application log directory."""
+    return _get_output_directory(config, "path_logs", "logs", get_run_id_from_config(config))
+
+
 def get_path_run(config: Config) -> Path:
     """Get the current runs results_path for storing run results and logs."""
-    return _get_shared_wg_path() / "results" / get_run_id_from_config(config)
+    return _get_output_directory(config, "path_results", "results", get_run_id_from_config(config))
 
 
 def get_path_model(config: Config | None = None, run_id: str | None = None) -> Path:
@@ -710,14 +727,16 @@ def get_path_model(config: Config | None = None, run_id: str | None = None) -> P
     else:
         msg = f"Missing run_id and cannot infer it from config: {config}"
         raise ValueError(msg)
-    return _get_shared_wg_path() / "models" / run_id
+    config = config if config is not None else _load_private_conf()
+    return _get_output_directory(config, "path_model", "models", run_id)
 
 
 def get_path_results(config: Config, mini_epoch: int) -> Path:
     """Get the path to validation results for a specific mini_epoch and rank."""
     ext = StoreType(config.zarr_store).value  # validate extension
     base_path = get_path_run(config)
     fname = f"validation_chkpt{mini_epoch:05d}_rank{config.rank:04d}.{ext}"
+    fname = config.get("output_name") or fname
 
     return base_path / fname
 

packages/common/src/weathergen/common/logger.py

@@ -96,7 +96,7 @@ def format(self, record, *args, **kwargs):
 
 
 @cache
-def init_loggers(run_id=None, logging_config=None):
+def init_loggers(run_id=None, logging_config=None, log_path=None):
     """
     Initialize the logger for the package and set output streams/files.
 
@@ -123,7 +123,7 @@ def init_loggers(run_id=None, logging_config=None):
     # output_dir = f"./output/{timestamp}-{run_id}"
     output_dir = ""
     if run_id is not None:
-        output_dir = f"./logs/{run_id}"
+        output_dir = str(log_path) if log_path is not None else f"./logs/{run_id}"
 
     # load the structure for logging config
     if logging_config is None:
@@ -142,7 +142,7 @@ def init_loggers(run_id=None, logging_config=None):
                 filename = f"{output_dir}/{v}"
                 ofile = pathlib.Path(filename)
                 # make sure the path is independent of path where job is launched
-                if not ofile.is_absolute():
+                if not ofile.is_absolute() and log_path is None:
                     work_dir = pathlib.Path(_load_private_conf().get("path_shared_working_dir"))
                     ofile = work_dir / ofile
                 pathlib.Path(ofile.parent).mkdir(parents=True, exist_ok=True)

src/weathergen/model/model_interface.py

@@ -181,7 +181,7 @@ def load_model(cf, model, device, run_id: str, mini_epoch=-1):
         mini_epoch : The mini_epoch to load. Default (-1) is the latest mini_epoch
     """
 
-    path_run = get_path_model(run_id=run_id)
+    path_run = get_path_model(cf, run_id=run_id)
     mini_epoch_id = (
         f"chkpt{mini_epoch:05d}" if mini_epoch != -1 and mini_epoch is not None else "latest"
     )

src/weathergen/run_train.py

@@ -100,7 +100,7 @@ def run_inference(args):
     devices = Trainer.init_torch()
     cf = Trainer.init_ddp(cf)
 
-    init_loggers(cf.general.run_id)
+    init_loggers(cf.general.run_id, log_path=config.get_path_logs(cf))
 
     logger.info(f"DDP initialization: rank={cf.rank}, world_size={cf.world_size}")
 
@@ -139,7 +139,7 @@ def run_continue(args):
     devices = Trainer.init_torch(multiprocessing_method=mp_method)
     cf = Trainer.init_ddp(cf)
 
-    init_loggers(cf.general.run_id)
+    init_loggers(cf.general.run_id, log_path=config.get_path_logs(cf))
 
     # track history of run to ensure traceability of results
     cf.general.run_history += [(args.from_run_id, cf.general.istep)]
@@ -176,7 +176,7 @@ def run_train(args):
 
     # this line should probably come after the processes have been sorted out else we get lots
     # of duplication due to multiple process in the multiGPU case
-    init_loggers(cf.general.run_id)
+    init_loggers(cf.general.run_id, log_path=config.get_path_logs(cf))
 
     logger.info(f"DDP initialization: rank={cf.rank}, world_size={cf.world_size}")