From 5c2392565d7e89b3634618a44e47f9dff2847165 Mon Sep 17 00:00:00 2001 From: Akshay Shah Date: Fri, 9 Oct 2026 12:57:10 -0500 Subject: [PATCH] fix(mlops): stop waiting on CANCELLED Athena queries in wait_for_athena_query wait_for_athena_query only stopped polling on SUCCEEDED or FAILED. CANCELLED is also a terminal Athena state, so waiting on a cancelled query (from the console, a workgroup limit, or StopQueryExecution) looped forever, hanging AthenaQuery.wait(), run_athena_query() and DatasetBuilder flows. Treat CANCELLED as terminal. Callers already raise for any state other than SUCCEEDED. --- .../mlops/feature_store/feature_utils.py | 2 +- .../mlops/feature_store/test_feature_utils.py | 17 +++++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py b/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py index 4da8e8dc8e..b1a7611dcd 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py +++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py @@ -194,7 +194,7 @@ def wait_for_athena_query(session: Session, query_execution_id: str, poll: int = state = get_query_execution(session, query_execution_id)["QueryExecution"]["Status"][ "State" ] - if state in ("SUCCEEDED", "FAILED"): + if state in ("SUCCEEDED", "FAILED", "CANCELLED"): logger.info("Query %s %s.", query_execution_id, state.lower()) break logger.info("Query %s is being executed.", query_execution_id) diff --git a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py index fc18b76adc..32b9ae66e9 100644 --- a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py +++ b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py @@ -13,6 +13,7 @@ ingest_dataframe, get_session_from_role, _is_collection_column, + wait_for_athena_query, ) @@ -201,6 +202,22 @@ def test_raises_when_no_metastore(self, mock_fg_class): create_athena_query("my-fg", session) +class TestWaitForAthenaQuery: + @pytest.mark.parametrize("terminal_state", ["SUCCEEDED", "FAILED", "CANCELLED"]) + @patch("sagemaker.mlops.feature_store.feature_utils.time.sleep") + @patch("sagemaker.mlops.feature_store.feature_utils.get_query_execution") + def test_returns_on_terminal_state(self, mock_get, mock_sleep, terminal_state): + mock_get.side_effect = [ + {"QueryExecution": {"Status": {"State": "RUNNING"}}}, + {"QueryExecution": {"Status": {"State": terminal_state}}}, + ] + + wait_for_athena_query(Mock(), "query-id", poll=1) + + assert mock_get.call_count == 2 + mock_sleep.assert_called_once_with(1) + + class TestIngestDataframe: @patch("sagemaker.mlops.feature_store.feature_utils.IngestionManagerPandas") @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup")