feat: cache index of configs automatically instead of generating

Author: jmtoepperwien

carps/analysis/gather_data.py

@@ -21,6 +21,7 @@
 
 from carps.analysis.calc_hypervolume import add_hypervolume_to_df
 from carps.analysis.utils import convert_mixed_types_to_str, get_ids_mo
+from carps.utils.index_configs import get_index_config
 from carps.utils.loggingutils import get_logger, setup_logging
 from carps.utils.task import Task
 from carps.utils.trials import TrialInfo
@@ -404,9 +405,7 @@ def maybe_postadd_task(logs: pd.DataFrame, overwrite: bool = False) -> pd.DataFr
         logger.debug("No task_id in logs. Can't add task info.")
         return logs
     index_fn = Path(__file__).parent.parent / "configs/task/index.csv"
-    if not index_fn.is_file():
-        raise ValueError("ObjectiveFunction ids have not been indexed. Run `python -m carps.utils.index_configs`.")
-    task_index = pd.read_csv(index_fn)
+    task_index = get_index_config(index_fn)
 
     new_logs = []
     for gid, gdf in logs.groupby(by=["task_id", "seed"]):

carps/analysis/process_data.py

@@ -9,6 +9,7 @@
 import pandas as pd
 from omegaconf import DictConfig, ListConfig, OmegaConf
 
+from carps.utils.index_configs import get_index_config
 from carps.utils.loggingutils import get_logger, setup_logging
 
 setup_logging()
@@ -69,9 +70,7 @@ def maybe_postadd_task(logs: pd.DataFrame) -> pd.DataFrame:
         Logs with task columns.
     """
     index_fn = Path(__file__).parent.parent / "configs/task/index.csv"
-    if not index_fn.is_file():
-        raise ValueError("Task ids have not been indexed. Run `python -m carps.utils.index_configs`.")
-    task_index = pd.read_csv(index_fn)
+    task_index = get_index_config(index_fn)
 
     def load_task_cfg(task_id: str) -> DictConfig:
         config_fn = task_index["config_fn"][task_index["task_id"] == task_id].iloc[0]

carps/utils/index_configs.py

@@ -2,21 +2,30 @@
 
 from __future__ import annotations
 
+import hashlib
 from pathlib import Path
 
 import fire
 import pandas as pd
+from deprecated import deprecated
 from omegaconf import OmegaConf
+from platformdirs import user_cache_dir
 from rich.progress import track
 
 from carps.utils.loggingutils import get_logger
 
 logger = get_logger("ConfigIndexer")
 
+
 config_folder = Path(__file__).parent.parent / "configs"
 config_folder_task = config_folder / "task"
 config_folder_optimizer = config_folder / "optimizer"
 
+PATH_KEY_ZIP = {
+    config_folder_task: "task_id",
+    config_folder_optimizer: "optimizer_id",
+}
+
 
 def index_configs(extra_task_paths: list[str] | None = None, extra_optimizer_paths: list[str] | None = None) -> None:
     """Index all task and optimizer configs.
