continuation.py

Go to the documentation of this file.

"""Portable search-space identity and read-only continuation preflight."""
 
from __future__ import annotations
 
from .error import HpoError
 
from dataclasses import dataclass, field, replace
import hashlib
import json
import math
import re
from pathlib import Path
import struct
import subprocess
from typing import Any
 
from .study_spec import (
    StudySpec,
    _finite_parameter_cardinality,
    _fused_grid_value,
    _reject_duplicate_keys,
    _reject_json_constant,
    _parse_objective,
    _parse_search_space,
    load_study_spec,
)
 
SPACE_HASH_VERSION = 1
WARM_START_EXIT = 4
SPACE_EXHAUSTED_EXIT = 5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
class WarmStartError(HpoError, ValueError):
    """A parent cannot safely seed the requested study (exit 4)."""
 
    code = "hpo_warm_start_rejected"
 
 
class SpaceExhaustedError(HpoError, ValueError):
    """Every finite parameter vector has already been tried (exit 5)."""
 
    code = "hpo_space_exhausted"
 
 
@dataclass(frozen=True)
class WarmHistory:
    """Validated parent history; no trial execution or billing occurs here."""
 
    source_sha256: str
    trials: tuple[dict[str, Any], ...]
    tried: frozenset[str]
    completed: int
    feasible: int
    next_id: int
 
 
def recorded_space(study: StudySpec) -> dict[str, Any]:
    """Return the version-independent declarative space recorded in results."""
    parameters = {}
    for name, parameter in study.search_space.items():
        value: dict[str, Any] = {"kind": parameter.kind}
        if parameter.kind in {"integer", "real"}:
            value.update(
                low=parameter.low,
                high=parameter.high,
                step=parameter.step,
                log=parameter.log,
            )
        elif parameter.kind == "categorical":
            value["choices"] = list(parameter.choices)
        parameters[name] = value
    if study.objective.kind != "expression":
        raise ValueError("space identity requires an expression objective")
    return {
        "parameters": parameters,
        "objective": {
            "expression": study.objective.expression,
            "direction": study.objective.direction,
            "constraints": list(study.objective.constraints),
        },
    }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
def canonical_space(space: dict[str, Any]) -> str:
    """Return canonical v1 JSON; numeric identities are independent of formatting."""
    parameters = {}
    for name, source in space["parameters"].items():
        kind = source["kind"]
        parameter = {"kind": kind}
        if kind in {"integer", "real"}:
            for key in ("low", "high", "step"):
                value = source[key]
                if value is None:
                    parameter[key] = None
                elif kind == "integer":
                    if type(value) is not int:
                        raise ValueError("expected integer bound/step")
                    parameter[key] = str(value)
                else:
                    if type(value) not in {int, float}:
                        raise ValueError("expected real bound/step")
                    parameter[key] = _real_bits(value)
            if type(source["log"]) is not bool:
                raise ValueError("log must be boolean")
            parameter["log"] = source["log"]
        elif kind == "categorical":
            parameter["choices"] = [_scalar(choice) for choice in source["choices"]]
        elif kind != "boolean":
            raise ValueError("unknown parameter kind")
        parameters[name] = parameter
    objective = space["objective"]
    return _dump(
        {
            "space_hash_version": SPACE_HASH_VERSION,
            "parameters": parameters,
            "objective": {
                "expression": objective["expression"],
                "direction": objective["direction"],
                "constraints": sorted(objective["constraints"]),
            },
        }
    )
 
 
def space_hash(study: StudySpec) -> str:
    """SHA-256 of canonical v1 JSON, independent of package versions and pins."""
    return hashlib.sha256(canonical_space(recorded_space(study)).encode()).hexdigest()
 
 
