try to fix

Author: voorhs

autointent/context/_context.py

@@ -88,14 +88,19 @@ def dump(self) -> None:
             yaml.dump(inference_config, file)
 
     def load(self) -> None:
-        """Load all information about optimization process from disk."""
+        """Restore the context state to resume the optimization process.
+
+        Raises:
+            RuntimeError: If the modules artifacts are not found.
+        """
         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"
+                "start again with dump_modules=True settings if you want to resume the run."
+                "To load optimization info only, use Context.optimization_info.load(logs_dir)."
             )
             raise RuntimeError(msg)
 

autointent/nodes/_node_optimizer.py

@@ -2,6 +2,7 @@
 
 import gc
 import itertools as it
+import json
 import logging
 from copy import deepcopy
 from functools import partial
@@ -106,16 +107,13 @@ def fit(self, context: Context, sampler: SamplerType = "brute", n_jobs: int = 1)
 
             study, finished_trials, n_trials = load_or_create_study(
                 study_name=f"{self.node_info.node_type}_{module_name}",
-                storage_dir=context.get_dump_dir(),
+                context=context,
                 direction="maximize",
                 sampler=sampler_instance,
                 n_trials=n_trials,
             )
             self._counter = max(self._counter, finished_trials)
 
-            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)
 
@@ -143,7 +141,7 @@ def objective(
         """
         config = self.suggest(trial, search_space)
 
-        self._logger.debug("Initializing %s module...", module_name)
+        self._logger.debug("Initializing %s module with config: %s", module_name, json.dumps(config))
         module = self.node_info.modules_available[module_name].from_context(context, **config)
 
         embedder_config = module.get_embedder_config()
@@ -338,7 +336,7 @@ def get_storage_url(study_name: str, storage_dir: Path | None) -> str | None:
 
 def load_or_create_study(
     study_name: str,
-    storage_dir: Path | None,
+    context: Context,
     sampler: optuna.samplers.BaseSampler,
     direction: str = "maximize",
     n_trials: int = 10,
@@ -347,7 +345,7 @@ def load_or_create_study(
 
     Args:
         study_name: Name of the study
-        storage_dir: Directory where study databases are stored
+        context: Context object
         direction: Optimization direction (maximize or minimize)
         sampler: Optuna sampler instance
         n_trials: n_trials
@@ -358,7 +356,7 @@ def load_or_create_study(
     remaining_trials = n_trials
     finished_trials = 0
 
-    storage_url = get_storage_url(study_name, storage_dir)
+    storage_url = get_storage_url(study_name, context.get_dump_dir())
 
     try:
         # will catch exception if study does not exist
@@ -373,6 +371,8 @@ def load_or_create_study(
             finished_trials = max(t.number for t in study.trials) + 1
             # Calculate remaining trials if n_trials is specified
             remaining_trials = n_trials if n_trials is None else max(0, n_trials - len(study.trials))
+        if remaining_trials == 0:
+            context.load()
         return study, finished_trials, remaining_trials  # noqa: TRY300
     except Exception:  # noqa: BLE001
         # Create a new study if none exists