@@ -31,22 +40,13 @@ def index_configs(extra_task_paths: list[str] | None = None, extra_optimizer_pat
     extra_optimizer_paths : list[str], optional
         Extra paths to custom optimizers, must be a folder containing only optimizer configs.
     """
-    config_folder_tasks = [config_folder_task] if extra_task_paths is None else [config_folder_task, *extra_task_paths]  # type: ignore[list-item]
-    config_folder_tasks = [Path(p) for p in config_folder_tasks]
-    config_folder_optimizers = (
-        [config_folder_optimizer]
-        if extra_optimizer_paths is None
-        else [config_folder_optimizer, *extra_optimizer_paths]  # type: ignore[list-item]
-    )
-    config_folder_optimizers = [Path(p) for p in config_folder_optimizers]
-    for key, paths in zip(["task_id", "optimizer_id"], [config_folder_tasks, config_folder_optimizers], strict=False):
-        logger.info(f"Search configs for {key} from {paths}...")
-        filenames = []
-        for path in paths:
-            filenames.extend(list(path.glob("**/*.yaml")))
+    register_extra_paths(extra_task_paths, extra_optimizer_paths)
+
+    for path, key in PATH_KEY_ZIP.items():
+        paths = list(path.glob("**/*.yaml"))
 
         table_list = []
-        for fn in track(filenames, total=len(filenames), description=f"Gathering for {key}..."):
+        for fn in track(paths, total=len(paths), description=f"Gathering for {key}..."):
             cfg = OmegaConf.load(fn)
             value = cfg.get(key)
             table_list.append(
@@ -59,5 +59,78 @@ def index_configs(extra_task_paths: list[str] | None = None, extra_optimizer_pat
         table.to_csv(paths[0] / "index.csv", index=False)
 
 
+def create_table(key, paths: list[Path], target: Path) -> None:
+    """Create index table."""
+    table_list = []
+    for p in paths:
+        cfg = OmegaConf.load(p)
+        value = cfg.get(key)
+        table_list.append(
+            {
+                "config_fn": str(p),
+                key: value,
+            }
+        )
+    table = pd.DataFrame(table_list)
+    table.to_csv(target, index=False)
+
+
+def hash_inputs(paths: list[Path]) -> str:
+    """Hash inputs so that index file can be cached."""
+    hasher = hashlib.sha256()
+    for path in sorted(paths):
+        with open(path, "rb") as f:
+            while chunk := f.read(16 * 1024 * 1024):
+                hasher.update(chunk)
+    return hasher.hexdigest()
+
+
+def register_extra_paths(extra_task_paths: list[str] | None, extra_optimizer_paths: list[str] | None) -> None:
+    """Register extra task and optimizer paths.
+
+    Parameters
+    ----------
+    extra_task_paths : list[str]
+        Extra paths to custom tasks, must be a folder containing only task configs.
+    extra_optimizer_paths : list[str]
+        Extra paths to custom optimizers, must be a folder containing only optimizer configs.
+    """
+    if not extra_task_paths:
+        extra_task_paths = []
+    if not extra_optimizer_paths:
+        extra_optimizer_paths = []
+
+    for optimizer_path_str in extra_optimizer_paths:
+        PATH_KEY_ZIP[Path(optimizer_path_str)] = "optimizer_id"
+    for task_path_str in extra_task_paths:
+        PATH_KEY_ZIP[Path(task_path_str)] = "task_id"
+
+
+def get_index_config(path: Path) -> pd.DataFrame:
+    """Index all task and optimizer configs.
+
+    Create `index.csv` containing the config filename `config_fn` and the
+    `task_id` or `optimizer_id` for all task and optimizer configs.
+    Replaces old indexing api by using caching directory
+
+    Parameters:
+    ----------
+    path: path the old index file would have been
+
+    returns: pd.DataFrame containing the index
+    """
+    path_dashed = str(path.parent).replace("/", "-")
+
+    paths = list(path.parent.glob("**/*.yaml"))
+    paths_hash = hash_inputs(paths)[:12]
+
+    cache_path = Path(user_cache_dir("carps")) / f"index-{path_dashed}-{paths_hash}.csv"
+    if not cache_path.is_file():
+        cache_path.parent.mkdir(exist_ok=True, parents=True)
+        create_table(PATH_KEY_ZIP[path.parent], paths, cache_path)
+
+    return pd.read_csv(cache_path)
+
+
 if __name__ == "__main__":
     fire.Fire(index_configs)

carps/utils/overridefinder.py

@@ -6,9 +6,8 @@
 from pathlib import Path
 
 import fire
-import pandas as pd
 
-from carps.utils.index_configs import index_configs
+from carps.utils.index_configs import get_index_config
 from carps.utils.loggingutils import get_logger
 
 logger = get_logger(__file__)
@@ -45,9 +44,7 @@ def find_override(task_id: str | None = None, optimizer_id: str | None = None) -
         raise ValueError("Please specify either `task_id` or `optimizer_id`.")
 
     index_fn = path / "index.csv"
-    if not index_fn.is_file():
-        index_configs()
-    table = pd.read_csv(index_fn)
+    table = get_index_config(index_fn)
 
     try:
         config_fn = table["config_fn"][table[key] == to_find].to_numpy()[0]

pyproject.toml

@@ -66,6 +66,7 @@ dev = [
   "pre-commit",
   "ruff",
   "mypy",
+  "types-Deprecated",
   "mkdocs",
   "mkdocs-material",
   "mkdocs-autorefs",

subselection/create_subset_configs.py

@@ -22,9 +22,7 @@ def write_subsets(subset_fn: str, identifier: str):
         task_ids = subset["task_id"].to_list()
 
         index_fn = config_target_path.parent.parent / "index.csv"
-        if not index_fn.is_file():
-            raise ValueError(f"Could not find {index_fn}. ObjectiveFunction ids have not been indexed. Run `python -m carps.utils.index_configs`.")
-        task_index = pd.read_csv(index_fn)
+        task_index = get_index_config(index_fn)
         print(task_index.head())
         print(task_ids)
         not_found = [pid for pid in task_ids if pid not in task_index["task_id"].to_list()]