try dumping

Author: Samoed

autointent/_dump_tools.py

@@ -81,7 +81,7 @@ def dump(obj: Any, path: Path) -> None:  # noqa: ANN401, C901
                 joblib.dump(val, path / Dumper.estimators / key)
             elif isinstance(val, Ranker):
                 val.save(str(path / Dumper.cross_encoders / key))
-            elif isinstance(val, CrossEncoderConfig | EmbedderConfig):
+            elif isinstance(val, BaseModel):
                 try:
                     pydantic_path = path / Dumper.pydantic_models / f"{key}.json"
                     with pydantic_path.open("w", encoding="utf-8") as file:

autointent/context/_context.py

@@ -77,15 +77,9 @@ def dump(self) -> None:
         Save metrics, hyperparameters, inference, configurations, and datasets to disk.
         """
         self._logger.debug("dumping logs...")
-        optimization_results = self.optimization_info.dump_evaluation_results()
-
         logs_dir = self.logging_config.dirpath
-        logs_dir.mkdir(parents=True, exist_ok=True)
-
-        logs_path = logs_dir / "logs.json"
-        with logs_path.open("w") as file:
-            json.dump(optimization_results, file, indent=4, ensure_ascii=False, cls=NumpyEncoder)
 
+        self.optimization_info.dump(logs_dir)
         self.data_handler.dataset.to_json(logs_dir / "dataset.json")
 
         self._logger.info("logs and other assets are saved to %s", logs_dir)
@@ -95,6 +89,18 @@ def dump(self) -> None:
         with inference_config_path.open("w") as file:
             yaml.dump(inference_config, file)
 
+    def load(self) -> None:
+        """Load all information about optimization process from disk."""
+        self._logger.debug("loading logs...")
+        logs_dir = self.logging_config.dirpath
+        self.optimization_info.load(logs_dir)
+        if not self.optimization_info.artifacts.has_artifacts():
+            msg = (
+                "It is impossible to continue from the previous point, "
+                "start again with dump_modules=True settings if you want to resume the run"
+            )
+            raise RuntimeError(msg)
+
     def get_dump_dir(self) -> Path | None:
         """Get the directory for saving dumped modules.
 

autointent/context/optimization_info/_data_models.py

@@ -136,6 +136,15 @@ def get_best_artifact(self, node_type: str, idx: int) -> Artifact:
         """
         return self.get_artifacts(node_type)[idx]
 
+    def has_artifacts(self) -> bool:
+        """Check if any artifacts have been saved in RAM.
+
+        Returns:
+            True if any artifacts exist, False otherwise.
+        """
+        node_types = [NodeType.regex, NodeType.embedding, NodeType.scoring, NodeType.decision]
+        return any(len(self.get_artifacts(nt)) > 0 for nt in node_types)
+
 
 class Trial(BaseModel):
     """Representation of an individual optimization trial.

autointent/context/optimization_info/_optimization_info.py

@@ -10,11 +10,13 @@
 
 import numpy as np
 from numpy.typing import NDArray
+from pathlib import Path
 
 from autointent.configs import EmbedderConfig, InferenceNodeConfig
 from autointent.custom_types import NodeType
 
 from ._data_models import Artifact, Artifacts, EmbeddingArtifact, ScorerArtifact, Trial, Trials, TrialsIds
+from ..._dump_tools import Dumper
 
 if TYPE_CHECKING:
     from autointent.modules.base import BaseModule
@@ -56,6 +58,19 @@ def add_module(self, node_type: str, module: "BaseModule") -> None:
         """
         self.get(node_type).append(module)
 
+    def model_dump(self) -> dict[str, list["BaseModule"]]:
+        """Dump the modules to a dictionary format.
+
+        Returns:
+            Dictionary representation of the modules.
+        """
+        return {
+            "regex": self.regex,
+            "embedding": self.embedding,
+            "scoring": self.scoring,
+            "decision": self.decision,
+        }
+
 
 class OptimizationInfo:
     """Tracks optimization results, including trials, artifacts, and modules.
@@ -225,8 +240,18 @@ def dump_evaluation_results(self) -> dict[str, Any]:
             "pipeline_metrics": self.pipeline_metrics,
             "metrics": node_wise_metrics,
             "configs": self.trials.model_dump(),
+            "artifacts": self.artifacts.model_dump(),
+            "modules": self.modules.model_dump(),
         }
 
+    def dump(self, path: Path) -> None:
+        """Dump the optimization information to a file."""
+        Dumper.dump(self, path / "optimization_info")
+
+    def load(self, path: Path) -> None:
+        """Load the optimization information from a file."""
+        Dumper.load(self, path / "optimization_info")
+
     def get_inference_nodes_config(self, asdict: bool = False) -> list[InferenceNodeConfig]:
         """Generate configuration for inference nodes based on the best trials.
 

autointent/nodes/_node_optimizer.py

@@ -113,17 +113,14 @@ def fit(self, context: Context, sampler: SamplerType = "brute", n_jobs: int = 1)
             )
             self._counter = max(self._counter, finished_trials)
 
-            # if n_trials == 0:
-            #     config = self.suggest(Trial(), search_space)
-            #
-            #     module = self.node_info.modules_available[module_name]
-            #     module.load()
+            if n_trials == 0:
+                context.load()
 
             optuna.logging.set_verbosity(optuna.logging.WARNING)
             obj = partial(self.objective, module_name=module_name, search_space=search_space, context=context)
 
             study.optimize(obj, n_trials=n_trials, n_jobs=n_jobs)
-
+        context.dump()
         self._logger.info("%s node optimization is finished!", self.node_info.node_type)
 
     def objective(