def cardinality(study: StudySpec) -> int | None:
    """Exact uint64 cardinality, or None for continuous/overflowing spaces."""
    total = 1
    for name, parameter in study.search_space.items():
        issues = []
        count = _finite_parameter_cardinality(parameter, name, issues)
        if count is None or issues or total > _UINT64_MAX // count:
            return None
        total *= count
    return total
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
@dataclass(frozen=True)
class ContinuationDataContext:
    """Internal compatibility contract made from an already admitted descriptor."""
 
    canonical_descriptor: str
    data_context_sha256: str = field(init=False)
    dataset_ids: tuple[str, ...] = field(init=False)
 
 
    def __post_init__(self) -> None:
        try:
            descriptor = _loads(self.canonical_descriptor)
            if (
                descriptor["schema"] != "pineforge-hpo-data-context/v1"
                or not isinstance(descriptor["datasets"], dict)
                or not descriptor["datasets"]
                or not isinstance(descriptor["bindings"], dict)
            ):
                raise ValueError("invalid admitted data context")
            encoded = _dump(descriptor)
            object.__setattr__(self, "canonical_descriptor", encoded)
            object.__setattr__(self, "_descriptor", descriptor)
            object.__setattr__(
                self, "dataset_ids", tuple(sorted(descriptor["datasets"]))
            )
            object.__setattr__(
                self,
                "data_context_sha256",
                hashlib.sha256(
                    b"pineforge:hpo:data-context:v1\0" + encoded.encode()
                ).hexdigest(),
            )
        except (ValueError, TypeError, KeyError) as error:
            raise WarmStartError(f"warm-start incompatible: {error}") from error
 
    def validate_owner(
        self, owner: Any, *, required: bool = False, complete: bool = False
    ) -> None:
        if not isinstance(owner, dict):
            raise ValueError("data context owner must be an object")
        if required or "data_context_sha256" in owner:
            if owner.get("data_context_sha256") != self.data_context_sha256data_context_sha256:
                raise ValueError("missing or differing data_context_sha256")
        if complete or "dataset_ids" in owner:
            identifiers = owner.get("dataset_ids")
            if (
                not isinstance(identifiers, list)
                or any(not isinstance(value, str) for value in identifiers)
                or tuple(sorted(identifiers)) != self.dataset_ids
            ):
                raise ValueError("missing or differing dataset_ids")
        runtime = owner.get("applied_runtime", {})
        if complete or "data_context" in runtime:
            if _dump(runtime.get("data_context")) != self.canonical_descriptor:
                raise ValueError("missing or differing data context descriptor")
        if (
            "data_context" in owner
            and _dump(owner["data_context"]) != self.canonical_descriptor
        ):
            raise ValueError("differing data context descriptor")
        if complete and owner.get("dataset_id") is not None:
            raise ValueError("bound result cannot declare one global dataset_id")
 
    def row_identity(self, parameters: Any) -> dict[str, str]:
        charts = [
            (title, binding)
            for title, binding in self._descriptor["bindings"].items()
            if binding["role"] == "chart"
        ]
        if len(charts) != 1:
            raise ValueError("data context requires one chart binding")
        title, binding = charts[0]
        symbol = parameters[title]
        if not isinstance(symbol, str) or symbol not in binding["choices"]:
            raise ValueError("trial symbol is outside the admitted binding")
        dataset_id = binding["datasets"][symbol]
        if self._descriptor["datasets"][dataset_id]["symbol"] != symbol:
            raise ValueError("trial binding contradicts dataset symbol")
        return {
            "symbol": symbol,
            "dataset_id": dataset_id,
            "data_context_sha256": self.data_context_sha256data_context_sha256,
        }
 
    def validate_row(self, row: Any, parameters: Any) -> None:
        self.validate_owner(row, required=True)
        if any(
            row.get(name) != value
            for name, value in self.row_identity(parameters).items()
        ):
            raise ValueError(
                "trial symbol or dataset_id contradicts its parameters and binding"
            )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
