fix: remaining specialization computation only in --no-opt mode, and at the end of the opt run

Author: Elwardi

src/foambo/api_server.py

@@ -2231,6 +2231,14 @@ def get_specialization_cost():
     if hasattr(_state, '_specialization_cache') and _state._specialization_cache:
         return SafeJSONResponse(content=_state._specialization_cache)
 
+    # Heavy GP simulation — only run in --no-opt dashboard mode.
+    # During an active optimization, this would compete with trial scheduling.
+    if not getattr(_state, 'no_opt', False):
+        return SafeJSONResponse(content={
+            "estimates": {},
+            "error": "Specialization cost is only computed in --no-opt mode",
+        })
+
     import numpy as np
     raw_cfg = _state.raw_cfg
     if not raw_cfg or not raw_cfg.get("robust_optimization"):
@@ -3188,7 +3196,7 @@ def get_trial_visualization(trial_index: int):
 # Server lifecycle
 
 def start_api_server(client, raw_cfg, orch_cfg, host: str = "127.0.0.1",
-                     port: int = 8098) -> threading.Thread | None:
+                     port: int = 8098, no_opt: bool = False) -> threading.Thread | None:
     """Start the API server in a background daemon thread.
 
     Args:
@@ -3208,6 +3216,7 @@ def start_api_server(client, raw_cfg, orch_cfg, host: str = "127.0.0.1",
     _state.client = client
     _state.raw_cfg = raw_cfg
     _state.orch_cfg = orch_cfg
+    _state.no_opt = no_opt
     _state.start_time = time.time()
     _state.last_callback = time.time()
 

src/foambo/optimize.py

@@ -1029,7 +1029,7 @@ def _warm_fit(self, *a, state_dict=None, refit=True, **kw):
         plain_cfg = OmegaConf.to_container(cfg, resolve=True) if hasattr(cfg, '_metadata') else cfg
         thread = start_api_server(
             client=client, raw_cfg=plain_cfg, orch_cfg=orch_cfg,
-            host=orch_cfg.api_host, port=orch_cfg.api_port,
+            host=orch_cfg.api_host, port=orch_cfg.api_port, no_opt=True,
         )
         if thread is None:
             log.error("API server failed to start")

src/foambo/robustness.py

@@ -187,6 +187,19 @@ def _construct_botorch_acquisition(self, botorch_acqf_class, botorch_acqf_option
             if issubclass(botorch_acqf_class, qLogNoisyExpectedHypervolumeImprovement):
                 botorch_acqf_class = qLogNoisyExpectedImprovement
                 log.debug("MARS: swapped acqf to qLogNoisyExpectedImprovement")
+            # Temporarily clear outcome_constraints for the parent construction:
+            # with MARS scalarization, the constraint transforms produce shape
+            # mismatches in prune_inferior_points.
+            _saved_oc = self._outcome_constraints
+            self._outcome_constraints = None
+            try:
+                return super()._construct_botorch_acquisition(
+                    botorch_acqf_class=botorch_acqf_class,
+                    botorch_acqf_options=botorch_acqf_options,
+                    model=model,
+                )
+            finally:
+                self._outcome_constraints = _saved_oc
         return super()._construct_botorch_acquisition(
             botorch_acqf_class=botorch_acqf_class,
             botorch_acqf_options=botorch_acqf_options,
@@ -203,6 +216,13 @@ def get_botorch_objective_and_transform(self, botorch_acqf_class, model,
         if is_moo:
             from botorch.acquisition.logei import qLogNoisyExpectedImprovement
             botorch_acqf_class = qLogNoisyExpectedImprovement
+            # MARS + constraint transforms produce shape mismatches in
+            # prune_inferior_points. Drop constraints for MARS — the
+            # GP-learned posterior already accounts for constraint violations
+            # via the training data.
+            if outcome_constraints is not None:
+                log.debug("MARS: dropping outcome_constraints to avoid shape mismatch")
+                outcome_constraints = None
 
         from ax.generators.torch.utils import get_botorch_objective_and_transform as _get_obj
         objective, posterior_transform = _get_obj(