Support kwargs in predict() and predict_proba()

.github/workflows/continuous_integration.yml

@@ -98,7 +98,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        AG_VERSION: ["source", "0.7.0"]
+        AG_VERSION: ["source", "1.1.0"]
     needs: cloud_lint_check
     runs-on: ubuntu-latest
     steps:
@@ -130,7 +130,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        AG_VERSION: ["source", "0.7.0"]
+        AG_VERSION: ["source", "1.1.0"]
     needs: cloud_lint_check
     runs-on: ubuntu-latest
     steps:
@@ -162,7 +162,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        AG_VERSION: ["source", "0.7.0"]
+        AG_VERSION: ["source", "1.1.0"]
     needs: cloud_lint_check
     runs-on: ubuntu-latest
     steps:
@@ -194,7 +194,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        AG_VERSION: ["source", "0.7.0"]
+        AG_VERSION: ["source", "1.1.0"]
     needs: cloud_lint_check
     runs-on: ubuntu-latest
     steps:
@@ -226,7 +226,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        AG_VERSION: ["source", "0.7.0"]
+        AG_VERSION: ["source", "1.1.0"]
     needs: cloud_lint_check
     runs-on: ubuntu-latest
     steps:

src/autogluon/cloud/backend/ray_backend.py

@@ -384,7 +384,7 @@ def _get_image_uri(self, framework_version: str, instance_type: str, custom_imag
             logger.log(20, f"Training with custom_image_uri=={custom_image_uri}")
         else:
             framework_version, py_version = parse_framework_version(
-                framework_version, "training", minimum_version="0.7.0"
+                framework_version, "training", minimum_version="1.0.0"
             )
             logger.log(20, f"Training with framework_version=={framework_version}")
             image_uri = image_uris.retrieve(

src/autogluon/cloud/backend/sagemaker_backend.py

@@ -558,7 +558,6 @@ def predict_real_time(
         test_data_image_column: Optional[str] = None,
         accept: str = "application/x-parquet",
         inference_kwargs: Optional[Dict[str, Any]] = None,
-        **kwargs,
     ) -> Union[pd.DataFrame, pd.Series]:
         """
         Predict with the deployed SageMaker endpoint. A deployed SageMaker endpoint is required.
@@ -595,7 +594,6 @@ def predict_proba_real_time(
         test_data_image_column: Optional[str] = None,
         accept: str = "application/x-parquet",
         inference_kwargs: Optional[Dict[str, Any]] = None,
-        **kwargs,
     ) -> Union[pd.DataFrame, pd.Series]:
         """
         Predict probability with the deployed SageMaker endpoint. A deployed SageMaker endpoint is required.
@@ -704,6 +702,7 @@ def predict(
         instance_count: int = 1,
         custom_image_uri: Optional[str] = None,
         wait: bool = True,
+        inference_kwargs: Optional[Dict[str, Any]] = None,
         download: bool = True,
         persist: bool = True,
         save_path: Optional[str] = None,
@@ -783,6 +782,7 @@ def predict(
             instance_count=instance_count,
             custom_image_uri=custom_image_uri,
             wait=wait,
+            inference_kwargs=inference_kwargs,
             download=download,
             persist=persist,
             save_path=save_path,
@@ -805,6 +805,7 @@ def predict_proba(
         instance_count: int = 1,
         custom_image_uri: Optional[str] = None,
         wait: bool = True,
+        inference_kwargs: Optional[Dict[str, Any]] = None,
         download: bool = True,
         persist: bool = True,
         save_path: Optional[str] = None,
@@ -889,6 +890,7 @@ def predict_proba(
             instance_count=instance_count,
             custom_image_uri=custom_image_uri,
             wait=wait,
+            inference_kwargs=inference_kwargs,
             download=download,
             persist=persist,
             save_path=save_path,
@@ -1133,6 +1135,7 @@ def _predict(
         instance_count=1,
         custom_image_uri=None,
         wait=True,
+        inference_kwargs=None,
         download=True,
         persist=True,
         save_path=None,
@@ -1256,6 +1259,7 @@ def _predict(
             transformer_kwargs=transformer_kwargs,
             model_kwargs=model_kwargs,
             repack_model=repack_model,
+            inference_kwargs=inference_kwargs,
             **transform_kwargs,
         )
         self._batch_transform_jobs[job_name] = batch_transform_job

src/autogluon/cloud/backend/timeseries_sagemaker_backend.py

@@ -15,9 +15,9 @@ class TimeSeriesSagemakerBackend(SagemakerBackend):
     def _preprocess_data(
         self,
         data: Union[pd.DataFrame, str],
-        id_column: str,
-        timestamp_column: str,
-        target: str,
+        id_column: Optional[str] = None,
+        timestamp_column: Optional[str] = None,
+        target: Optional[str] = None,
         static_features: Optional[Union[pd.DataFrame, str]] = None,
     ) -> pd.DataFrame:
         if isinstance(data, str):
@@ -27,12 +27,15 @@ def _preprocess_data(
         cols = data.columns.to_list()
         # Make sure id and timestamp columns are the first two columns, and target column is in the end
         # This is to ensure in the container we know how to find id and timestamp columns, and whether there are static features being merged
-        timestamp_index = cols.index(timestamp_column)
-        cols.insert(0, cols.pop(timestamp_index))
-        id_index = cols.index(id_column)
-        cols.insert(0, cols.pop(id_index))
-        target_index = cols.index(target)
-        cols.append(cols.pop(target_index))
+        if timestamp_column is not None:
+            timestamp_index = cols.index(timestamp_column)
+            cols.insert(0, cols.pop(timestamp_index))
+        if id_column is not None:
+            id_index = cols.index(id_column)
+            cols.insert(0, cols.pop(id_index))
+        if target is not None:
+            target_index = cols.index(target)
+            cols.append(cols.pop(target_index))
         data = data[cols]
 
         if static_features is not None:
@@ -48,8 +51,8 @@ def fit(
         *,
         predictor_init_args: Dict[str, Any],
         predictor_fit_args: Dict[str, Any],
-        id_column: str,
-        timestamp_column: str,
+        id_column: Optional[str] = None,
+        timestamp_column: Optional[str] = None,
         static_features: Optional[Union[str, pd.DataFrame]] = None,
         framework_version: str = "latest",
         job_name: Optional[str] = None,
@@ -199,9 +202,9 @@ def predict_proba_real_time(self, **kwargs) -> pd.DataFrame:
     def predict(
         self,
         test_data: Union[str, pd.DataFrame],
-        id_column: str,
-        timestamp_column: str,
-        target: str,
+        id_column: Optional[str] = None,
+        timestamp_column: Optional[str] = None,
+        target: Optional[str] = None,
         static_features: Optional[Union[str, pd.DataFrame]] = None,
         **kwargs,
     ) -> Optional[pd.DataFrame]:

src/autogluon/cloud/default_cluster_configs/ray_aws_default_cluster_config.yaml

@@ -6,7 +6,7 @@ max_workers: 1
 
 docker:
     # The image uri will be dynamically replaced by cloud predictor
-    image: "763104351884.dkr.ecr.us-east-1.amazonaws.com/autogluon-training:0.7.0-cpu-py39-ubuntu20.04"
+    image: "763104351884.dkr.ecr.us-east-1.amazonaws.com/autogluon-training:1.0.0-cpu-py39-ubuntu20.04"
     container_name: "ag_dlc"
 
 # Cloud-provider specific configuration.

src/autogluon/cloud/job/sagemaker_job.py

@@ -10,6 +10,7 @@
     AutoGluonSagemakerEstimator,
 )
 from ..utils.constants import LOCAL_MODE, LOCAL_MODE_GPU, MODEL_ARTIFACT_NAME
+from ..utils.utils import serialize_kwargs
 from .remote_job import RemoteJob
 
 logger = logging.getLogger(__name__)
@@ -257,6 +258,7 @@ def run(
         model_kwargs,
         transformer_kwargs,
         repack_model=False,
+        inference_kwargs=None,
         **kwargs,
     ):
         self._local_mode = instance_type in (LOCAL_MODE, LOCAL_MODE_GPU)
@@ -265,6 +267,10 @@ def run(
         else:
             model_cls = AutoGluonNonRepackInferenceModel
         logger.log(20, "Creating inference model...")
+        inference_kwargs_str = serialize_kwargs(inference_kwargs) if inference_kwargs is not None else None
+        env = {}
+        if len(inference_kwargs_str) > 0:
+            env["inference_kwargs"] = inference_kwargs_str
         model = model_cls(
             model_data=model_data,
             role=role,
@@ -275,6 +281,7 @@ def run(
             custom_image_uri=custom_image_uri,
             entry_point=entry_point,
             predictor_cls=predictor_cls,
+            env=env,
             **model_kwargs,
         )
         logger.log(20, "Inference model created successfully")

src/autogluon/cloud/predictor/cloud_predictor.py

@@ -541,7 +541,7 @@ def predict_proba_real_time(
         """
         self._validate_inference_kwargs(inference_kwargs=kwargs)
         return self.backend.predict_proba_real_time(
-            test_data=test_data, test_data_image_column=test_data_image_column, accept=accept
+            test_data=test_data, test_data_image_column=test_data_image_column, accept=accept, inference_kwargs=kwargs
         )
 
     def predict(
@@ -556,6 +556,7 @@ def predict(
         custom_image_uri: Optional[str] = None,
         wait: bool = True,
         backend_kwargs: Optional[Dict] = None,
+        **kwargs,
     ) -> Optional[pd.Series]:
         """
         Batch inference.
@@ -632,6 +633,7 @@ def predict(
             instance_count=instance_count,
             custom_image_uri=custom_image_uri,
             wait=wait,
+            inference_kwargs=kwargs,
             **backend_kwargs,
         )
 
@@ -648,6 +650,7 @@ def predict_proba(
         custom_image_uri: Optional[str] = None,
         wait: bool = True,
         backend_kwargs: Optional[Dict] = None,
+        **kwargs,
     ) -> Optional[Union[Tuple[pd.Series, Union[pd.DataFrame, pd.Series]], Union[pd.DataFrame, pd.Series]]]:
         """
         Batch inference
@@ -730,6 +733,7 @@ def predict_proba(
             instance_count=instance_count,
             custom_image_uri=custom_image_uri,
             wait=wait,
+            inference_kwargs=kwargs,
             **backend_kwargs,
         )
 

src/autogluon/cloud/predictor/timeseries_cloud_predictor.py

@@ -50,8 +50,8 @@ def fit(
         *,
         predictor_init_args: Dict[str, Any],
         predictor_fit_args: Dict[str, Any],
-        id_column: str = "item_id",
-        timestamp_column: str = "timestamp",
+        id_column: Optional[str] = None,
+        timestamp_column: Optional[str] = None,
         static_features: Optional[Union[str, pd.DataFrame]] = None,
         framework_version: str = "latest",
         job_name: Optional[str] = None,
@@ -120,7 +120,7 @@ def fit(
         if backend_kwargs is None:
             backend_kwargs = {}
 
-        self.target_column = predictor_init_args.get("target", "target")
+        self.target_column = predictor_init_args.get("target")
         self.id_column = id_column
         self.timestamp_column = timestamp_column
 
@@ -146,6 +146,9 @@ def fit(
     def predict_real_time(
         self,
         test_data: Union[str, pd.DataFrame],
+        id_column: Optional[str] = None,
+        timestamp_column: Optional[str] = None,
+        target: Optional[str] = None,
         static_features: Optional[Union[str, pd.DataFrame]] = None,
         accept: str = "application/x-parquet",
         **kwargs,
@@ -175,13 +178,18 @@ def predict_real_time(
         Pandas.DataFrame
         Predict results in DataFrame
         """
+        self.id_column = id_column or self.id_column
+        self.timestamp_column = timestamp_column or self.timestamp_column
+        self.target_column = target or self.target_column
+
         return self.backend.predict_real_time(
             test_data=test_data,
             id_column=self.id_column,
             timestamp_column=self.timestamp_column,
             target=self.target_column,
             static_features=static_features,
             accept=accept,
+            inference_kwargs=kwargs,
         )
 
     def predict_proba_real_time(self, **kwargs) -> pd.DataFrame:
@@ -190,6 +198,9 @@ def predict_proba_real_time(self, **kwargs) -> pd.DataFrame:
     def predict(
         self,
         test_data: Union[str, pd.DataFrame],
+        id_column: Optional[str] = None,
+        timestamp_column: Optional[str] = None,
+        target: Optional[str] = None,
         static_features: Optional[Union[str, pd.DataFrame]] = None,
         predictor_path: Optional[str] = None,
         framework_version: str = "latest",
@@ -199,6 +210,7 @@ def predict(
         custom_image_uri: Optional[str] = None,
         wait: bool = True,
         backend_kwargs: Optional[Dict] = None,
+        **kwargs,
     ) -> Optional[pd.DataFrame]:
         """
         Predict using SageMaker batch transform.
@@ -263,6 +275,10 @@ def predict(
                     Please refer to
                     https://sagemaker.readthedocs.io/en/stable/api/inference/transformer.html#sagemaker.transformer.Transformer.transform for all options.
         """
+        self.id_column = id_column or self.id_column
+        self.timestamp_column = timestamp_column or self.timestamp_column
+        self.target_column = target or self.target_column
+
         if backend_kwargs is None:
             backend_kwargs = {}
         backend_kwargs = self.backend.parse_backend_predict_kwargs(backend_kwargs)
@@ -279,6 +295,7 @@ def predict(
             instance_count=instance_count,
             custom_image_uri=custom_image_uri,
             wait=wait,
+            inference_kwargs=kwargs,
             **backend_kwargs,
         )