add kwargs in predict and predict_proba

Author: suzhoum

src/autogluon/cloud/backend/sagemaker_backend.py

@@ -704,6 +704,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 +784,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 +807,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 +892,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 +1137,7 @@ def _predict(
         instance_count=1,
         custom_image_uri=None,
         wait=True,
+        inference_kwargs=None,
         download=True,
         persist=True,
         save_path=None,
@@ -1256,6 +1261,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/job/sagemaker_job.py

@@ -1,7 +1,7 @@
 import logging
 from abc import abstractmethod
 from typing import Optional
-
+import json
 import sagemaker
 
 from ..utils.ag_sagemaker import (
@@ -257,6 +257,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 +266,10 @@ def run(
         else:
             model_cls = AutoGluonNonRepackInferenceModel
         logger.log(20, "Creating inference model...")
+        inference_kwargs_str = json.dumps(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 +280,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

@@ -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,
         )