Ship ag_args as JSON instead of pickle

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,34 @@
 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 as e:
+        original_error = e
+    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.
+    raise TypeError(
+        f"The provided arguments are not JSON-serializable. Original error: {original_error}"
+    ) from original_error
+
+
 class Backend(ABC):
     name = "backend"
 
@@ -83,18 +111,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

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
         )

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(

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.
         """

src/autogluon/cloud/predictor/cloud_predictor.py

@@ -194,9 +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.
         leaderboard: bool, default = True
             Whether to include the leaderboard in the output artifact
         framework_version: str, default = `latest`
@@ -263,6 +260,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,

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()}"

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)

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)

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