From fb0fc597ae8a428a6b25c2f1ef8faa464c31c304 Mon Sep 17 00:00:00 2001 From: Nastaran Alipour Date: Sun, 5 Jul 2026 15:43:25 +0200 Subject: [PATCH] add support of artifacts --- CHANGELOG.md | 2 + neps/optimizers/optimizer.py | 78 +++++++++++++++- neps/runtime.py | 88 ++++++++++++++++--- neps/status/status.py | 8 +- neps/utils/files.py | 166 ++++++++++++++++++++++++++++++++++- 5 files changed, 324 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8d77063fb..9e72e29cf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - Added 3 tutorials and embended existing colab tutorial in it. +- Support artifacts for each optimizer based on evaluated trials. +- Add a generic structure for the filesystem writers. ### Fixed - fix device mismatch in the tensors created in gp acqisition process diff --git a/neps/optimizers/optimizer.py b/neps/optimizers/optimizer.py index 4e5cf270b..28cda5263 100644 --- a/neps/optimizers/optimizer.py +++ b/neps/optimizers/optimizer.py @@ -26,7 +26,8 @@ def __call__( from abc import abstractmethod from collections.abc import Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field +from enum import Enum from typing import TYPE_CHECKING, Any, Protocol, TypedDict if TYPE_CHECKING: @@ -48,6 +49,33 @@ class OptimizerInfo(TypedDict): """ +class ArtifactType(Enum): + """Supported artifact types for optimizer results.""" + + TEXT = "text" + FIGURE = "figure" + JSON = "json" + PICKLE = "pickle" + BYTES = "bytes" + + +@dataclass +class Artifact: + """Single artifact produced by optimizer. + + Attributes: + name: Unique identifier for this artifact. + content: The actual artifact content (string, Figure, dict, bytes, etc). + artifact_type: Type of artifact, determines how it will be persisted. + metadata: Optional metadata about the artifact (e.g., format, dimensions). + """ + + name: str + content: Any + artifact_type: ArtifactType + metadata: dict[str, Any] = field(default_factory=dict) + + @dataclass class SampledConfig: id: str @@ -101,3 +129,51 @@ def import_trials( trials: All of the trials that are known about. """ ... + + @abstractmethod + def get_trial_artifacts( + self, + trials: Mapping[str, Trial] | None = None, + ) -> list[Artifact] | None: + """Return artifacts for the runtime to persist. + + This is an optional method that optimizers can implement to return artifacts + (figures, logs, metadata, etc) that the neps runtime will handle for persistence. + The runtime takes responsibility for all I/O operations, enabling: + + - Clean separation between optimizer logic and I/O concerns + - Uniform handling of different artifact types (figures, text, JSON, etc) + - Future extensibility without optimizer changes + - Easy testing and mocking + + Args: + trials: All evaluated trials, passed by runtime for context. + Optional - optimizers can ignore if not needed. + + Returns: + List of Artifact objects to persist, or None if no artifacts should be + persisted. + + Note: + This method is optional. Optimizers that don't need to persist artifacts + should return None (the default behavior). + + Example: + >>> from neps.optimizers.optimizer import Artifact, ArtifactType + >>> def get_trial_artifacts(self, trials=None) -> list[Artifact] | None: + ... import matplotlib.pyplot as plt + ... fig, ax = plt.subplots() + ... + ... # Can use trials to generate more meaningful artifacts + ... if trials: + ... results = [ + ... t.report.objective_to_minimize for t in trials.values() + ... ] + ... ax.plot(results) + ... + ... return [ + ... Artifact("loss_curve", fig, ArtifactType.FIGURE), + ... Artifact("best_config", self.best_config, ArtifactType.JSON), + ... ] + """ + return None diff --git a/neps/runtime.py b/neps/runtime.py index 012d36748..c9e21090b 100644 --- a/neps/runtime.py +++ b/neps/runtime.py @@ -7,6 +7,7 @@ from __future__ import annotations +import inspect import logging import os import shutil @@ -18,7 +19,7 @@ from pathlib import Path from typing import TYPE_CHECKING, ClassVar, Literal -from filelock import FileLock +from filelock import BaseFileLock, FileLock from portalocker import portalocker from neps.env import ( @@ -56,6 +57,7 @@ status, ) from neps.utils.common import gc_disabled +from neps.utils.files import get_file_writer if TYPE_CHECKING: from neps import SearchSpace @@ -537,7 +539,9 @@ def _write_trajectory_files( best_config_path: Path, final_stopping_criteria: ResourceUsage | None = None, ) -> None: - """Writes the trajectory and best config files safely.""" + """Writes the trajectory and best config files safely using generic file + writer. + """ trace_text = _build_incumbent_content(incumbent_configs) best_config_text = _build_optimal_set_content(optimal_configs) @@ -549,12 +553,50 @@ def _write_trajectory_files( best_config_text += f"\n{metric}: {value}" with trace_lock: + text_writer = get_file_writer("text") if incumbent_configs: - with improvement_trace_path.open(mode="w") as f: - f.write(trace_text) + try: + text_writer.write(trace_text, improvement_trace_path) + except Exception as e: # noqa: BLE001 + logger.error(f"Failed to write improvement trace: {e}") if optimal_configs: - with best_config_path.open(mode="w") as f: - f.write(best_config_text) + try: + text_writer.write(best_config_text, best_config_path) + except Exception as e: # noqa: BLE001 + logger.error(f"Failed to write best config: {e}") + + def _save_optimizer_artifacts(self, artifacts: list, summary_dir: Path) -> None: + """Save optimizer artifacts to summary directory. + + Args: + artifacts: List of Artifact objects to persist. + summary_dir: Summary directory where artifacts will be saved. + """ + logger.info("saving artifacts...") + + for artifact in artifacts: + try: + # Map ArtifactType enum to string for writer lookup + content_type = artifact.artifact_type.value + writer = get_file_writer(content_type) + file_path = summary_dir / artifact.name + + accepted = set(inspect.signature(writer.write).parameters) + unknown = set(artifact.metadata) - accepted + if unknown: + raise TypeError( + f"metadata key(s) {sorted(unknown)} are not accepted by " + f"{type(writer).__name__}.write(); valid keys: {sorted(accepted)}" + ) + + writer.write(artifact.content, file_path, **artifact.metadata) + except Exception as e: # noqa: BLE001 + logger.error( + f"Failed to save artifact '{artifact.name}' " + f"(type={artifact.artifact_type.value}): {e}" + ) + # Allow optimization to continue even if artifact save fails + continue def _get_next_trial(self) -> Trial | Literal["break"]: # If there are no global stopping criterion, we can no just return early. @@ -682,13 +724,13 @@ def run(self) -> None: # noqa: C901, PLR0912, PLR0915 improvement_trace_path.touch(exist_ok=True) best_config_path = summary_dir / "best_config.txt" best_config_path.touch(exist_ok=True) - _trace_lock = FileLock(".trace.lock") + _trace_lock = FileLock(str(main_dir / ".trace.lock")) _trace_lock_path = Path(str(_trace_lock.lock_file)) _trace_lock_path.touch(exist_ok=True) - full_df_path, short_path, csv_locker = _initiate_summary_csv(main_dir) + full_df_path, short_path, summary_locker = _initiate_summary_csv(main_dir) # Create empty CSV files - with csv_locker.lock(): + with summary_locker.lock(): full_df_path.parent.mkdir(parents=True, exist_ok=True) full_df_path.touch(exist_ok=True) short_path.touch(exist_ok=True) @@ -707,6 +749,15 @@ def run(self) -> None: # noqa: C901, PLR0912, PLR0915 _repeated_fail_get_next_trial_count = 0 n_repeated_failed_check_should_stop = 0 + + evaluated_trials = self.state._trial_repo.get_valid_evaluated_trials() + self.load_incumbent_trace( + evaluated_trials, + _trace_lock, + improvement_trace_path, + best_config_path, + ) + while True: try: # First check local worker settings @@ -839,9 +890,22 @@ def run(self) -> None: # noqa: C901, PLR0912, PLR0915 improvement_trace_path, best_config_path, ) + # Persist optimizer artifacts if available + if hasattr(self.optimizer, "get_trial_artifacts"): + try: + artifacts = self.optimizer.get_trial_artifacts( + trials=evaluated_trials + ) + if artifacts is not None: + with _trace_lock: + self._save_optimizer_artifacts(artifacts, summary_dir) + except Exception as e: + logger.error( + f"Failed to persist optimizer artifacts: {e}", exc_info=True + ) full_df, short = status(main_dir) - with csv_locker.lock(): + with summary_locker.lock(): full_df.to_csv(full_df_path) short.to_frame().to_csv(short_path) @@ -855,7 +919,7 @@ def run(self) -> None: # noqa: C901, PLR0912, PLR0915 def load_incumbent_trace( self, trials: dict[str, Trial], - _trace_lock: FileLock, + _trace_lock: BaseFileLock, improvement_trace_path: Path, best_config_path: Path, ) -> None: @@ -865,7 +929,7 @@ def load_incumbent_trace( Args: trials (dict): A dictionary of the evaluated trials which have a valid report. - _trace_lock (FileLock): A file lock to ensure thread-safe writing. + _trace_lock (BaseFileLock): A file lock to ensure thread-safe writing. improvement_trace_path (Path): Path to the improvement trace file. best_config_path (Path): Path to the best configuration file. """ diff --git a/neps/status/status.py b/neps/status/status.py index f5ab360b3..04393a492 100644 --- a/neps/status/status.py +++ b/neps/status/status.py @@ -407,9 +407,9 @@ def _initiate_summary_csv(root_directory: str | Path) -> tuple[Path, Path, FileL full_path = summary_path / "full.csv" short_path = summary_path / "short.csv" - csv_locker = FileLocker(summary_path / ".csv_lock", poll=0.1, timeout=10) + summary_locker = FileLocker(summary_path / ".summary.lock", poll=0.1, timeout=10) - return (full_path, short_path, csv_locker) + return (full_path, short_path, summary_locker) def post_run_csv(root_directory: str | Path) -> tuple[Path, Path]: @@ -422,9 +422,9 @@ def post_run_csv(root_directory: str | Path) -> tuple[Path, Path]: The paths to the configuration data CSV and the run data CSV. """ full_df, short = status(root_directory, print_summary=False) - full_df_path, short_path, csv_locker = _initiate_summary_csv(root_directory) + full_df_path, short_path, summary_locker = _initiate_summary_csv(root_directory) - with csv_locker.lock(): + with summary_locker.lock(): full_df.to_csv(full_df_path) short.to_frame().to_csv(short_path) diff --git a/neps/utils/files.py b/neps/utils/files.py index d79e43047..72985aac9 100644 --- a/neps/utils/files.py +++ b/neps/utils/files.py @@ -4,15 +4,19 @@ import dataclasses import io +import json +import logging import os from collections.abc import Iterable, Iterator, Mapping from contextlib import contextmanager from enum import Enum from pathlib import Path -from typing import IO, Any, Literal +from typing import IO, Any, Literal, Protocol import yaml +logger = logging.getLogger(__name__) + try: from yaml import ( CDumper as YamlDumper, # type: ignore @@ -148,3 +152,163 @@ def load_and_merge_yamls(*paths: str | Path | IO[str]) -> dict[str, Any]: config.update(read_config) return config + + +class FileWriter(Protocol): + """Protocol for writing content to disk.""" + + def write(self, content: Any, file_path: Path | str) -> None: + """Write content to disk. + + Args: + content: Content to write. + file_path: Path where content should be saved. + """ + ... + + +class TextWriter: + """Write text content to disk.""" + + def write(self, content: str, file_path: Path | str) -> None: + """Write text to file. + + Args: + content: Text content to save. + file_path: Path to save to (will use .txt extension). + """ + file_path = Path(file_path).with_suffix(".txt") + try: + file_path.parent.mkdir(parents=True, exist_ok=True) + file_path.write_text(content) + logger.debug(f"Wrote text to {file_path}") + except Exception as e: + logger.error(f"Failed to write text to {file_path}: {e}") + raise + + +class JsonWriter: + """Write JSON-serializable content to disk.""" + + def write(self, content: Any, file_path: Path | str) -> None: + """Write JSON to file. + + Args: + content: JSON-serializable object to save. + file_path: Path to save to (will use .json extension). + """ + file_path = Path(file_path).with_suffix(".json") + try: + file_path.parent.mkdir(parents=True, exist_ok=True) + with file_path.open("w") as f: + json.dump(content, f, indent=2, default=str) + logger.debug(f"Wrote JSON to {file_path}") + except Exception as e: + logger.error(f"Failed to write JSON to {file_path}: {e}") + raise + + +class FigureWriter: + """Write matplotlib Figure objects to disk.""" + + def write( + self, + content: Any, + file_path: Path | str, + dpi: int = 100, + fmt: str = "png", + ) -> None: + """Write matplotlib figure to file. + + Args: + content: matplotlib Figure object. + file_path: Path to save to (extension added based on fmt). + dpi: Resolution for raster formats (default 100). + fmt: Format to save as (default 'png'). Can also be 'pdf', 'svg', etc. + """ + file_path = Path(file_path).with_suffix(f".{fmt}") + try: + file_path.parent.mkdir(parents=True, exist_ok=True) + content.savefig(file_path, dpi=dpi, bbox_inches="tight") + logger.debug(f"Wrote figure to {file_path}") + # Clean up matplotlib resource + try: + import matplotlib.pyplot as plt + + plt.close(content) + except ImportError: + pass + except Exception as e: + logger.error(f"Failed to write figure to {file_path}: {e}") + raise + + +class PickleWriter: + """Write Python objects using pickle.""" + + def write(self, content: Any, file_path: Path | str) -> None: + """Write object using pickle. + + Args: + content: Python object to pickle. + file_path: Path to save to (will use .pkl extension). + """ + import pickle + + file_path = Path(file_path).with_suffix(".pkl") + try: + file_path.parent.mkdir(parents=True, exist_ok=True) + with file_path.open("wb") as f: + pickle.dump(content, f) + logger.debug(f"Wrote pickle to {file_path}") + except Exception as e: + logger.error(f"Failed to write pickle to {file_path}: {e}") + raise + + +class BytesWriter: + """Write raw bytes to disk.""" + + def write(self, content: bytes, file_path: Path | str) -> None: + """Write bytes to file. + + Args: + content: Bytes content to save. + file_path: Path to save to (will use .bin extension). + """ + file_path = Path(file_path).with_suffix(".bin") + try: + file_path.parent.mkdir(parents=True, exist_ok=True) + file_path.write_bytes(content) + logger.debug(f"Wrote bytes to {file_path}") + except Exception as e: + logger.error(f"Failed to write bytes to {file_path}: {e}") + raise + + +def get_file_writer(content_type: str) -> FileWriter: + """Get the appropriate writer for a content type. + + Args: + content_type: Type of content ('text', 'json', 'figure', 'pickle', 'bytes'). + + Returns: + Writer instance for the given type. + + Raises: + ValueError: If content type is not supported. + """ + # Writer registry - maps content type to writer instance + _file_writers: dict[str, FileWriter] = { + "text": TextWriter(), + "json": JsonWriter(), + "figure": FigureWriter(), + "pickle": PickleWriter(), + "bytes": BytesWriter(), + } + if content_type not in _file_writers: + raise ValueError( + f"Unsupported content type: {content_type}. " + f"Supported types: {list(_file_writers.keys())}" + ) + return _file_writers[content_type]