From a3b2f0af19929fbbc85245e84c193171bd43d716 Mon Sep 17 00:00:00 2001 From: Oleksandr Shchur Date: Thu, 6 Aug 2026 15:00:22 +0000 Subject: [PATCH 1/3] Ship ag_args as JSON instead of pickle The remote-training config (predictor init/fit args) was serialized with pickle and unpickled inside the SageMaker/Ray container, so a tampered ag_args file would be arbitrary-code execution on load. Switch to JSON. Non-JSON-serializable inputs (search spaces, custom metric objects, classes/callables in hyperparameters) now raise an informative error that names the offending user argument. A future 'extra' cloudpickle channel can carry such objects opt-in. - backend.dumps_ag_args(): serialize-or-raise helper, used by prepare_args - flip the three read sites to json.load; rename ag_args.pkl -> ag_args.json - add unit tests for the serialization error paths --- src/autogluon/cloud/backend/backend.py | 40 ++++++++++++++--- src/autogluon/cloud/backend/ray_backend.py | 2 +- .../cloud/backend/sagemaker_backend.py | 15 +++---- .../backend/timeseries_sagemaker_backend.py | 2 +- .../cloud/scripts/ray_scripts/train.py | 6 +-- .../cloud/scripts/sagemaker_scripts/train.py | 5 +-- .../general/test_ag_args_serialization.py | 44 +++++++++++++++++++ 7 files changed, 93 insertions(+), 21 deletions(-) create mode 100644 tests/unittests/general/test_ag_args_serialization.py diff --git a/src/autogluon/cloud/backend/backend.py b/src/autogluon/cloud/backend/backend.py index 3817b70c..e122f835 100644 --- a/src/autogluon/cloud/backend/backend.py +++ b/src/autogluon/cloud/backend/backend.py @@ -1,7 +1,7 @@ from __future__ import annotations +import json import os -import pickle from abc import ABC, abstractmethod from typing import Any, Dict, List, Optional, Union @@ -10,6 +10,35 @@ from ..endpoint.endpoint import Endpoint +def dumps_ag_args(config: Dict[str, Any]) -> str: + """Serialize the remote-training config to JSON, raising a user-facing error on failure. + + The config carries the user's predictor init/fit arguments. Some objects (search spaces, + custom metric objects, classes/callables in ``hyperparameters``) are not JSON-serializable; + when that happens we pinpoint the offending argument so the error names what the user passed. + """ + try: + return json.dumps(config) + except TypeError: + pass + for group in ("predictor_init_args", "predictor_fit_args"): + for key, value in (config.get(group) or {}).items(): + try: + json.dumps(value) + except TypeError as e: + raise TypeError( + f"The value passed for `{key}` is not JSON-serializable. This can happen with " + f"search spaces (Real/Categorical/Int), custom metric objects, or classes/callables " + f"in `hyperparameters`. Please pass these as plain values (e.g. metric names as strings). " + f"Original error: {e}" + ) from e + # Culprit is outside the known arg groups; re-raise the original error with generic guidance. + try: + return json.dumps(config) + except TypeError as e: + raise TypeError(f"The provided arguments are not JSON-serializable. Original error: {e}") from e + + class Backend(ABC): name = "backend" @@ -83,18 +112,19 @@ def get_fit_job_info(self) -> Dict[str, Any]: def prepare_args(self, path: str, **kwargs): """ prepare parameter args required to be passed to remote AG, i.e. init args and fit args - The args will be saved as a pickle object + The args will be saved as a JSON file Parameters ---------- path: str - Path to save the pickle file + Path to save the JSON file """ assert self.predictor_type is not None config = self._construct_ag_args(**kwargs) + payload = dumps_ag_args(config) os.makedirs(os.path.dirname(path), exist_ok=True) - with open(path, "wb") as f: - pickle.dump(config, f) + with open(path, "w") as f: + f.write(payload) def _construct_ag_args(**kwargs): raise NotImplementedError diff --git a/src/autogluon/cloud/backend/ray_backend.py b/src/autogluon/cloud/backend/ray_backend.py index 2490d378..f9f03a02 100644 --- a/src/autogluon/cloud/backend/ray_backend.py +++ b/src/autogluon/cloud/backend/ray_backend.py @@ -211,7 +211,7 @@ def fit( image_uri = self._get_image_uri( framework_version=framework_version, custom_image_uri=custom_image_uri, instance_type=instance_type ) - ag_args_path = os.path.join(self.local_output_path, "job", "ag_args.pkl") + ag_args_path = os.path.join(self.local_output_path, "job", "ag_args.json") self.prepare_args( path=ag_args_path, predictor_init_args=predictor_init_args, predictor_fit_args=predictor_fit_args ) diff --git a/src/autogluon/cloud/backend/sagemaker_backend.py b/src/autogluon/cloud/backend/sagemaker_backend.py index 82346327..94a1c96a 100644 --- a/src/autogluon/cloud/backend/sagemaker_backend.py +++ b/src/autogluon/cloud/backend/sagemaker_backend.py @@ -2,7 +2,6 @@ import json import logging import os -import pickle import shutil import tarfile import tempfile @@ -229,7 +228,7 @@ def fit( Any extra arguments needed to pass to fit. Please refer to https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html#sagemaker.estimator.Framework.fit for all options extra_ag_args: Optional[Dict[str, Any]], default = None - Additional entries to merge into ``ag_args.pkl``. Use this to ship caller-specific metadata to the + Additional entries to merge into ``ag_args.json``. Use this to ship caller-specific metadata to the train script (e.g. ``predict_after_fit``, or ``id_column`` / ``timestamp_column`` for time series). """ if data_channels.get("train_data") is None: @@ -314,7 +313,7 @@ def fit( f"`predictions_path` must be a full S3 URL ending in '.csv' or '.parquet' " f"(e.g. 's3://bucket/key/predictions.parquet'), got {predictions_path!r}." ) - ag_args_path = os.path.join(self.local_output_path, "utils", "ag_args.pkl") + ag_args_path = os.path.join(self.local_output_path, "utils", "ag_args.json") self.prepare_args(path=ag_args_path, **ag_args) inputs = self._upload_fit_artifact( data_channels=data_channels, @@ -1015,9 +1014,9 @@ def get_fit_predict_results(self) -> pd.DataFrame: return load_pd.load(local_path) def _download_ag_args_from_job(self) -> Dict[str, Any]: - """Fetch and unpickle the ``ag_args.pkl`` that was uploaded as the ``ag_args`` channel. + """Fetch and parse the ``ag_args.json`` that was uploaded as the ``ag_args`` channel. - Each training job carries the exact pickle it was launched with as an input channel, + Each training job carries the exact config it was launched with as an input channel, making this the authoritative source — independent of local-disk lifetime. """ job_name = self._fit_job.job_name @@ -1029,12 +1028,12 @@ def _download_ag_args_from_job(self) -> Dict[str, Any]: f"Training job {job_name!r} has no `ag_args` input channel — cannot recover predictions_path." ) bucket, key = s3_path_to_bucket_prefix(ag_args_uri) - assert key.endswith(".pkl"), f"Expected ag_args channel to point to a .pkl file, got {ag_args_uri!r}" + assert key.endswith(".json"), f"Expected ag_args channel to point to a .json file, got {ag_args_uri!r}" with tempfile.TemporaryDirectory(prefix="ag_args_") as tmpdir: local_path = os.path.join(tmpdir, os.path.basename(key)) self.sagemaker_session.boto_session.client("s3").download_file(bucket, key, local_path) - with open(local_path, "rb") as f: - return pickle.load(f) + with open(local_path, "r") as f: + return json.load(f) def _construct_ag_args(self, predictor_init_args, predictor_fit_args, leaderboard, **kwargs): config = dict( diff --git a/src/autogluon/cloud/backend/timeseries_sagemaker_backend.py b/src/autogluon/cloud/backend/timeseries_sagemaker_backend.py index 37663873..66f32bdf 100644 --- a/src/autogluon/cloud/backend/timeseries_sagemaker_backend.py +++ b/src/autogluon/cloud/backend/timeseries_sagemaker_backend.py @@ -38,7 +38,7 @@ def fit( ) -> None: """Fit a TimeSeriesPredictor in SageMaker. - ``id_column`` / ``timestamp_column`` are forwarded to the training script via ``ag_args.pkl``. + ``id_column`` / ``timestamp_column`` are forwarded to the training script via ``ag_args.json``. ``known_covariates`` (if present in ``data_channels``) is only honored when ``extra_ag_args["predict_after_fit"]`` is True. """ diff --git a/src/autogluon/cloud/scripts/ray_scripts/train.py b/src/autogluon/cloud/scripts/ray_scripts/train.py index d546c67a..d1c3c115 100644 --- a/src/autogluon/cloud/scripts/ray_scripts/train.py +++ b/src/autogluon/cloud/scripts/ray_scripts/train.py @@ -1,6 +1,6 @@ import argparse +import json import os -import pickle import shutil import time from datetime import datetime, timezone @@ -91,8 +91,8 @@ def upload_file(file_name: str, bucket: str, prefix: Optional[str] = None): tune_data = None if args.tune_data is not None: tune_data = TabularDataset(args.tune_data) - with open(args.ag_args_path, "rb") as f: - ag_args = pickle.load(f) + with open(args.ag_args_path, "r") as f: + ag_args = json.load(f) predictor_init_args = ag_args["predictor_init_args"] predictor_fit_args = ag_args["predictor_fit_args"] save_path = f"ag_distributed_training_{get_utc_timestamp_now()}" diff --git a/src/autogluon/cloud/scripts/sagemaker_scripts/train.py b/src/autogluon/cloud/scripts/sagemaker_scripts/train.py index b0b17977..56748632 100644 --- a/src/autogluon/cloud/scripts/sagemaker_scripts/train.py +++ b/src/autogluon/cloud/scripts/sagemaker_scripts/train.py @@ -10,7 +10,6 @@ import boto3 import pandas as pd -import pickle from autogluon.common.loaders import load_pd from autogluon.common.savers import save_pd @@ -85,8 +84,8 @@ def prepare_data(data_file, predictor_type, ag_args, static_features_df=None): os.makedirs(args.output_data_dir, mode=0o777, exist_ok=True) ag_args_file = get_input_path(args.ag_args) - with open(ag_args_file, "rb") as f: - ag_args = pickle.load(f) # AutoGluon-specific args + with open(ag_args_file, "r") as f: + ag_args = json.load(f) # AutoGluon-specific args if args.n_gpus: ag_args["num_gpus"] = int(args.n_gpus) diff --git a/tests/unittests/general/test_ag_args_serialization.py b/tests/unittests/general/test_ag_args_serialization.py new file mode 100644 index 00000000..b9d7b848 --- /dev/null +++ b/tests/unittests/general/test_ag_args_serialization.py @@ -0,0 +1,44 @@ +"""Unit tests for the JSON serialization of remote-training args (see backend.dumps_ag_args).""" + +import json + +import pytest + +from autogluon.cloud.backend.backend import dumps_ag_args + + +def test_when_config_is_json_native_then_roundtrips(): + config = { + "predictor_type": "tabular", + "predictor_init_args": {"label": "y", "eval_metric": "roc_auc"}, + "predictor_fit_args": {"presets": "best_quality", "time_limit": 60}, + "leaderboard": True, + } + assert json.loads(dumps_ag_args(config)) == config + + +def test_when_init_arg_not_serializable_then_error_names_that_arg(): + config = { + "predictor_init_args": {"label": "y", "eval_metric": object()}, + "predictor_fit_args": {}, + } + with pytest.raises(TypeError) as exc: + dumps_ag_args(config) + assert "`eval_metric`" in str(exc.value) + + +def test_when_fit_arg_not_serializable_then_error_names_that_arg(): + config = { + "predictor_init_args": {"label": "y"}, + "predictor_fit_args": {"presets": "best_quality", "hyperparameters": {"GBM": object}}, + } + with pytest.raises(TypeError) as exc: + dumps_ag_args(config) + assert "`hyperparameters`" in str(exc.value) + + +def test_error_message_does_not_leak_internal_ag_args_name(): + config = {"predictor_init_args": {"eval_metric": object()}, "predictor_fit_args": {}} + with pytest.raises(TypeError) as exc: + dumps_ag_args(config) + assert "ag_args" not in str(exc.value) From 8ca25993c4ffb88488b7d0442c5aa418e7b6aba2 Mon Sep 17 00:00:00 2001 From: Oleksandr Shchur Date: Fri, 7 Aug 2026 11:18:00 +0000 Subject: [PATCH 2/3] Simplify JSON error path and block image_column for tabular - dumps_ag_args: reuse the original TypeError on the fallback path instead of re-serializing the config just to re-raise. - Raise an informative error when image_column is passed for a tabular predictor; tabular image support requires autogluon.multimodal, which is being deprecated. MultiModal predictors (string image_column) are unaffected. --- src/autogluon/cloud/backend/backend.py | 11 +++++------ src/autogluon/cloud/predictor/cloud_predictor.py | 7 +++++++ 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/src/autogluon/cloud/backend/backend.py b/src/autogluon/cloud/backend/backend.py index e122f835..4aa63321 100644 --- a/src/autogluon/cloud/backend/backend.py +++ b/src/autogluon/cloud/backend/backend.py @@ -19,8 +19,8 @@ def dumps_ag_args(config: Dict[str, Any]) -> str: """ try: return json.dumps(config) - except TypeError: - pass + except TypeError as e: + original_error = e for group in ("predictor_init_args", "predictor_fit_args"): for key, value in (config.get(group) or {}).items(): try: @@ -33,10 +33,9 @@ def dumps_ag_args(config: Dict[str, Any]) -> str: f"Original error: {e}" ) from e # Culprit is outside the known arg groups; re-raise the original error with generic guidance. - try: - return json.dumps(config) - except TypeError as e: - raise TypeError(f"The provided arguments are not JSON-serializable. Original error: {e}") from e + raise TypeError( + f"The provided arguments are not JSON-serializable. Original error: {original_error}" + ) from original_error class Backend(ABC): diff --git a/src/autogluon/cloud/predictor/cloud_predictor.py b/src/autogluon/cloud/predictor/cloud_predictor.py index 6ac73e94..cd5c1646 100644 --- a/src/autogluon/cloud/predictor/cloud_predictor.py +++ b/src/autogluon/cloud/predictor/cloud_predictor.py @@ -197,6 +197,7 @@ def fit( image_column: str, default = None The column name in the training/tuning data that contains the image paths. The image paths MUST be absolute paths to you local system. + Not supported for tabular predictors (raises an error) — use `MultiModalCloudPredictor` for image data. leaderboard: bool, default = True Whether to include the leaderboard in the output artifact framework_version: str, default = `latest` @@ -263,6 +264,12 @@ def fit( ) if data_channels["train_data"] is None: raise TypeError("fit() missing required argument: 'train_data'") + if image_column is not None and self.predictor_type == "tabular": + raise ValueError( + "`image_column` is no longer supported for tabular predictors: image models in " + "AutoGluon-Tabular require autogluon.multimodal, which is being deprecated. " + "Use `MultiModalCloudPredictor` for image data." + ) backend_kwargs = self.backend.parse_backend_fit_kwargs(backend_kwargs) self.backend.fit( predictor_init_args=predictor_init_args, From cf997747a60f6b37a62ca3e5925c11bd19877e7c Mon Sep 17 00:00:00 2001 From: Oleksandr Shchur Date: Fri, 7 Aug 2026 11:25:40 +0000 Subject: [PATCH 3/3] Remove dead tabular image test and hide image_column from fit docstring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit test_tabular_tabular_text_image trained a tabular predictor with an image column (AG_IMAGE_NN/AG_TEXT_NN + a live FeatureMetadata object) — a workflow that requires autogluon.multimodal and no longer works over JSON ag_args. Drop it along with the now-deprecated image_column docstring entry. --- .../cloud/predictor/cloud_predictor.py | 4 -- tests/unittests/tabular/test_tabular.py | 70 ------------------- 2 files changed, 74 deletions(-) diff --git a/src/autogluon/cloud/predictor/cloud_predictor.py b/src/autogluon/cloud/predictor/cloud_predictor.py index cd5c1646..656ba0d7 100644 --- a/src/autogluon/cloud/predictor/cloud_predictor.py +++ b/src/autogluon/cloud/predictor/cloud_predictor.py @@ -194,10 +194,6 @@ def fit( predictor_fit_args: Optional[dict], default = None Additional fit args forwarded to the underlying predictor's ``fit()``. Must NOT contain ``train_data`` or ``tuning_data`` — pass those as explicit arguments above. - image_column: str, default = None - The column name in the training/tuning data that contains the image paths. - The image paths MUST be absolute paths to you local system. - Not supported for tabular predictors (raises an error) — use `MultiModalCloudPredictor` for image data. leaderboard: bool, default = True Whether to include the leaderboard in the output artifact framework_version: str, default = `latest` diff --git a/tests/unittests/tabular/test_tabular.py b/tests/unittests/tabular/test_tabular.py index 8cd585cb..783a9eeb 100644 --- a/tests/unittests/tabular/test_tabular.py +++ b/tests/unittests/tabular/test_tabular.py @@ -3,76 +3,6 @@ import pandas as pd -from autogluon.cloud import TabularCloudPredictor -from autogluon.common.features.feature_metadata import FeatureMetadata - - -def test_tabular_tabular_text_image(test_helper, framework_version): - train_data = "tabular_text_image_train.csv" - test_data = "tabular_text_image_test.csv" - images = "tabular_text_image_images.zip" - image_column = "Images" - timestamp = test_helper.get_utc_timestamp_now() - with tempfile.TemporaryDirectory() as temp_dir: - os.chdir(temp_dir) - test_helper.prepare_data(train_data, test_data, images) - test_helper.extract_images(images) - train_data = test_helper.replace_image_abspath(train_data, image_column) - test_data = test_helper.replace_image_abspath(test_data, image_column) - feature_metadata = FeatureMetadata.from_df(train_data) - feature_metadata = feature_metadata.add_special_types({"Images": ["image_path"]}) - - time_limit = 600 - - predictor_init_args = dict( - label="AdoptionSpeed", - ) - text_model = "AG_TEXT_NN" - image_model = "AG_IMAGE_NN" - predictor_fit_args = dict( - time_limit=time_limit, - hyperparameters={ - "XGB": {}, - text_model: {"presets": "medium_quality_faster_train"}, - image_model: {}, - }, - feature_metadata=feature_metadata, - ) - cloud_predictor = TabularCloudPredictor( - cloud_output_path=f"s3://autogluon-cloud-ci/test-tabular-tabular-text-image/{framework_version}/{timestamp}", - local_output_path="test_tabular_tabular_text_image_cloud_predictor", - ) - training_custom_image_uri = test_helper.get_custom_image_uri(framework_version, type="training", gpu=True) - inference_custom_image_uri = test_helper.get_custom_image_uri(framework_version, type="inference", gpu=False) - test_helper.test_basic_functionality( - cloud_predictor, - train_data, - predictor_init_args, - predictor_fit_args, - test_data, - fit_kwargs=dict( - instance_type="ml.g4dn.2xlarge", - volume_size=100, - image_column=image_column, - framework_version=framework_version, - custom_image_uri=training_custom_image_uri, - ), - deploy_kwargs=dict(framework_version=framework_version, custom_image_uri=inference_custom_image_uri), - predict_real_time_kwargs=dict( - test_data_image_column=image_column, - ), - predict_kwargs=dict( - test_data_image_column=image_column, - framework_version=framework_version, - custom_image_uri=inference_custom_image_uri, - ), - ) - local_predictor = cloud_predictor.to_local_predictor( - require_version_match=False, require_py_version_match=False - ) - models = local_predictor.model_names() - assert "ImagePredictor" in models - def test_tabular_foundation_model_predict(test_helper, framework_version): import boto3