[569] Load eagerly the stream content in order (ecmwf#585)

* changes

* change

* changes

* Remove loading of streams also from inference.

---------

Co-authored-by: Christian Lessig <christian.lessig@ecmwf.int>

Author: tjhunter

src/weathergen/run_train.py

@@ -61,7 +61,6 @@ def inference_from_args(argl: list[str]):
     cf = config.set_run_id(cf, args.run_id, args.reuse_run_id)
 
     cf.run_history += [(args.from_run_id, cf.istep)]
-    cf.streams = config.load_streams(Path(cf.streams_directory))
     cf = config.set_paths(cf)
 
     trainer = Trainer()
@@ -115,7 +114,6 @@ def train_continue() -> None:
 
     # track history of run to ensure traceability of results
     cf.run_history += [(args.from_run_id, cf.istep)]
-    cf.streams = config.load_streams(Path(cf.streams_directory))
     cf = config.set_paths(cf)
 
     if args.finetune_forecast:

src/weathergen/utils/config.py

@@ -121,10 +121,14 @@ def load_config(
             # all the paths may be concatenated with ":"
             p = str(overwrite).split(":")
             for path in p:
-                overwrite_configs.append(_load_overwrite_conf(Path(path)))
+                c = _load_overwrite_conf(Path(path))
+                c = _load_streams_in_config(c)
+                overwrite_configs.append(c)
         else:
             # If it is a dict or DictConfig, we can directly use it
-            overwrite_configs.append(_load_overwrite_conf(overwrite))
+            c = _load_overwrite_conf(overwrite)
+            c = _load_streams_in_config(c)
+            overwrite_configs.append(c)
 
     if from_run_id is None:
         base_config = _load_default_conf()
@@ -135,6 +139,22 @@ def load_config(
     return OmegaConf.merge(base_config, private_config, *overwrite_configs)
 
 
+def _load_streams_in_config(config: Config) -> Config:
+    """If the config contains a streams_directory, loads the streams and returns the config with
+    the streams set."""
+    streams_directory = config.get("streams_directory", None)
+    config = config.copy()
+    if streams_directory is not None:
+        streams_directory = Path(streams_directory)
+        if not streams_directory.is_dir():
+            msg = f"Streams directory {streams_directory} does not exist."
+            raise FileNotFoundError(msg)
+
+        _logger.info(f"Loading streams from {streams_directory}")
+        config.streams = load_streams(streams_directory)
+    return config
+
+
 def set_run_id(config: Config, run_id: str | None, reuse_run_id: bool) -> Config:
     """
     Determine and set run_id of current run.