diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py index 1f506fffb..ccc543cb1 100644 --- a/packages/common/src/weathergen/common/config.py +++ b/packages/common/src/weathergen/common/config.py @@ -208,7 +208,12 @@ 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. @@ -216,7 +221,9 @@ def load_run_config(run_id: str, mini_epoch: int | None, model_path: str | None) 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 diff --git a/packages/common/src/weathergen/common/logger.py b/packages/common/src/weathergen/common/logger.py index a32d44ba3..f543b93c4 100644 --- a/packages/common/src/weathergen/common/logger.py +++ b/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.") diff --git a/packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py b/packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py index 8e4ef8ca4..2d1d8cdf1 100644 --- a/packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py +++ b/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}") diff --git a/src/weathergen/model/model_interface.py b/src/weathergen/model/model_interface.py index c0b475b15..7617eaf68 100644 --- a/src/weathergen/model/model_interface.py +++ b/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" ) diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py index 7995b5864..5af219f7c 100644 --- a/src/weathergen/run_train.py +++ b/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}") diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py index 276da0bd6..13c60edc5 100644 --- a/src/weathergen/train/trainer.py +++ b/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", {}) diff --git a/src/weathergen/utils/train_logger.py b/src/weathergen/utils/train_logger.py index d7501a1cd..8d391e1bd 100644 --- a/src/weathergen/utils/train_logger.py +++ b/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"] diff --git a/src/weathergen/utils/validation_io.py b/src/weathergen/utils/validation_io.py index 6f989b2bb..9216ef635 100644 --- a/src/weathergen/utils/validation_io.py +++ b/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():