address review comments

Author: grassesi

src/weathergen/__init__.py

@@ -54,7 +54,7 @@ def evaluate_from_args(argl: list[str]):
         evaluate_overwrite,
         cli_overwrite,
     )
-    config.set_run_id(cf, args.run_id, args.reuse_run_id)
+    cf = config.set_run_id(cf, args.run_id, args.reuse_run_id)
 
     cf.run_history += [(args.from_run_id, cf.istep)]
 
@@ -103,7 +103,7 @@ def train_continue() -> None:
         *args.config,
         cli_overwrite,
     )
-    config.set_run_id(cf, args.run_id, args.reuse_run_id)
+    cf = config.set_run_id(cf, args.run_id, args.reuse_run_id)
 
     # track history of run to ensure traceability of results
     cf.run_history += [(args.from_run_id, cf.istep)]
@@ -144,7 +144,7 @@ def train_with_args(argl: list[str], stream_dir: str | None):
 
     cli_overwrite = config.from_cli_arglist(args.options)
     cf = config.load_config(args.private_config, None, None, *args.config, cli_overwrite)
-    config.set_run_id(cf, args.run_id, False)
+    cf = config.set_run_id(cf, args.run_id, False)
 
     if cf.with_flash_attention:
         assert cf.with_mixed_precision

src/weathergen/utils/config.py

@@ -129,20 +129,37 @@ def load_config(
 
 def set_run_id(config: Config, run_id: str | None, reuse_run_id: bool):
     """
-    Determine run_id of current run.
+    Determine and set run_id of current run.
+
+    Determining the run id should follow the following logic:
+
+    1. (default case): run train, train_continue or evaluate without any flags => generate a new run_id for this run.
+    2. (assign run_id): run train, train_continue or evaluate with --run_id <RUNID> flag => assign a run_id manually to this run
+    3. (reuse run_id -> only for train_continue and evaluate): reuse the run_id from the run specified by --from_run_id <RUNID>. Since the run_id correct run_id is already loaded in the config nothing has to be assigned. This case will happen if --reuse_run_id is specified.
+
 
     Args:
         config: Base configuration loaded from previous run or default.
         run_id: Id assigned to this run. If None a new one will be generated.
         reuse_run_id: Reuse run_id from base configuration instead.
+
+    Returns:
+        config object with the run_id attribute properly set.
     """
-    if not reuse_run_id:
+    config = config.copy()
+    if reuse_run_id:
+        assert config.run_id is not None, "run_id loaded from previous run should not be None."
+        _logger.info(f"reusing run_id from previous run: {config.run_id}")
+    else:
         if run_id is None:
-            run_id = get_run_id()
-
-        config.run_id = run_id
-
-    assert config.run_id is not None
+            # generate new id if run_id is None
+            config.run_id = run_id or get_run_id()
+            _logger.info(f"using generated run_id: {config.run_id}")
+        else:
+            config.run_id = run_id
+            _logger.info(f"using assigned run_id: {config.run_id}")
+
+    return config
 
 
 def from_cli_arglist(arg_list: list[str]) -> Config:

tests/test_config.py

@@ -145,7 +145,7 @@ def overwrite_file(overwrite_config):
 @pytest.fixture
 def config_fresh(private_config_file):
     cf = config.load_config(private_config_file, None, None)
-    config.set_run_id(cf, TEST_RUN_ID, False)
+    cf = config.set_run_id(cf, TEST_RUN_ID, False)
     cf.data_loader_rng_seed = 42
 
     return cf
@@ -205,14 +205,25 @@ def test_from_cli(options, cf):
     parsed_config = config.from_cli_arglist(options)
 
     assert parsed_config == OmegaConf.create(cf)
-@pytest.mark.parametrize("run_id,reuse,expected", [(None, False, "generated"), ("new_id", False, "new_id"), (None, True, TEST_RUN_ID), ("new_id", True, TEST_RUN_ID)])
+
+
+@pytest.mark.parametrize(
+    "run_id,reuse,expected",
+    [
+        (None, False, "generated"),
+        ("new_id", False, "new_id"),
+        (None, True, TEST_RUN_ID),
+        ("new_id", True, TEST_RUN_ID),
+    ],
+)
 def test_set_run_id(config_fresh, run_id, reuse, expected, mocker):
-    patch = mocker.patch("weathergen.utils.config.get_run_id", return_value="generated")
-    
-    config.set_run_id(config_fresh, run_id, reuse)    
-    
+    mocker.patch("weathergen.utils.config.get_run_id", return_value="generated")
+
+    config_fresh = config.set_run_id(config_fresh, run_id, reuse)
+
     assert config_fresh.run_id == expected
 
+
 def test_print_cf_no_secrets(config_fresh):
     output = config._format_cf(config_fresh)
 
@@ -252,4 +263,4 @@ def test_save(epoch, config_fresh):
     config.save(config_fresh, epoch)
 
     cf = config.load_model_config(config_fresh.run_id, epoch, config_fresh.model_path)
-    assert is_equal(cf, config_fresh)
+    assert is_equal(cf, config_fresh)