From 6b4905d3bfa752fa804c7efbeae10e59f286c2c1 Mon Sep 17 00:00:00 2001 From: Jerry Lin Date: Thu, 6 Jun 2024 22:44:03 -0700 Subject: [PATCH] Passing current_node_name when gs is node based (#2505) Summary: Pull Request resolved: https://github.com/facebook/Ax/pull/2505 As titled. Reviewed By: mgarrard Differential Revision: D58260119 fbshipit-source-id: f338e2d86167467d9f66f3100f3b64a0073c40f7 --- ax/storage/sqa_store/save.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/ax/storage/sqa_store/save.py b/ax/storage/sqa_store/save.py index dd8798497a2..19c662699e3 100644 --- a/ax/storage/sqa_store/save.py +++ b/ax/storage/sqa_store/save.py @@ -343,10 +343,15 @@ def _update_generation_strategy( "should be saved before generation strategy." ) + curr_index = ( + generation_strategy.current_node_name + if generation_strategy.is_node_based + else generation_strategy.current_step_index + ) with session_scope() as session: session.query(gs_sqa_class).filter_by(id=gs_id).update( { - "curr_index": generation_strategy.current_step_index, + "curr_index": curr_index, "experiment_id": experiment_id, } )