Refactor/artifacts (#199)

* get rid of trials_ids storing

* bug fix

* fix typing errors

* adjust context loading to minimize errors

* optimize memory occupied by modules' dumps

* bug fix

* add proper modules dumping

* bug fix

* remove nodes tests

Author: voorhs

.github/workflows/test-nodes.yaml

@@ -69,6 +69,10 @@ def set_dataset(self, dataset: Dataset, config: DataConfig) -> None:
         """
         self.data_handler = DataHandler(dataset=dataset, random_seed=self.seed, config=config)
 
+    def dump_optimization_info(self) -> None:
+        """Save optimization info to disk."""
+        self.optimization_info.dump(self.logging_config.dirpath)
+
     def dump(self) -> None:
         """Save all information about optimization process to disk.
 
@@ -77,7 +81,7 @@ def dump(self) -> None:
         self._logger.debug("dumping logs...")
         logs_dir = self.logging_config.dirpath
 
-        self.optimization_info.dump(logs_dir)
+        self.dump_optimization_info()
         self.data_handler.dataset.to_json(logs_dir / "dataset.json")
 
         self._logger.info("logs and other assets are saved to %s", logs_dir)
@@ -87,7 +91,7 @@ def dump(self) -> None:
         with inference_config_path.open("w") as file:
             yaml.dump(inference_config, file)
 
-    def load(self) -> None:
+    def load_optimization_info(self) -> None:
         """Restore the context state to resume the optimization process.
 
         Raises:
@@ -124,7 +128,7 @@ def is_ram_to_clear(self) -> bool:
     def has_saved_modules(self) -> bool:
         """Check if any modules have been saved in RAM."""
         node_types = ["regex", "embedding", "scoring", "decision"]
-        return any(len(self.optimization_info.modules.get(nt)) > 0 for nt in node_types)
+        return any(self.optimization_info.modules.get(nt) is not None for nt in node_types)
 
     def resolve_embedder(self) -> EmbedderConfig:
         """Resolve the embedder configuration.

autointent/context/_context.py

@@ -114,78 +114,87 @@ def validate_node_name(value: str) -> str:
 class Artifacts(BaseModel):
     """Container for storing and managing artifacts generated by pipeline nodes.
 
-    Modules hyperparams and outputs. The best ones are transmitted between nodes of the pipeline.
+    Only stores the best artifact for each node type to optimize memory usage.
+    The best ones are transmitted between nodes of the pipeline.
 
     Attributes:
-        regex: List of artifacts from the regex node.
-        embedding: List of artifacts from the embedding node.
-        scoring: List of artifacts from the scoring node.
-        decision: List of artifacts from the decision node.
+        regex: Best artifact from the regex node.
+        embedding: Best artifact from the embedding node.
+        scoring: Best artifact from the scoring node.
+        decision: Best artifact from the decision node.
     """
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    regex: list[RegexArtifact] = []
-    embedding: list[EmbeddingArtifact] = []
-    scoring: list[ScorerArtifact] = []
-    decision: list[DecisionArtifact] = []
+    regex: RegexArtifact | None = None
+    embedding: EmbeddingArtifact | None = None
+    scoring: ScorerArtifact | None = None
+    decision: DecisionArtifact | None = None
 
     def model_dump(self, **kwargs: Any) -> dict[str, Any]:  # noqa: ANN401
         """Convert the model to a dictionary, ensuring nested artifacts are properly serialized."""
         data = super().model_dump(**kwargs)
         for node_type in [NodeType.regex, NodeType.embedding, NodeType.scoring, NodeType.decision]:
-            artifacts = getattr(self, node_type.value)
-            data[node_type.value] = [artifact.model_dump(**kwargs) for artifact in artifacts]
+            artifact = getattr(self, node_type.value)
+            if artifact is not None:
+                data[node_type.value] = artifact.model_dump(**kwargs)
+            else:
+                data[node_type.value] = None
         return data
 
     @classmethod
     def model_validate(cls, obj: dict[str, Any]) -> "Artifacts":
         """Convert the dictionary back to an Artifacts instance, ensuring nested artifacts are properly deserialized."""
         # First convert the lists back to numpy arrays in the scoring artifacts
-        if "scoring" in obj:
-            for artifact in obj["scoring"]:
-                if artifact.get("train_scores") is not None:
-                    artifact["train_scores"] = np.array(artifact["train_scores"])
-                if artifact.get("validation_scores") is not None:
-                    artifact["validation_scores"] = np.array(artifact["validation_scores"])
-                if artifact.get("test_scores") is not None:
-                    artifact["test_scores"] = np.array(artifact["test_scores"])
-                if artifact.get("folded_scores") is not None:
-                    artifact["folded_scores"] = [np.array(arr) for arr in artifact["folded_scores"]]
+        if "scoring" in obj and obj["scoring"] is not None:
+            if obj["scoring"].get("train_scores") is not None:
+                obj["scoring"]["train_scores"] = np.array(obj["scoring"]["train_scores"])
+            if obj["scoring"].get("validation_scores") is not None:
+                obj["scoring"]["validation_scores"] = np.array(obj["scoring"]["validation_scores"])
+            if obj["scoring"].get("test_scores") is not None:
+                obj["scoring"]["test_scores"] = np.array(obj["scoring"]["test_scores"])
+            if obj["scoring"].get("folded_scores") is not None:
+                obj["scoring"]["folded_scores"] = [np.array(arr) for arr in obj["scoring"]["folded_scores"]]
 
         return super().model_validate(obj)
 
     def add_artifact(self, node_type: str, artifact: Artifact) -> None:
-        """Add an artifact to the specified node type.
+        """Add an artifact to the specified node type, replacing any existing artifact.
 
         Args:
             node_type: Node type as a string.
             artifact: The artifact to add.
         """
-        self.get_artifacts(node_type).append(artifact)
+        setattr(self, validate_node_name(node_type), artifact)
 
-    def get_artifacts(self, node_type: str) -> list[Artifact]:
-        """Retrieve all artifacts for a specified node type.
+    def get_artifact(self, node_type: str) -> Artifact | None:
+        """Retrieve the artifact for a specified node type.
 
         Args:
             node_type: Node type as a string.
 
         Returns:
-            A list of artifacts for the node type.
+            The artifact for the node type, or None if no artifact exists.
         """
         return getattr(self, validate_node_name(node_type))  # type: ignore[no-any-return]
 
-    def get_best_artifact(self, node_type: str, idx: int) -> Artifact:
-        """Retrieve the best artifact for a specified node type and index.
+    def get_best_artifact(self, node_type: str) -> Artifact:
+        """Retrieve the artifact for a specified node type.
 
         Args:
             node_type: Node type as a string.
-            idx: Index of the artifact.
 
         Returns:
-            The best artifact.
+            The artifact for the node type.
+
+        Raises:
+            ValueError: If no artifact exists for the node type.
         """
-        return self.get_artifacts(node_type)[idx]
+        artifact = self.get_artifact(node_type)
+        if artifact is None:
+            msg = f"No artifact for {node_type}"
+            raise ValueError(msg)
+        return artifact
 
     def has_artifacts(self) -> bool:
         """Check if any artifacts have been saved in RAM.
@@ -194,7 +203,7 @@ def has_artifacts(self) -> bool:
             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)
