ft: registries now handle isolated input I/O for future concurrency-safe execution

Author: pabloitu

floatcsep/experiment.py

@@ -76,6 +76,7 @@ class Experiment:
 
         model_config (str): Path to the models' configuration file
         test_config (str): Path to the evaluations' configuration file
+        run_mode (str): 'sequential' or 'parallel'
         default_test_kwargs (dict): Default values for the testing
          (seed, number of simulations, etc.)
         postprocess (dict): Contains the instruction for postprocessing
@@ -100,6 +101,7 @@ def __init__(
         postprocess: str = None,
         default_test_kwargs: dict = None,
         rundir: str = "results",
+        run_mode: str = "sequential",
         report_hook: dict = None,
         **kwargs,
     ) -> None:
@@ -127,6 +129,7 @@ def __init__(
         self.original_config = kwargs.get("original_config", None)
         self.original_run_dir = kwargs.get("original_rundir", None)
         self.run_dir = rundir
+        self.run_mode = run_mode
         self.seed = kwargs.get("seed", None)
         self.time_config = read_time_cfg(time_config, **kwargs)
         self.region_config = read_region_cfg(region_config, **kwargs)
@@ -296,7 +299,7 @@ def stage_models(self) -> None:
         """
         log.info("Staging models")
         for i in self.models:
-            i.stage(self.time_windows)
+            i.stage(self.time_windows, run_mode=self.run_mode, run_dir=self.run_dir)
             self.registry.add_model_registry(i)
 
     def set_tests(self, test_config: Union[str, Dict, List]) -> list:
@@ -377,7 +380,7 @@ def set_tasks(self) -> None:
         """
 
         # Set the file path structure
-        self.registry.build_tree(self.time_windows, self.models, self.tests)
+        self.registry.build_tree(self.time_windows, self.models, self.tests, self.run_mode)
 
         log.debug("Pre-run forecast summary")
         log_models_tree(log, self.registry, self.time_windows)

floatcsep/infrastructure/registries.py

@@ -3,7 +3,7 @@
 from abc import ABC, abstractmethod
 from datetime import datetime
 from os.path import join, abspath, relpath, normpath, dirname, exists
-from typing import Sequence, Union, TYPE_CHECKING, Any
+from typing import Sequence, Union, TYPE_CHECKING, Any, Optional
 
 from floatcsep.utils.helpers import timewindow2str
 
@@ -149,6 +149,7 @@ def factory(cls, registry_type: str = 'file', **kwargs) -> "ModelRegistry":
 class ModelFileRegistry(ModelRegistry, FilepathMixin):
     def __init__(
         self,
+        model_name: str,
         workdir: str,
         path: str,
         database: str = None,
@@ -159,19 +160,25 @@ def __init__(
         """
 
         Args:
+            model_name (str): Model's identifier string
             workdir (str): The current working directory of the experiment.
             path (str): The path of the model working directory (or model filepath).
             database (str): The path of the database, in case forecasts are stored therein.
             args_file (str): The path of the arguments file (only for TimeDependentModel).
             input_cat (str): : The path of the arguments file (only for TimeDependentModel).
         """
 
+        self.model_name = model_name
         self.workdir = workdir
         self.path = path
         self.database = database
         self.args_file = args_file
         self.input_cat = input_cat
+
         self.forecasts = {}
+        self.input_args = {}
+        self.input_cats = {}
+        self.input_store = None
         self._fmt = fmt
 
     @property
@@ -227,7 +234,7 @@ def get_input_catalog_key(self, *args: Sequence[str]) -> str:
         Returns:
            The input catalog registry key from a sequence of key values
         """
-        return self.get_attr("input_cat", *args)
+        return self.get_attr("input_cats", *args)
 
     def get_forecast_key(self, *args: Sequence[str]) -> str:
         """
@@ -253,15 +260,15 @@ def get_args_key(self, *args: Sequence[str]) -> str:
         Returns:
            The argument file's key(s) from a sequence of key values
         """
-        return self.get_attr("args_file", *args)
+        return self.get_attr("input_args", *args)
 
     def build_tree(
         self,
         time_windows: Sequence[Sequence[datetime]] = None,
         model_class: str = "TimeIndependentModel",
         prefix: str = None,
-        args_file: str = None,
-        input_cat: str = None
+        run_mode: str = 'sequential',
+        run_dir: Optional[str] = None
     ) -> None:
         """
         Creates the run directory, and reads the file structure inside.
@@ -270,33 +277,43 @@ def build_tree(
             time_windows (list(str)): List of time windows or strings.
             model_class (str): Model's class name
             prefix (str): prefix of the model forecast filenames if TD
-            args_file (str, bool): input arguments path of the model if TD
-            input_cat (str, bool): input catalog path of the model if TD
-            fmt (str, bool): for time dependent mdoels
-
+            run_mode (str): if run mode is sequential, input data (args and cat) will be
+                dynamically overwritten in 'model/input/`  through time_windows. If 'parallel',
+                input data is dynamically writing anew in
+                'results/{time_window}/input/{model_name}/'.
+            run_dir (str): Where experiment's results are stored.
         """
 
         windows = timewindow2str(time_windows)
-
         if model_class == "TimeIndependentModel":
             fname = self.database if self.database else self.path
             self.forecasts = {win: fname for win in windows}
 
         elif model_class == "TimeDependentModel":
 
-            args = args_file if args_file else join("input", "args.txt")
-            self.args_file = join(self.path, args)
-            input_cat = input_cat if input_cat else join("input", "catalog.csv")
-            self.input_cat = join(self.path, input_cat)
-            # grab names for creating directories
+            # grab names for creating model directories
             subfolders = ["input", "forecasts"]
             dirtree = {folder: self.abs(self.path, folder) for folder in subfolders}
-
-            # create directories if they don't exist
             for _, folder_ in dirtree.items():
                 os.makedirs(folder_, exist_ok=True)
 
-            # set forecast names
+            if run_mode == 'sequential':
+                self.input_args = {
+                    win: join(self.path, 'input', self.args_file) for win in windows
+                }
+                self.input_cats = {
+                    win: join(self.path, 'input', self.input_cat) for win in windows
+                }
+            elif run_mode == 'parallel':
+                self.input_args = {
+                    win: join(run_dir, win,  'input', self.model_name, self.args_file)
+                    for win in windows
+                }
+                self.input_cats = {
+                    win: join(run_dir, win,  'input', self.model_name, self.input_cat)
+                    for win in windows
+                }
+
             self.forecasts = {
                 win: join(dirtree["forecasts"], f"{prefix}_{win}.{self.fmt}") for win in windows
             }
@@ -492,6 +509,7 @@ def build_tree(
         time_windows: Sequence[Sequence[datetime]],
         models: Sequence["Model"],
         tests: Sequence["Evaluation"],
+        run_mode: str = 'sequential'
     ) -> None:
         """
         Creates the run directory and reads the file structure inside.
@@ -500,6 +518,7 @@ def build_tree(
             time_windows: List of time windows, or representing string.
             models: List of models or model names
             tests: List of tests or test names
+            run_mode: 'parallel' or 'sequential'
 
         """
         windows = timewindow2str(time_windows)
@@ -509,6 +528,8 @@ def build_tree(
 
         run_folder = self.run_dir
         subfolders = ["catalog", "evaluations", "figures"]
+        if run_mode == 'parallel':
+            subfolders.append('input')
         dirtree = {
             win: {folder: self.abs(run_folder, win, folder) for folder in subfolders}
             for win in windows
@@ -518,7 +539,11 @@ def build_tree(
         for tw, tw_folder in dirtree.items():
             for _, folder_ in tw_folder.items():
                 os.makedirs(folder_, exist_ok=True)
-
+                print(folder_)
+                if run_mode == 'parallel' and folder_.endswith('input'):
+                    print('a')
+                    for model in models:
+                        os.makedirs(join(folder_, model), exist_ok=True)
         results = {
             win: {
                 test: {

floatcsep/infrastructure/repositories.py

@@ -202,7 +202,7 @@ def set_input_cat(self, tstring: str, model: "Model", fmt: str ='ascii') -> None
             fmt (str): Output catalog format
         """
         start, end = str2timewindow(tstring)
-        input_cat_name = model.registry.get_input_catalog_key()
+        input_cat_name = model.registry.get_input_catalog_key(tstring)
         sub_cat = self.catalog.filter([f"origin_time < {start.timestamp() * 1000}"])
 
         writer = getattr(CatalogSerializer, fmt)

floatcsep/model.py

@@ -210,13 +210,14 @@ def __init__(self, name: str, model_path: str, forecast_unit=1, store_db=False,
 
         self.forecast_unit = forecast_unit
         self.store_db = store_db
-        self.registry = ModelRegistry.factory(workdir=kwargs.get("workdir", os.getcwd()),
-                                          path=model_path)
+        self.registry = ModelRegistry.factory(model_name=name,
+                                              workdir=kwargs.get("workdir", os.getcwd()),
+                                              path=model_path)
         self.repository = ForecastRepository.factory(
             self.registry, model_class=self.__class__.__name__, **kwargs
         )
 
-    def stage(self, time_windows: Sequence[Sequence[datetime]] = None) -> None:
+    def stage(self, time_windows: Sequence[Sequence[datetime]] = None, **kwargs) -> None:
         """
         Acquire the forecast data if it is not in the file system. Sets the paths internally
         (or database pointers) to the forecast data.
@@ -298,6 +299,8 @@ def __init__(
         model_path: str,
         func: Union[str, Callable] = None,
         func_kwargs: dict = None,
+        args_file: str = "args.txt",
+        input_cat: str = "catalog.csv",
         fmt: str = 'csv',
         **kwargs,
     ) -> None:
@@ -310,6 +313,8 @@ def __init__(
             func: A function/command that runs the model.
             func_kwargs: The keyword arguments to run the model. They are usually (over)written
                 into the file `{model_path}/input/{args_file}`
+            args_file: Name of the arguments file that will be used to create forecasts
+            input_cat: Name of the file that will be used as input catalog to create forecasts
             **kwargs: Additional keyword parameters, such as a ``prefix`` (str) for the
                 resulting forecast file paths, ``args_file`` (str) as the path for the model
                 arguments file or ``input_cat`` that indicates where the input catalog will be
@@ -321,9 +326,12 @@ def __init__(
         self.func = func
         self.func_kwargs = func_kwargs or {}
 
-        self.registry = ModelRegistry.factory(workdir=kwargs.get("workdir", os.getcwd()),
-                                          path=model_path,
-                                          fmt=fmt)
+        self.registry = ModelRegistry.factory(model_name=name,
+                                              workdir=kwargs.get("workdir", os.getcwd()),
+                                              path=model_path,
+                                              fmt=fmt,
+                                              args_file=args_file,
+                                              input_cat=input_cat)
         self.repository = ForecastRepository.factory(
             self.registry, model_class=self.__class__.__name__, **kwargs
         )
@@ -334,7 +342,7 @@ def __init__(
                 self.build, self.name, self.registry.abs(model_path)
             )
 
-    def stage(self, time_windows=None) -> None:
+    def stage(self, time_windows=None, run_mode='sequential', run_dir='') -> None:
         """
         Core method to interface a model with the experiment.
 
@@ -355,8 +363,8 @@ def stage(self, time_windows=None) -> None:
             time_windows=time_windows,
             model_class=self.__class__.__name__,
             prefix=self.__dict__.get("prefix", self.name),
-            args_file=self.__dict__.get("args_file", None),
-            input_cat=self.__dict__.get("input_cat", None),
+            run_mode=run_mode,
+            run_dir=run_dir
         )
 
     def get_forecast(
@@ -407,7 +415,7 @@ def create_forecast(self, tstring: str, **kwargs) -> None:
             f"Running {self.name} using {self.environment.__class__.__name__}:"
             f" {timewindow2str([start_date, end_date])}"
         )
-        self.environment.run_command(f"{self.func} {self.registry.get_args_key()}")
+        self.environment.run_command(f"{self.func} {self.registry.get_args_key(tstring)}")
 
     def prepare_args(self, start: datetime, end: datetime, **kwargs) -> None:
         """
@@ -422,7 +430,9 @@ def prepare_args(self, start: datetime, end: datetime, **kwargs) -> None:
             **kwargs: represents additional model arguments (name/value pair)
 
         """
-        filepath = self.registry.get_args_key()
+        window_str = timewindow2str([start, end])
+
+        filepath = self.registry.get_args_key(window_str)
         fmt = os.path.splitext(filepath)[1]
 
         if fmt == ".txt":
@@ -477,8 +487,14 @@ def nested_update(dest: dict, src: dict, max_depth: int = 3, _level: int = 1):
                     else:
                         dest[key] = val
 
+            if not os.path.exists(filepath):
+                template_file = os.path.join(self.registry.path,
+                                             'input',
+                                             self.registry.args_file)
+            else:
+                template_file = filepath
 
-            with open(filepath, "r") as file_:
+            with open(template_file, "r") as file_:
                 args = yaml.safe_load(file_)
             args["start_date"] = start.isoformat()
             args["end_date"] = end.isoformat()

tests/unit/test_model.py

@@ -198,14 +198,15 @@ def setUp(self):
         self.model = TimeDependentModel(
             name=self.name, model_path=self.model_path, func=self.func
         )
-
     def tearDown(self):
         patch.stopall()
 
     def test_init(self):
         # Assertions to check if the components were instantiated correctly
         self.mock_registry_factory.assert_called_once_with(
-            workdir=os.getcwd(), path=self.model_path, fmt='csv'
+            model_name=self.name,
+            workdir=os.getcwd(), path=self.model_path, fmt='csv',
+            args_file='args.txt', input_cat='catalog.csv'
         )  # Ensure the registry is initialized correctly
         self.mock_repository_factory.assert_called_once_with(
             self.mock_registry_instance, model_class="TimeDependentModel"
@@ -233,8 +234,8 @@ def test_stage(self, mk):
             time_windows=["2020-01-01_2020-12-31"],
             model_class="TimeDependentModel",
             prefix=self.model.__dict__.get("prefix", self.name),
-            args_file=self.model.__dict__.get("args_file", None),
-            input_cat=self.model.__dict__.get("input_cat", None),
+            run_mode='sequential',
+            run_dir=''
         )
         self.mock_environment_instance.create_environment.assert_called_once()
 

tests/unit/test_registry.py

@@ -9,10 +9,16 @@ class TestModelFileRegistry(unittest.TestCase):
 
     def setUp(self):
         self.registry_for_filebased_model = ModelFileRegistry(
-            workdir="/test/workdir", path="/test/workdir/model.txt"
+            model_name='test',
+            workdir="/test/workdir",
+            path="/test/workdir/model.txt"
         )
         self.registry_for_folderbased_model = ModelFileRegistry(
-            workdir="/test/workdir", path="/test/workdir/model"
+            model_name='test',
+            workdir="/test/workdir",
+            path="/test/workdir/model",
+            args_file="args.txt",
+            input_cat="catalog.csv"
         )
 
     def test_call(self):
@@ -83,6 +89,7 @@ def test_build_tree_time_dependent(self, mock_listdir, mock_makedirs):
             [datetime(2023, 1, 1), datetime(2023, 1, 2)],
             [datetime(2023, 1, 2), datetime(2023, 1, 3)],
         ]
+        print(self.registry_for_folderbased_model.__dict__)
         self.registry_for_folderbased_model.build_tree(
             time_windows=time_windows, model_class="TimeDependentModel", prefix="forecast"
         )

tutorials/case_h/models.yml

@@ -2,7 +2,7 @@
     giturl: https://git.gfz-potsdam.de/csep/it_experiment/models/vetas.git
     repo_hash: v3.2
     path: models/etas
-    args_file: input/args.json
+    args_file: args.json
     func: etas-run
     func_kwargs:
       n_sims: 100

tutorials/case_i/config.yml

@@ -14,6 +14,7 @@ region_config:
   depth_min: 0
   depth_max: 70
 
+run_mode: parallel
 force_rerun: True
 catalog: catalog.csep
 model_config: models.yml