def load_warm_start(
    study: StudySpec,
    path: str | Path,
    *,
    data_context: ContinuationDataContext | None = None,
) -> WarmHistory:
    """Auto-detect fd-3 JSONL, a complete result object, or its trials array."""
    try:
        source = Path(path).expanduser().resolve()
        content = source.read_bytes()
        if len(content) > 256 * 1024 * 1024:
            raise ValueError("warm-start document exceeds 256 MiB")
        text = content.decode("utf-8")
        try:
            document = _loads(text)
        except json.JSONDecodeError:
            document = None
        parent_space = None
        if isinstance(document, dict) and "trials" in document:
            _validate_data_context(
                document, data_context, required=True, complete=data_context is not None
            )
            if "tpe_sampler_state" in document:
                checkpoint = document["tpe_sampler_state"]
                _validate_sampler_checkpoint(checkpoint)
            if document.get("trials_out", "all") != "all":
                raise ValueError("summary/none result is not a complete trial history")
            own_trials = document["trials"]
            inherited = document.get("warm_start_trials", [])
            if "trials_completed" in document and (
                type(document["trials_completed"]) is not int
                or document["trials_completed"] != len(own_trials)
            ):
                raise ValueError("result does not contain all completed trials")
            if "warm_start" in document and (
                type(document["warm_start"]["trials"]) is not int
                or document["warm_start"]["trials"] != len(inherited)
            ):
                raise ValueError("result is missing ancestor trials")
            trials = inherited + own_trials
            parent_space = document.get("space")
            if parent_space is None and "study_spec" in document:
                embedded = document["study_spec"]
                issues = []
                parameters = _parse_search_space(
                    embedded["strategies"][0]["search_space"],
                    "$.strategies[0].search_space",
                    issues,
                )
                objective = _parse_objective(embedded["objective"], issues)
                if issues:
                    raise ValueError("invalid legacy study_spec space or objective")
                parent_space = recorded_space(
                    replace(
                        study,
                        strategy=replace(study.strategy, search_space=parameters),
                        objective=objective,
                    )
                )
            if parent_space is None and "study" in document:
                parent_spec = Path(document["study"])
                if not parent_spec.is_absolute():
                    parent_spec = source.parent / parent_spec
                parent_space = recorded_space(
                    load_study_spec(parent_spec, continuation=True)
                )
        elif isinstance(document, list):
            trials = document
        else:
            trials = [_loads(line) for line in text.splitlines() if line.strip()]
        if not isinstance(trials, list) or not trials:
            raise ValueError("parent has no trials")
        expected = canonical_space(recorded_space(study))
        expected_hash = hashlib.sha256(expected.encode()).hexdigest()
 
        def verify_hash(owner):
            version = owner.get("space_hash_version")
            if version is not None and type(version) is not int:
                raise ValueError("space_hash_version must be an integer")
            if (
                version == SPACE_HASH_VERSION
                and "space_hash" in owner
                and owner["space_hash"] != expected_hash
            ):
                raise ValueError("recorded space_hash does not match space")
 
        if parent_space is not None and canonical_space(parent_space) != expected:
            raise ValueError("search space or objective differs (space_hash)")
        if parent_space is not None:
            verify_hash(document)
        tried = set()
        identifiers = set()
        completed = feasible = next_id = 0
        for trial in trials:
            space = trial.get("space", parent_space)
            if space is None:
                raise ValueError(
                    "missing recorded space; compatibility cannot be proven"
                )
            if canonical_space(space) != expected:
                raise ValueError("search space or objective differs (space_hash)")
            verify_hash(trial)
            status = trial["status"]
            if status not in _STATUSES:
                raise ValueError(f"unknown trial status: {status}")
            identifier = trial["trial_id"]
            if type(identifier) is not int or not 0 <= identifier < _UINT64_MAX:
                raise ValueError("expected uint64 trial_id with a continuation ID")
            if identifier in identifiers:
                raise ValueError("duplicate trial_id")
            identifiers.add(identifier)
            next_id = max(next_id, identifier + 1)
            values = _parameters(study, trial["parameters"])
            _validate_data_context(trial, data_context, required=True)
            if data_context is not None:
                data_context.validate_row(trial, values)
            tried.add(_dump({name: _scalar(value) for name, value in values.items()}))
            is_feasible = trial["feasible"]
            objective = trial["objective"]
            if type(is_feasible) is not bool or is_feasible != (status == "ok"):
                raise ValueError("inconsistent status/feasibility/objective")
            if objective is not None and (
                type(objective) not in {int, float} or not math.isfinite(objective)
            ):
                raise ValueError("objective must be finite or null")
            for score in trial.get("pruning", {}).get("rung_scores", []):
                if score is not None and (
                    type(score) not in {int, float} or not math.isfinite(score)
                ):
                    raise ValueError("invalid pruning rung score")
            completed += status in {"ok", "constraint_violation"}
            feasible += is_feasible
        return WarmHistory(
            hashlib.sha256(content).hexdigest(),
            tuple(sorted(trials, key=lambda trial: trial["trial_id"])),
            frozenset(tried),
            completed,
            feasible,
            next_id,
        )
    except (
        OSError,
        ValueError,
        TypeError,
        KeyError,
        AttributeError,
        OverflowError,
    ) as error:
        raise WarmStartError(f"warm-start incompatible: {error}") from error
 
 
def is_binary_warm(path: str | Path) -> bool:
    """Detect binary input without reading its payload or interpreting JSON."""
    with Path(path).expanduser().open("rb") as source:
        return source.read(1) == b"P"
 
 
def warm_start_metadata(
    study: StudySpec,
    path: str | Path,
    *,
    native: str | Path | None = None,
    defer_symbol_feeds: bool = False,
) -> dict[str, Any]:
    """Inspect sampler history; binary validation stays in the native mmap loader."""
    try:
        if is_binary_warm(path):
            from .cli import _resolve_native
 
            process = subprocess.run(
                [
                    str(_resolve_native(native)),
                    "space-info",
                    "--spec",
                    str(study.spec_path),
                    "--warm-start",
                    str(Path(path).expanduser().resolve()),
                    "--warm-details",
                    *(["--defer-symbol-feeds"] if defer_symbol_feeds else []),
                ],
                capture_output=True,
                text=True,
                check=False,
            )
            if process.returncode:
                message = process.stderr.strip().split("warm-start incompatible: ")[-1]
                raise WarmStartError(f"warm-start incompatible: {message}")
            return json.loads(process.stdout)
        history = load_warm_start(study, path)
        count = cardinality(study)
        return {
            "cardinality": count,
            "tried": len(history.tried),
            "remaining": None if count is None else count - len(history.tried),
            "space_hash": space_hash(study),
            "warm_trials": len(history.trials),
            "next_id": history.next_id,
            "completed": history.completed,
            "feasible": history.feasible,
            "source_sha256": history.source_sha256,
        }
    except OSError as error:
        raise WarmStartError(f"warm-start incompatible: {error}") from error
 
 
def space_info(
    study: StudySpec,
    warm_start: str | Path | None = None,
    *,
    native: str | Path | None = None,
) -> dict[str, Any]:
    """Inspect coverage without compiling a strategy, loading data, or drawing candidates."""
    count = cardinality(study)
    tried = (
        0
        if warm_start is None
        else warm_start_metadata(study, warm_start, native=native)["tried"]
    )
    return {
        "cardinality": count,
        "tried": tried,
        "remaining": None if count is None else count - tried,
        "space_hash": space_hash(study),
    }