From 4147fabbd5f623fc14ed54e15b7ef1c16841acb9 Mon Sep 17 00:00:00 2001 From: Dmytro Soltysiuk Date: Mon, 19 May 2025 17:36:11 -0700 Subject: [PATCH] Improve Sagemaker serve script error logs --- .../sagemaker_scripts/tabular_serve.py | 22 +++++++++++++------ 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/src/autogluon/cloud/scripts/sagemaker_scripts/tabular_serve.py b/src/autogluon/cloud/scripts/sagemaker_scripts/tabular_serve.py index e22a163..c6c24ab 100644 --- a/src/autogluon/cloud/scripts/sagemaker_scripts/tabular_serve.py +++ b/src/autogluon/cloud/scripts/sagemaker_scripts/tabular_serve.py @@ -135,15 +135,23 @@ def _read_with_fallback(read_func, buf, expected_columns): """ # Attempt to read with headers data = read_func(buf) - if set(data.columns) != set(expected_columns): - # Reset buffer and read without headers - buf.seek(0) - data = read_func(buf, header=None) - # Assign expected column names - data.columns = expected_columns - else: + data_cols, model_cols = set(data.columns), set(expected_columns) + required_fields_missing = model_cols.difference(data_cols) + + if required_fields_missing: + raise ValueError("Missing required columns", required_fields_missing) + + if data_cols == model_cols: # Reorder columns to match expected_columns data = data[expected_columns] + return data + + # Reset buffer and read without headers + buf.seek(0) + data = read_func(buf, header=None) + # Assign expected column names + data.columns = expected_columns + return data