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

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

@@ -208,15 +208,22 @@ 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,
+    private_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.
+        private_config: Configuration whose full_model_path specifies the exact checkpoint
+            directory (without appending run_id), used when model_path is None.
 
     Returns:
         Configuration object loaded from the specified run and mini_epoch.
@@ -228,7 +235,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(private_config, run_id=run_id)
         else:
             path = Path(model_path) / run_id
 
@@ -451,7 +458,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, private_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,26 +705,57 @@ def load_streams(streams_directory: Path) -> Config:
     return OmegaConf.create(streams)
 
 
-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)
+def _get_path_output(path_config: Config, key: str, folder: str, run_id: str) -> Path:
+    """Use an exact directory override, or the existing shared per-run location."""
+    if path_config.get(key) is not None:
+        return Path(path_config[key])
+    working_dir = path_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_path_output(config, "path_logs", "logs", get_run_id_from_config(config))
 
 
-def get_path_model(config: Config | None = None, run_id: str | None = None) -> Path:
-    """Get the current runs model_path for storing model checkpoints."""
-    if config or run_id:
-        run_id = run_id if run_id else get_run_id_from_config(config)
+def get_path_model(model_config: Config | None = None, run_id: str | None = None) -> Path:
+    """Get full_model_path if set, otherwise the shared per-run checkpoint directory."""
+    if model_config or run_id:
+        run_id = run_id if run_id else get_run_id_from_config(model_config)
     else:
-        msg = f"Missing run_id and cannot infer it from config: {config}"
+        msg = f"Missing run_id and cannot infer it from config: {model_config}"
         raise ValueError(msg)
