"""Problem type definitions and registry for EngiBench integration.
This is the SINGLE FILE to update when adding a new problem type:
1. Add to ProblemId Literal type
2. Add the engibench import path to _PROBLEM_IMPORTS
Everything else (SUPPORTED_PROBLEMS, PROBLEM_CLASSES, the ``Beams2D`` /
``Photonics2D`` re-exports) is derived automatically.
The engibench problem classes are imported *lazily*, on first lookup. Importing
``engibench.problems.beams2d.v1`` eagerly costs ~1.7 s (it drags in seaborn ->
scipy.stats -> IPython/ipywidgets, and datasets -> pandas/pyarrow), and this
module sits on the import path of every agent via ``src.utils.prompts`` ->
``src.tools`` -- including agents that never touch engibench (arxiv, rag,
search, cli). Deferring the class import keeps ``SUPPORTED_PROBLEMS`` (derived
from the Literal, not from the classes) free.
"""
import logging
import os
from collections.abc import Iterator, Mapping
from importlib import import_module
from importlib.util import find_spec
from typing import TYPE_CHECKING, Literal, get_args
if TYPE_CHECKING:
from engibench.problems.beams2d.v1 import Beams2D # type: ignore[import-untyped]
from engibench.problems.photonics2d.v1 import (
Photonics2D, # type: ignore[import-untyped]
)
logger = logging.getLogger(__name__)
# engibench itself is checked here even though its problem classes are resolved
# lazily below. Every tool that resolves one wraps its body in a broad
# ``except`` that returns ``{"success": False}`` -- create_problem even carries
# an ``except ImportError`` branch that was unreachable while this module
# imported engibench eagerly. The benchmark scorer reads that same flag as the
# agent's score, so a degraded install would otherwise be recorded as agent
# failure rather than a broken environment, and create_problem gates every
# downstream step of the engineering workflow. find_spec does not execute the
# package, so this keeps the fail-loud property at no import cost.
if find_spec("engibench") is None: # pragma: no cover - install-time guard
_msg = (
"engibench is required by src.tools.problems but is not installed. It "
"is declared in pyproject.toml; the problem classes are imported "
"lazily, so without this check the failure would surface as a silent "
"tool error rather than a broken install."
)
raise ImportError(_msg)
# Type alias for problem IDs - this is the single source of truth for supported problems
ProblemId = Literal["beams2d", "photonics2d"]
# Where each problem class lives: {problem_id: (module, class name)}.
# The class name doubles as the lazily re-exported module attribute
# (``from src.tools.problems import Beams2D`` keeps working).
_PROBLEM_IMPORTS: dict[str, tuple[str, str]] = {
"beams2d": ("engibench.problems.beams2d.v1", "Beams2D"),
"photonics2d": ("engibench.problems.photonics2d.v1", "Photonics2D"),
}
class _LazyProblemClasses(Mapping[str, type]):
"""``{problem_id: problem class}`` that imports engibench on first lookup.
A drop-in replacement for the eager dict: membership tests, ``keys()``,
``list()`` and iteration only need the ids (no engibench import), while
``PROBLEM_CLASSES[problem_id]`` imports and returns the real class.
"""
def __init__(self) -> None:
self._resolved: dict[str, type] = {}
def __getitem__(self, key: str) -> type:
cached = self._resolved.get(key)
if cached is not None:
return cached
# KeyError for an unknown id is raised here, before import_module, so an
# invalid lookup never imports engibench.
module_path, class_name = _PROBLEM_IMPORTS[key]
resolved: type = getattr(import_module(module_path), class_name)
self._resolved[key] = resolved
return resolved
def __contains__(self, key: object) -> bool:
# Override Mapping's default (which calls __getitem__ and would import
# engibench just to answer a membership test).
return key in _PROBLEM_IMPORTS
def __iter__(self) -> Iterator[str]:
return iter(_PROBLEM_IMPORTS)
def __len__(self) -> int:
return len(_PROBLEM_IMPORTS)
def __repr__(self) -> str:
return f"_LazyProblemClasses({sorted(_PROBLEM_IMPORTS)})"
# Map problem IDs to their class implementations (lazy - see above)
PROBLEM_CLASSES: Mapping[str, type] = _LazyProblemClasses()
_CLASS_EXPORTS = {
class_name: module_path for module_path, class_name in _PROBLEM_IMPORTS.values()
}
def __getattr__(name: str) -> type:
"""Lazily resolve the re-exported problem classes (PEP 562)."""
module_path = _CLASS_EXPORTS.get(name)
if module_path is None:
msg = f"module {__name__!r} has no attribute {name!r}"
raise AttributeError(msg)
return getattr(import_module(module_path), name) # type: ignore[no-any-return]
# Supported problem types - automatically derived from ProblemId
SUPPORTED_PROBLEMS = list(get_args(ProblemId))
# Env var the eval harness sets per training-code run to bind the whole pipeline
# to the registered benchmark's problem (see ``harness_problem_id``).
PROBLEM_ID_ENV_VAR = "TRAINING_CODE_PROBLEM_ID"
[docs]
def harness_problem_id(problem_id: str) -> str:
"""Return the harness-injected ``problem_id`` if set, else ``problem_id``.
The evaluation harness sets ``TRAINING_CODE_PROBLEM_ID`` per training-code run
(from the registered benchmark, e.g. ``training_code_photonics2d`` ->
``photonics2d``) so the whole pipeline -- SLURM training, the local dry-run,
contract validation, and model evaluation -- targets that problem without the
prompt having to name it. Unset -> the tool/agent-supplied value (default
``beams2d``), so behaviour is unchanged when the harness does not set it.
The override is stripped and validated against :data:`SUPPORTED_PROBLEMS`; an
empty/whitespace or unsupported value (e.g. a typo in a manual export) is
ignored with a warning and the supplied ``problem_id`` is kept, rather than
silently binding the pipeline to a bogus problem.
"""
override = os.getenv(PROBLEM_ID_ENV_VAR)
override = override.strip() if override else ""
if not override:
return problem_id
if override not in SUPPORTED_PROBLEMS:
logger.warning(
"Ignoring %s=%r (not a supported problem: %s); using %r.",
PROBLEM_ID_ENV_VAR,
override,
", ".join(SUPPORTED_PROBLEMS),
problem_id,
)
return problem_id
if override != problem_id:
logger.debug(
"%s overriding problem_id %r -> %r",
PROBLEM_ID_ENV_VAR,
problem_id,
override,
)
return override
# Export problem classes for type hints
__all__ = [
"PROBLEM_CLASSES",
"PROBLEM_ID_ENV_VAR",
"SUPPORTED_PROBLEMS",
"Beams2D",
"Photonics2D",
"ProblemId",
"harness_problem_id",
]
