Sgrasse/develop/issue 185 (ecmwf#221)

* disable prefix parsing for CLI

* remove run_id flag fresh training

* separate cli definition into its own module

* enable multiple config overwrites from different sources

* add method to parse cli arglist into config

* pack extra overwrites for forecasting into dict and supply as overwrite

* add `--config` arg to train_continue to harmonize cli interfaces

* pass overwrite constructed from additional cli args to load_config

* make print_cf testable, filter secrets

* change: raise ValueError for unknown type argument

* add tests for the configuration

* test config.load_streams, config.save

* improve logging/error handling of load_streams

* refactor: centralize naming of run config files

* align semantics and default value of `--epoch`

* unify of `--eval_run_id` and `--run_id_new`.

- `--eval_run_id` (for evaluation)
- `--run_id_new` (for continue)
- `--run_id` (for train)
are coalesced into `--run_id` which now indicates
always the resulting `run_id` of a run.

the previous `--run_id` flag used to specify a
model to load is renamed into `--from_run_id`.

* factor out shared cli argument definitions

* remove evaluation args that can be overwritten using new mechanism

* move liminting of workers by samples into evaluation method

* make `--analysis_streams_output` and `--config` accept multiple items

* clean up evaluate_from_args

* change implementation of how to handle direct option overwrites: use argument --options

* move date formatting to cli

* make flags store_true and assure nargs='+' always returns list

* improve typehints and cli help text

Author: grassesi

integration_tests/small1.py

@@ -57,9 +57,9 @@ def test_train(setup, test_run_id):
     evaluate_from_args(
         "-start 2022-10-10 -end 2022-10-11 --samples 10 --epoch 0".split()
         + [
-            "--run_id",
+            "--from_run_id",
             test_run_id,
-            "--eval_run_id",
+            "--run_id",
             test_run_id,
             "--config",
             f"{weathergen_home}/integration_tests/small1.yaml",

src/weathergen/__init__.py

@@ -7,15 +7,12 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-import argparse
 import pdb
 import sys
 import time
 import traceback
-from pathlib import Path
-
-import pandas as pd
 
+import weathergen.utils.cli as cli
 import weathergen.utils.config as config
 from weathergen.train.trainer import Trainer
 from weathergen.utils.logger import init_loggers
@@ -33,189 +30,90 @@ def evaluate_from_args(argl: list[str]):
 
     When running integration tests, the arguments are directly provided.
     """
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-        "--run_id",
-        type=str,
-        required=True,
-        help="Run/model id of pretrained WeatherGenerator model.",
-    )
-    parser.add_argument(
-        "--start_date",
-        "-start",
-        type=str,
-        required=False,
-        default="2022-10-01",
-        help="Start date for evaluation. Format must be parsable with pd.to_datetime.",
-    )
-    parser.add_argument(
-        "--end_date",
-        "-end",
-        type=str,
-        required=False,
-        default="2022-12-01",
-        help="End date for evaluation. Format must be parsable with pd.to_datetime.",
-    )
-    parser.add_argument(
-        "--epoch",
-        type=int,
-        default=None,
-        help="Epoch of pretrained WeatherGenerator model used for evaluation (Default None corresponds to the last checkpoint).",
-    )
-    parser.add_argument(
-        "--forecast_steps",
-        type=int,
-        default=None,
-        help="Number of forecast steps for evaluation. Uses attribute from config when None is set.",
-    )
-    parser.add_argument(
-        "--samples", type=int, default=10000000, help="Number of evaluation samples."
-    )
-    parser.add_argument(
-        "--shuffle", type=bool, default=False, help="Shuffle samples from evaluation."
-    )
-    parser.add_argument(
-        "--save_samples", type=bool, default=True, help="Save samples from evaluation."
-    )
-    parser.add_argument(
-        "--analysis_streams_output",
-        type=list,
-        default=["ERA5"],
-        help="Analysis output streams during evaluation.",
-    )
-    parser.add_argument(
-        "--private_config",
-        type=Path,
-        default=None,
-        help="Path to private configuration file for paths.",
-    )
-    parser.add_argument(
-        "--eval_run_id",
-        type=str,
-        required=False,
-        dest="eval_run_id",
-        help="(optional) if specified, uses the provided run id to store the evaluation results",
-    )
-    parser.add_argument(
-        "--config",
-        type=Path,
-        default=None,
-        help="Optional experiment specfic configuration file",
-    )
-
+    parser = cli.get_evaluate_parser()
     args = parser.parse_args(argl)
 
     # TODO: move somewhere else
     init_loggers()
 
-    cf = config.load_config(args.private_config, args.run_id, args.epoch, args.config)
-
-    cf.run_history += [(cf.run_id, cf.istep)]
-
-    cf.samples_per_validation = args.samples
-    cf.log_validation = args.samples if args.save_samples else 0
-    start_date, end_date = pd.to_datetime(args.start_date), pd.to_datetime(args.end_date)
-
-    cf.start_date_val = start_date.strftime("%Y%m%d%H%M")
-    cf.end_date_val = end_date.strftime("%Y%m%d%H%M")
-
-    cf.shuffle = args.shuffle
-
-    cf.forecast_steps = args.forecast_steps if args.forecast_steps else cf.forecast_steps
-    # cf.forecast_policy = 'fixed'
+    evaluate_overwrite = dict(
+        shuffle=False,
+        start_date_val=args.start_date,
+        end_date_val=args.end_date,
+        samples_per_validation=args.samples,
+        log_validation=args.samples if args.save_samples else 0,
+        analysis_streams_output=args.analysis_streams_output,
+    )
 
-    # cf.analysis_streams_output = ['Surface', 'Air', 'METEOSAT', 'ATMS', 'IASI', 'AMSR2']
-    cf.analysis_streams_output = args.analysis_streams_output
+    cli_overwrite = config.from_cli_arglist(args.options)
+    cf = config.load_config(
+        args.private_config,
+        args.from_run_id,
+        args.epoch,
+        *args.config,
+        evaluate_overwrite,
+        cli_overwrite,
+    )
 
-    # make sure number of loaders does not exceed requested samples
-    cf.loader_num_workers = min(cf.loader_num_workers, args.samples)
+    cf.run_history += [(args.from_run_id, cf.istep)]
 
     trainer = Trainer()
-    trainer.evaluate(cf, args.run_id, args.epoch, run_id_new=args.eval_run_id)
+    trainer.evaluate(cf, args.from_run_id, args.epoch, run_id_new=args.run_id)
 
 
 ####################################################################################################
 def train_continue() -> None:
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-        "-id",
-        "--run_id",
-        type=str,
-        required=True,
-        help="run id of to be continued",
-    )
-    parser.add_argument(
-        "-e",
-        "--epoch",
-        type=int,
-        required=False,
-        default=-1,
-        help="epoch where to continue run",
-    )
-    parser.add_argument(
-        "-n",
-        "--run_id_new",
-        type=bool,
-        required=False,
-        default=False,
-        help="create new run id for cont'd run",
-    )
-    parser.add_argument(
-        "--private_config",
-        type=Path,
-        default=None,
-        help="Path to private configuration file for paths.",
-    )
-    parser.add_argument(
-        "--finetune_forecast",
-        action="store_true",
-        help="Fine tune for forecasting. It overwrites some of the Config settings.",
-    )
-
+    parser = cli.get_continue_parser()
     args = parser.parse_args()
 
-    cf = config.load_config(args.private_config, args.run_id, args.epoch, None)
+    if args.finetune_forecast:
+        finetune_overwrite = dict(
+            training_mode="forecast",
+            forecast_delta_hrs=0,  # 12
+            forecast_steps=1,  # [j for j in range(1,9) for i in range(4)]
+            forecast_policy="fixed",  # 'sequential_random' # 'fixed' #'sequential' #_random'
+            forecast_freeze_model=True,
+            forecast_att_dense_rate=1.0,  # 0.25
+            fe_num_blocks=8,
+            fe_num_heads=16,
+            fe_dropout_rate=0.1,
+            fe_with_qk_lnorm=True,
+            lr_start=0.000001,
+            lr_max=0.00003,
+            lr_final_decay=0.00003,
+            lr_final=0.0,
+            lr_steps_warmup=1024,
+            lr_steps_cooldown=4096,
+            lr_policy_warmup="cosine",
+            lr_policy_decay="linear",
+            lr_policy_cooldown="linear",
+            num_epochs=12,  # len(cf.forecast_steps) + 4
+            istep=0,
+        )
+    else:
+        finetune_overwrite = dict()
+
+    cli_overwrite = config.from_cli_arglist(args.options)
+    cf = config.load_config(
+        args.private_config,
+        args.from_run_id,
+        args.epoch,
+        finetune_overwrite,
+        *args.config,
+        cli_overwrite,
+    )
 
     # track history of run to ensure traceability of results
-    cf.run_history += [(cf.run_id, cf.istep)]
+    cf.run_history += [(args.from_run_id, cf.istep)]
 
-    #########################
     if args.finetune_forecast:
-        cf.training_mode = "forecast"
-        cf.forecast_delta_hrs = 0  # 12
-        cf.forecast_steps = 1  # [j for j in range(1,9) for i in range(4)]
-        cf.forecast_policy = "fixed"  # 'sequential_random' # 'fixed' #'sequential' #_random'
-        cf.forecast_freeze_model = True
-        cf.forecast_att_dense_rate = 1.0  # 0.25
-
         if cf.forecast_freeze_model:
             cf.with_fsdp = False
             import torch
 
             torch._dynamo.config.optimize_ddp = False
-
-        cf.fe_num_blocks = 8
-        cf.fe_num_heads = 16
-        cf.fe_dropout_rate = 0.1
-        cf.fe_with_qk_lnorm = True
-
-        cf.lr_start = 0.000001
-        cf.lr_max = 0.00003
-        cf.lr_final_decay = 0.00003
-        cf.lr_final = 0.0
-        cf.lr_steps_warmup = 1024
-        cf.lr_steps_cooldown = 4096
-        cf.lr_policy_warmup = "cosine"
-        cf.lr_policy_decay = "linear"
-        cf.lr_policy_cooldown = "linear"
-
-        cf.num_epochs = 12  # len(cf.forecast_steps) + 4
-        cf.istep = 0
-
     trainer = Trainer()
-    trainer.run(cf, args.run_id, args.epoch, args.run_id_new)
+    trainer.run(cf, args.from_run_id, args.epoch, run_id_new=args.run_id)
 
 
 ####################################################################################################
@@ -236,34 +134,14 @@ def train() -> None:
 def train_with_args(argl: list[str], stream_dir: str | None):
     """
     Training function for WeatherGenerator model."""
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-        "--run_id",
-        type=str,
-        default=None,
-        help="Run id",
-    )
-    parser.add_argument(
-        "--private_config",
-        type=Path,
-        default=None,
-        help="Path to private configuration file for paths",
-    )
-    parser.add_argument(
-        "--config",
-        type=Path,
-        default=None,
-        help="Optional experiment specfic configuration file",
-    )
-
+    parser = cli.get_train_parser()
     args = parser.parse_args(argl)
 
     # TODO: move somewhere else
     init_loggers()
 
-    cf = config.load_config(args.private_config, None, None, args.config)
-    cf.run_id = args.run_id
+    cli_overwrite = config.from_cli_arglist(args.options)
+    cf = config.load_config(args.private_config, None, None, *args.config, cli_overwrite)
 
     if cf.with_flash_attention:
         assert cf.with_mixed_precision

src/weathergen/train/trainer.py

@@ -100,11 +100,13 @@ def evaluate(self, cf, run_id_trained, epoch, run_id_new=False):
             shuffle=cf.shuffle,
         )
 
+        # make sure number of loaders does not exceed requested samples
+        loader_num_workers = min(cf.samples_per_validation, cf.loader_num_workers)
         loader_params = {
             "batch_size": None,
             "batch_sampler": None,
             "shuffle": False,
-            "num_workers": cf.loader_num_workers,
+            "num_workers": loader_num_workers,
             "pin_memory": True,
         }
         self.data_loader_validation = torch.utils.data.DataLoader(