From e30f9da377fce5c78d0929ab8bf85e4ff02f3298 Mon Sep 17 00:00:00 2001 From: nikitakuklev Date: Fri, 21 Aug 2026 05:01:45 -0500 Subject: [PATCH 1/3] rework serialization around typed pydantic codecs in xopt.types --- docs/index.md | 2 +- environment.yml | 4 +- pyproject.toml | 2 +- xopt/base.py | 119 +- xopt/entrypoint.py | 4 - xopt/evaluator.py | 19 +- xopt/generator.py | 33 +- xopt/generators/bayesian/bax/algorithms.py | 6 +- xopt/generators/bayesian/bax_generator.py | 7 +- .../bayesian/bayesian_exploration.py | 9 +- .../generators/bayesian/bayesian_generator.py | 82 +- .../bayesian/expected_improvement.py | 7 +- xopt/generators/bayesian/mggpo.py | 7 +- xopt/generators/bayesian/mobo.py | 5 +- xopt/generators/bayesian/models/standard.py | 71 +- xopt/generators/bayesian/multi_fidelity.py | 7 +- xopt/generators/bayesian/time_dependent.py | 5 +- .../bayesian/upper_confidence_bound.py | 7 +- xopt/generators/deduplicated.py | 16 +- xopt/generators/ga/cnsga.py | 28 +- xopt/generators/ga/nsga2.py | 7 +- xopt/generators/ga/operators.py | 21 +- xopt/generators/random.py | 12 +- xopt/generators/scipy/latin_hypercube.py | 10 +- xopt/generators/sequential/extremumseeking.py | 3 +- xopt/generators/sequential/neldermead.py | 44 +- xopt/generators/sequential/rcds.py | 8 +- xopt/pydantic.py | 553 +++----- xopt/resources/testing.py | 3 +- xopt/stopping_conditions.py | 21 +- .../bayesian/test_bayesian_generator.py | 63 +- xopt/tests/test_entrypoint.py | 6 +- xopt/tests/test_io.py | 7 - xopt/tests/test_pydantic.py | 1110 ++++++++--------- xopt/types.py | 709 +++++++++++ xopt/utils.py | 36 +- 36 files changed, 1756 insertions(+), 1297 deletions(-) create mode 100644 xopt/types.py diff --git a/docs/index.md b/docs/index.md index 58fcb9514..916691716 100644 --- a/docs/index.md +++ b/docs/index.md @@ -77,7 +77,7 @@ generator: output_path: . evaluator: - function: my_function + function: my_module.my_function function_kwargs: my_arguments: 42 diff --git a/environment.yml b/environment.yml index b44bd3332..c4061701c 100644 --- a/environment.yml +++ b/environment.yml @@ -6,14 +6,14 @@ dependencies: - python>=3.10 - deap - numpy - - pydantic>=2.3 + - pydantic>=2.12 - pyyaml - botorch>=0.13.0 - scipy>=1.10.1 - pandas - ipywidgets - tqdm - - orjson + - zstandard - matplotlib # parallel - mpi4py diff --git a/pyproject.toml b/pyproject.toml index 0fbb7b02e..94e607fae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,7 @@ dependencies = [ "pandas", "ipywidgets", "tqdm", - "orjson", + "zstandard", "matplotlib", "gest-api>=0.2" ] diff --git a/xopt/base.py b/xopt/base.py index 485070f5d..551fa7678 100644 --- a/xopt/base.py +++ b/xopt/base.py @@ -1,41 +1,42 @@ import json import logging import os +import warnings from copy import deepcopy from typing import Any, Optional, Union import numpy as np import pandas as pd import yaml +from gest_api.vocs import VOCS from pandas import DataFrame from pydantic import ( Field, SerializeAsAny, ValidationInfo, + field_serializer, field_validator, model_validator, ) -import warnings from xopt.errors import VOCSError from xopt.evaluator import Evaluator, validate_outputs from xopt.generator import Generator, StateOwner from xopt.generators import get_generator from xopt.generators.sequential import SequentialGenerator -from xopt.pydantic import XoptBaseModel +from xopt.pydantic import XoptBaseModel, serialization_fallback +from xopt.stopping_conditions import ( + MaxEvaluationsCondition, + StoppingConditionUnion, +) +from xopt.types import DataFrameCodec, XDataFrame from xopt.utils import explode_all_columns, get_generator_name from xopt.vocs import ( ContextualVariable, - validate_input_data, - random_inputs, grid_inputs, + random_inputs, + validate_input_data, ) -from gest_api.vocs import VOCS -from xopt.stopping_conditions import ( - MaxEvaluationsCondition, - StoppingConditionUnion, -) - from .errors import DataError @@ -133,7 +134,7 @@ class Xopt(XoptBaseModel): data_dump_file: Optional[str] = Field( None, description="file to dump the evaluation data to as CSV" ) - data: Optional[DataFrame] = Field(None, description="internal DataFrame object") + data: Optional[XDataFrame] = Field(None, description="internal DataFrame object") serialize_torch: bool = Field( False, description="flag to indicate that torch models should be serialized " @@ -151,7 +152,7 @@ class Xopt(XoptBaseModel): @model_validator(mode="before") @classmethod - def validate_generator_and_legacy_vocs(cls, data: Any): + def validate_generator_and_legacy_vocs(cls, data: Any, info: ValidationInfo): """ Validate the Xopt model by checking the generator and evaluator. """ @@ -195,13 +196,45 @@ def validate_generator_and_legacy_vocs(cls, data: Any): if isinstance(data["generator"], dict): name = data["generator"].pop("name") generator_class = get_generator(name) - data["generator"] = generator_class.model_validate(data["generator"]) + data["generator"] = generator_class.model_validate( + data["generator"], context=info.context + ) # make a copy of the generator / vocs objects to avoid modifying the original data["generator"] = deepcopy(data["generator"]) return data + @field_serializer("generator", mode="wrap", when_used="always") + def serialize_generator(self, generator, handler, info): + # Pydantic currently narrows a field to its annotation when any field + # serializer is present, even if the annotation is SerializeAsAny, so + # bypass the supplied handler and serialize the runtime model directly + # with the same options so subclass fields survive. + if not isinstance(generator, Generator): + name = getattr(generator, "name", None) or ( + f"{type(generator).__module__}.{type(generator).__qualname__}" + ) + return {"name": name} + serializer = generator.__pydantic_serializer__ + result = serializer.to_python( + generator, + mode=info.mode, + include=info.include, + exclude=info.exclude, + by_alias=info.by_alias, + exclude_unset=info.exclude_unset, + exclude_defaults=info.exclude_defaults, + exclude_none=info.exclude_none, + round_trip=info.round_trip, + fallback=serialization_fallback, + serialize_as_any=False, + context=info.context, + ) + if not isinstance(result, dict): + result = {"name": result} + return {"name": get_generator_name(generator)} | result + @field_validator("evaluator", mode="before") def validate_evaluator(cls, value): if isinstance(value, dict): @@ -211,6 +244,13 @@ def validate_evaluator(cls, value): @field_validator("data", mode="before") def validate_data(cls, v, info: ValidationInfo): + if v is None: + # explicit null (dump of an un-evaluated Xopt) is not a frame to propagate + return None + if isinstance(v, str): + # decode encoded forms (e.g. "b64df:") before the generator + # propagation below, which needs a real DataFrame + v = DataFrameCodec.validate(v, info) if isinstance(v, dict): try: v = pd.DataFrame(v) @@ -253,9 +293,11 @@ def max_evaluations_legacy(cls, data: Any): "Use 'stopping_condition' with MaxEvaluationsCondition instead." ) max_evals = data.pop("max_evaluations") - data["stopping_condition"] = MaxEvaluationsCondition( - max_evaluations=max_evals - ) + # 2.x dumps carry an explicit `max_evaluations: null`; treat it as absent + if max_evals is not None: + data["stopping_condition"] = MaxEvaluationsCondition( + max_evaluations=max_evals + ) return data @model_validator(mode="before") @@ -659,14 +701,9 @@ def yaml(self, **kwargs): The Xopt configuration serialized as a YAML string. """ - output = json.loads( - self.json( - serialize_torch=self.serialize_torch, - serialize_inline=self.serialize_inline, - **kwargs, - ) - ) - return yaml.dump(output) + kwargs.setdefault("serialize_torch", self.serialize_torch) + kwargs.setdefault("serialize_inline", self.serialize_inline) + return super().yaml(**kwargs) def dump(self, file: str = None, **kwargs): """ @@ -697,6 +734,8 @@ def dump(self, file: str = None, **kwargs): ) fname = os.path.expanduser(os.path.expandvars(fname)) + # write torch sidecar files next to the dump file by default + kwargs.setdefault("file_dir", os.path.dirname(os.path.abspath(fname))) with open(fname, "w") as f: f.write(self.yaml(**kwargs)) logger.debug(f"Dumped state to YAML file: {fname}") @@ -738,13 +777,7 @@ def dict(self, **kwargs) -> dict: A dictionary representation of the Xopt configuration. """ - result = super().model_dump(**kwargs) - if not isinstance(result["generator"], dict): # may return as module.path - result["generator"] = {"name": result["generator"]} - result["generator"] = {"name": get_generator_name(self.generator)} | result[ - "generator" - ] - return result + return self.model_dump(**kwargs) def json(self, **kwargs) -> str: """ @@ -761,27 +794,9 @@ def json(self, **kwargs) -> str: The Xopt configuration serialized as a JSON string. """ - result = super().to_json(**kwargs) - dict_result = json.loads(result) - if not isinstance(dict_result["generator"], dict): # may return as module.path - dict_result["generator"] = {"name": dict_result["generator"]} - dict_result["generator"] = { - "name": get_generator_name(self.generator) - } | dict_result["generator"] - dict_result["data"] = ( - json.loads(self.data.to_json()) if self.data is not None else None - ) - - if "stopping_condition" in dict_result: - if dict_result["stopping_condition"] is not None: - dict_result["stopping_condition"] = { - "name": self.stopping_condition.__class__.__name__ - } | dict_result["stopping_condition"] - - # TODO: implement version checking - # dict_result["xopt_version"] = __version__ - - return json.dumps(dict_result) + kwargs.setdefault("serialize_torch", self.serialize_torch) + kwargs.setdefault("serialize_inline", self.serialize_inline) + return super().to_json(**kwargs) def __repr__(self): """ diff --git a/xopt/entrypoint.py b/xopt/entrypoint.py index e8687243d..ed9388d61 100644 --- a/xopt/entrypoint.py +++ b/xopt/entrypoint.py @@ -1,6 +1,5 @@ from .base import Xopt from .evaluator import DummyExecutor -from .pydantic import remove_none_values from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor from contextlib import contextmanager import argparse @@ -205,9 +204,6 @@ def main(): # Open file config = yaml.safe_load(f) - # Clean up (replicate behavior of Xopt.from_file) - config = remove_none_values(config) - # Apply the overrides to the config dict if args.override: logger.info("Applying config file overrides:") diff --git a/xopt/evaluator.py b/xopt/evaluator.py index ed7768237..dfa43a3d3 100644 --- a/xopt/evaluator.py +++ b/xopt/evaluator.py @@ -11,7 +11,8 @@ from xopt.errors import XoptError from xopt.pydantic import NormalExecutor, XoptBaseModel -from xopt.utils import get_function, get_function_defaults, safe_call +from xopt.types import CallableRef +from xopt.utils import get_function_defaults, safe_call logger = logging.getLogger(__name__) @@ -42,7 +43,7 @@ class Evaluator(XoptBaseModel): mapping. """ - function: Callable + function: CallableRef max_workers: int = Field(1, ge=1) executor: NormalExecutor = Field(exclude=True) # Do not serialize function_kwargs: dict = Field({}) @@ -65,12 +66,6 @@ def validate_all(cls, values: Dict) -> Dict: dict The validated input values. """ - f = get_function(values["function"]) - kwargs = values.get("function_kwargs", {}) - kwargs = {**get_function_defaults(f), **kwargs} - values["function"] = f - values["function_kwargs"] = kwargs - max_workers = values.pop("max_workers", 1) executor = values.pop("executor", None) @@ -87,6 +82,14 @@ def validate_all(cls, values: Dict) -> Dict: return values + @model_validator(mode="after") + def fill_function_default_kwargs(self): + """Fills in place so revalidation (e.g. nesting into Xopt) keeps the + dict identity and stays idempotent.""" + for key, val in get_function_defaults(self.function).items(): + self.function_kwargs.setdefault(key, val) + return self + def evaluate(self, input: Dict, **kwargs) -> Dict: """ Evaluate a single input dict using Evaluator.function with diff --git a/xopt/generator.py b/xopt/generator.py index 62e044044..fa30e1192 100644 --- a/xopt/generator.py +++ b/xopt/generator.py @@ -1,22 +1,31 @@ import logging from abc import ABC, abstractmethod -from typing import Any, ClassVar, Optional, List, Hashable - +from typing import Any, ClassVar, Hashable, List, Optional import pandas as pd +from gest_api.generator import Generator as BaseGenerator +from gest_api.vocs import VOCS, DiscreteVariable from pydantic import ConfigDict, Field, field_validator from pydantic_core.core_schema import ValidationInfo from xopt.errors import VOCSError from xopt.pydantic import XoptBaseModel +from xopt.types import XDataFrame from xopt.vocs import ContextualVariable -from gest_api.vocs import VOCS, DiscreteVariable -from gest_api.generator import Generator as BaseGenerator - logger = logging.getLogger(__name__) +def support_flag(default: bool): + """Field spec for generator capability flags. + + Subclasses overriding a ``supports_*`` default must use this instead of a + bare ``bool = True`` annotation, which would silently drop the base field's + ``frozen``/``exclude`` settings (pydantic replaces the whole FieldInfo). + """ + return Field(default=default, frozen=True, exclude=True) + + class Generator(XoptBaseModel, BaseGenerator, ABC): """ Base class for Generators. @@ -92,9 +101,7 @@ class Generator(XoptBaseModel, BaseGenerator, ABC): ) vocs: VOCS = Field(description="generator VOCS") - data: Optional[pd.DataFrame] = Field( - None, description="generator data", exclude=True - ) + data: Optional[XDataFrame] = Field(None, description="generator data", exclude=True) model_config = ConfigDict(validate_assignment=True) @@ -200,16 +207,6 @@ def add_data(self, new_data: pd.DataFrame): else: self.data = new_data - def model_dump(self, *args: Any, **kwargs: Any) -> dict[str, Any]: - """overwrite model dump to remove faux class attrs""" - - res = super().model_dump(*args, **kwargs) - - res.pop("supports_batch_generation", None) - res.pop("supports_multi_objective", None) - - return res - class StateOwner: """ diff --git a/xopt/generators/bayesian/bax/algorithms.py b/xopt/generators/bayesian/bax/algorithms.py index 42731f139..2957a8b4f 100644 --- a/xopt/generators/bayesian/bax/algorithms.py +++ b/xopt/generators/bayesian/bax/algorithms.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import Any +from typing import Any, Optional import torch from botorch.models.model import Model, ModelList @@ -25,11 +25,11 @@ class OptimizationAlgorithmResult(AlgorithmResult): best_objective: Tensor = Field( description="The optimal objective values from the sample-wise optimization of the virtual objective." ) - solution_center: Tensor = Field( + solution_center: Optional[Tensor] = Field( None, description="The mean of the distribution of optimal inputs from the sample-wise optimization of the virtual objective.", ) - solution_entropy: float = Field( + solution_entropy: Optional[float] = Field( None, description="The entropy of the distribution of optimal inputs from the sample-wise optimization of the virtual objective.", ) diff --git a/xopt/generators/bayesian/bax_generator.py b/xopt/generators/bayesian/bax_generator.py index 9db94e755..4e0600cae 100644 --- a/xopt/generators/bayesian/bax_generator.py +++ b/xopt/generators/bayesian/bax_generator.py @@ -14,6 +14,7 @@ model_validator, ) from pydantic.fields import ModelPrivateAttr, PrivateAttr +from xopt.generator import support_flag from xopt.errors import VOCSError from xopt.generators.bayesian.bax.acquisition import ModelListExpectedInformationGain from xopt.generators.bayesian.bax.algorithms import Algorithm, GridOptimize @@ -59,9 +60,9 @@ class BaxGenerator(BayesianGenerator): """ name = "bax" - supports_constraints: bool = True - supports_no_objective: bool = True - supports_discrete_variables: bool = False + supports_constraints: bool = support_flag(True) + supports_no_objective: bool = support_flag(True) + supports_discrete_variables: bool = support_flag(False) algorithm: SerializeAsAny[Algorithm] = Field( default=GridOptimize(observable_names_ordered=[]), description="algorithm evaluated in the BAX process", diff --git a/xopt/generators/bayesian/bayesian_exploration.py b/xopt/generators/bayesian/bayesian_exploration.py index ebf934597..0de336ed1 100644 --- a/xopt/generators/bayesian/bayesian_exploration.py +++ b/xopt/generators/bayesian/bayesian_exploration.py @@ -9,6 +9,7 @@ from torch import Tensor from gest_api.vocs import ExploreObjective +from xopt.generator import support_flag from xopt.errors import VOCSError from xopt.generators.bayesian.bayesian_generator import ( BayesianGenerator, @@ -23,10 +24,10 @@ class BayesianExplorationGenerator(BayesianGenerator): """ name = "bayesian_exploration" - supports_batch_generation: bool = True - supports_constraints: bool = True - supports_multi_objective: bool = True - supports_single_objective: bool = True + supports_batch_generation: bool = support_flag(True) + supports_constraints: bool = support_flag(True) + supports_multi_objective: bool = support_flag(True) + supports_single_objective: bool = support_flag(True) __doc__ = "Bayesian exploration generator\n" + formatted_base_docstring() diff --git a/xopt/generators/bayesian/bayesian_generator.py b/xopt/generators/bayesian/bayesian_generator.py index 1390997c7..9db689553 100644 --- a/xopt/generators/bayesian/bayesian_generator.py +++ b/xopt/generators/bayesian/bayesian_generator.py @@ -1,12 +1,11 @@ import logging -import os import time import warnings from abc import ABC, abstractmethod from copy import deepcopy from itertools import islice, product from math import prod -from typing import Any, Dict, Hashable, List, Optional, Union, cast +from typing import Annotated, Any, Dict, Hashable, List, Optional, Union, cast import numpy as np import pandas as pd @@ -26,13 +25,15 @@ PositiveInt, SerializeAsAny, field_validator, + model_serializer, model_validator, ) from pydantic.fields import ModelPrivateAttr, PrivateAttr from pydantic_core.core_schema import ValidationInfo from torch import Tensor + from xopt.errors import FeasibilityError, VOCSError, XoptError -from xopt.generator import Generator +from xopt.generator import Generator, support_flag from xopt.generators.bayesian.base_model import ModelConstructor from xopt.generators.bayesian.custom_botorch.constrained_acquisition import ( ConstrainedMCAcquisitionFunction, @@ -64,7 +65,12 @@ ) from xopt.generators.bayesian.visualize import visualize_generator_model from xopt.numerical_optimizer import GridOptimizer, LBFGSOptimizer, NumericalOptimizer -from xopt.pydantic import decode_torch_module +from xopt.types import ( + TorchModuleCodec, + XDataFrame, + get_serialization_options, + save_module_sidecar, +) from xopt.vocs import ( ContextualVariable, convert_numpy_to_inputs, @@ -152,12 +158,11 @@ class BayesianGenerator(Generator, ABC): """ name = "base_bayesian_generator" - supports_discrete_variables: bool = True - supports_contextual_variables: bool = True - supports_no_objective: bool = ( - True # note: only supports if custom objective is provided - ) - model: Optional[Model] = Field( + supports_discrete_variables: bool = support_flag(True) + supports_contextual_variables: bool = support_flag(True) + # note: no-objective mode is only supported if a custom objective is provided + supports_no_objective: bool = support_flag(True) + model: Optional[Annotated[Model, TorchModuleCodec()]] = Field( None, description="botorch model used by the generator to perform optimization" ) n_monte_carlo_samples: int = Field( @@ -181,13 +186,15 @@ class BayesianGenerator(Generator, ABC): fixed_features: Optional[Dict[str, float]] = Field( None, description="fixed features used in Bayesian optimization" ) - computation_time: Optional[pd.DataFrame] = Field( + computation_time: Optional[XDataFrame] = Field( None, description="data frame tracking computation time in seconds", ) - custom_objective: Optional[CustomXoptObjective] = Field( - None, - description="custom objective for optimization, replaces objective specified by VOCS", + custom_objective: Optional[Annotated[CustomXoptObjective, TorchModuleCodec()]] = ( + Field( + None, + description="custom objective for optimization, replaces objective specified by VOCS", + ) ) n_interpolate_points: Optional[PositiveInt] = None @@ -250,21 +257,28 @@ def get_compatible_numerical_optimizers( compatible = cast(ModelPrivateAttr, cls._compatible_numerical_optimizers) return compatible.get_default() - @field_validator("model", mode="before") - @classmethod - def validate_torch_modules(cls, value: Any) -> Any: - if isinstance(value, str): - if value.startswith("base64:"): - value = decode_torch_module(value) - elif os.path.exists(value): - value = torch.load(value, weights_only=False) - else: - raise XoptError(f"cannot load torch module from {value}") - return value + @model_serializer(mode="wrap", when_used="json") + def serialize_torch_modules(self, handler, info): + result = handler(self) + options = get_serialization_options(info.context) + # only touch fields the handler kept (respect include/exclude) + for field_name in ("model", "custom_objective"): + if field_name not in result: + continue + module = getattr(self, field_name) + if module is None: + continue + if options.module_mode == "drop": + del result[field_name] + elif options.module_mode == "file": + result[field_name] = save_module_sidecar( + module, options, f"generator_{field_name}.pt" + ) + return result @field_validator("gp_constructor", mode="before") @classmethod - def validate_gp_constructor(cls, value: Any) -> Any: + def validate_gp_constructor(cls, value: Any, info: ValidationInfo): constructor_dict = { "standard": StandardModelConstructor, "batched": BatchedModelConstructor, @@ -281,10 +295,12 @@ def validate_gp_constructor(cls, value: Any) -> Any: else: raise ValueError(f"{value} not found") elif isinstance(value, dict): - _value = cast(dict[str, Any], value) + _value = dict(cast(dict[str, Any], value)) name = _value.pop("name", "") if name in constructor_dict: - value = constructor_dict[name](**_value) + value = constructor_dict[name].model_validate( + _value, context=info.context if info is not None else None + ) else: raise ValueError(f"{value} not found") @@ -307,7 +323,7 @@ def validate_numerical_optimizer(cls, value: Any) -> Any: else: raise ValueError(f"{value} not found") elif isinstance(value, dict): - _value = cast(dict[str, Any], value) + _value = dict(cast(dict[str, Any], value)) name: str = _value.pop("name", "") if name in optimizer_dict: value = optimizer_dict[name](**_value) @@ -346,6 +362,8 @@ def validate_computation_time(cls, value: Any) -> Any: return value elif isinstance(value, dict): value = pd.DataFrame(value) + elif isinstance(value, str): + return value else: raise ValueError( "computation_time must be a pandas DataFrame, dict, or None" @@ -1182,17 +1200,19 @@ class MultiObjectiveBayesianGenerator(BayesianGenerator, ABC): description="dict specifying reference point for multi-objective optimization", # validate_default=True, ) - pareto_front_history: Optional[pd.DataFrame] = Field( + pareto_front_history: Optional[XDataFrame] = Field( None, description="history of pareto front statistics every time points are added to the generator", exclude=True, ) - supports_multi_objective: bool = True + supports_multi_objective: bool = support_flag(True) @field_validator("pareto_front_history", mode="before") @classmethod def validate_pareto_front_history(cls, value: Any): + if isinstance(value, str): + return value return pd.DataFrame(value) if value is not None else None @model_validator(mode="after") diff --git a/xopt/generators/bayesian/expected_improvement.py b/xopt/generators/bayesian/expected_improvement.py index 4a197e43e..c2a305366 100644 --- a/xopt/generators/bayesian/expected_improvement.py +++ b/xopt/generators/bayesian/expected_improvement.py @@ -7,6 +7,7 @@ ) from gest_api.vocs import MinimizeObjective +from xopt.generator import support_flag from xopt.generators.bayesian.bayesian_generator import ( BayesianGenerator, formatted_base_docstring, @@ -28,9 +29,9 @@ class ExpectedImprovementGenerator(BayesianGenerator): """ name = "expected_improvement" - supports_batch_generation: bool = True - supports_single_objective: bool = True - supports_constraints: bool = True + supports_batch_generation: bool = support_flag(True) + supports_single_objective: bool = support_flag(True) + supports_constraints: bool = support_flag(True) __doc__ = ( "Bayesian optimization generator using Expected improvement\n" diff --git a/xopt/generators/bayesian/mggpo.py b/xopt/generators/bayesian/mggpo.py index cfc6f1433..b80ca4d75 100644 --- a/xopt/generators/bayesian/mggpo.py +++ b/xopt/generators/bayesian/mggpo.py @@ -8,6 +8,7 @@ ) from pydantic import Field +from xopt.generator import support_flag from xopt.generators.bayesian.objectives import create_mobo_objective from xopt.generators.ga.cnsga import CNSGAGenerator from .bayesian_generator import MultiObjectiveBayesianGenerator @@ -46,9 +47,9 @@ class MGGPOGenerator(MultiObjectiveBayesianGenerator): name = "mggpo" population_size: int = Field(64, description="population size for ga") - supports_batch_generation: bool = True - supports_constraints: bool = True - supports_discrete_variables: bool = False + supports_batch_generation: bool = support_flag(True) + supports_constraints: bool = support_flag(True) + supports_discrete_variables: bool = support_flag(False) ga_generator: Optional[CNSGAGenerator] = Field( None, description="CNSGA generator used to generate candidates" diff --git a/xopt/generators/bayesian/mobo.py b/xopt/generators/bayesian/mobo.py index d737a4f34..65b595b2b 100644 --- a/xopt/generators/bayesian/mobo.py +++ b/xopt/generators/bayesian/mobo.py @@ -9,6 +9,7 @@ from pydantic import Field, field_validator from torch import Tensor +from xopt.generator import support_flag from xopt.generators.bayesian.bayesian_generator import MultiObjectiveBayesianGenerator from xopt.generators.bayesian.objectives import create_mobo_objective from xopt.generators.bayesian.turbo import SafetyTurboController @@ -33,8 +34,8 @@ class MOBOGenerator(MultiObjectiveBayesianGenerator): """ name = "mobo" - supports_batch_generation: bool = True - supports_constraints: bool = True + supports_batch_generation: bool = support_flag(True) + supports_constraints: bool = support_flag(True) use_pf_as_initial_points: bool = Field( False, description="flag to specify if pareto front points are to be used during " diff --git a/xopt/generators/bayesian/models/standard.py b/xopt/generators/bayesian/models/standard.py index 8a91d0ffe..5c96d1145 100644 --- a/xopt/generators/bayesian/models/standard.py +++ b/xopt/generators/bayesian/models/standard.py @@ -1,14 +1,13 @@ -import os.path import warnings from copy import deepcopy from functools import partial -from typing import Any, Dict, List, Literal, Optional, Union, cast +from typing import Annotated, Any, Dict, List, Literal, Optional, Union -from botorch.exceptions import ModelFittingError import botorch.settings import pandas as pd import torch from botorch import fit_gpytorch_mll +from botorch.exceptions import ModelFittingError from botorch.models import ModelListGP, SingleTaskGP from botorch.models.gpytorch import BatchedMultiOutputGPyTorchModel from botorch.models.transforms import Normalize, Standardize @@ -22,7 +21,13 @@ from gpytorch.likelihoods import GaussianLikelihood, Likelihood from gpytorch.likelihoods.gaussian_likelihood import FixedNoiseGaussianLikelihood from gpytorch.priors import GammaPrior, Prior -from pydantic import ConfigDict, Field, field_validator +from pydantic import ( + ConfigDict, + Field, + SerializeAsAny, + field_validator, + model_serializer, +) from pydantic_core.core_schema import ValidationInfo from torch.nn import Module from torch.optim import Adam @@ -30,9 +35,9 @@ from xopt.generators.bayesian.base_model import ModelConstructor from xopt.generators.bayesian.models.prior_mean import CustomMean from xopt.generators.bayesian.utils import get_training_data, get_training_data_batched -from xopt.pydantic import XoptBaseModel, decode_torch_module +from xopt.pydantic import XoptBaseModel +from xopt.types import TorchModuleCodec, get_serialization_options, save_module_sidecar -DECODERS = {"torch.float32": torch.float32, "torch.float64": torch.float64} MIN_INFERRED_NOISE_LEVEL = 1e-4 # TODO: make custom stopping criterion that checks lengthscales @@ -150,10 +155,10 @@ class StandardModelConstructor(ModelConstructor): use_low_noise_prior: bool = Field( False, description="specify if model should assume a low noise environment" ) - covar_modules: Dict[str, Kernel] = Field( + covar_modules: Dict[str, Annotated[Kernel, TorchModuleCodec()]] = Field( {}, description="covariance modules for GP models" ) - mean_modules: Dict[str, Module] = Field( + mean_modules: Dict[str, Annotated[Module, TorchModuleCodec()]] = Field( {}, description="prior mean modules for GP models" ) trainable_mean_keys: List[str] = Field( @@ -164,7 +169,7 @@ class StandardModelConstructor(ModelConstructor): description="specify if inputs should be transformed inside the gp " "model, can optionally specify a dict of specifications", ) - custom_noise_prior: Optional[Prior] = Field( + custom_noise_prior: Optional[Annotated[Prior, TorchModuleCodec()]] = Field( None, description="specify custom noise prior for the GP likelihood, " "overwrites value specified by use_low_noise_prior", @@ -181,7 +186,7 @@ class StandardModelConstructor(ModelConstructor): True, description="flag to specify if the model should be trained (fitted to data)", ) - train_config: NumericalOptimizerConfig | None = Field( + train_config: SerializeAsAny[NumericalOptimizerConfig] | None = Field( None, description="configuration of the numerical optimizer - see fit_gpytorch_mll_scipy" " and fit_gpytorch_mll_torch", @@ -225,16 +230,22 @@ def validate_train_kwargs(cls, train_kwargs, info: ValidationInfo): ) return train_kwargs - @field_validator("train_config") + @field_validator("train_config", mode="before") def validate_train_config(cls, v, info: ValidationInfo): if v is None: return v if info.data["train_method"] == "adam": + if isinstance(v, dict): + v = AdamNumericalOptimizerConfig.model_validate(v, context=info.context) if not isinstance(v, AdamNumericalOptimizerConfig): raise ValueError( "train_config must be of type AdamOptimizerConfig when method is 'adam'" ) elif info.data["train_method"] == "lbfgs": + if isinstance(v, dict): + v = LBFGSNumericalOptimizerConfig.model_validate( + v, context=info.context + ) if not isinstance(v, LBFGSNumericalOptimizerConfig): raise ValueError( "train_config must be of type LBFGSOptimizerConfig when method is 'lbfgs'" @@ -243,20 +254,30 @@ def validate_train_config(cls, v, info: ValidationInfo): raise ValueError("method must be either 'adam' or 'lbfgs'") return v - @field_validator("covar_modules", "mean_modules", mode="before") - def validate_torch_modules(cls, value: Any): - if not isinstance(value, dict): - raise ValueError("must be dict") - else: - value = cast(dict[str, Any], value) - for key, val in value.items(): - if isinstance(val, str): - if val.startswith("base64:"): - value[key] = decode_torch_module(val) - elif os.path.exists(val): - value[key] = torch.load(val, weights_only=False) - - return value + @model_serializer(mode="wrap", when_used="json") + def serialize_torch_modules(self, handler, info): + result = handler(self) + options = get_serialization_options(info.context) + # only touch fields the handler kept (respect include/exclude) + for field_name in ("covar_modules", "mean_modules"): + if field_name not in result: + continue + modules = getattr(self, field_name) + if options.module_mode == "drop": + result[field_name] = {} + elif options.module_mode == "file": + result[field_name] = { + key: save_module_sidecar(module, options, f"{field_name}_{key}.pt") + for key, module in modules.items() + } + if "custom_noise_prior" in result and self.custom_noise_prior is not None: + if options.module_mode == "drop": + del result["custom_noise_prior"] + elif options.module_mode == "file": + result["custom_noise_prior"] = save_module_sidecar( + self.custom_noise_prior, options, "custom_noise_prior.pt" + ) + return result @field_validator("trainable_mean_keys") def validate_trainable_mean_keys(cls, value: Any, info: ValidationInfo): diff --git a/xopt/generators/bayesian/multi_fidelity.py b/xopt/generators/bayesian/multi_fidelity.py index dafddac94..6590d29c5 100644 --- a/xopt/generators/bayesian/multi_fidelity.py +++ b/xopt/generators/bayesian/multi_fidelity.py @@ -11,6 +11,7 @@ ) from pydantic import Field, field_validator +from xopt.generator import support_flag from xopt.generators.bayesian.custom_botorch.constrained_acquisition import ( ConstrainedMCAcquisitionFunction, ) @@ -72,9 +73,9 @@ class MultiFidelityGenerator(MOBOGenerator): exclude=True, ) reference_point: Optional[Dict[str, float]] = None - supports_multi_objective: bool = True - supports_batch_generation: bool = True - supports_constraints: bool = True + supports_multi_objective: bool = support_flag(True) + supports_batch_generation: bool = support_flag(True) + supports_constraints: bool = support_flag(True) __doc__ = """Implements Multi-fidelity Bayesian optimization Assumes a fidelity parameter [0,1] diff --git a/xopt/generators/bayesian/time_dependent.py b/xopt/generators/bayesian/time_dependent.py index c3ae72bdc..42bb8ebfc 100644 --- a/xopt/generators/bayesian/time_dependent.py +++ b/xopt/generators/bayesian/time_dependent.py @@ -9,6 +9,7 @@ from botorch.acquisition import FixedFeatureAcquisitionFunction from pydantic import Field, ValidationInfo, field_validator, PositiveFloat +from xopt.generator import support_flag from xopt.generators.bayesian.bayesian_generator import BayesianGenerator from xopt.generators.bayesian.models.time_dependent import TimeDependentModelConstructor @@ -45,8 +46,8 @@ class TimeDependentBayesianGenerator(BayesianGenerator, ABC): """ name = "time_dependent_bayesian_generator" - supports_single_objective: bool = True - supports_constraints: bool = True + supports_single_objective: bool = support_flag(True) + supports_constraints: bool = support_flag(True) target_prediction_time: Optional[PositiveFloat] = Field(None) added_time: PositiveFloat = Field( 0.1, diff --git a/xopt/generators/bayesian/upper_confidence_bound.py b/xopt/generators/bayesian/upper_confidence_bound.py index f9ce3553e..84ebe1205 100644 --- a/xopt/generators/bayesian/upper_confidence_bound.py +++ b/xopt/generators/bayesian/upper_confidence_bound.py @@ -9,6 +9,7 @@ from pydantic import Field import torch +from xopt.generator import support_flag from xopt.errors import GeneratorWarning from xopt.generators.bayesian.bayesian_generator import ( BayesianGenerator, @@ -52,9 +53,9 @@ class UpperConfidenceBoundGenerator(BayesianGenerator): 0.0, description="Vertical shift applied to the UCB acquisition function for use with constraints", ) - supports_batch_generation: bool = True - supports_single_objective: bool = True - supports_constraints: bool = True + supports_batch_generation: bool = support_flag(True) + supports_single_objective: bool = support_flag(True) + supports_constraints: bool = support_flag(True) _compatible_turbo_controllers = [OptimizeTurboController, SafetyTurboController] __doc__ = """Bayesian optimization generator using Upper Confidence Bound diff --git a/xopt/generators/deduplicated.py b/xopt/generators/deduplicated.py index c94fba4c3..3bfb2e804 100644 --- a/xopt/generators/deduplicated.py +++ b/xopt/generators/deduplicated.py @@ -1,10 +1,11 @@ -import numpy as np -from pydantic import field_validator -from typing import Optional import logging import time +from typing import Optional + +import numpy as np from xopt.generator import Generator +from xopt.types import NDArray from xopt.vocs import get_variable_data @@ -35,7 +36,7 @@ class DeduplicatedGeneratorBase(Generator): deduplicate_output: bool = True # The decision vars seen so far - decision_vars_seen: Optional[np.ndarray] = None + decision_vars_seen: Optional[NDArray] = None # For per-object log output in child objects (see eg NSGA2Generator) _logger: Optional[logging.Logger] = None @@ -46,13 +47,6 @@ def model_post_init(self, context): f"{__name__}.DeduplicatedGeneratorBase.{id(self)}" ) - @field_validator("decision_vars_seen", mode="before") - @classmethod - def cast_arr(cls, value): - if isinstance(value, list): - return np.array(value) - return value - def generate(self, n_candidates: int) -> list[dict]: """ Generate the unique candidates. diff --git a/xopt/generators/ga/cnsga.py b/xopt/generators/ga/cnsga.py index 4a7a12596..6ce0d4a9a 100644 --- a/xopt/generators/ga/cnsga.py +++ b/xopt/generators/ga/cnsga.py @@ -2,16 +2,19 @@ import logging import os import random -from typing import Dict, List, Optional +from typing import Any, Dict, List, Optional import pandas as pd -from deap import algorithms as deap_algorithms, base as deap_base, tools as deap_tools -from pydantic import ConfigDict, confloat, Field, PrivateAttr +from deap import algorithms as deap_algorithms +from deap import base as deap_base +from deap import tools as deap_tools +from pydantic import ConfigDict, Field, PrivateAttr, confloat import xopt.utils -from xopt.generator import Generator +from xopt.generator import Generator, support_flag from xopt.generators.ga import deap_creator from xopt.generators.ga.deap_fitness_with_constraints import FitnessWithConstraints +from xopt.types import XDataFrame from xopt.vocs import ( VOCS, convert_dataframe_to_inputs, @@ -70,9 +73,9 @@ class CNSGAGenerator(Generator): """ name = "cnsga" - supports_multi_objective: bool = True - supports_constraints: bool = True - supports_single_objective: bool = True + supports_multi_objective: bool = support_flag(True) + supports_constraints: bool = support_flag(True) + supports_single_objective: bool = support_flag(True) population_size: int = Field(64, description="Population size") crossover_probability: confloat(ge=0, le=1) = Field( 0.9, description="Crossover probability" @@ -88,18 +91,17 @@ class CNSGAGenerator(Generator): ) _children: List[Dict] = PrivateAttr([]) _offspring: Optional[pd.DataFrame] = PrivateAttr(None) - population: Optional[pd.DataFrame] = Field(None) + # use to generate children until the first pop is made + _loaded_population: Optional[pd.DataFrame] = PrivateAttr(None) + # DEAP toolbox; must be a declared PrivateAttr or extra="allow" serializes it + _toolbox: Any = PrivateAttr(None) + population: Optional[XDataFrame] = Field(None) model_config = ConfigDict(extra="allow") def __init__(self, **kwargs): super().__init__(**kwargs) - self._loaded_population = ( - None # use these to generate children until the first pop is made - ) - - # DEAP toolbox (internal) self._toolbox = cnsga_toolbox(self.vocs, selection="auto") if self.population_file is not None: diff --git a/xopt/generators/ga/nsga2.py b/xopt/generators/ga/nsga2.py index ba716cccb..a583e7181 100644 --- a/xopt/generators/ga/nsga2.py +++ b/xopt/generators/ga/nsga2.py @@ -10,6 +10,7 @@ import time import warnings +from xopt.generator import support_flag from xopt.vocs import get_constraint_data, get_objective_data, get_variable_data from ...errors import DataError from ...generator import StateOwner @@ -369,9 +370,9 @@ class NSGA2Generator(DeduplicatedGeneratorBase, StateOwner): """ name = "nsga2" - supports_multi_objective: bool = True - supports_constraints: bool = True - supports_single_objective: bool = True + supports_multi_objective: bool = support_flag(True) + supports_constraints: bool = support_flag(True) + supports_single_objective: bool = support_flag(True) # Checkpoint loading checkpoint_file: str | None = Field( diff --git a/xopt/generators/ga/operators.py b/xopt/generators/ga/operators.py index 9d23ece04..a90be1a95 100644 --- a/xopt/generators/ga/operators.py +++ b/xopt/generators/ga/operators.py @@ -1,6 +1,7 @@ +from typing import Annotated, Literal, Optional, Tuple + import numpy as np -from pydantic import Field, field_validator -from typing import Optional, Literal, Annotated, Tuple +from pydantic import Field from ...pydantic import XoptBaseModel @@ -8,14 +9,6 @@ class MutationOperator(XoptBaseModel): name: Literal["abstract"] = "abstract" - @field_validator("name", mode="after") - def validate_files(cls, value, info): - """ - Hack to override the wildcard before validator in `XoptBaseModel` for - the discriminator field. Before validators are dissallowed in this case. - """ - return value - def __call__(self, parent: np.ndarray, bounds: np.ndarray) -> np.ndarray: raise NotImplementedError @@ -131,14 +124,6 @@ def __call__(self, parent: np.ndarray, bounds: np.ndarray) -> np.ndarray: class CrossoverOperator(XoptBaseModel): name: Literal["abstract"] = "abstract" - @field_validator("name", mode="after") - def validate_files(cls, value, info): - """ - Hack to override the wildcard before validator in `XoptBaseModel` for - the discriminator field. Before validators are dissallowed in this case. - """ - return value - def __call__( self, parent_a: np.ndarray, parent_b: np.ndarray, bounds: np.ndarray ) -> Tuple[np.ndarray, np.ndarray]: diff --git a/xopt/generators/random.py b/xopt/generators/random.py index a9b31ad56..a15cbfbe6 100644 --- a/xopt/generators/random.py +++ b/xopt/generators/random.py @@ -1,4 +1,4 @@ -from xopt.generator import Generator +from xopt.generator import Generator, support_flag from xopt.vocs import random_inputs @@ -8,11 +8,11 @@ class RandomGenerator(Generator): """ name = "random" - supports_batch_generation: bool = True - supports_multi_objective: bool = True - supports_single_objective: bool = True - supports_constraints: bool = True - supports_discrete_variables: bool = True + supports_batch_generation: bool = support_flag(True) + supports_multi_objective: bool = support_flag(True) + supports_single_objective: bool = support_flag(True) + supports_constraints: bool = support_flag(True) + supports_discrete_variables: bool = support_flag(True) def generate(self, n_candidates) -> list[dict]: """generate uniform random data points""" diff --git a/xopt/generators/scipy/latin_hypercube.py b/xopt/generators/scipy/latin_hypercube.py index 660a676b3..7225e62ef 100644 --- a/xopt/generators/scipy/latin_hypercube.py +++ b/xopt/generators/scipy/latin_hypercube.py @@ -4,7 +4,7 @@ from scipy.stats import qmc from typing_extensions import Annotated -from xopt.generator import Generator +from xopt.generator import Generator, support_flag from gest_api.vocs import ExploreObjective from xopt.errors import VOCSError @@ -42,10 +42,10 @@ class LatinHypercubeGenerator(Generator): """ name = "latin_hypercube" - supports_batch_generation: bool = True - supports_multi_objective: bool = True - supports_single_objective: bool = True - supports_constraints: bool = True + supports_batch_generation: bool = support_flag(True) + supports_multi_objective: bool = support_flag(True) + supports_single_objective: bool = support_flag(True) + supports_constraints: bool = support_flag(True) batch_size: Annotated[ Optional[int], diff --git a/xopt/generators/sequential/extremumseeking.py b/xopt/generators/sequential/extremumseeking.py index f48dcb6c5..77b06e196 100644 --- a/xopt/generators/sequential/extremumseeking.py +++ b/xopt/generators/sequential/extremumseeking.py @@ -4,6 +4,7 @@ import pandas as pd from pydantic import Field, PositiveFloat +from xopt.generator import support_flag from xopt.vocs import get_variable_data, get_objective_data from xopt.generators.sequential.sequential_generator import SequentialGenerator @@ -67,7 +68,7 @@ class ExtremumSeekingGenerator(SequentialGenerator): k: PositiveFloat = Field(2.0, description="feedback gain") oscillation_size: PositiveFloat = Field(0.1, description="oscillation size") decay_rate: PositiveFloat = Field(1.0, description="decay rate") - supports_single_objective: bool = True + supports_single_objective: bool = support_flag(True) _nES: int = 0 _wES: np.ndarray = np.array([]) diff --git a/xopt/generators/sequential/neldermead.py b/xopt/generators/sequential/neldermead.py index 447eee6ac..26711d0c1 100644 --- a/xopt/generators/sequential/neldermead.py +++ b/xopt/generators/sequential/neldermead.py @@ -1,14 +1,16 @@ import logging import warnings from copy import deepcopy -from typing import Dict, List, Optional, Union, Tuple +from typing import Dict, List, Optional, Tuple, Union import numpy as np import pandas as pd from pydantic import ConfigDict, Field, field_validator +from xopt.generator import support_flag from xopt.generators.sequential.sequential_generator import SequentialGenerator from xopt.pydantic import XoptBaseModel +from xopt.types import NDArray from xopt.vocs import ( VOCS, get_objective_data, @@ -26,16 +28,16 @@ class SimplexState(XoptBaseModel): N: Optional[int] = None kend: int = 0 jend: int = 0 - ind: Optional[np.ndarray] = None - sim: Optional[np.ndarray] = None - fsim: Optional[np.ndarray] = None + ind: Optional[NDArray] = None + sim: Optional[NDArray] = None + fsim: Optional[NDArray] = None fxr: Optional[float] = None - x: Optional[np.ndarray] = None - xr: Optional[np.ndarray] = None - xe: Optional[np.ndarray] = None - xc: Optional[np.ndarray] = None - xcc: Optional[np.ndarray] = None - xbar: Optional[np.ndarray] = None + x: Optional[NDArray] = None + xr: Optional[NDArray] = None + xe: Optional[NDArray] = None + xc: Optional[NDArray] = None + xcc: Optional[NDArray] = None + xbar: Optional[NDArray] = None doshrink: int = 0 ngen: int = 0 model_config = ConfigDict(arbitrary_types_allowed=True) @@ -44,6 +46,9 @@ class SimplexState(XoptBaseModel): "ind", "fsim", "sim", "x", "xr", "xe", "xc", "xcc", "xbar", mode="before" ) def to_numpy(cls, v): + # strings are decoded downstream by the NDArray codec + if v is None or isinstance(v, str): + return v return np.array(v, dtype=np.float64) @@ -104,12 +109,25 @@ class NelderMeadGenerator(SequentialGenerator): """ name = "neldermead" - supports_single_objective: bool = True + supports_single_objective: bool = support_flag(True) initial_point: Optional[Dict[str, float]] = None # replaces x0 argument - initial_simplex: Optional[Dict[str, Union[List[float], np.ndarray]]] = ( + initial_simplex: Optional[Dict[str, Union[List[float], NDArray]]] = ( None # This overrides the use of initial_point ) + + @field_validator("initial_simplex", mode="before") + def check_simplex_entries(cls, v): + # the NDArray codec accepts bare scalars (0-d arrays), but a simplex + # entry must be one point per vertex + if isinstance(v, dict): + for key, entry in v.items(): + if isinstance(entry, (int, float)): + raise ValueError( + f"initial_simplex entry {key!r} must be a list of floats" + ) + return v + # Same as scipy.optimize._optimize._minimize_neldermead adaptive: bool = Field( True, description="Change hyperparameters based on dimensionality" @@ -118,7 +136,7 @@ class NelderMeadGenerator(SequentialGenerator): future_state: Optional[SimplexState] = None # Internal data structures - x: Optional[np.ndarray] = None + x: Optional[NDArray] = None y: Optional[float] = None manual_data_cnt: int = Field( 0, description="How many points are considered manual/not part of simplex run" diff --git a/xopt/generators/sequential/rcds.py b/xopt/generators/sequential/rcds.py index 4e04c240d..3faeddbd3 100644 --- a/xopt/generators/sequential/rcds.py +++ b/xopt/generators/sequential/rcds.py @@ -4,11 +4,13 @@ import numpy as np import pandas as pd +from gest_api.vocs import MaximizeObjective, MinimizeObjective from pydantic import ConfigDict, Field from pydantic.types import PositiveFloat -from gest_api.vocs import MinimizeObjective, MaximizeObjective +from xopt.generator import support_flag from xopt.generators.sequential.sequential_generator import SequentialGenerator +from xopt.types import NDArray logger = logging.getLogger(__name__) @@ -746,8 +748,8 @@ class RCDSGenerator(SequentialGenerator): """ name = "rcds" - supports_single_objective: bool = True - init_mat: Optional[np.ndarray] = Field(None) + supports_single_objective: bool = support_flag(True) + init_mat: Optional[NDArray] = Field(None) noise: PositiveFloat = Field(1e-5) step: PositiveFloat = Field(1e-2) diff --git a/xopt/pydantic.py b/xopt/pydantic.py index b1a7544b3..6b075f72f 100644 --- a/xopt/pydantic.py +++ b/xopt/pydantic.py @@ -1,14 +1,11 @@ import copy import inspect -import io import json import logging import os.path import typing from concurrent.futures import Future -from functools import partial -from importlib import import_module -from types import FunctionType, MethodType +from types import BuiltinFunctionType, FunctionType, MethodType from typing import ( Any, Callable, @@ -18,274 +15,152 @@ Optional, TextIO, TypeVar, - cast, ) import numpy as np -import orjson import pandas as pd -import torch.nn +import torch import yaml from pydantic import ( BaseModel, ConfigDict, Field, + SerializeAsAny, create_model, - field_serializer, field_validator, - model_serializer, model_validator, ) -from pydantic.v1.json import custom_pydantic_encoder -from pydantic_core.core_schema import SerializationInfo, ValidationInfo +from pydantic_core.core_schema import ValidationInfo + +from xopt.types import ( + CallableRef, + SerializationOptions, + TypeRef, + maybe_decompress, + module_load_base_dir, + normalize_serialization_context, + object_from_qualified_name, + qualified_name, + resolve_callable, +) ObjType = TypeVar("ObjType") logger = logging.getLogger(__name__) -JSON_ENCODERS = { - # function/method type distinguished for class members - # and not recognized as callables - FunctionType: lambda x: f"{x.__module__}.{x.__qualname__}", - MethodType: lambda x: f"{x.__module__}.{x.__qualname__}", - Callable: lambda x: f"{x.__module__}.{x.__qualname__}", - type: lambda x: f"{x.__module__}.{x.__name__}", - # for encoding instances of the ObjType} - # ObjType: lambda x: f"{x.__module__}.{x.__class__.__qualname__}", - np.ndarray: lambda x: x.tolist(), - np.int64: lambda x: int(x), - np.float64: lambda x: float(x), - # torch.nn.Module: lambda x: process_torch_module(x), - # torch.Tensor: lambda x: x.detach().cpu().numpy().tolist(), -} - - -def _serialize_non_finite_float(value: float | np.floating) -> str: - value = float(value) - if np.isnan(value): - return "nan" - if value > 0: - return "inf" - return "-inf" - - -def _serialize_list(values, base_key="", serialize_torch=False, serialize_inline=False): - serialized_values = [] - for i, item in enumerate(values): - list_key = f"{base_key}_{i}" if base_key else str(i) - - if isinstance(item, dict): - item = recursive_serialize( - item, list_key, serialize_torch, serialize_inline - ) - elif isinstance(item, list): - item = _serialize_list(item, list_key, serialize_torch, serialize_inline) - elif isinstance(item, (float, np.floating)) and not np.isfinite(float(item)): - item = _serialize_non_finite_float(item) - else: - for _type, func in JSON_ENCODERS.items(): - if isinstance(item, _type): - item = func(item) - - if isinstance(item, (float, np.floating)) and not np.isfinite(float(item)): - item = _serialize_non_finite_float(item) - - try: - json.dumps(item) - except (TypeError, OverflowError): - item = f"{item.__module__}.{item.__class__.__qualname__}" - - serialized_values.append(item) - - return serialized_values - - -# The problem with v2 serialization is that model_serialize_json() does not accept kwargs -# meaning whichever model method is decorated with @model_serializer cant adjust for 'base_key' -# and other similar options - it renders native whole v2 scheme quite useless. We can still try -# to use @field_serializer, but there is a lack of documentation on how to call these -# handlers from a custom function. -# -# So, we implement two serialization options for now. -# First is the native one, with no customization, under serialize_json() method. It needs to -# invoked as xopt.model_dump_json(), the standard pydantic v2 syntax. -# Second method bypasses pydantic completely. It is invoked via 'xopt.json()' -# or '.to_json()' - -# Pydantic v2 will by default serialize submodels as annotated types, dropping subclass attributes - - -def recursive_serialize( - v, base_key="", serialize_torch=False, serialize_inline: bool = False -) -> dict: - for key in list(v): - if isinstance(v[key], dict): - v[key] = recursive_serialize(v[key], key, serialize_torch, serialize_inline) - elif isinstance(v[key], list): - v[key] = _serialize_list(v[key], key, serialize_torch, serialize_inline) - elif isinstance(v[key], torch.nn.Module): - if serialize_torch: - if serialize_inline: - v[key] = "base64:" + encode_torch_module(v[key]) - else: - v[key] = process_torch_module( - module=v[key], name="_".join((base_key, key)) - ) - else: - del v[key] - elif isinstance(v[key], torch.dtype): - v[key] = str(v[key]) - elif isinstance(v[key], pd.DataFrame): - v[key] = json.loads(v[key].to_json()) - elif isinstance(v[key], set): - v[key] = list(v[key]) - elif isinstance(v[key], (float, np.floating)) and not np.isfinite( - float(v[key]) - ): - v[key] = _serialize_non_finite_float(v[key]) - else: - for _type, func in JSON_ENCODERS.items(): - if isinstance(v[key], _type): - v[key] = func(v[key]) - - if isinstance(v[key], (float, np.floating)) and not np.isfinite( - float(v[key]) - ): - v[key] = _serialize_non_finite_float(v[key]) - # check to make sure object has been serialized, - # if not use a generic serializer - try: - # handle case when key is (not) deleted - if key in v: - json.dumps(v[key]) - except (TypeError, OverflowError): - v[key] = f"{v[key].__module__}.{v[key].__class__.__qualname__}" +def serialization_fallback(value: Any) -> Any: + """Return JSON-native representations for values below untyped fields.""" + if isinstance(value, np.generic): + return value.item() + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, torch.dtype): + return str(value) + if isinstance(value, torch.Tensor): + return value.detach().cpu().tolist() + if isinstance(value, pd.DataFrame): + return json.loads(value.to_json()) + if isinstance(value, (type, FunctionType, MethodType, BuiltinFunctionType)): + # deliberately narrow: callable *instances* (nn.Module, partial, ...) + # fall through to the generic class-name form below + return qualified_name(value) + if isinstance(value, Exception): + return str(value) + return f"{type(value).__module__}.{type(value).__qualname__}" - return v - -DECODERS = {"torch.float32": torch.float32, "torch.float64": torch.float64} - - -def recursive_deserialize(v: dict) -> dict: - """deserialize strings from xopt outputs""" - for key, value in v.items(): - # process dicts - if isinstance(value, dict): - v[key] = recursive_deserialize(value) - - elif isinstance(value, str): - if value in DECODERS: - v[key] = DECODERS[value] - - return v - - -def orjson_dumps( - v: BaseModel, *, base_key="", serialize_torch=False, serialize_inline=False -) -> str: - # TODO: move away from borrowing pydantic v1 encoder preset - json_encoder = partial(custom_pydantic_encoder, JSON_ENCODERS) - return orjson_dumps_custom( - v, - default=json_encoder, - base_key=base_key, - serialize_torch=serialize_torch, - serialize_inline=serialize_inline, +class XoptBaseModel(BaseModel): + model_config = ConfigDict( + arbitrary_types_allowed=True, extra="forbid", ser_json_inf_nan="strings" ) + @classmethod + def model_validate(cls, *args, **kwargs): + # models with custom __init__ methods re-enter validation without the + # caller's context, so a context-supplied base_dir must also travel via + # the contextvar to reach nested file-loading codecs + context = kwargs.get("context") + base_dir = context.get("base_dir") if isinstance(context, dict) else None + if base_dir is not None: + with module_load_base_dir(base_dir): + return super().model_validate(*args, **kwargs) + return super().model_validate(*args, **kwargs) + + def model_dump_json(self, *args, **kwargs) -> str: + kwargs["context"] = normalize_serialization_context(kwargs.get("context")) + kwargs.setdefault("fallback", serialization_fallback) + return super().model_dump_json(*args, **kwargs) + + def model_dump(self, *args, **kwargs) -> dict[str, Any]: + if kwargs.get("mode") == "json": + kwargs["context"] = normalize_serialization_context(kwargs.get("context")) + kwargs.setdefault("fallback", serialization_fallback) + return super().model_dump(*args, **kwargs) + + def to_json(self, **kwargs) -> str: + context = kwargs.pop("context", None) + option_values = { + name: kwargs.pop(name) + for name in ( + "array_mode", + "module_mode", + "df_mode", + "compress", + "level", + "file_dir", + ) + if name in kwargs + } -def orjson_dumps_custom(v: BaseModel, *, default, base_key="", **kwargs) -> str: - v = recursive_serialize(v.model_dump(), base_key=base_key, **kwargs) - return orjson.dumps(v, default=default).decode() - - -def orjson_dumps_except_root(v: BaseModel, *, base_key="", **kwargs) -> dict: - """Same as above but start at fields of root model, instead of model itself""" - dump = v.model_dump() - encoded_dump = recursive_serialize(dump, base_key=base_key, **kwargs) - return encoded_dump - - -def orjson_loads(v, default=None) -> dict: - v = orjson.loads(v) - v = recursive_deserialize(v) - return v - - -def process_torch_module(module, name): - """save module to file based on module name and return file path to json""" - # module_name = "".join(random.choices(string.ascii_uppercase + string.digits, - # k=7)) + ".pt" - module_name = f"{name}.pt" - torch.save(module, module_name) - return module_name - - -def encode_torch_module(module): - import base64 - import gzip - - buffer = io.BytesIO() - # 5 supported since 3.8 - torch.save(module, buffer, pickle_protocol=5) - module_bytes = buffer.getbuffer().tobytes() - cb = gzip.compress(module_bytes, compresslevel=9) - encoded_bytes = base64.standard_b64encode(cb) - return encoded_bytes.decode("ascii") - - -def decode_torch_module(modulestr: str): - import base64 - import gzip - - assert modulestr.startswith("base64:") - base64str = modulestr.split("base64:", 1)[1] - decoded = base64.standard_b64decode(base64str) - decompressed = gzip.decompress(decoded) - bytestream = io.BytesIO(decompressed) - module = torch.load(bytestream, weights_only=False) - return module - - -class XoptBaseModel(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + serialize_torch = kwargs.pop("serialize_torch", None) + serialize_inline = kwargs.pop("serialize_inline", None) + if "module_mode" not in option_values and ( + serialize_torch is not None or serialize_inline is not None + ): + option_values["module_mode"] = ( + "inline" + if serialize_torch and serialize_inline + else "file" + if serialize_torch + else "drop" + ) - @model_validator(mode="before") - @classmethod - def validate_files(cls, data: Any) -> Any: - if not isinstance(data, dict): - return data - for key, value in data.items(): - # Exclude field `name`` from before validator for use in discriminated fields - if key == "name": - continue - if isinstance(value, str): - if os.path.exists(value): - extension = value.split(".")[-1] - if extension == "pt": - data[key] = torch.load(value, weights_only=False) - return data - - # Note that this function still returns a dict, NOT a string. Pydantic will handle - # final serialization of basic types in Rust. - @model_serializer(mode="plain", when_used="json") - def serialize_json(self, sinfo: SerializationInfo) -> dict: - return orjson_dumps_except_root(self) - - def to_json(self, **kwargs: Any) -> str: - return orjson_dumps(self, **kwargs) + if isinstance(context, dict): + normalized_context = normalize_serialization_context(context) + context_options = normalized_context["serialization_options"] + merged_values = dict(option_values) + merged_values.update( + { + name: getattr(context_options, name) + for name in context_options._explicit + } + ) + merged_context = { + key: value + for key, value in normalized_context.items() + if key not in {"serialization_options", *option_values} + } + merged_context["serialization_options"] = SerializationOptions( + **merged_values + ) + elif isinstance(context, SerializationOptions): + merged_values = dict(option_values) + merged_values.update( + {name: getattr(context, name) for name in context._explicit} + ) + merged_context = SerializationOptions(**merged_values) + elif context is None: + merged_context = option_values + else: + raise TypeError("context must be a mapping or SerializationOptions") + return self.model_dump_json(context=merged_context, **kwargs) def json(self, **kwargs: Any) -> str: return self.to_json(**kwargs) def yaml(self, **kwargs: Any) -> str: """serialize first then dump to yaml string""" - output = json.loads( - self.to_json( - **kwargs, - ) - ) + output = json.loads(self.to_json(**kwargs)) return yaml.dump(output) @classmethod @@ -293,28 +168,19 @@ def from_file(cls, filename: str) -> "XoptBaseModel": if not os.path.exists(filename): raise OSError(f"file {filename} is not found") - with open(filename, "r") as file: - return cls.from_yaml(file) + with open(filename, "rb") as file: + raw = maybe_decompress(file.read()) + data = yaml.safe_load(raw.decode("utf-8")) + base_dir = str(os.path.dirname(os.path.abspath(filename))) + return cls.model_validate(data, context={"base_dir": base_dir}) @classmethod def from_yaml(cls, yaml_obj: str | TextIO) -> "XoptBaseModel": - return cls.model_validate(remove_none_values(yaml.safe_load(yaml_obj))) + return cls.model_validate(yaml.safe_load(yaml_obj)) @classmethod def from_dict(cls, config: dict) -> "XoptBaseModel": - return cls.model_validate(remove_none_values(config)) - - -def remove_none_values(d: Any) -> Any: - if isinstance(d, dict): - d = cast(dict[str, Any], d) - # Create a copy of the dictionary to avoid modifying the original while iterating - d = {k: remove_none_values(v) for k, v in d.items() if v is not None} - elif isinstance(d, list): - d = cast(list[Any], d) - # If it's a list, recursively process each item in the list - d = [remove_none_values(item) for item in d if item is not None] - return d + return cls.model_validate(config) def get_descriptions_defaults(model: XoptBaseModel): @@ -327,25 +193,17 @@ def get_descriptions_defaults(model: XoptBaseModel): if isinstance(value, XoptBaseModel): description_dict[name] = get_descriptions_defaults(value) else: - try: - description_dict[name] = [val.description, val.default] - except TypeError: - # if the val is an object or callable type - description_dict[name] = val.description + description_dict[name] = [val.description, val.default] return description_dict -class CallableModel(BaseModel): - callable: Callable - signature: BaseModel +class CallableModel(XoptBaseModel): + callable: CallableRef + signature: SerializeAsAny[BaseModel] model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - @model_serializer(mode="plain", when_used="json", return_type="str") - def serialize(self): - return orjson_dumps(self) - @model_validator(mode="before") def validate_all(cls, values): callable = values.pop("callable") @@ -409,17 +267,13 @@ def __call__(self, *args, **kwargs): class ObjLoader( - BaseModel, + XoptBaseModel, Generic[ObjType], ): model_config = ConfigDict(arbitrary_types_allowed=True) object: Optional[ObjType] = None - loader: CallableModel = None - object_type: Optional[type] = None - - @model_serializer(mode="plain", when_used="json", return_type="str") - def serialize_json(self) -> str: - return orjson_dumps(self) + loader: Optional[CallableModel] = None + object_type: Optional[TypeRef] = None @model_validator(mode="before") def validate_all(cls, values): @@ -482,39 +336,9 @@ def load(self, store: bool = False): return self.loader() -# For testing -class ObjLoaderMinimal( - BaseModel, - Generic[ObjType], -): - model_config = ConfigDict(arbitrary_types_allowed=True) - object: Optional[ObjType] = None - object_type: Optional[type] = None - - @model_validator(mode="before") - def validate_all(cls, values): - print("model validator before: ", values) - annotation = cls.model_fields["object"].annotation - inner_types = typing.get_args(annotation) - obj_type = inner_types[0] - print(f"{obj_type=}") - return {"object_type": obj_type} - - @model_validator(mode="after") - def validate_print(self, values): - print("model validator after: ", values) - return values - - @field_serializer("object_type", when_used="json") - def serialize_object_type(self, x): - if x is None: - return x - return f"{x.__module__}.{x.__name__}" - - # COMMON BASE FOR EXECUTORS class BaseExecutor( - BaseModel, + XoptBaseModel, Generic[ObjType], ): model_config = ConfigDict(arbitrary_types_allowed=True) @@ -525,18 +349,14 @@ class BaseExecutor( # This is a utility field not included in reps. The typing lib has opened # issues on access of generic type within class. # This tracks for if-necessary future use. - executor_type: Optional[type] = Field(None, exclude=True, validate_default=True) + executor_type: Optional[TypeRef] = Field(None, exclude=True, validate_default=True) submit_callable: str = "submit" map_callable: str = "map" shutdown_callable: str = "shutdown" # executor will not be explicitly serialized, but loaded using loader with class # and kwargs - executor: Optional[ObjType] = None - - @model_serializer(mode="plain", when_used="json", return_type="str") - def serialize_json(self) -> str: - return orjson_dumps(self) + executor: Optional[ObjType] = Field(None, exclude=True) @model_validator(mode="before") def validate_all(cls, values): @@ -666,98 +486,37 @@ def map(self, fn, *iter: Iterable, **kwargs) -> Iterable[Future]: def get_callable_from_string(callable: str, bind: Any = None) -> Callable: - """Get callable from a string. In the case that the callable points to a bound method, - the function returns a callable taking the bind instance as the first arg. + """Get callable from its fully qualified name, e.g. ``module.func`` or + ``module.Class.method``. - Parameters - ---------- - callable: String representation of callable abiding convention - __module__:callable - bind: Class to bind as self + Parameters + ---------- + callable : str + Fully qualified name of the callable. + bind : Any, optional + Instance to bind to when the name points to a method of the instance's + class; the bound method is returned. Returns ------- - Callable + Callable """ - callable_split = callable.rsplit(".", 1) + fn = resolve_callable(callable) - if len(callable_split) != 2: - raise ValueError(f"Improperly formatted callable string: {callable_split}") - - module_name, callable_name = callable_split + if bind is None: + return fn + owner_name, _, attr_name = callable.rpartition(".") try: - module = import_module(module_name) - - except ModuleNotFoundError: - try: - module_split = module_name.rsplit(".", 1) - - if len(module_split) != 2: - raise ValueError(f"Unable to access: {callable}") - - module_name, class_name = module_split - - module = import_module(module_name) - callable_name = f"{class_name}.{callable_name}" - - except ModuleNotFoundError as err: - logger.error("Unable to import module %s", module_name) - raise err - - except ValueError as err: - logger.error(err) - raise err - - # construct partial in case of bound method - if "." in callable_name: - bound_class, callable_name = callable_name.rsplit(".") - - try: - bound_class = getattr(module, bound_class) - except Exception as e: - logger.error("Unable to get %s from %s", bound_class, module_name) - raise e - - # require right partial for assembly of callable - # https://funcy.readthedocs.io/en/stable/funcs.html#rpartial - def rpartial(func, *args): - return lambda *a: func(*(a + args)) - - callable = getattr(bound_class, callable_name) - params = inspect.signature(callable).parameters - - # check bindings - is_bound = params.get("self", None) is not None - if not is_bound and bind is not None: - raise ValueError("Cannot bind %s to %s.", callable_name, bind) - - # bound, return partial - if bind is not None: - if not isinstance(bind, (bound_class,)): - raise ValueError( - "Provided bind %s is not instance of %s", - bind, - bound_class.__qualname__, - ) - - if is_bound and isinstance(callable, (FunctionType,)) and bind is None: - callable = rpartial(getattr, callable_name) - - elif is_bound and isinstance(callable, (FunctionType,)) and bind is not None: - callable = getattr(bind, callable_name) - - else: - if bind is not None: - raise ValueError("Cannot bind %s to %s.", callable_name, type(bind)) - - try: - callable = getattr(module, callable_name) - except Exception as e: - logger.error("Unable to get %s from %s", callable_name, module_name) - raise e + owner = object_from_qualified_name(owner_name) if owner_name else None + except ValueError: + owner = None + if not isinstance(owner, type) or not isinstance(bind, owner): + raise ValueError( + f"Cannot bind {callable!r} to instance of {type(bind).__name__}" + ) - return callable + return getattr(bind, attr_name) class SignatureModel(BaseModel): diff --git a/xopt/resources/testing.py b/xopt/resources/testing.py index a9ddb58ab..89a2c6cf6 100644 --- a/xopt/resources/testing.py +++ b/xopt/resources/testing.py @@ -10,7 +10,6 @@ from torch import nn from xopt import Generator -from xopt.pydantic import remove_none_values from xopt.vocs import VOCS # TODO: make a config module like gpytorch has @@ -246,7 +245,7 @@ def reload_gen_from_json(gen): def reload_gen_from_yaml(gen): assert isinstance(gen, Generator) gen_class = gen.__class__ - gen_new = gen_class(**remove_none_values(yaml.safe_load(gen.yaml()))) + gen_new = gen_class(**yaml.safe_load(gen.yaml())) gen_new.add_data(gen.data.copy()) return gen_new diff --git a/xopt/stopping_conditions.py b/xopt/stopping_conditions.py index 7b6e0cf48..bc01915d8 100644 --- a/xopt/stopping_conditions.py +++ b/xopt/stopping_conditions.py @@ -7,23 +7,22 @@ """ from abc import ABC, abstractmethod -from typing import List, Literal, Annotated, Union +from typing import Annotated, List, Literal, Union + import pandas as pd +from gest_api.vocs import VOCS, MinimizeObjective from pydantic import ( ConfigDict, + Discriminator, Field, PositiveFloat, PositiveInt, - field_serializer, - field_validator, - Discriminator, TypeAdapter, + field_validator, ) - from xopt.pydantic import XoptBaseModel from xopt.vocs import get_feasibility_data -from gest_api.vocs import MinimizeObjective, VOCS class StoppingCondition(XoptBaseModel, ABC): @@ -316,16 +315,6 @@ class CompositeCondition(StoppingCondition): default="or", description="Logic to combine conditions: 'and' or 'or'" ) - @field_serializer("conditions") - @classmethod - def serialize_conditions(cls, v): - serialized_conditions = [] - for condition in v: - serialized_conditions.append( - condition.model_dump() | {"name": condition.__class__.__name__} - ) - return serialized_conditions - @field_validator("logic") @classmethod def validate_logic(cls, v): diff --git a/xopt/tests/generators/bayesian/test_bayesian_generator.py b/xopt/tests/generators/bayesian/test_bayesian_generator.py index 8518045f3..f817264d1 100644 --- a/xopt/tests/generators/bayesian/test_bayesian_generator.py +++ b/xopt/tests/generators/bayesian/test_bayesian_generator.py @@ -1,36 +1,37 @@ -from copy import deepcopy import os +import tempfile +from copy import deepcopy from unittest import TestCase from unittest.mock import MagicMock, patch import numpy as np import pandas as pd -from pydantic import ValidationInfo import pytest import torch -from torch.nn import Module +from botorch.models import SingleTaskGP from botorch.models.gpytorch import GPyTorchModel from botorch.models.transforms import Normalize, Standardize from gpytorch.kernels import PeriodicKernel - +from pydantic import ValidationInfo from xopt import VOCS from xopt.base import Xopt -from xopt.errors import VOCSError, XoptError +from xopt.generator import support_flag +from xopt.errors import VOCSError from xopt.evaluator import Evaluator -from xopt.generators.bayesian.models.standard import StandardModelConstructor from xopt.generators.bayesian.base_model import ModelConstructor from xopt.generators.bayesian.bayesian_generator import ( BayesianGenerator, MultiObjectiveBayesianGenerator, ) +from xopt.generators.bayesian.models.standard import StandardModelConstructor from xopt.generators.bayesian.turbo import ( OptimizeTurboController, ) from xopt.numerical_optimizer import GridOptimizer, LBFGSOptimizer -from xopt.pydantic import encode_torch_module from xopt.resources.test_functions.sinusoid_1d import evaluate_sinusoid, sinusoid_vocs from xopt.resources.testing import TEST_VOCS_BASE, TEST_VOCS_DATA +from xopt.types import encode_torch_module from xopt.vocs import random_inputs @@ -87,14 +88,19 @@ def test_init(self): gen.get_acquisition(gen.model) # test asking for batch generation when not supported - gen.supports_batch_generation = False + class NoBatchGenerator(PatchBayesianGenerator): + supports_batch_generation: bool = support_flag(False) + + gen = NoBatchGenerator(vocs=TEST_VOCS_BASE) with pytest.raises(NotImplementedError): gen.generate(2) # test with n_interpolate_points but mutiple candidates - gen = PatchBayesianGenerator(vocs=TEST_VOCS_BASE) + class BatchGenerator(PatchBayesianGenerator): + supports_batch_generation: bool = support_flag(True) + + gen = BatchGenerator(vocs=TEST_VOCS_BASE) gen.n_interpolate_points = 5 - gen.supports_batch_generation = True with pytest.raises(RuntimeError): gen.generate(2) @@ -162,18 +168,19 @@ class CustomBayesianGenerator(BayesianGenerator): assert CustomBayesianGenerator.get_compatible_turbo_controllers() == [None] def test_torch_module_validation(self): - # test validate torch modules - encoded_module = encode_torch_module(torch.nn.Linear(5, 2)) - exit_val = BayesianGenerator.validate_torch_modules("base64: " + encoded_module) - assert isinstance(exit_val, Module) - - torch.save(torch.nn.Linear(3, 1), "test_module.pt") - exit_val = BayesianGenerator.validate_torch_modules("test_module.pt") - assert isinstance(exit_val, torch.nn.Linear) - os.remove("test_module.pt") + model = SingleTaskGP(torch.rand(3, 1), torch.rand(3, 1)) + encoded_module = encode_torch_module(model) + with patch.multiple(PatchBayesianGenerator, __abstractmethods__=set()): + generator = PatchBayesianGenerator( + vocs=TEST_VOCS_BASE, model=encoded_module + ) + assert isinstance(generator.model, SingleTaskGP) - with pytest.raises(XoptError): - BayesianGenerator.validate_torch_modules("invalid_string") + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "test_module.pt") + torch.save(model, path) + generator = PatchBayesianGenerator(vocs=TEST_VOCS_BASE, model=path) + assert isinstance(generator.model, SingleTaskGP) def test_numerical_optimizer_validation(self): # test with None @@ -541,27 +548,29 @@ def test_bad_mo_vocs(self): def test_validate_gp_constructor_none(self): # Should return StandardModelConstructor instance - result = BayesianGenerator.validate_gp_constructor(None) + result = BayesianGenerator.validate_gp_constructor(None, None) assert isinstance(result, StandardModelConstructor) def test_validate_gp_constructor_instance(self): dummy = DummyModelConstructor() - result = BayesianGenerator.validate_gp_constructor(dummy) + result = BayesianGenerator.validate_gp_constructor(dummy, None) assert result is dummy def test_validate_gp_constructor_str(self): - result = BayesianGenerator.validate_gp_constructor("standard") + result = BayesianGenerator.validate_gp_constructor("standard", None) assert isinstance(result, StandardModelConstructor) with pytest.raises(ValueError): - BayesianGenerator.validate_gp_constructor("not_a_constructor") + BayesianGenerator.validate_gp_constructor("not_a_constructor", None) def test_validate_gp_constructor_dict(self): # Valid dict - result = BayesianGenerator.validate_gp_constructor({"name": "standard"}) + result = BayesianGenerator.validate_gp_constructor({"name": "standard"}, None) assert isinstance(result, StandardModelConstructor) # Invalid dict with pytest.raises(ValueError): - BayesianGenerator.validate_gp_constructor({"name": "not_a_constructor"}) + BayesianGenerator.validate_gp_constructor( + {"name": "not_a_constructor"}, None + ) def test_validate_turbo_controller(self): # Should return None diff --git a/xopt/tests/test_entrypoint.py b/xopt/tests/test_entrypoint.py index d907216c9..802a0d80b 100644 --- a/xopt/tests/test_entrypoint.py +++ b/xopt/tests/test_entrypoint.py @@ -30,9 +30,8 @@ def make_config(self, tmp_path): config_path.write_text(yaml.dump(config)) return str(config_path), config - @mock.patch("xopt.entrypoint.remove_none_values", side_effect=lambda x: x) @mock.patch("xopt.entrypoint.Xopt") - def test_main_basic(self, mock_Xopt, mock_remove_none, tmp_path): + def test_main_basic(self, mock_Xopt, tmp_path): config_path, config = self.make_config(tmp_path) sys_argv = [ "entrypoint.py", @@ -102,10 +101,9 @@ def test_normalize_initial_data_keeps_xopt_error_str(self): @mock.patch("xopt.entrypoint.normalize_initial_data") @mock.patch("xopt.entrypoint.pd.read_csv") - @mock.patch("xopt.entrypoint.remove_none_values", side_effect=lambda x: x) @mock.patch("xopt.entrypoint.Xopt") def test_main_with_initial_data( - self, mock_Xopt, mock_remove_none, mock_read_csv, mock_normalize, tmp_path + self, mock_Xopt, mock_read_csv, mock_normalize, tmp_path ): config_path, config = self.make_config(tmp_path) csv_path = str(tmp_path / "initial.csv") diff --git a/xopt/tests/test_io.py b/xopt/tests/test_io.py index c2d4df52d..17bd6dcd1 100644 --- a/xopt/tests/test_io.py +++ b/xopt/tests/test_io.py @@ -11,13 +11,6 @@ def dummy(): class Test_IO: - def test_options_to_dict(self): - evaluator = Evaluator(function=dummy) - generator = RandomGenerator(vocs=TEST_VOCS_BASE) - X = Xopt(generator=generator, evaluator=evaluator) - print(X.model_dump_json()) - print(X.to_json(base_key="bk")) - def test_state_to_dict(self): evaluator = Evaluator(function=dummy) generator = RandomGenerator(vocs=TEST_VOCS_BASE) diff --git a/xopt/tests/test_pydantic.py b/xopt/tests/test_pydantic.py index 56625f25f..ef22e3109 100644 --- a/xopt/tests/test_pydantic.py +++ b/xopt/tests/test_pydantic.py @@ -1,682 +1,636 @@ -import inspect +import base64 +import gzip import io import json -import os -import tempfile -from functools import partial -from types import FunctionType, MethodType -from typing import Callable, Optional, Union +from concurrent.futures import Executor +from typing import Annotated, Any import numpy as np import pandas as pd import pytest import torch import yaml -from pydantic import BaseModel, ConfigDict, Field, SerializeAsAny, field_validator -from pydantic.json import custom_pydantic_encoder - +from gpytorch.kernels import RBFKernel +from pydantic import Field + +from xopt.base import Xopt +from xopt.evaluator import Evaluator +from xopt.generators.bayesian.models.standard import StandardModelConstructor +from xopt.generators.bayesian.objectives import CustomXoptObjective +from xopt.generators.random import RandomGenerator from xopt.pydantic import ( - JSON_ENCODERS, CallableModel, NormalExecutor, ObjLoader, - ObjLoaderMinimal, - SignatureModel, XoptBaseModel, - decode_torch_module, - encode_torch_module, - get_callable_from_string, get_descriptions_defaults, - orjson_dumps, - orjson_dumps_custom, - orjson_dumps_except_root, - orjson_loads, - process_torch_module, - recursive_deserialize, - recursive_serialize, - remove_none_values, validate_and_compose_signature, ) +from xopt.resources.testing import TEST_VOCS_BASE +from xopt.types import ( + CallableRef, + DataFrameCodec, + NDArray, + NDArrayCodec, + SerializationOptions, + TorchDType, + TorchModuleCodec, + TorchTensor, + XDataFrame, + encode_torch_module, + maybe_decompress, +) -def misc_fn(x, y=1, *args, **kwargs): - pass +def misc_fn(x=1, y=2): + return x + y class MiscClass: - @staticmethod - def misc_static_method(x, y=1, *args, **kwargs): - return - - @classmethod - def misc_cls_method(cls, x, y=1, *args, **kwargs): - return cls - - def misc_method(self, x, y=1, *args, **kwargs): - return - - -class TestJsonEncoders: - misc_class = MiscClass() - - @pytest.mark.parametrize( - ("fn",), - [ - (misc_fn,), - pytest.param(misc_class.misc_method, marks=pytest.mark.xfail(strict=True)), - (misc_class.misc_static_method,), - pytest.param( - misc_class.misc_cls_method, marks=pytest.mark.xfail(strict=True) - ), - ], - ) - def test_function_type(self, fn): - encoder = {FunctionType: JSON_ENCODERS[FunctionType]} - json_encoder = partial(custom_pydantic_encoder, encoder) - - serialized = json.dumps(fn, default=json_encoder) - loaded = json.loads(serialized) - callable_from_str = get_callable_from_string(loaded) - - assert fn == callable_from_str - - @pytest.mark.parametrize( - ("fn",), - [ - pytest.param( - misc_class.misc_static_method, marks=pytest.mark.xfail(strict=True) - ), - pytest.param(misc_fn, marks=pytest.mark.xfail(strict=True)), - (misc_class.misc_method,), - pytest.param( - misc_class.misc_cls_method, marks=pytest.mark.xfail(strict=True) - ), - ], - ) - def test_method_type(self, fn): - encoder = {MethodType: JSON_ENCODERS[MethodType]} - json_encoder = partial(custom_pydantic_encoder, encoder) - - serialized = json.dumps(fn, default=json_encoder) - loaded = json.loads(serialized) - callable = get_callable_from_string(loaded, bind=self.misc_class) - - assert fn == callable - - @pytest.mark.parametrize( - ("fn",), - [ - (misc_class.misc_static_method,), - (misc_fn,), - (misc_class.misc_method,), - (misc_class.misc_cls_method,), - ], - ) - def test_full_encoder(self, fn): - json_encoder = partial(custom_pydantic_encoder, JSON_ENCODERS) - serialized = json.dumps(fn, default=json_encoder) - loaded = json.loads(serialized) - - get_callable_from_string(loaded) - - -class TestSignatureValidateAndCompose: - misc_class = MiscClass() - - @pytest.mark.parametrize( - ("args", "kwargs"), - [ - pytest.param((5, 2, 1), {"x": 2}, marks=pytest.mark.xfail(strict=True)), - pytest.param((), ({"y": 2}), marks=pytest.mark.xfail(strict=True)), - pytest.param((2,), ({"x": 2}), marks=pytest.mark.xfail(strict=True)), - ((), ({"x": 2})), - ((), {}), - ], - ) - def test_validate_kwarg_only(self, args, kwargs): - def run(*, x: int = 4): - pass - - signature_model = validate_and_compose_signature(run, *args, **kwargs) - assert all( - [kwargs[kwarg] == getattr(signature_model, kwarg) for kwarg in kwargs] - ) - # run - - args, kwargs = signature_model.build() - - run(*args, **kwargs) - - @pytest.mark.parametrize( - ("args", "kwargs"), - [ - pytest.param( - ( - 5, - 3, - 2, - ), - {"x": 1}, - marks=pytest.mark.xfail(strict=True), - ), - ((2, 1, 0), {}), - ((), {}), - ], - ) - def test_validate_var_positional(self, args, kwargs): - def run(*args): - pass - - signature_model = validate_and_compose_signature(run, *args, **kwargs) - args, kwargs = signature_model.build() - assert len(kwargs) == 0 - assert len(args) == len(args) - - # run - run(*args) - - @pytest.mark.parametrize( - ("args", "kwargs"), - [ - pytest.param((5,), {"x": 2}, marks=pytest.mark.xfail(strict=True)), - ((), {"x": 2, "y": 3}), - pytest.param((), {}, marks=pytest.mark.xfail(strict=True)), - ( - ( - 2, - 4, - ), - {}, - ), - ((2,), {"y": 4, "extra": True}), - ((2,), {"y": 4}), - ((2,), {"y": 4, "z": 3}), - ], - ) - def test_validate_full_sig(self, args, kwargs): - def run(x, y, z=4, *args, **kwargs): - pass - - signature_model = validate_and_compose_signature(run, *args, **kwargs) - args, kwargs = signature_model.build() - - # run - run(*args, **kwargs) - - @pytest.mark.parametrize( - ("args", "kwargs"), - [ - pytest.param((5, 1), {"y": 2}, marks=pytest.mark.xfail(strict=True)), - ( - ( - 2, - 4, - ), - {}, - ), - ((5,), {"y": 2}), - ], - ) - def test_validate_classmethod(self, args, kwargs): - signature_model = validate_and_compose_signature( - self.misc_class.misc_cls_method, *args, **kwargs - ) - args, kwargs = signature_model.build() - self.misc_class.misc_cls_method(*args, **kwargs) - - @pytest.mark.parametrize( - ("args", "kwargs"), - [ - pytest.param((5, 1), {"y": 2}, marks=pytest.mark.xfail(strict=True)), - ( - ( - 2, - 4, - ), - {}, - ), - ((5,), {"y": 2}), - ], - ) - def test_validate_staticmethod(self, args, kwargs): - signature_model = validate_and_compose_signature( - self.misc_class.misc_static_method, *args, **kwargs - ) - args, kwargs = signature_model.build() - self.misc_class.misc_static_method(*args, **kwargs) - - @pytest.mark.parametrize( - ("args", "kwargs"), - [ - pytest.param((5, 1), {"y": 2}, marks=pytest.mark.xfail(strict=True)), - ( - ( - 2, - 4, - ), - {}, - ), - ((5,), {"y": 2}), - ], - ) - def test_validate_bound_method(self, args, kwargs): - signature_model = validate_and_compose_signature( - self.misc_class.misc_method, *args, **kwargs - ) + def __init__(self, value=1): + self.value = value - args, kwargs = signature_model.build() + def misc_method(self, x=1): + return self.value + x - self.misc_class.misc_method(*args, **kwargs) +class DummyExecutor(Executor): + def __init__(self, tag="default"): + self.tag = tag -class TestCallableModel: - misc_class = MiscClass() + def submit(self, fn, *args, **kwargs): + return fn(*args, **kwargs) - @pytest.mark.parametrize( - ("fn", "args", "kwargs"), - [ - (misc_fn, (5,), {"y": 2}), - (misc_class.misc_cls_method, (5,), {"y": 2}), - (misc_class.misc_static_method, (5,), {"y": 2}), - pytest.param( - misc_class.misc_method, - (5,), - {"y": 2}, - marks=pytest.mark.xfail(strict=True), - ), - ], - ) - def test_construct_callable(self, fn, args, kwargs): - json_encoder = partial(custom_pydantic_encoder, JSON_ENCODERS) - serialized = json.dumps(fn, default=json_encoder) - loaded = json.loads(serialized) - - callable = CallableModel(callable=loaded) - callable(*args, **kwargs) - - @pytest.mark.parametrize( - ("fn", "args", "kwargs"), - [ - pytest.param(misc_fn, (5,), {"y": 2}, marks=pytest.mark.xfail(strict=True)), - pytest.param( - misc_class.misc_cls_method, - (5,), - {"y": 2}, - marks=pytest.mark.xfail(strict=True), - ), - pytest.param( - misc_class.misc_static_method, - (5,), - {"y": 2}, - marks=pytest.mark.xfail(strict=True), - ), - (misc_class.misc_method, (5,), {"y": 2}), - ], - ) - def test_bound_callables(self, fn, args, kwargs): - json_encoder = partial(custom_pydantic_encoder, JSON_ENCODERS) - serialized = json.dumps(fn, default=json_encoder) - loaded = json.loads(serialized) - - callable = CallableModel(callable=loaded, bind=self.misc_class) - callable(*args, **kwargs) - - -class TestObjLoader: - misc_class_loader_type = ObjLoader[MiscClass] + def map(self, fn, *iterables, timeout=None, chunksize=1): + return list(map(fn, *iterables)) - def test_class_loader(self): - loader = self.misc_class_loader_type() - assert loader.object_type == MiscClass + def shutdown(self, wait=True, *, cancel_futures=False): + self.was_shutdown = True - def test_load_model(self): - loader = self.misc_class_loader_type() - misc_obj = loader.load() - assert isinstance(misc_obj, (MiscClass,)) - def test_serialize_loader(self): - loader = self.misc_class_loader_type() +class CodecModel(XoptBaseModel): + array: NDArray + tensor: TorchTensor + dtype: TorchDType + dataframe: XDataFrame - json_encoder = partial(custom_pydantic_encoder, JSON_ENCODERS) - serialized = json.dumps(loader, default=json_encoder) - # self.misc_class_loader_type.parse_raw(serialized) - # This works in 2.2+ as it should - self.misc_class_loader_type.model_validate_json(serialized) - - -# tests to verify v2 behavior remains same (for things that changed from v1) +@pytest.fixture +def codec_model(): + return CodecModel( + array=np.array([[1.25, 2.5]], dtype=np.float64), + tensor=torch.tensor([[3.0, 4.0]], dtype=torch.float64), + dtype=torch.float16, + dataframe=pd.DataFrame( + {"x": [1.1234567890123, np.nan], "label": ["a", "b"]}, + index=[3, 7], + ), + ) -class DummyObj: - pass +def test_codec_default_shapes_and_python_mode(codec_model): + dumped = json.loads(codec_model.model_dump_json()) + assert dumped["array"] == [[1.25, 2.5]] + assert dumped["tensor"] == [[3.0, 4.0]] + assert dumped["dtype"] == "torch.float16" + # to_json's default 10-digit precision applies + assert dumped["dataframe"] == { + "x": {"3": 1.123456789, "7": None}, + "label": {"3": "a", "7": "b"}, + } + python_dump = codec_model.model_dump(mode="python") + assert python_dump["array"] is codec_model.array + assert python_dump["tensor"] is codec_model.tensor + assert python_dump["dataframe"] is codec_model.dataframe + + loaded = CodecModel.model_validate_json(codec_model.model_dump_json()) + np.testing.assert_array_equal(loaded.array, codec_model.array) + torch.testing.assert_close(loaded.tensor, codec_model.tensor) + assert loaded.dtype is torch.float16 + pd.testing.assert_frame_equal(loaded.dataframe, codec_model.dataframe) + + +@pytest.mark.parametrize("compression", [None, "gzip", "zstd"]) +def test_binary_codec_round_trips(codec_model, compression): + dumped = codec_model.model_dump_json( + context={ + "array_mode": "b64", + "df_mode": "b64", + "compress": compression, + "level": 3, + } + ) + data = json.loads(dumped) + assert data["array"].startswith("b64np:") + assert data["tensor"].startswith("b64pt:") + assert data["dataframe"].startswith("b64df:") + + raw = base64.b64decode(data["array"].removeprefix("b64np:")) + expected_magic = { + None: b"\x93NUMPY", + "gzip": b"\x1f\x8b", + "zstd": b"\x28\xb5\x2f\xfd", + } + assert raw.startswith(expected_magic[compression]) + + loaded = CodecModel.model_validate_json(dumped) + np.testing.assert_array_equal(loaded.array, codec_model.array) + torch.testing.assert_close(loaded.tensor, codec_model.tensor) + # b64df decoding matches dict mode: the integer index is restored + pd.testing.assert_frame_equal( + loaded.dataframe, codec_model.dataframe, check_exact=False, atol=1e-9 + ) -class Dummy(BaseModel): - default_obj: DummyObj = Field(DummyObj()) - model_config = ConfigDict(arbitrary_types_allowed=True) - @field_validator("default_obj") - def validate_obj(cls, value): - assert isinstance(value, DummyObj) - return value +def test_annotation_defaults_and_context_precedence(): + class ParameterizedModel(XoptBaseModel): + array: Annotated[np.ndarray, NDArrayCodec(array_mode="b64")] + dataframe: Annotated[pd.DataFrame, DataFrameCodec(df_mode="b64")] + model = ParameterizedModel( + array=np.array([1.0]), dataframe=pd.DataFrame({"x": [1.0]}) + ) + annotated = json.loads(model.model_dump_json()) + assert annotated["array"].startswith("b64np:") + assert annotated["dataframe"].startswith("b64df:") -# Test subclass model resolution order -# we want behavior like v1 had https://github.com/pydantic/pydantic/issues/1932 -class Parent(BaseModel): - a1: str = "a1" + overridden = json.loads( + model.model_dump_json(context={"array_mode": "list", "df_mode": "dict"}) + ) + assert overridden == {"array": [1.0], "dataframe": {"x": {"0": 1.0}}} + options = SerializationOptions(array_mode="b64", df_mode="b64") + object_context = json.loads(model.model_dump_json(context=options)) + assert object_context["array"].startswith("b64np:") -class Child1(Parent): - name: str = "child1" + wrapper_context = json.loads( + model.to_json( + array_mode="b64", + context={"serialization_options": SerializationOptions(array_mode="list")}, + ) + ) + assert wrapper_context["array"] == [1.0] + + +def test_torch_dtype_generalized_validation(): + class DTypeModel(XoptBaseModel): + dtype: TorchDType + + for dtype in ( + torch.float32, + torch.float64, + torch.float16, + torch.int64, + torch.bool, + ): + loaded = DTypeModel.model_validate({"dtype": str(dtype)}) + assert loaded.dtype is dtype + with pytest.raises(ValueError, match="invalid torch dtype"): + DTypeModel.model_validate({"dtype": "torch.not_a_dtype"}) + + +def test_module_modes_and_sidecar_sanitization(tmp_path, monkeypatch): + constructor = StandardModelConstructor(covar_modules={"y/bad key": RBFKernel()}) + + dropped = json.loads(constructor.to_json()) + assert dropped["covar_modules"] == {} + + inline = json.loads(constructor.to_json(module_mode="inline")) + assert inline["covar_modules"]["y/bad key"].startswith("b64pt:") + inline_loaded = StandardModelConstructor.model_validate(inline) + assert isinstance(inline_loaded.covar_modules["y/bad key"], RBFKernel) + + written = json.loads(constructor.to_json(module_mode="file", file_dir=tmp_path)) + assert written["covar_modules"] == {"y/bad key": "covar_modules_y_bad_key.pt"} + assert (tmp_path / "covar_modules_y_bad_key.pt").is_file() + monkeypatch.chdir(tmp_path) + file_loaded = StandardModelConstructor.model_validate(written) + assert isinstance(file_loaded.covar_modules["y/bad key"], RBFKernel) + + +def test_module_inline_legacy_and_prefix_collision_errors(): + class ModuleModel(XoptBaseModel): + module: Annotated[torch.nn.Module, TorchModuleCodec()] + + legacy = "base64:" + encode_torch_module(torch.nn.Linear(1, 1))[len("b64pt:") :] + loaded = ModuleModel.model_validate({"module": legacy}) + assert isinstance(loaded.module, torch.nn.Linear) + with pytest.raises(ValueError, match="invalid b64pt: torch module payload"): + ModuleModel.model_validate({"module": "base64:AAAA"}) + with pytest.raises(ValueError, match="malformed base64"): + ModuleModel.model_validate({"module": "b64pt:not-valid!"}) + with pytest.raises(ValueError, match="cannot load torch module"): + ModuleModel.model_validate({"module": "missing-module.pt"}) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("array", "b64np:@@@@", "malformed base64"), + ("tensor", "b64pt:@@@@", "malformed base64"), + ("dataframe", "b64df:@@@@", "malformed base64"), + ( + "array", + "b64np:" + base64.b64encode(b"not a numpy file").decode(), + "invalid b64np", + ), + ( + "dataframe", + "b64df:" + base64.b64encode(b"not json").decode(), + "invalid b64df", + ), + ], +) +def test_malformed_binary_payloads(codec_model, field, value, message): + data = codec_model.model_dump(mode="python") + data[field] = value + with pytest.raises(ValueError, match=message): + CodecModel.model_validate(data) -class Child2(Parent): - name: str = "child2" +def test_magic_decompression_and_size_cap(): + raw = b"payload" + assert maybe_decompress(raw) == raw + assert maybe_decompress(gzip.compress(raw)) == raw + import zstandard -class Container(BaseModel): - obj: SerializeAsAny[Optional[Parent]] = Field(None) - obj2: SerializeAsAny[Optional[Union[Child1, Child2, Parent]]] = Field(None) + compressed = zstandard.ZstdCompressor().compress(raw) + assert maybe_decompress(compressed) == raw + with pytest.raises(ValueError, match="size limit"): + maybe_decompress(gzip.compress(b"x" * 100), max_size=16) + with pytest.raises(ValueError, match="size limit"): + maybe_decompress(zstandard.ZstdCompressor().compress(b"x" * 100), max_size=16) + # the cap applies to decompression only, not raw passthrough + assert maybe_decompress(b"x" * 100, max_size=16) == b"x" * 100 + with pytest.raises(ValueError, match="invalid gzip"): + maybe_decompress(b"\x1f\x8btruncated") + with pytest.raises(ValueError, match="invalid zstd"): + maybe_decompress(b"\x28\xb5\x2f\xfdtruncated") -class TestPydanticInitialization: - def test_object_init(self): - d = Dummy() - assert isinstance(d.default_obj, DummyObj) - def test_subclass_init(self): - c1 = Container() - print("c1", c1.model_dump()) - c2 = Container(obj=Child2()) - print("c2", c2.model_dump()) - # doesn't resolve child1 - c3 = Container(**{"obj": {"a1": "a1", "name": "child1"}}) - print(type(c3.obj), type(c3.obj2), c3) - # works - c4 = Container(**{"obj2": {"a1": "a1", "name": "child1"}}) - print(type(c4.obj), type(c4.obj2), c4) +def test_fallback_is_recursive_and_json_native(): + class AnyModel(XoptBaseModel): + values: dict[str, Any] + class Unknown: + pass -class DummyModel(BaseModel): - a: int = 1 - b: str = "foo" - c: None = None + model = AnyModel( + values={ + "np_scalar": np.int64(2), + "array": np.array([1, 2]), + "dtype": torch.float32, + "tensor": torch.tensor([3, 4]), + "callable": misc_fn, + "type": MiscClass, + "exception": RuntimeError("boom"), + "unknown": Unknown(), + } + ) + dumped = json.loads(model.model_dump_json()) + assert dumped["values"] == { + "np_scalar": 2, + "array": [1, 2], + "dtype": "torch.float32", + "tensor": [3, 4], + "callable": f"{__name__}.misc_fn", + "type": f"{__name__}.MiscClass", + "exception": "boom", + "unknown": f"{Unknown.__module__}.{Unknown.__qualname__}", + } -class DummyTorchModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.linear = torch.nn.Linear(2, 2) +def test_non_finite_floats_round_trip(): + class NonFiniteModel(XoptBaseModel): + values: list[float] = [-float("inf"), float("inf"), float("nan")] + loaded = json.loads(NonFiniteModel().model_dump_json()) + assert loaded["values"][:2] == ["-Infinity", "Infinity"] + assert loaded["values"][2] == "NaN" + reloaded = NonFiniteModel.model_validate(loaded) + assert reloaded.values[0] == -float("inf") + assert reloaded.values[1] == float("inf") + assert reloaded.values[2] != reloaded.values[2] -def test_recursive_serialize_and_deserialize(): - d = { - "a": 1, - "b": {"c": 2}, - "d": np.array([1, 2]), - "e": set([1, 2]), - "f": pd.DataFrame({"x": [1, 2]}), - } - ser = recursive_serialize(d.copy()) - assert isinstance(ser["d"], list) - assert isinstance(ser["e"], list) - assert isinstance(ser["f"], dict) - # test recursive_deserialize - d2 = {"a": 1, "b": {"c": 2}, "dtype": "torch.float32"} - deser = recursive_deserialize(d2.copy()) - assert deser["dtype"] == torch.float32 - - ser = recursive_serialize({"float": torch.float32, "unserizable": DummyModel()}) - assert ser == { - "float": "torch.float32", - "unserizable": "xopt.tests.test_pydantic.DummyModel", - } +def test_xoptbase_public_io_and_whole_file_compression(tmp_path): + class Model(XoptBaseModel): + a: int = 1 -def test_orjson_dumps_and_loads(): - m = DummyModel() - s = orjson_dumps(m) - assert isinstance(s, str) - loaded = orjson_loads(s) - assert loaded["a"] == 1 - # test custom - s2 = orjson_dumps_custom(m, default=lambda x: str(x)) - assert isinstance(s2, str) - # except root - d = orjson_dumps_except_root(m) - assert isinstance(d, dict) + model = Model() + assert json.loads(model.to_json()) == {"a": 1} + assert yaml.safe_load(model.yaml()) == {"a": 1} + assert Model.from_dict({"a": 2}).a == 2 + assert Model.from_yaml(io.StringIO("a: 3\n")).a == 3 + + plain = tmp_path / "model.yaml" + plain.write_text("a: 4\n") + assert Model.from_file(str(plain)).a == 4 + compressed = tmp_path / "model.yaml.gz" + compressed.write_bytes(gzip.compress(b"a: 5\n")) + assert Model.from_file(str(compressed)).a == 5 + with pytest.raises(OSError): + Model.from_file(str(tmp_path / "missing.yaml")) -def test_orjson_dumps_preserves_non_finite_floats(): - class NonFiniteModel(XoptBaseModel): - values: list[float] = [-float("inf"), float("inf")] +def test_xopt_generator_name_and_serialize_as_any(): + class CustomRandomGenerator(RandomGenerator): + subclass_only: int = 17 - m = NonFiniteModel() - loaded = json.loads(orjson_dumps(m)) - assert loaded["values"] == ["-inf", "inf"] + generator = CustomRandomGenerator(vocs=TEST_VOCS_BASE) + xopt = Xopt(generator=generator, evaluator=Evaluator(function=misc_fn)) + python_dump = xopt.model_dump() + json_dump = json.loads(xopt.model_dump_json()) + assert python_dump["generator"]["name"] == generator.name + assert python_dump["generator"]["subclass_only"] == 17 + assert json_dump["generator"]["name"] == generator.name + assert json_dump["generator"]["subclass_only"] == 17 -def test_process_and_encode_decode_torch_module(): - mod = DummyTorchModule() - with tempfile.TemporaryDirectory() as tmpdir: - path = process_torch_module(mod, os.path.join(tmpdir, "testmod")) - assert os.path.exists(path) - # encode/decode - encoded = encode_torch_module(mod) - decoded = decode_torch_module("base64:" + encoded) - assert isinstance(decoded, torch.nn.Module) +def test_callable_serialization_warns_when_not_reloadable(monkeypatch): + import functools + import warnings -def test_xoptbasemodel_to_json_yaml(tmp_path): - class M(XoptBaseModel): - a: int = 1 + import xopt.types - m = M() - assert isinstance(m.to_json(), str) - assert isinstance(m.json(), str) - assert isinstance(m.yaml(), str) - # test from_dict - m2 = M.from_dict({"a": 2}) - assert m2.a == 2 - # test from_yaml - yaml_str = yaml.dump({"a": 3}) - m3 = M.from_yaml(io.StringIO(yaml_str)) - assert m3.a == 3 - # test from_file - file = tmp_path / "test.yaml" - file.write_text(yaml_str) - m4 = M.from_file(str(file)) - assert m4.a == 3 - # test file not found - with pytest.raises(OSError): - M.from_file("nonexistent.yaml") + monkeypatch.setattr(xopt.types, "_WARNED_CALLABLE_NAMES", set()) - # test torch load in XoptBaseModel - torch.save(torch.nn.Linear(2, 2), tmp_path / "model.pt") - M.validate_files(yaml.safe_load(str(tmp_path / "model.pt"))) + class FnModel(XoptBaseModel): + fn: CallableRef + with warnings.catch_warnings(): + warnings.simplefilter("error") + dump = json.loads(FnModel(fn=misc_fn).model_dump_json()) + assert dump["fn"] == f"{__name__}.misc_fn" -def test_remove_none_values(): - d = {"a": 1, "b": None, "c": {"d": None, "e": 2}, "f": [None, 3]} - cleaned = remove_none_values(d) - assert "b" not in cleaned - assert "d" not in cleaned["c"] - assert cleaned["f"] == [3] + model = FnModel(fn=functools.partial(misc_fn, y=3)) + with pytest.warns(UserWarning, match="will not reload"): + model.model_dump_json() + # warned once per callable per process, not once per dump + with warnings.catch_warnings(): + warnings.simplefilter("error") + model.model_dump_json() + # lambdas have no importable qualified name + with pytest.warns(UserWarning, match="will not reload"): + FnModel(fn=lambda x: x).model_dump_json() -def test_get_descriptions_defaults(): - class M(XoptBaseModel): - """desc""" - a: int = 1 - f: Callable = lambda x: x + 1 +@pytest.mark.parametrize( + ("bind_args", "bind_kwargs", "build_kwargs", "expected"), + [ + # tuple defaults degrade to None, None/empty defaults survive + ((), {}, {"a": 1}, (1, 2, None, None)), + # bound positional and keyword values override defaults + ((1,), {"b": 3}, {}, (1, 3, None, None)), + # build-time kwargs override stored values + ((1,), {"b": 3}, {"b": 4, "d": "x"}, (1, 4, None, "x")), + ], +) +def test_validate_and_compose_signature_defaults( + bind_args, bind_kwargs, build_kwargs, expected +): + def fn(a, b=2, c=(1, 2), d=None): + return (a, b, c, d) + + signature = validate_and_compose_signature(fn, *bind_args, **bind_kwargs) + args, kwargs = signature.build(**build_kwargs) + assert fn(*args, **kwargs) == expected + + +def test_validate_and_compose_signature_varargs_and_invalid(): + def fn(*args, x=1): + return args, x + + signature = validate_and_compose_signature(fn, 1, 2) + assert signature.model_dump() == {"args": [1, 2], "x": 1} + # partial positional replacement keeps the remaining stored args + args, kwargs = signature.build(4) + assert (args, kwargs) == ([4, 2], {"x": 1}) + + def plain(a, b=2): + return a, b + + with pytest.raises(TypeError, match="too many positional"): + validate_and_compose_signature(plain, 1, 2, 3) + with pytest.raises(TypeError, match="unexpected keyword"): + validate_and_compose_signature(plain, nope=1) + + +def test_callable_model_reload_and_bind(): + model = CallableModel(callable=misc_fn, kwargs={"y": 7}) + dumped = json.loads(model.model_dump_json()) + assert dumped["callable"] == f"{__name__}.misc_fn" + reloaded = CallableModel.model_validate(dumped) + assert reloaded(x=5) == 12 + # positional call args map through the stored kwarg_order + assert reloaded(3, 4) == 7 + + instance = MiscClass(value=10) + bound = CallableModel(callable=f"{__name__}.MiscClass.misc_method", bind=instance) + assert bound(x=3) == 13 + with pytest.raises(ValueError, match="Cannot bind"): + CallableModel(callable=f"{__name__}.MiscClass.misc_method", bind=object()) + with pytest.raises(ValueError, match="must be object or a string"): + CallableModel(callable=123) + + +def test_objloader_guards_and_store(): + # loader callable must match the parameterized type + with pytest.raises(ValueError): + ObjLoader[MiscClass].model_validate( + {"loader": {"callable": f"{__name__}.misc_fn"}} + ) - m = M() - desc = get_descriptions_defaults(m) - assert "a" in desc - assert "f" in desc + loader_dump = json.loads(ObjLoader[MiscClass]().model_dump_json()) + assert isinstance( + ObjLoader[MiscClass].model_validate(loader_dump).load(), MiscClass + ) - class Inner(XoptBaseModel): - """inner desc""" + loader = ObjLoader[MiscClass](kwargs={"value": 5}) + assert loader.object is None + obj = loader.load(store=True) + assert loader.object is obj and obj.value == 5 - x: int = 42 - class Outer(XoptBaseModel): - """outer desc""" +def test_normal_executor_reload_reconstructs_executor(): + executor = NormalExecutor[DummyExecutor](loader={"kwargs": {"tag": "loaded"}}) + dumped = json.loads(executor.model_dump_json()) + assert "executor" not in dumped - inner: Inner = Inner() - y: float = 3.14 + reloaded = NormalExecutor[DummyExecutor].model_validate(dumped) + assert isinstance(reloaded.executor, DummyExecutor) + assert reloaded.executor.tag == "loaded" + assert reloaded.submit(misc_fn, 1, 2) == 3 + reloaded.shutdown() + assert reloaded.executor.was_shutdown - o = Outer() - desc = get_descriptions_defaults(o) - # Should recurse into inner and get its description dict - assert "inner" in desc - assert isinstance(desc["inner"], dict) - assert "x" in desc["inner"] - assert "y" in desc + with pytest.raises(ValueError, match="instance of DummyExecutor"): + NormalExecutor[DummyExecutor](executor="not an executor") - class DummyCallable: - pass - class M(XoptBaseModel): - """desc""" +def test_dump_and_reload_sidecars_from_other_cwd(tmp_path, monkeypatch): + from xopt.generators.bayesian.upper_confidence_bound import ( + UpperConfidenceBoundGenerator, + ) - a: DummyCallable = Field(DummyCallable(), description="callable field") + generator = UpperConfidenceBoundGenerator( + vocs=TEST_VOCS_BASE, + gp_constructor=StandardModelConstructor(covar_modules={"y1": RBFKernel()}), + ) + X = Xopt( + generator=generator, + evaluator=Evaluator(function=misc_fn), + serialize_torch=True, + ) + run_dir = tmp_path / "run" + run_dir.mkdir() + dump_file = run_dir / "xopt.yaml" + X.dump(str(dump_file)) + assert (run_dir / "covar_modules_y1.pt").is_file() - m = M() - desc = get_descriptions_defaults(m) - # Should handle object/callable type - assert desc["a"][0] == "callable field" + elsewhere = tmp_path / "elsewhere" + elsewhere.mkdir() + monkeypatch.chdir(elsewhere) + # from_file resolves sidecars relative to the dump file, not the cwd + reloaded = Xopt.from_file(str(dump_file)) + assert isinstance(reloaded.generator.gp_constructor.covar_modules["y1"], RBFKernel) -def test_objloader_minimal(): - class Dummy: - pass + # an explicit base_dir context must work too, surviving the custom + # __init__ methods in the generator chain + config = yaml.safe_load(dump_file.read_text()) + reloaded = Xopt.model_validate(config, context={"base_dir": str(run_dir)}) + assert isinstance(reloaded.generator.gp_constructor.covar_modules["y1"], RBFKernel) - loader = ObjLoaderMinimal[Dummy]() - assert loader.object_type == Dummy - # test serialize object type - res = loader.serialize_object_type(None) - assert res is None - res = loader.serialize_object_type(Dummy) - assert res == f"{Dummy.__module__}.{Dummy.__name__}" +def test_retained_model_helpers(): + class Inner(XoptBaseModel): + x: int = Field(1, description="x field") + class Outer(XoptBaseModel): + inner: Inner = Inner() + fn: Any = Field(misc_fn, description="function") -def test_signaturemodel_build(): - class S(SignatureModel): - args: list = [1, 2] - kwarg_order: list = ["x"] - x: int = 3 + descriptions = get_descriptions_defaults(Outer()) + assert descriptions["inner"]["x"][0] == "x field" + assert descriptions["fn"][0] == "function" - s = S() - args, kwargs = s.build(4, x=5) - assert args == [4, 2] # positional overwrite - assert kwargs["x"] == 5 +def test_scalar_and_nonfinite_array_round_trips(): + class ArrayModel(XoptBaseModel): + array: NDArray + tensor: TorchTensor -def test_baseexecutor_and_normalexecutor(): - class DummyExec: - def submit(self, fn, *args, **kwargs): - return "submitted" + m = ArrayModel( + array=np.array([np.nan, 1.0, np.inf]), + tensor=torch.tensor(2.5, dtype=torch.float64), + ) + dumped = json.loads(m.model_dump_json()) + assert dumped["array"] == ["NaN", 1.0, "Infinity"] + assert dumped["tensor"] == 2.5 + + loaded = ArrayModel.model_validate(dumped) + assert loaded.array.dtype == np.float64 + assert np.isnan(loaded.array[0]) and loaded.array[1] == 1.0 + assert np.isinf(loaded.array[2]) + assert loaded.tensor.ndim == 0 and float(loaded.tensor) == 2.5 + + assert json.loads(loaded.model_dump_json())["array"] == dumped["array"] + + # legacy lowercase non-finite scalars still load, and bare finite scalars + # round-trip as 0-d values in both codecs + legacy = ArrayModel.model_validate({"array": "nan", "tensor": "-inf"}) + assert np.isnan(legacy.array) and float(legacy.tensor) == -np.inf + scalars = ArrayModel.model_validate({"array": 3.5, "tensor": 0.0}) + assert scalars.array.ndim == 0 and float(scalars.array) == 3.5 + assert scalars.tensor.ndim == 0 and float(scalars.tensor) == 0.0 + assert json.loads(scalars.model_dump_json()) == {"array": 3.5, "tensor": 0.0} + + +def test_xopt_data_b64_round_trip(): + X = Xopt( + generator=RandomGenerator(vocs=TEST_VOCS_BASE), + evaluator=Evaluator(function=misc_fn), + ) + X.add_data(pd.DataFrame({"x1": [0.1, 0.2], "x2": [0.3, 0.4], "y1": [1.0, 2.0]})) + reloaded = Xopt.from_yaml(X.yaml(df_mode="b64")) + assert list(reloaded.data.index) == [0, 1] + assert reloaded.data["x1"].tolist() == [0.1, 0.2] + + +def test_module_wrap_serializers_respect_exclude(tmp_path): + constructor = StandardModelConstructor(covar_modules={"y1": RBFKernel()}) + dumped = json.loads(constructor.model_dump_json(exclude={"covar_modules"})) + assert "covar_modules" not in dumped + dumped = json.loads( + constructor.model_dump_json( + exclude={"covar_modules"}, + context={"module_mode": "file", "file_dir": tmp_path}, + ) + ) + assert "covar_modules" not in dumped + assert not list(tmp_path.iterdir()), "excluded field must not write sidecars" - def map(self, fn, *args, **kwargs): - return ["mapped"] - def shutdown(self): - return "shutdown" +def test_custom_noise_prior_round_trips(tmp_path): + from gpytorch.priors import GammaPrior - loader = ObjLoader[DummyExec]() # <-- Use ObjLoader, not ObjLoaderMinimal - be = NormalExecutor[DummyExec](loader=loader, executor=DummyExec()) - assert be.submit(lambda x: x, 1) == "submitted" - assert be.map(lambda x: x, [1, 2]) == ["mapped"] - be.shutdown() + constructor = StandardModelConstructor(custom_noise_prior=GammaPrior(1.0, 100.0)) + # drop (default): key removed, like other module-valued fields + dumped = json.loads(constructor.model_dump_json()) + assert "custom_noise_prior" not in dumped + assert StandardModelConstructor.model_validate(dumped).custom_noise_prior is None -def test_validate_and_compose_signature_tuple_and_empty(): - def fn_tuple(x=(1, 2)): - pass # pragma: no cover + # inline round trip + dumped = json.loads(constructor.to_json(module_mode="inline")) + assert dumped["custom_noise_prior"].startswith("b64pt:") + loaded = StandardModelConstructor.model_validate(dumped) + assert isinstance(loaded.custom_noise_prior, GammaPrior) - model = validate_and_compose_signature(fn_tuple) - # Should create a field with type tuple and default None - assert hasattr(model, "x") - assert model.model_fields["x"].annotation is tuple - assert model.model_fields["x"].default is None + # file round trip + dumped = json.loads(constructor.to_json(module_mode="file", file_dir=str(tmp_path))) + assert dumped["custom_noise_prior"] == "custom_noise_prior.pt" + assert (tmp_path / "custom_noise_prior.pt").exists() - def fn_empty(x=inspect.Parameter.empty): - pass # pragma: no cover + # an unset prior stays as an explicit null in every mode + empty = json.loads(StandardModelConstructor().model_dump_json()) + assert empty["custom_noise_prior"] is None - model = validate_and_compose_signature(fn_empty) - # Should create a field with type inspect.Parameter.empty and default inspect.Parameter.empty - assert hasattr(model, "x") - assert model.model_fields["x"].annotation == inspect.Parameter.empty - assert model.model_fields["x"].default == inspect.Parameter.empty - def fn_none(x=None): - pass # pragma: no cover +class _FirstOutputObjective(CustomXoptObjective): + # module scope so torch pickling can resolve it + def forward(self, samples, X=None): + return samples[..., 0] - model = validate_and_compose_signature(fn_none) - # Should create a field with type inspect.Parameter.empty and default None - assert hasattr(model, "x") - assert model.model_fields["x"].annotation == inspect.Parameter.empty - assert model.model_fields["x"].default is None - def fn_int(x=5): - pass # pragma: no cover +def test_custom_objective_round_trips(): + from xopt.generators.bayesian.expected_improvement import ( + ExpectedImprovementGenerator, + ) - model = validate_and_compose_signature(fn_int) - # Should create a field with type int and default 5 - assert hasattr(model, "x") - assert model.model_fields["x"].annotation is int - assert model.model_fields["x"].default == 5 + gen = ExpectedImprovementGenerator( + vocs=TEST_VOCS_BASE, custom_objective=_FirstOutputObjective(TEST_VOCS_BASE) + ) + dumped = json.loads(gen.model_dump_json()) + assert "custom_objective" not in dumped # dropped by default, like model + inline = json.loads(gen.to_json(module_mode="inline")) + assert inline["custom_objective"].startswith("b64pt:") + loaded = ExpectedImprovementGenerator.model_validate(inline) + assert isinstance(loaded.custom_objective, CustomXoptObjective) -def test_objloader_validate_all_loader_variants(): - class Dummy: - pass - # Loader not in values: should create CallableModel with Dummy as callable - loader = ObjLoader[Dummy]() - assert loader.object_type == Dummy - assert isinstance(loader.loader, type(loader.loader)) - # Loader is already a CallableModel - loader2 = ObjLoader[Dummy](loader=loader.loader) - assert loader2.object_type == Dummy - assert isinstance(loader2.loader, type(loader.loader)) - # Loader is a dict with 'callable' key - loader3 = ObjLoader[Dummy](loader={"callable": Dummy}) - assert loader3.object_type == Dummy - assert isinstance(loader3.loader, type(loader.loader)) - # Loader is a dict without 'callable' key - loader4 = ObjLoader[Dummy](loader={}) - assert loader4.object_type == Dummy - assert isinstance(loader4.loader, type(loader.loader)) - - # test serialization of loader - for loader in [loader, loader2, loader3, loader4]: - loader.serialize_json() - - # Loader with wrong callable type should raise ValueError - class Other: - pass +def test_model_dump_json_mode_uses_fallback(): + class KwargsModel(XoptBaseModel): + kwargs: dict = {} - with pytest.raises(ValueError): - ObjLoader[Dummy](loader={"callable": Other}) - - -def test_objloader_load_store_and_no_store(): - class Dummy: - def __init__(self): - self.value = 42 - - loader = ObjLoader[Dummy]() - # Test store=False (should return a new Dummy instance, not store it) - result1 = loader.load(store=False) - assert isinstance(result1, Dummy) - assert loader.object is None # Should not store - # Test store=True (should store the Dummy instance) - result2 = loader.load(store=True) - assert isinstance(result2, Dummy) - assert loader.object is result2 # Should store + m = KwargsModel(kwargs={"a": np.float32(1.5), "df": pd.DataFrame({"x": [1]})}) + dumped = m.model_dump(mode="json") + assert dumped["kwargs"]["a"] == 1.5 + assert dumped["kwargs"]["df"] == {"x": {"0": 1}} diff --git a/xopt/types.py b/xopt/types.py new file mode 100644 index 000000000..a0f25a431 --- /dev/null +++ b/xopt/types.py @@ -0,0 +1,709 @@ +"""Typed, context-aware codecs used by Xopt's pydantic models. + +Binary prefixes identify the payload type, while compression is inferred from +magic bytes. Loading a serialized torch module is intentionally restricted to +module-typed fields and uses ``weights_only=False``; configuration files that +contain modules must therefore be treated as trusted input. +""" + +from __future__ import annotations + +import base64 +import gzip +import importlib +import io +import json +import os +import re +import warnings +import zlib +from contextlib import contextmanager +from types import MethodType +from contextvars import ContextVar +from dataclasses import dataclass, field +from pathlib import Path +from typing import Annotated, Any, Callable, Literal + +import numpy as np +import pandas as pd +import torch +import zstandard +from pydantic import GetCoreSchemaHandler +from pydantic_core import core_schema +from pydantic_core.core_schema import SerializationInfo, ValidationInfo + +MAX_DECOMPRESSED_SIZE = 256 * 1024 * 1024 +_GZIP_MAGIC = b"\x1f\x8b" +_ZSTD_MAGIC = b"\x28\xb5\x2f\xfd" +_UNSET = object() +_VALIDATION_BASE_DIR: ContextVar[str | os.PathLike[str] | None] = ContextVar( + "xopt_validation_base_dir", default=None +) + + +@dataclass(frozen=True, init=False) +class SerializationOptions: + """Normalized writer options carried through pydantic serialization context.""" + + array_mode: Literal["list", "b64"] = "list" + module_mode: Literal["drop", "file", "inline"] = "drop" + df_mode: Literal["dict", "b64"] = "dict" + compress: Literal["gzip", "zstd"] | None = None + level: int | None = None + file_dir: Path = field(default_factory=Path.cwd) + _explicit: frozenset[str] = field(default_factory=frozenset, repr=False) + + def __init__( + self, + *, + array_mode: Literal["list", "b64"] | object = _UNSET, + module_mode: Literal["drop", "file", "inline"] | object = _UNSET, + df_mode: Literal["dict", "b64"] | object = _UNSET, + compress: Literal["gzip", "zstd"] | None | object = _UNSET, + level: int | None | object = _UNSET, + file_dir: str | os.PathLike[str] | object = _UNSET, + ) -> None: + values = { + "array_mode": "list" if array_mode is _UNSET else array_mode, + "module_mode": "drop" if module_mode is _UNSET else module_mode, + "df_mode": "dict" if df_mode is _UNSET else df_mode, + "compress": None if compress is _UNSET else compress, + "level": None if level is _UNSET else level, + "file_dir": Path.cwd() if file_dir is _UNSET else Path(file_dir), + } + explicit = frozenset( + name + for name, value in { + "array_mode": array_mode, + "module_mode": module_mode, + "df_mode": df_mode, + "compress": compress, + "level": level, + "file_dir": file_dir, + }.items() + if value is not _UNSET + ) + + if values["array_mode"] not in ("list", "b64"): + raise ValueError("array_mode must be 'list' or 'b64'") + if values["module_mode"] not in ("drop", "file", "inline"): + raise ValueError("module_mode must be 'drop', 'file', or 'inline'") + if values["df_mode"] not in ("dict", "b64"): + raise ValueError("df_mode must be 'dict' or 'b64'") + if values["compress"] not in (None, "gzip", "zstd"): + raise ValueError("compress must be None, 'gzip', or 'zstd'") + + for name, value in values.items(): + object.__setattr__(self, name, value) + object.__setattr__(self, "_explicit", explicit) + + def resolve(self, name: str, annotation_default: Any = _UNSET) -> Any: + """Resolve context > annotation > library-default precedence.""" + if name in self._explicit or annotation_default is _UNSET: + return getattr(self, name) + return annotation_default + + +_OPTION_NAMES = { + "array_mode", + "module_mode", + "df_mode", + "compress", + "level", + "file_dir", +} + + +def normalize_serialization_context(context: Any = None) -> dict[str, Any]: + """Return a context dict containing one normalized options object.""" + if isinstance(context, SerializationOptions): + return {"serialization_options": context} + if context is None: + return {"serialization_options": SerializationOptions()} + if not isinstance(context, dict): + raise TypeError( + "serialization context must be a mapping or SerializationOptions" + ) + + result = dict(context) + existing = result.get("serialization_options", result.get("options")) + raw = {key: result[key] for key in _OPTION_NAMES if key in result} + if existing is not None and not isinstance(existing, SerializationOptions): + if not isinstance(existing, dict): + raise TypeError( + "serialization_options must be a mapping or SerializationOptions" + ) + raw = {**existing, **raw} + existing = None + + if existing is None: + options = SerializationOptions(**raw) + elif raw: + base = {name: getattr(existing, name) for name in existing._explicit} + options = SerializationOptions(**{**base, **raw}) + else: + options = existing + result["serialization_options"] = options + result.pop("options", None) + return result + + +def get_serialization_options(context: Any = None) -> SerializationOptions: + return normalize_serialization_context(context)["serialization_options"] + + +@contextmanager +def module_load_base_dir(base_dir: str | os.PathLike[str]): + """Preserve file-relative loading through models with custom ``__init__`` methods.""" + token = _VALIDATION_BASE_DIR.set(base_dir) + try: + yield + finally: + _VALIDATION_BASE_DIR.reset(token) + + +def _compress(raw: bytes, algorithm: str | None, level: int | None) -> bytes: + if algorithm is None: + return raw + if algorithm == "gzip": + return gzip.compress(raw, compresslevel=9 if level is None else level) + if algorithm == "zstd": + kwargs = {} if level is None else {"level": level} + return zstandard.ZstdCompressor(**kwargs).compress(raw) + raise ValueError(f"unknown compression algorithm: {algorithm}") + + +def maybe_decompress(raw: bytes, *, max_size: int = MAX_DECOMPRESSED_SIZE) -> bytes: + """Decompress gzip/zstd data inferred by magic bytes, enforcing a size cap.""" + if raw.startswith(_GZIP_MAGIC): + decompressor = zlib.decompressobj(wbits=31) + try: + result = decompressor.decompress(raw, max_size + 1) + if len(result) > max_size or decompressor.unconsumed_tail: + raise ValueError("decompressed payload exceeds the size limit") + result += decompressor.flush(max_size + 1 - len(result)) + except zlib.error as exc: + raise ValueError("invalid gzip-compressed payload") from exc + if len(result) > max_size: + raise ValueError("decompressed payload exceeds the size limit") + if not decompressor.eof or decompressor.unused_data: + raise ValueError("invalid gzip-compressed payload") + return result + + if raw.startswith(_ZSTD_MAGIC): + try: + result = zstandard.ZstdDecompressor().decompress( + raw, max_output_size=max_size + 1 + ) + except zstandard.ZstdError as exc: + raise ValueError("invalid zstd-compressed payload") from exc + if len(result) > max_size: + raise ValueError("decompressed payload exceeds the size limit") + return result + + # the cap guards decompression amplification only + return raw + + +def _b64encode(prefix: str, raw: bytes) -> str: + return prefix + base64.b64encode(raw).decode("ascii") + + +def _b64decode(value: str, prefix: str) -> bytes: + if not value.startswith(prefix): + raise ValueError(f"expected a {prefix} payload") + try: + return base64.b64decode(value[len(prefix) :], validate=True) + except (ValueError, base64.binascii.Error) as exc: + raise ValueError(f"malformed base64 data in {prefix} payload") from exc + + +def sanitize_sidecar_name(name: str) -> str: + return re.sub(r"[^A-Za-z0-9_.-]", "_", name) + + +def _is_nonfinite_token(value: Any) -> bool: + return isinstance(value, str) and value.lower().lstrip("+-") in ( + "nan", + "inf", + "infinity", + ) + + +def _coerce_nonfinite_tokens(value: Any) -> Any: + """Convert "NaN"/"Infinity"-style strings (ser_json_inf_nan output, and the + legacy walker's lowercase forms) back to floats inside nested lists.""" + if isinstance(value, list): + return [_coerce_nonfinite_tokens(item) for item in value] + if _is_nonfinite_token(value): + return float(value) + return value + + +def sidecar_path(options: SerializationOptions, name: str) -> tuple[Path, str]: + filename = sanitize_sidecar_name(name) + directory = Path(options.file_dir) + directory.mkdir(parents=True, exist_ok=True) + return directory / filename, filename + + +# sidecar files this process has written; replacing one of these is a routine +# checkpoint overwrite, replacing anything else clobbers a foreign file +_WRITTEN_SIDECARS: set[Path] = set() + + +def save_module_sidecar( + module: torch.nn.Module, options: SerializationOptions, name: str +) -> str: + """Atomically write a module sidecar file; returns the serialized filename.""" + # TODO: consider prefixing sidecar names with the dump-file stem so dumps + # sharing a directory (or keys that sanitize identically) cannot collide + path, filename = sidecar_path(options, name) + key = path.resolve() + if path.exists() and key not in _WRITTEN_SIDECARS: + warnings.warn( + f"overwriting existing sidecar file {path}", UserWarning, stacklevel=2 + ) + _WRITTEN_SIDECARS.add(key) + temp_path = path.with_name(path.name + ".tmp") + torch.save(module, temp_path) + os.replace(temp_path, path) + return filename + + +class NDArrayCodec: + def __init__( + self, + *, + array_mode: Literal["list", "b64"] | object = _UNSET, + compress: Literal["gzip", "zstd"] | None | object = _UNSET, + level: int | None | object = _UNSET, + ) -> None: + self.array_mode = array_mode + self.compress = compress + self.level = level + + @staticmethod + def validate(value: Any, _: ValidationInfo) -> np.ndarray: + if isinstance(value, np.ndarray): + return value + if isinstance(value, list): + return np.asarray(_coerce_nonfinite_tokens(value)) + if isinstance(value, (int, float)) and not isinstance(value, bool): + # 0-d arrays serialize as bare scalars via tolist() + return np.asarray(float(value)) + if isinstance(value, str) and value.startswith("b64np:"): + raw = maybe_decompress(_b64decode(value, "b64np:")) + try: + result = np.load(io.BytesIO(raw), allow_pickle=False) + except Exception as exc: + raise ValueError("invalid b64np: NumPy payload") from exc + if not isinstance(result, np.ndarray): + raise ValueError("b64np: payload did not contain an ndarray") + return result + if _is_nonfinite_token(value): + # non-finite scalars are dumped as strings (ser_json_inf_nan); + # legacy dumps also encoded 0-d nan/inf arrays this way + return np.asarray(float(value)) + raise ValueError("expected a NumPy array, list, scalar, or b64np: string") + + def serialize(self, value: np.ndarray, info: SerializationInfo) -> Any: + options = get_serialization_options(info.context) + mode = options.resolve("array_mode", self.array_mode) + if mode == "list": + return value.tolist() + buffer = io.BytesIO() + np.save(buffer, value, allow_pickle=False) + compress = options.resolve("compress", self.compress) + level = options.resolve("level", self.level) + return _b64encode("b64np:", _compress(buffer.getvalue(), compress, level)) + + def __get_pydantic_core_schema__( + self, source_type: Any, handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + return core_schema.with_info_plain_validator_function( + self.validate, + json_schema_input_schema=core_schema.union_schema( + [core_schema.list_schema(), core_schema.str_schema()] + ), + serialization=core_schema.plain_serializer_function_ser_schema( + self.serialize, info_arg=True, when_used="json" + ), + ) + + +class TorchTensorCodec: + def __init__( + self, + *, + array_mode: Literal["list", "b64"] | object = _UNSET, + compress: Literal["gzip", "zstd"] | None | object = _UNSET, + level: int | None | object = _UNSET, + ) -> None: + self.array_mode = array_mode + self.compress = compress + self.level = level + + @staticmethod + def validate(value: Any, _: ValidationInfo) -> torch.Tensor: + if isinstance(value, torch.Tensor): + return value + if isinstance(value, list): + return torch.tensor(_coerce_nonfinite_tokens(value)) + if isinstance(value, (int, float)): + # 0-d tensors serialize as bare scalars via tolist() + return torch.tensor(value) + if _is_nonfinite_token(value): + return torch.tensor(float(value)) + if isinstance(value, str) and value.startswith("b64pt:"): + raw = maybe_decompress(_b64decode(value, "b64pt:")) + try: + result = torch.load(io.BytesIO(raw), weights_only=True) + except Exception as exc: + raise ValueError("invalid b64pt: tensor payload") from exc + if not isinstance(result, torch.Tensor): + raise ValueError("b64pt: tensor payload did not contain a tensor") + return result + raise ValueError("expected a torch tensor, list, or b64pt: string") + + def serialize(self, value: torch.Tensor, info: SerializationInfo) -> Any: + options = get_serialization_options(info.context) + mode = options.resolve("array_mode", self.array_mode) + if mode == "list": + return value.detach().cpu().tolist() + buffer = io.BytesIO() + torch.save(value.detach().cpu(), buffer) + compress = options.resolve("compress", self.compress) + level = options.resolve("level", self.level) + return _b64encode("b64pt:", _compress(buffer.getvalue(), compress, level)) + + def __get_pydantic_core_schema__( + self, source_type: Any, handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + return core_schema.with_info_plain_validator_function( + self.validate, + json_schema_input_schema=core_schema.union_schema( + [core_schema.list_schema(), core_schema.str_schema()] + ), + serialization=core_schema.plain_serializer_function_ser_schema( + self.serialize, info_arg=True, when_used="json" + ), + ) + + +def encode_torch_module( + module: torch.nn.Module, + *, + compress: Literal["gzip", "zstd"] | None = None, + level: int | None = None, +) -> str: + """Encode a trusted torch module as a typed inline payload.""" + buffer = io.BytesIO() + torch.save(module, buffer, pickle_protocol=5) + return _b64encode("b64pt:", _compress(buffer.getvalue(), compress, level)) + + +def decode_torch_module(value: str, *, base_dir: str | os.PathLike[str] | None = None): + """Decode a trusted inline/path torch module (uses ``weights_only=False``).""" + if value.startswith("base64:"): + # legacy inline prefix written by pre-3.3 releases; the payload is the + # same base64 of torch.save as b64pt: + value = "b64pt:" + value[len("base64:") :] + if value.startswith("b64pt:"): + raw = maybe_decompress(_b64decode(value, "b64pt:")) + try: + return torch.load(io.BytesIO(raw), weights_only=False) + except Exception as exc: + raise ValueError("invalid b64pt: torch module payload") from exc + + if base_dir is None: + base_dir = _VALIDATION_BASE_DIR.get() + candidates = [Path(value)] + if base_dir is not None and not Path(value).is_absolute(): + candidates.append(Path(base_dir) / value) + for candidate in candidates: + if candidate.exists(): + return torch.load(candidate, weights_only=False) + raise ValueError(f"cannot load torch module from {value}") + + +class TorchModuleCodec: + def __init__( + self, + *, + compress: Literal["gzip", "zstd"] | None | object = _UNSET, + level: int | None | object = _UNSET, + ) -> None: + self.compress = compress + self.level = level + + def __get_pydantic_core_schema__( + self, source_type: Any, handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + def validate(value: Any, info: ValidationInfo) -> Any: + if isinstance(value, str): + context = info.context if isinstance(info.context, dict) else {} + value = decode_torch_module(value, base_dir=context.get("base_dir")) + if not isinstance(value, source_type): + raise ValueError(f"expected a {source_type.__qualname__} torch module") + return value + + def serialize(value: torch.nn.Module, info: SerializationInfo) -> Any: + options = get_serialization_options(info.context) + if options.module_mode != "inline": + # Owning-model wrap serializers remove or replace this placeholder. + return None + compress = options.resolve("compress", self.compress) + level = options.resolve("level", self.level) + return encode_torch_module(value, compress=compress, level=level) + + return core_schema.with_info_plain_validator_function( + validate, + json_schema_input_schema=core_schema.str_schema(), + serialization=core_schema.plain_serializer_function_ser_schema( + serialize, info_arg=True, when_used="json" + ), + ) + + +class TorchDTypeCodec: + @staticmethod + def validate(value: Any, _: ValidationInfo) -> torch.dtype: + if isinstance(value, torch.dtype): + return value + if isinstance(value, str) and value.startswith("torch."): + result = getattr(torch, value.removeprefix("torch."), None) + if isinstance(result, torch.dtype): + return result + raise ValueError(f"invalid torch dtype: {value!r}") + + @staticmethod + def serialize(value: torch.dtype, _: SerializationInfo) -> str: + return str(value) + + def __get_pydantic_core_schema__( + self, source_type: Any, handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + return core_schema.with_info_plain_validator_function( + self.validate, + json_schema_input_schema=core_schema.str_schema(), + serialization=core_schema.plain_serializer_function_ser_schema( + self.serialize, info_arg=True, when_used="json" + ), + ) + + +class DataFrameCodec: + def __init__( + self, + *, + df_mode: Literal["dict", "b64"] | object = _UNSET, + compress: Literal["gzip", "zstd"] | None | object = _UNSET, + level: int | None | object = _UNSET, + ) -> None: + self.df_mode = df_mode + self.compress = compress + self.level = level + + @staticmethod + def _restore_index(df: pd.DataFrame) -> pd.DataFrame: + # JSON object keys are always strings; recover the integer index a + # round trip started with so reloads match the original frame + # (pandas >= 3 uses the dedicated "str" dtype, older versions "object") + if pd.api.types.is_string_dtype(df.index) or df.index.dtype == object: + try: + df.index = df.index.astype(np.int64) + except (TypeError, ValueError): + pass + return df + + @staticmethod + def validate(value: Any, _: ValidationInfo) -> pd.DataFrame: + if isinstance(value, pd.DataFrame): + return value + if isinstance(value, dict): + return DataFrameCodec._restore_index(pd.DataFrame(value)) + if isinstance(value, str) and value.startswith("b64df:"): + raw = maybe_decompress(_b64decode(value, "b64df:")) + try: + # conversion heuristics disabled so decoding matches the plain + # dict mode (no date sniffing) + result = pd.read_json( + io.StringIO(raw.decode("utf-8")), + orient="columns", + convert_axes=False, + convert_dates=False, + ) + except Exception as exc: + raise ValueError("invalid b64df: DataFrame payload") from exc + return DataFrameCodec._restore_index(result) + raise ValueError("expected a DataFrame, dict, or b64df: string") + + def serialize(self, value: pd.DataFrame, info: SerializationInfo) -> Any: + options = get_serialization_options(info.context) + mode = options.resolve("df_mode", self.df_mode) + text = value.to_json() + if mode == "dict": + return json.loads(text) + compress = options.resolve("compress", self.compress) + level = options.resolve("level", self.level) + return _b64encode("b64df:", _compress(text.encode(), compress, level)) + + def __get_pydantic_core_schema__( + self, source_type: Any, handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + return core_schema.with_info_plain_validator_function( + self.validate, + json_schema_input_schema=core_schema.union_schema( + [core_schema.dict_schema(), core_schema.str_schema()] + ), + serialization=core_schema.plain_serializer_function_ser_schema( + self.serialize, info_arg=True, when_used="json" + ), + ) + + +def qualified_name(value: Callable[..., Any] | type) -> str: + qualname = getattr(value, "__qualname__", None) or getattr(value, "__name__", None) + module = getattr(value, "__module__", None) + if qualname is None or module is None: + # callable instances (functools.partial, nn.Module, ...) have no + # qualname of their own + cls = type(value) + return f"{cls.__module__}.{cls.__qualname__}" + return f"{module}.{qualname}" + + +def object_from_qualified_name(value: str) -> Any: + parts = value.split(".") + for split_at in range(len(parts) - 1, 0, -1): + module_name = ".".join(parts[:split_at]) + try: + result = importlib.import_module(module_name) + except ModuleNotFoundError as error: + # only fall back to a shorter prefix when the prefix itself is not + # a module; a missing transitive dependency must propagate, or an + # unrelated attribute could shadow the intended submodule + missing = error.name or "" + if missing == module_name or module_name.startswith(missing + "."): + continue + raise + for attribute in parts[split_at:]: + try: + result = getattr(result, attribute) + except AttributeError as error: + raise ValueError(f"cannot import object from {value!r}") from error + return result + raise ValueError(f"cannot import object from {value!r}") + + +def resolve_callable(value: Any) -> Callable[..., Any]: + """Return ``value`` if it is already callable, otherwise import it from its + fully qualified name (e.g. ``"math.sqrt"`` or ``"module.Class.method"``).""" + if callable(value): + return value + if not isinstance(value, str): + raise ValueError(f"{value!r} must be a callable or a qualified-name string") + result = object_from_qualified_name(value) + if not callable(result): + raise ValueError(f"{value!r} does not name a callable") + return result + + +def _serialized_name_round_trips(value: Any, name: str) -> bool: + """True if resolving ``name`` recovers ``value`` (lossless serialization).""" + try: + resolved = object_from_qualified_name(name) + except (ValueError, ModuleNotFoundError): + return False + if resolved is value: + return True + # bound-method objects are created anew on each attribute access, so + # identity fails even for a lossless round trip; compare the parts + if isinstance(value, MethodType) and isinstance(resolved, MethodType): + return ( + resolved.__func__ is value.__func__ and resolved.__self__ is value.__self__ + ) + return False + + +_WARNED_CALLABLE_NAMES: set[str] = set() + + +def _warn_callable_once(name: str, message: str) -> None: + # periodic dump_file checkpointing serializes after every batch; warn once + # per callable per process instead of once per dump + if name not in _WARNED_CALLABLE_NAMES: + _WARNED_CALLABLE_NAMES.add(name) + warnings.warn(message, UserWarning) + + +class CallableCodec: + @staticmethod + def validate(value: Any, _: ValidationInfo) -> Callable[..., Any]: + return resolve_callable(value) + + @staticmethod + def serialize(value: Callable[..., Any], _: SerializationInfo) -> str: + name = qualified_name(value) + # TODO: change these warnings to raise ValueError in a future release + if not _serialized_name_round_trips(value, name): + _warn_callable_once( + name, + f"serialized callable {name!r} will not reload as the same object " + "(lambda, functools.partial, bound method, or closure); pass a " + "module-level callable instead", + ) + elif name.split(".", 1)[0] == "__main__": + _warn_callable_once( + name, + f"serialized callable {name!r} is defined in __main__ and will " + "only reload from a process that defines it (e.g. re-running " + "the same script); define it in an importable module for " + "portable dumps", + ) + return name + + def __get_pydantic_core_schema__( + self, source_type: Any, handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + return core_schema.with_info_plain_validator_function( + self.validate, + json_schema_input_schema=core_schema.str_schema(), + serialization=core_schema.plain_serializer_function_ser_schema( + self.serialize, info_arg=True, when_used="json" + ), + ) + + +class TypeCodec: + @staticmethod + def validate(value: Any, _: ValidationInfo) -> type: + if isinstance(value, str): + value = object_from_qualified_name(value) + if not isinstance(value, type): + raise ValueError("expected a type or qualified-name string") + return value + + @staticmethod + def serialize(value: type, _: SerializationInfo) -> str: + return qualified_name(value) + + def __get_pydantic_core_schema__( + self, source_type: Any, handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + return core_schema.with_info_plain_validator_function( + self.validate, + json_schema_input_schema=core_schema.str_schema(), + serialization=core_schema.plain_serializer_function_ser_schema( + self.serialize, info_arg=True, when_used="json" + ), + ) + + +NDArray = Annotated[np.ndarray, NDArrayCodec()] +TorchTensor = Annotated[torch.Tensor, TorchTensorCodec()] +TorchDType = Annotated[torch.dtype, TorchDTypeCodec()] +XDataFrame = Annotated[pd.DataFrame, DataFrameCodec()] +CallableRef = Annotated[Callable[..., Any], CallableCodec()] +TypeRef = Annotated[type, TypeCodec()] diff --git a/xopt/utils.py b/xopt/utils.py index bb6a4b0fa..4c9602aff 100644 --- a/xopt/utils.py +++ b/xopt/utils.py @@ -2,7 +2,6 @@ from pydantic import BaseModel from typing import List, Tuple import datetime -import importlib import inspect import logging import numpy as np @@ -12,12 +11,14 @@ import time import torch import traceback +import warnings import yaml from gest_api.vocs import VOCS, GreaterThanConstraint from .generator import Generator from .pydantic import get_descriptions_defaults +from .types import resolve_callable # Grab the logger logger = logging.getLogger(__name__) @@ -66,31 +67,16 @@ def get_generator_name(generator): def get_function(name): """ - Returns a function from a fully qualified name or global name. - """ - - # Check if already a function - if callable(name): - return name - - if not isinstance(name, str): - raise ValueError(f"{name} must be callable or a string.") + Deprecated alias of :func:`xopt.types.resolve_callable`. - if name in globals(): - if callable(globals()[name]): - f = globals()[name] - else: - raise ValueError(f"global {name} is not callable") - else: - if "." in name: - # try to import - m_name, f_name = name.rsplit(".", 1) - module = importlib.import_module(m_name) - f = getattr(module, f_name) - else: - raise Exception(f"function {name} does not exist") - - return f + Returns a callable as-is, or imports it from a fully qualified name. + """ + warnings.warn( + "get_function is deprecated, use xopt.types.resolve_callable instead", + DeprecationWarning, + stacklevel=2, + ) + return resolve_callable(name) def get_function_defaults(f): From f47f8ba51b1bb358e3b0879dda0c923beb59b4fb Mon Sep 17 00:00:00 2001 From: nikitakuklev Date: Fri, 21 Aug 2026 05:18:52 -0500 Subject: [PATCH 2/3] remove stale x0 key from rcds example notebook --- docs/examples/sequential/rcds.ipynb | 2 -- 1 file changed, 2 deletions(-) diff --git a/docs/examples/sequential/rcds.ipynb b/docs/examples/sequential/rcds.ipynb index 1c2890b29..1975964cd 100644 --- a/docs/examples/sequential/rcds.ipynb +++ b/docs/examples/sequential/rcds.ipynb @@ -129,7 +129,6 @@ " max_evaluations: 100\n", "generator:\n", " name: rcds\n", - " x0: null\n", " init_mat: null\n", " noise: 0.00001\n", " step: 0.01\n", @@ -199,7 +198,6 @@ " max_evaluations: 400\n", "generator:\n", " name: rcds\n", - " x0: null\n", " init_mat: null\n", " noise: 1e-8\n", " step: 0.01\n", From c9fc86683a33195b1391c88c6f3b7b09495a474b Mon Sep 17 00:00:00 2001 From: nikitakuklev Date: Fri, 21 Aug 2026 05:38:19 -0500 Subject: [PATCH 3/3] add serialization options example notebook --- docs/examples/basic/xopt_serialization.ipynb | 159 +++++++++++++++++++ mkdocs.yml | 1 + 2 files changed, 160 insertions(+) create mode 100644 docs/examples/basic/xopt_serialization.ipynb diff --git a/docs/examples/basic/xopt_serialization.ipynb b/docs/examples/basic/xopt_serialization.ipynb new file mode 100644 index 000000000..6095c1f51 --- /dev/null +++ b/docs/examples/basic/xopt_serialization.ipynb @@ -0,0 +1,159 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "aab681aa", + "metadata": {}, + "source": [ + "# Serialization Options\n", + "Xopt objects serialize to YAML/JSON for checkpointing and restarts. What goes\n", + "into a dump is controlled by a few writer options, passed to `yaml()`, `json()`\n", + "or `dump()`:\n", + "\n", + "- `module_mode`: `\"drop\"` (default), `\"file\"` (write torch modules as `.pt`\n", + " sidecar files next to the dump file) or `\"inline\"` (embed as base64 strings)\n", + "- `array_mode`: `\"list\"` (default) or `\"b64\"` (binary numpy/torch payloads)\n", + "- `df_mode`: `\"dict\"` (default) or `\"b64\"` for dataframes\n", + "- `compress`: `None` (default), `\"gzip\"` or `\"zstd\"`, with an optional `level`\n", + "\n", + "This example compares dump sizes for a sizeable Xopt instance." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "797033b0", + "metadata": {}, + "outputs": [], + "source": [ + "import math\n", + "\n", + "from xopt import Xopt, Evaluator, VOCS\n", + "from xopt.generators.bayesian import UpperConfidenceBoundGenerator\n", + "from xopt.resources.test_functions.sinusoid_1d import evaluate_sinusoid\n", + "\n", + "vocs = VOCS(variables={\"x1\": [0, 1.75 * math.pi]}, objectives={\"y1\": \"MINIMIZE\"})\n", + "X = Xopt(\n", + " generator=UpperConfidenceBoundGenerator(vocs=vocs),\n", + " evaluator=Evaluator(function=evaluate_sinusoid),\n", + ")\n", + "\n", + "# evaluate 500 random points and train the GP model\n", + "X.random_evaluate(500)\n", + "X.generator.train_model()\n", + "X.generator.model" + ] + }, + { + "cell_type": "markdown", + "id": "8292f847", + "metadata": {}, + "source": [ + "## Default dump\n", + "By default torch modules (like the trained GP model above) are dropped and\n", + "arrays/dataframes are written as plain lists and dicts." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6291e520", + "metadata": {}, + "outputs": [], + "source": [ + "print(f\"default: {len(X.yaml()) / 1024:.0f} KiB\")" + ] + }, + { + "cell_type": "markdown", + "id": "56e2f2c2", + "metadata": {}, + "source": [ + "## Inline torch modules\n", + "With `module_mode=\"inline\"` the trained model is embedded in the dump as a\n", + "base64 payload, so the file is fully self-contained. Compression shrinks the\n", + "binary payloads considerably; combining it with `b64` array and dataframe\n", + "modes compresses the evaluation data as well." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "944cd42d", + "metadata": {}, + "outputs": [], + "source": [ + "import time\n", + "\n", + "b64 = dict(module_mode=\"inline\", array_mode=\"b64\", df_mode=\"b64\")\n", + "variants = {\n", + " \"inline raw (lists)\": dict(module_mode=\"inline\"),\n", + " \"inline b64\": b64,\n", + " \"inline b64 + gzip 9\": dict(**b64, compress=\"gzip\", level=9),\n", + " \"inline b64 + zstd 3\": dict(**b64, compress=\"zstd\", level=3),\n", + " \"inline b64 + zstd 22 (max)\": dict(**b64, compress=\"zstd\", level=22),\n", + "}\n", + "for name, kwargs in variants.items():\n", + " start = time.perf_counter()\n", + " size = len(X.yaml(**kwargs))\n", + " elapsed = time.perf_counter() - start\n", + " print(f\"{name:28s} {size / 1024:4.0f} KiB {elapsed * 1e3:6.1f} ms\")" + ] + }, + { + "cell_type": "markdown", + "id": "a89d5733", + "metadata": {}, + "source": [ + "When `level` is not given, gzip defaults to level 9 and zstd to the\n", + "zstandard default (level 3). Higher zstd levels trade dump time for\n", + "size." + ] + }, + { + "cell_type": "markdown", + "id": "2626cce4", + "metadata": {}, + "source": [ + "## Reloading\n", + "Every variant reloads with `from_yaml`/`from_file`, including the trained\n", + "model." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "419226f4", + "metadata": {}, + "outputs": [], + "source": [ + "X2 = Xopt.from_yaml(X.yaml(module_mode=\"inline\", compress=\"zstd\"))\n", + "X2.generator.model" + ] + }, + { + "cell_type": "markdown", + "id": "620b96b0", + "metadata": {}, + "source": [ + "For long optimization runs prefer `module_mode=\"file\"` (the default used by\n", + "`X.dump()` when `serialize_torch=True`): modules are written as `.pt` files\n", + "next to the dump file, and `Xopt.from_file` finds them from any working\n", + "directory." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/mkdocs.yml b/mkdocs.yml index 91a7af1fe..2cdbd72d5 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -16,6 +16,7 @@ nav: - Working with generators: examples/basic/xopt_generator.ipynb - Stopping conditions: examples/basic/xopt_stopping_condition.ipynb - Checkpointing and restarts: examples/basic/checkpointing_and_restarts.ipynb + - Serialization options: examples/basic/xopt_serialization.ipynb - Bayesian: - Gaussian Process Model Creation: - Basic example: examples/gp_model_creation/model_creation.ipynb