Support custom kwargs for timeseries

Author: suzhoum

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

src/autogluon/cloud/scripts/sagemaker_scripts/timeseries_serve.py

@@ -5,9 +5,12 @@
 from io import BytesIO, StringIO
 
 import pandas as pd
+import logging
+import sys
 
 from autogluon.timeseries import TimeSeriesDataFrame, TimeSeriesPredictor
-
+logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 def model_fn(model_dir):
     """loads model from previously saved artifact"""
@@ -31,12 +34,15 @@ def model_fn(model_dir):
 def prepare_timeseries_dataframe(df, predictor):
     target = predictor.target
     cols = df.columns.to_list()
+    logger.info(f"COLUMN {cols}")
     id_column = cols[0]
     timestamp_column = cols[1]
     df[timestamp_column] = pd.to_datetime(df[timestamp_column])
     static_features = None
     if target != cols[-1]:
         # target is not the last column, then there are static features being merged in
+        logger.info(f"Inside condition: {cols}, {target}")
+        logger.info(f"Inside condition: {cols}, {target}")
         target_index = cols.index(target)
         static_columns = cols[target_index + 1 :]
         static_features = df[[id_column] + static_columns].groupby([id_column], sort=False).head(1)
@@ -56,6 +62,7 @@ def transform_fn(model, request_body, input_content_type, output_content_type="a
 
     elif input_content_type == "text/csv":
         buf = StringIO(request_body)
+        logger.info(f"request body data path: {buf}")
         data = pd.read_csv(buf)
 
     elif input_content_type == "application/json":
@@ -77,6 +84,8 @@ def transform_fn(model, request_body, input_content_type, output_content_type="a
     else:
         raise ValueError(f"{input_content_type} input content type not supported.")
 
+    logger.info(f"Model is: {model}")
+    logger.info(f"Columns are: {data.columns}")
     data = prepare_timeseries_dataframe(data, model)
     prediction = model.predict(data, **inference_kwargs)
     prediction = pd.DataFrame(prediction)

src/autogluon/cloud/scripts/sagemaker_scripts/train.py

@@ -31,13 +31,13 @@ def get_env_if_present(name):
 
 
 def prepare_timeseries_dataframe(df, predictor_init_args):
-    target = predictor_init_args["target"]
+    target = predictor_init_args.get("target")
     cols = df.columns.to_list()
     id_column = cols[0]
     timestamp_column = cols[1]
     df[timestamp_column] = pd.to_datetime(df[timestamp_column])
     static_features = None
-    if target != cols[-1]:
+    if target is not None and target != cols[-1]:
         # target is not the last column, then there are static features being merged in
         target_index = cols.index(target)
         static_columns = cols[target_index + 1 :]
@@ -46,7 +46,6 @@ def prepare_timeseries_dataframe(df, predictor_init_args):
         df.drop(columns=static_columns, inplace=True)
     df = TimeSeriesDataFrame.from_data_frame(df, id_column=id_column, timestamp_column=timestamp_column)
     if static_features is not None:
-        print(static_features)
         df.static_features = static_features
     return df