+        return any(self.get_artifact(nt) is not None for nt in node_types)
 
 
 class Trial(BaseModel):
@@ -263,39 +272,3 @@ def add_trial(self, node_type: str, trial: Trial) -> None:
             trial: The trial to add.
         """
         self.get_trials(node_type).append(trial)
-
-
-class TrialsIds(BaseModel):
-    """Representation of the best trial IDs for each pipeline node.
-
-    Attributes:
-        regex: Best trial index for the regex node.
-        embedding: Best trial index for the embedding node.
-        scoring: Best trial index for the scoring node.
-        decision: Best trial index for the decision node.
-    """
-
-    regex: int | None = None
-    embedding: int | None = None
-    scoring: int | None = None
-    decision: int | None = None
-
-    def get_best_trial_idx(self, node_type: str) -> int | None:
-        """Retrieve the best trial index for a specified node type.
-
-        Args:
-            node_type: Node type as a string.
-
-        Returns:
-            The index of the best trial, or None if not set.
-        """
-        return getattr(self, validate_node_name(node_type))  # type: ignore[no-any-return]
-
-    def set_best_trial_idx(self, node_type: str, idx: int) -> None:
-        """Set the best trial index for a specified node type.
-
-        Args:
-            node_type: Node type as a string.
-            idx: Index of the best trial.
-        """
-        setattr(self, validate_node_name(node_type), idx)