Simplify CloudPredictor / FoundationModel signatures

docs/conf.py

@@ -31,8 +31,9 @@
     "sm-runtime-api": ("https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_runtime_%s.html", "%s"),
 }
 
-# Docstring section listing the SageMaker requests a method sends (rendered like "Notes").
-napoleon_custom_sections = ["SageMaker API"]
+# "SageMaker API" lists the SageMaker requests a method sends (rendered like "Notes"). "Other Parameters" holds rarely
+# used arguments; by default napoleon merges it into the "Parameters" field list, so render it as a separate one.
+napoleon_custom_sections = ["SageMaker API", ("Other Parameters", "params_style")]
 
 # Render the return type inline in the "Returns" field instead of a separate "Return type" field.
 napoleon_use_rtype = False

src/autogluon/cloud/backend/timeseries_sagemaker_backend.py

@@ -31,6 +31,7 @@ def fit(
         instance_count: int = 1,
         volume_size: int = DEFAULT_VOLUME_SIZE,
         custom_image_uri: str | None = None,
+        timeout: int = 24 * 60 * 60,
         wait: bool = True,
         backend_overrides: dict[str, dict[str, Any]] | None = None,
         extra_ag_args: dict[str, Any] | None = None,
@@ -62,6 +63,7 @@ def fit(
             instance_count=instance_count,
             volume_size=volume_size,
             custom_image_uri=custom_image_uri,
+            timeout=timeout,
             wait=wait,
             backend_overrides=backend_overrides,
             extra_ag_args=extra_ag_args,

src/autogluon/cloud/endpoint/timeseries_endpoint.py

@@ -26,6 +26,7 @@ class TimeSeriesEndpoint(Endpoint):
     def predict(
         self,
         data: str | Path | pd.DataFrame,
+        *,
         known_covariates: str | Path | pd.DataFrame | None = None,
         static_features: str | Path | pd.DataFrame | None = None,
         prediction_length: int | None = None,

src/autogluon/cloud/predictor/multimodal_cloud_predictor.py

@@ -1,8 +1,23 @@
+from __future__ import annotations
+
 import logging
 import warnings
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+from typing_extensions import Self, Unpack
 
 from ..backend.constant import MULTIMODL_SAGEMAKER, SAGEMAKER
 from ..endpoint.multimodal_endpoint import MultiModalEndpoint
+from ..utils.constants import DEFAULT_FRAMEWORK_VERSION
+from ..utils.sagemaker_api import (
+    IGNORED_TRAINING_KWARGS,
+    BatchTransformKwargs,
+    TrainingJobKwargs,
+    check_backend_kwargs,
+    reject_legacy_kwargs,
+)
 from .cloud_predictor import CloudPredictor
 
 logger = logging.getLogger(__name__)
@@ -40,3 +55,118 @@ def _get_local_predictor_cls(self):
 
         predictor_cls = MultiModalPredictor
         return predictor_cls
+
+    @reject_legacy_kwargs
+    def fit(
+        self,
+        train_data: str | Path | pd.DataFrame | None = None,
+        *,
+        tuning_data: str | Path | pd.DataFrame | None = None,
+        predictor_init_args: dict[str, Any],
+        predictor_fit_args: dict[str, Any] | None = None,
+        image_column: str | None = None,
+        framework_version: str = DEFAULT_FRAMEWORK_VERSION,
+        instance_type: str = "ml.m5.2xlarge",
+        wait: bool = True,
+        backend_overrides: dict[str, dict[str, Any]] | None = None,
+        **kwargs: Unpack[TrainingJobKwargs],
+    ) -> Self:
+        """
+        Fit the predictor in a SageMaker training job.
+
+        Same as :meth:`CloudPredictor.fit`, with one additional argument:
+
+        Parameters
+        ----------
+        image_column: str | None, default = None
+            Name of the column containing absolute local paths to images, if the data contains images.
+        """
+        kwargs = check_backend_kwargs(kwargs, TrainingJobKwargs, "fit", IGNORED_TRAINING_KWARGS)
+        self._fit(
+            data_channels={"train_data": train_data, "tuning_data": tuning_data},
+            predictor_init_args=predictor_init_args,
+            predictor_fit_args=predictor_fit_args,
+            image_column=image_column,
+            framework_version=framework_version,
+            instance_type=instance_type,
+            wait=wait,
+            backend_overrides=backend_overrides,
+            **kwargs,
+        )
+        return self
+
+    @reject_legacy_kwargs
+    def predict(
+        self,
+        test_data: str | pd.DataFrame,
+        *,
+        test_data_image_column: str | None = None,
+        predictor_path: str | None = None,
+        framework_version: str | None = None,
+        instance_type: str = "ml.m5.2xlarge",
+        wait: bool = True,
+        predictions_path: str | None = None,
+        backend_overrides: dict[str, dict[str, Any]] | None = None,
+        **kwargs: Unpack[BatchTransformKwargs],
+    ) -> pd.Series | None:
+        """
+        Predict with a SageMaker batch transform job.
+
+        Same as :meth:`CloudPredictor.predict`, with one additional argument:
+
+        Parameters
+        ----------
+        test_data_image_column: str | None, default = None
+            Name of the column containing absolute paths to images, if the data contains images.
+        """
+        kwargs = check_backend_kwargs(kwargs, BatchTransformKwargs, "predict")
+        return self.backend.predict(
+            test_data=test_data,
+            test_data_image_column=test_data_image_column,
+            predictor_path=predictor_path,
+            framework_version=framework_version,
+            instance_type=instance_type,
+            wait=wait,
+            predictions_path=predictions_path,
+            backend_overrides=backend_overrides,
+            **kwargs,
+        )
+
+    @reject_legacy_kwargs
+    def predict_proba(
+        self,
+        test_data: str | pd.DataFrame,
+        *,
+        test_data_image_column: str | None = None,
+        include_predict: bool = True,
+        predictor_path: str | None = None,
+        framework_version: str | None = None,
+        instance_type: str = "ml.m5.2xlarge",
+        wait: bool = True,
+        predictions_path: str | None = None,
+        backend_overrides: dict[str, dict[str, Any]] | None = None,
+        **kwargs: Unpack[BatchTransformKwargs],
+    ) -> tuple[pd.Series, pd.DataFrame | pd.Series] | pd.DataFrame | pd.Series | None:
+        """
+        Predict class probabilities with a SageMaker batch transform job.
+
+        Same as :meth:`CloudPredictor.predict_proba`, with one additional argument:
+
+        Parameters
+        ----------
+        test_data_image_column: str | None, default = None
+            Name of the column containing absolute paths to images, if the data contains images.
+        """
+        kwargs = check_backend_kwargs(kwargs, BatchTransformKwargs, "predict_proba")
+        return self.backend.predict_proba(
+            test_data=test_data,
+            test_data_image_column=test_data_image_column,
+            include_predict=include_predict,
+            predictor_path=predictor_path,
+            framework_version=framework_version,
+            instance_type=instance_type,
+            wait=wait,
+            predictions_path=predictions_path,
+            backend_overrides=backend_overrides,
+            **kwargs,
+        )