-    return _get_shared_wg_path() / "models" / run_id
+    private_config = model_config if model_config is not None else _load_private_conf()
+    return _get_path_output(private_config, "full_model_path", "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}"
+def get_path_results(
+    model_config: Config,
+    mini_epoch: int | None = None,
+    step: int | None = None,
+) -> Path:
+    """Get the path for run results. Returns the results directory when mini_epoch is None."""
+    base_path = _get_path_output(
+        model_config, "path_results", "results", get_run_id_from_config(model_config)
+    )
+    if mini_epoch is None:
+        return base_path
+
+    ext = StoreType(model_config.zarr_store).value  # validate extension
+    default_name = f"validation_chkpt{mini_epoch:05d}_rank{model_config.rank:04d}.{ext}"
+    fname_template = model_config.get("output_name")
+    if fname_template is None:
+        fname = default_name
+    else:
+        try:
+            fname = fname_template.format(
+                epoch=mini_epoch,
+                step=step,
+                rank=model_config.rank,
+            )
+        except (IndexError, KeyError, ValueError):
+            fname = default_name
 
     return base_path / fname
 

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

@@ -14,8 +14,6 @@
 import pathlib
 from functools import cache
 
-from weathergen.common.config import _load_private_conf
-
 LOGGING_CONFIG = """
 {
     "version": 1,
@@ -96,10 +94,13 @@ def format(self, record, *args, **kwargs):
 
 
 @cache
-def init_loggers(run_id=None, logging_config=None):
+def init_loggers(log_path=None, logging_config=None):
     """
     Initialize the logger for the package and set output streams/files.
 
+    log_path is the directory resolved by config.get_path_logs(cf).
+    When omitted, the default logging configuration writes only to the console.
+
     WARNING: this function resets all the logging handlers.
 
     This function follows a singleton pattern, it will only operate once per process
@@ -115,21 +116,11 @@ def init_loggers(run_id=None, logging_config=None):
                 not supported
     """
 
-    # Get current time
-    # Shelved until decided how to change logging directory structure
-    # now = datetime.now()
-    # timestamp = now.strftime("%Y-%m-%d-%H%M")
-
-    # output_dir = f"./output/{timestamp}-{run_id}"
-    output_dir = ""
-    if run_id is not None:
-        output_dir = f"./logs/{run_id}"
-
     # load the structure for logging config
     if logging_config is None:
         logging_config = json.loads(LOGGING_CONFIG)
 
-        if run_id is None:
+        if log_path is None:
             del logging_config["handlers"]["logfile"]
             del logging_config["handlers"]["errorfile"]
             del logging_config["root"]["handlers"][2:]
@@ -139,12 +130,7 @@ def init_loggers(run_id=None, logging_config=None):
             if k == "formatter":
                 handler[k] = v
             elif k == "filename":
-                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():
-                    work_dir = pathlib.Path(_load_private_conf().get("path_shared_working_dir"))
-                    ofile = work_dir / ofile
+                ofile = pathlib.Path(log_path or ".") / v
                 pathlib.Path(ofile.parent).mkdir(parents=True, exist_ok=True)
                 handler[k] = ofile
             else:
@@ -153,7 +139,7 @@ def init_loggers(run_id=None, logging_config=None):
     # make sure the parent directory exists
     logging.config.dictConfig(logging_config)
 
-    if output_dir:
-        logging.info(f"Logging set up. Logs are in {output_dir}")
+    if log_path is not None:
+        logging.info(f"Logging set up. Logs are in {log_path}")
     else:
         logging.info("Logging set up. No log files created.")

packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py

@@ -21,7 +21,7 @@
 
 # Local application / package
 from weathergen.common.config import (
-    get_path_run,
+    get_path_results,
     load_merge_configs,
     load_run_config,
 )
@@ -52,7 +52,7 @@ def __init__(self, eval_cfg: dict, run_id: str, private_paths: dict | None = Non
         self.inference_cfg = self.get_inference_config()
 
         if not self.results_base_dir:
-            self.results_base_dir = get_path_run(self.inference_cfg)
+            self.results_base_dir = get_path_results(self.inference_cfg)
             _logger.info(f"Results directory obtained from private config: {self.results_base_dir}")
         else:
             _logger.info(f"Results directory parsed: {self.results_base_dir}")

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(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(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(log_path=config.get_path_logs(cf))
 
     logger.info(f"DDP initialization: rank={cf.rank}, world_size={cf.world_size}")
 

src/weathergen/train/trainer.py

@@ -156,10 +156,10 @@ def init(self, cf: Config, devices):
 
         # create output directory
         if is_root():
-            config.get_path_run(cf).mkdir(exist_ok=True, parents=True)
+            config.get_path_results(cf).mkdir(exist_ok=True, parents=True)
             config.get_path_model(cf).mkdir(exist_ok=True, parents=True)
 
-        self.train_logger = TrainLogger(cf, config.get_path_run(self.cf))
+        self.train_logger = TrainLogger(cf, config.get_path_results(self.cf))
 
         # Initialize collapse monitor for SSL training
         collapse_config = cf.train_logging.get("collapse_monitoring", {})

src/weathergen/utils/train_logger.py

@@ -58,9 +58,9 @@ def by_mode(self, s: str) -> pl.DataFrame:
 
 class TrainLogger:
     #######################################
-    def __init__(self, cf, path_run: Path) -> None:
+    def __init__(self, cf, path_results: Path) -> None:
         self.cf = cf
-        self.path_run = path_run
+        self.path_results = path_results
 
     def log_metrics(self, stage: Stage, metrics: dict[str, float], step: int | None = None) -> None:
         """
@@ -86,7 +86,7 @@ def log_metrics(self, stage: Stage, metrics: dict[str, float], step: int | None
         # but we can probably do better and rely for example on the logging module.
 
         metrics_path = get_train_metrics_path(
-            base_path=config.get_path_run(self.cf), run_id=self.cf.general.run_id
+            base_path=self.path_results, run_id=self.cf.general.run_id
         )
         with open(metrics_path, "ab") as f:
             s = json.dumps(clean_metrics) + "\n"
@@ -152,7 +152,7 @@ def read(
             )
         run_id = cf.general.run_id
 
-        result_dir_base = config.get_path_run(cf)
+        result_dir_base = config.get_path_results(cf)
 
         # define cols for training
         cols1 = [_weathergen_timestamp, "num_samples", "loss_avg_mean", "learning_rate"]

src/weathergen/utils/validation_io.py

@@ -187,7 +187,7 @@ def write_output(
         forecast_offset=forecast_offset,
     )
 
-    store_path = config.get_path_results(cf, mini_epoch)
+    store_path = config.get_path_results(cf, mini_epoch, batch_idx)
 
     with zarrio_writer(store_path) as zio:
         for subset in data.items():