From 3a59e88675f6048616b6447c71f996447f5951fd Mon Sep 17 00:00:00 2001
From: thomashopkins32 <thopkins1@bnl.gov>
Date: Mon, 20 Jul 2026 15:42:53 -0400
Subject: [PATCH] Set checkpoint path after init on Ax agent and optimizer

---
 src/blop/ax/agent.py                |  4 ++++
 src/blop/ax/optimizer.py            |  5 +++++
 src/blop/tests/ax/test_agent.py     | 29 ++++++++++++++++++++++++
 src/blop/tests/ax/test_optimizer.py | 35 +++++++++++++++++++++++++++++
 4 files changed, 73 insertions(+)

diff --git a/src/blop/ax/agent.py b/src/blop/ax/agent.py
index b5a6ad8a..bbcdd71f 100644
--- a/src/blop/ax/agent.py
+++ b/src/blop/ax/agent.py
@@ -48,6 +48,10 @@ def ax_client(self) -> Client:
     def checkpoint_path(self) -> str | None:
         return self._optimizer.checkpoint_path
 
+    @checkpoint_path.setter
+    def checkpoint_path(self, value: str) -> None:
+        self._optimizer.checkpoint_path = value
+
     @property
     def fixed_dofs(self) -> dict[str, Any] | None:
         return self._optimizer.fixed_parameters
diff --git a/src/blop/ax/optimizer.py b/src/blop/ax/optimizer.py
index 5fd4c357..a21831c2 100644
--- a/src/blop/ax/optimizer.py
+++ b/src/blop/ax/optimizer.py
@@ -93,6 +93,11 @@ def checkpoint_path(self) -> str | None:
         """The file path for saving and restoring optimizer state, or ``None`` if disabled."""
         return self._checkpoint_path
 
+    @checkpoint_path.setter
+    def checkpoint_path(self, value: str) -> None:
+        """Set file path for saving and restoring optimizer state, or ``None`` if disabled."""
+        self._checkpoint_path = value
+
     @property
     def ax_client(self) -> Client:
         """The underlying Ax ``Client`` used for experiment management."""
diff --git a/src/blop/tests/ax/test_agent.py b/src/blop/tests/ax/test_agent.py
index b5b1ee57..573ff0ed 100644
--- a/src/blop/tests/ax/test_agent.py
+++ b/src/blop/tests/ax/test_agent.py
@@ -92,6 +92,35 @@ def test_agent_checkpoint(mock_evaluation_function, mock_acquisition_plan, tmp_p
     assert len(agent.ax_client.summarize()) == 1
 
 
+def test_agent_set_checkpoint_after_init(mock_evaluation_function, mock_acquisition_plan, tmp_path):
+    checkpoint_path = tmp_path / "checkpoint.json"
+    readable = ReadableSignal(name="test_readable")
+    agent = Agent(
+        sensors=[ReadableSignal(name="test_readable")],
+        dofs=[RangeDOF(name="x1", bounds=(0, 10), parameter_type="float")],
+        objectives=[Objective(name="test_objective", minimize=False)],
+        evaluation_function=mock_evaluation_function,
+        acquisition_plan=mock_acquisition_plan,
+    )
+
+    assert agent.checkpoint_path is None
+    agent.ingest([{"x1": 0.1, "test_objective": 0.2}])
+    agent.ax_client.configure_generation_strategy()
+    agent.checkpoint_path = str(checkpoint_path)
+    assert not checkpoint_path.exists()
+    agent.checkpoint()
+    assert checkpoint_path.exists()
+
+    agent = Agent.from_checkpoint(
+        str(checkpoint_path),
+        sensors=[readable],
+        actuators=[],
+        evaluation_function=mock_evaluation_function,
+        acquisition_plan=mock_acquisition_plan,
+    )
+    assert len(agent.ax_client.summarize()) == 1
+
+
 def test_agent_to_optimization_problem(mock_evaluation_function):
     """Test that the agent can be converted to an optimization problem."""
     movable1 = MovableSignal(name="test_movable1")
diff --git a/src/blop/tests/ax/test_optimizer.py b/src/blop/tests/ax/test_optimizer.py
index 81b54af5..c1cf9218 100644
--- a/src/blop/tests/ax/test_optimizer.py
+++ b/src/blop/tests/ax/test_optimizer.py
@@ -172,6 +172,41 @@ def test_ax_optimizer_checkpoint(tmp_path):
     assert optimizer.checkpoint_path == str(checkpoint_path)
 
 
+def test_ax_optimizer_set_checkpoint_after_init(tmp_path):
+    checkpoint_path = tmp_path / "checkpoint.json"
+
+    # Save to checkpoint
+    optimizer = AxOptimizer(
+        parameters=[
+            RangeParameterConfig(name="x1", bounds=(-5.0, 5.0), parameter_type="float"),
+        ],
+        objective="y1",
+    )
+    assert optimizer.checkpoint_path is None
+    suggestions = optimizer.suggest(num_points=2)
+    outcomes = [
+        {"_id": suggestions[0]["_id"], "y1": 1.0},
+        {"_id": suggestions[1]["_id"], "y1": 3.0},
+    ]
+    optimizer.ingest(outcomes)
+
+    with pytest.raises(ValueError, match="path is not set"):
+        optimizer.checkpoint()
+
+    optimizer.checkpoint_path = str(checkpoint_path)
+    assert not checkpoint_path.exists()
+    optimizer.checkpoint()
+    assert checkpoint_path.exists()
+
+    # Load from checkpoint
+    optimizer = AxOptimizer.from_checkpoint(str(checkpoint_path))
+    summary_df = optimizer.ax_client.summarize()
+    assert "x1" in summary_df.columns
+    assert "y1" in summary_df.columns
+    assert len(summary_df) == 2
+    assert optimizer.checkpoint_path == str(checkpoint_path)
+
+
 def test_ax_optimizer_register_failures():
     optimizer = AxOptimizer(
         parameters=[