diff --git a/docs/conf.py b/docs/conf.py index 712ac76..a78c0d3 100644 --- a/docs/conf.py +++ b/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 diff --git a/src/autogluon/cloud/backend/timeseries_sagemaker_backend.py b/src/autogluon/cloud/backend/timeseries_sagemaker_backend.py index 62b64e2..b5f8b94 100644 --- a/src/autogluon/cloud/backend/timeseries_sagemaker_backend.py +++ b/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, diff --git a/src/autogluon/cloud/endpoint/timeseries_endpoint.py b/src/autogluon/cloud/endpoint/timeseries_endpoint.py index c724983..3b420c3 100644 --- a/src/autogluon/cloud/endpoint/timeseries_endpoint.py +++ b/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, diff --git a/src/autogluon/cloud/model/foundation_model.py b/src/autogluon/cloud/model/foundation_model.py index 0fa3828..fa86372 100644 --- a/src/autogluon/cloud/model/foundation_model.py +++ b/src/autogluon/cloud/model/foundation_model.py @@ -12,7 +12,7 @@ import pandas as pd from packaging.version import InvalidVersion, Version -from typing_extensions import Self +from typing_extensions import Self, Unpack from autogluon.common.loaders import load_pd from autogluon.common.utils.s3_utils import s3_path_to_bucket_prefix @@ -25,7 +25,13 @@ from ..scripts.script_manager import ScriptManager from ..utils.aws_utils import resolve_cloud_output_path from ..utils.constants import DEFAULT_FRAMEWORK_VERSION -from ..utils.sagemaker_api import reject_legacy_kwargs +from ..utils.sagemaker_api import ( + IGNORED_FM_JOB_KWARGS, + DeployKwargs, + TrainingJobKwargs, + check_backend_kwargs, + reject_legacy_kwargs, +) from ..utils.utils import split_pred_and_pred_proba from ..version import __version__ from .registry import FOUNDATION_MODEL_REGISTRY, get_model_config @@ -231,14 +237,15 @@ def _deploy_backend( endpoint_name: str | None = None, hyperparameters: dict[str, Any] | None = None, framework_version: str = DEFAULT_FRAMEWORK_VERSION, - custom_image_uri: str | None = None, wait: bool = True, inference_mode: Literal["realtime", "serverless"] = "realtime", inference_config: dict[str, Any] | None = None, - **backend_kwargs, + backend_overrides: dict[str, dict[str, Any]] | None = None, + **kwargs, ) -> None: """Shared deploy logic. Subclasses call this then wrap the endpoint.""" - self._check_framework_version(framework_version, custom_image_uri) + kwargs = check_backend_kwargs(kwargs, DeployKwargs, "deploy") + self._check_framework_version(framework_version, kwargs.get("custom_image_uri")) if inference_mode == "serverless" and instance_type is not None: raise ValueError("`instance_type` must not be set when `inference_mode='serverless'`.") if instance_type is None and inference_mode == "realtime": @@ -268,15 +275,15 @@ def _deploy_backend( endpoint_name=endpoint_name, framework_version=framework_version, instance_type=instance_type, - custom_image_uri=custom_image_uri, wait=wait, + backend_overrides=backend_overrides, entry_point=self._serve_script_path, fm_serve_config=fm_serve_config, inference_mode=inference_mode, inference_config=inference_config, repack=False, extra_tags=[{"Key": "autogluon-cloud-model-id", "Value": self.model_id}], - **backend_kwargs, + **kwargs, ) assert self._backend.endpoint_name is not None @@ -467,15 +474,16 @@ def _serve_script_path(self) -> str: @reject_legacy_kwargs def deploy( self, + *, instance_type: str | None = None, endpoint_name: str | None = None, hyperparameters: dict[str, Any] | None = None, framework_version: str = DEFAULT_FRAMEWORK_VERSION, - custom_image_uri: str | None = None, wait: bool = True, inference_mode: Literal["realtime", "serverless"] = "realtime", inference_config: dict[str, Any] | None = None, - **backend_kwargs, + backend_overrides: dict[str, dict[str, Any]] | None = None, + **kwargs: Unpack[DeployKwargs], ) -> TimeSeriesEndpoint: """ Deploy the model to a real-time or serverless endpoint. @@ -491,8 +499,6 @@ def deploy( Model hyperparameters for inference. Overrides values passed to the constructor. framework_version: str, optional AutoGluon version, e.g. "1.6". Uses the official AutoGluon DLC image for this version. - custom_image_uri: str | None, default = None - Custom Docker image URI for the inference container. wait: bool, default = True Whether to block until the endpoint is ready. inference_mode: Literal["realtime", "serverless"], default = "realtime" @@ -501,24 +507,28 @@ def deploy( inference_config: dict[str, Any] | None, default = None Serverless settings: ``memory_size_in_mb`` (default 4096), ``max_concurrency`` (default 5), and ``provisioned_concurrency``. - **backend_kwargs: Any - Additional SageMaker arguments: - - * ``initial_instance_count``: Number of instances for the endpoint. Defaults to 1. Ignored when - ``inference_mode="serverless"``. - * ``volume_size``: Size in GB of the EBS volume to use for the endpoint. Ignored for GPU instances. - * ``backend_overrides``: raw SageMaker request fields for settings without a dedicated argument. - - * Keys: request names from the *SageMaker API* section below. - * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request - built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. - * Example: ``{"ProductionVariant": {"ModelDataDownloadTimeoutInSeconds": 1200}}`` Returns ------- TimeSeriesEndpoint Handle to the deployed endpoint. + Other Parameters + ---------------- + backend_overrides: dict[str, dict[str, Any]] | None, default = None + Raw SageMaker request fields for settings without a dedicated argument. + + * Keys: request names from the *SageMaker API* section below. + * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request + built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. + * Example: ``{"ProductionVariant": {"ModelDataDownloadTimeoutInSeconds": 1200}}`` + initial_instance_count: int, default = 1 + Number of instances for the endpoint. Ignored when ``inference_mode="serverless"``. + volume_size: int | None, default = None + Size in GB of the EBS volume to use for the endpoint. Ignored for GPU instances. + custom_image_uri: str | None, default = None + Custom inference container image URI. If set, ``framework_version`` is ignored. + SageMaker API ------------- * :sm-api:`CreateModel`: registers the model artifact and inference image as a SageMaker model. @@ -533,11 +543,11 @@ def deploy( endpoint_name=endpoint_name, hyperparameters=hyperparameters, framework_version=framework_version, - custom_image_uri=custom_image_uri, wait=wait, inference_mode=inference_mode, inference_config=inference_config, - **backend_kwargs, + backend_overrides=backend_overrides, + **kwargs, ) return TimeSeriesEndpoint( endpoint_name=self._backend.endpoint_name, @@ -571,6 +581,7 @@ def _build_predictor_init_args( def predict( self, data: str | Path | pd.DataFrame, + *, target: str = "target", id_column: str = "item_id", timestamp_column: str = "timestamp", @@ -582,9 +593,9 @@ def predict( hyperparameters: dict[str, Any] | None = None, instance_type: str | None = None, framework_version: str = DEFAULT_FRAMEWORK_VERSION, - custom_image_uri: str | None = None, wait: bool = True, - **backend_kwargs, + backend_overrides: dict[str, dict[str, Any]] | None = None, + **kwargs: Unpack[TrainingJobKwargs], ) -> pd.DataFrame | JobPredictionFuture: """ Forecast the future values of ``data`` in a one-off SageMaker job. @@ -621,23 +632,10 @@ def predict( Instance type for the prediction job. Defaults to the model registry value. framework_version: str, optional AutoGluon version, e.g. "1.6". Uses the official AutoGluon DLC image for this version. - custom_image_uri: str | None, default = None - Custom Docker image URI for the container. wait: bool, default = True If True, block until the job completes and return the forecasts. If False, return a :class:`~autogluon.cloud.endpoint.prediction_future.JobPredictionFuture` immediately; call its ``.status()`` to check progress and ``.result()`` to get the forecasts. - **backend_kwargs: Any - Additional SageMaker arguments: - - * ``job_name``: Name of the training job that runs the prediction. Auto-generated if not set. - * ``volume_size``: Size in GB of the storage volume to use for the job. Defaults to 100. - * ``backend_overrides``: raw SageMaker request fields for settings without a dedicated argument. - - * Keys: request names from the *SageMaker API* section below. - * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request - built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. - * Example: ``{"CreateTrainingJob": {"RetryStrategy": {"MaximumRetryAttempts": 2}}}`` Returns ------- @@ -646,12 +644,31 @@ def predict( level if ``wait=True``; a :class:`~autogluon.cloud.endpoint.prediction_future.JobPredictionFuture` otherwise. + Other Parameters + ---------------- + backend_overrides: dict[str, dict[str, Any]] | None, default = None + Raw SageMaker request fields for settings without a dedicated argument. + + * Keys: request names from the *SageMaker API* section below. + * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request + built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. + * Example: ``{"CreateTrainingJob": {"RetryStrategy": {"MaximumRetryAttempts": 2}}}`` + job_name: str | None, default = None + Name of the training job that runs the prediction. Auto-generated if not set. + volume_size: int, default = 100 + Size in GB of the storage volume to use for the job. + custom_image_uri: str | None, default = None + Custom container image URI. If set, ``framework_version`` is ignored. + timeout: int, default = 86400 + Maximum job runtime in seconds. Defaults to 24 hours. + SageMaker API ------------- * :sm-api:`CreateTrainingJob`: runs the prediction as a training job (not a batch transform job) on ``instance_type``. Predictions are written to ``predictions_path``. """ - self._check_framework_version(framework_version, custom_image_uri) + kwargs = check_backend_kwargs(kwargs, TrainingJobKwargs, "predict", IGNORED_FM_JOB_KWARGS) + self._check_framework_version(framework_version, kwargs.get("custom_image_uri")) if instance_type is None: instance_type = self._config.predict_instance_type @@ -680,11 +697,11 @@ def predict( timestamp_column=timestamp_column, framework_version=framework_version, instance_type=instance_type, - custom_image_uri=custom_image_uri, wait=wait, + backend_overrides=backend_overrides, extra_ag_args=extra_ag_args, extra_tags=[{"Key": "autogluon-cloud-model-id", "Value": self.model_id}], - **backend_kwargs, + **kwargs, ) if not wait: @@ -719,15 +736,16 @@ def _serve_script_path(self) -> str: @reject_legacy_kwargs def deploy( self, + *, instance_type: str | None = None, endpoint_name: str | None = None, hyperparameters: dict[str, Any] | None = None, framework_version: str = DEFAULT_FRAMEWORK_VERSION, - custom_image_uri: str | None = None, wait: bool = True, inference_mode: Literal["realtime"] = "realtime", inference_config: dict[str, Any] | None = None, - **backend_kwargs, + backend_overrides: dict[str, dict[str, Any]] | None = None, + **kwargs: Unpack[DeployKwargs], ) -> TabularEndpoint: """Deploy the model to a real-time endpoint. @@ -744,31 +762,34 @@ def deploy( Model hyperparameters for inference. Overrides values passed to the constructor. framework_version: str, optional AutoGluon version, e.g. "1.6". Uses the official AutoGluon DLC image for this version. - custom_image_uri: str | None, default = None - Custom Docker image URI for the inference container. wait: bool, default = True Whether to block until the endpoint is ready. inference_mode: Literal["realtime"], default = "realtime" Endpoint type. Only ``"realtime"`` is supported. inference_config: dict[str, Any] | None, default = None Not supported; must be None. - **backend_kwargs: Any - Additional SageMaker arguments: - - * ``initial_instance_count``: Number of instances for the endpoint. Defaults to 1. - * ``volume_size``: Size in GB of the EBS volume to use for the endpoint. Ignored for GPU instances. - * ``backend_overrides``: raw SageMaker request fields for settings without a dedicated argument. - - * Keys: request names from the *SageMaker API* section below. - * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request - built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. - * Example: ``{"ProductionVariant": {"ModelDataDownloadTimeoutInSeconds": 1200}}`` Returns ------- TabularEndpoint Handle to the deployed endpoint. + Other Parameters + ---------------- + backend_overrides: dict[str, dict[str, Any]] | None, default = None + Raw SageMaker request fields for settings without a dedicated argument. + + * Keys: request names from the *SageMaker API* section below. + * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request + built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. + * Example: ``{"ProductionVariant": {"ModelDataDownloadTimeoutInSeconds": 1200}}`` + initial_instance_count: int, default = 1 + Number of instances for the endpoint. + volume_size: int | None, default = None + Size in GB of the EBS volume to use for the endpoint. Ignored for GPU instances. + custom_image_uri: str | None, default = None + Custom inference container image URI. If set, ``framework_version`` is ignored. + SageMaker API ------------- * :sm-api:`CreateModel`: registers the model artifact and inference image as a SageMaker model. @@ -793,10 +814,10 @@ def deploy( endpoint_name=endpoint_name, hyperparameters=hyperparameters, framework_version=framework_version, - custom_image_uri=custom_image_uri, wait=wait, inference_mode="realtime", - **backend_kwargs, + backend_overrides=backend_overrides, + **kwargs, ) return TabularEndpoint( endpoint_name=self._backend.endpoint_name, @@ -840,9 +861,9 @@ def predict( hyperparameters: dict[str, Any] | None = None, instance_type: str | None = None, framework_version: str = DEFAULT_FRAMEWORK_VERSION, - custom_image_uri: str | None = None, wait: bool = True, - **backend_kwargs, + backend_overrides: dict[str, dict[str, Any]] | None = None, + **kwargs: Unpack[TrainingJobKwargs], ) -> pd.Series | JobPredictionFuture: """ Predict ``test_data`` in a one-off SageMaker job, using the labeled ``train_data`` as in-context examples. @@ -865,23 +886,10 @@ def predict( Instance type for the prediction job. Defaults to the model registry value. framework_version: str, optional AutoGluon version, e.g. "1.6". Uses the official AutoGluon DLC image for this version. - custom_image_uri: str | None, default = None - Custom Docker image URI for the container. wait: bool, default = True If True, block until the job completes and return the predictions. If False, return a :class:`~autogluon.cloud.endpoint.prediction_future.JobPredictionFuture` immediately; call its ``.status()`` to check progress and ``.result()`` to get the predictions. - **backend_kwargs: Any - Additional SageMaker arguments: - - * ``job_name``: Name of the training job that runs the prediction. Auto-generated if not set. - * ``volume_size``: Size in GB of the storage volume to use for the job. Defaults to 100. - * ``backend_overrides``: raw SageMaker request fields for settings without a dedicated argument. - - * Keys: request names from the *SageMaker API* section below. - * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request - built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. - * Example: ``{"CreateTrainingJob": {"RetryStrategy": {"MaximumRetryAttempts": 2}}}`` Returns ------- @@ -889,11 +897,31 @@ def predict( Predictions if ``wait=True``; a :class:`~autogluon.cloud.endpoint.prediction_future.JobPredictionFuture` otherwise. + Other Parameters + ---------------- + backend_overrides: dict[str, dict[str, Any]] | None, default = None + Raw SageMaker request fields for settings without a dedicated argument. + + * Keys: request names from the *SageMaker API* section below. + * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request + built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. + * Example: ``{"CreateTrainingJob": {"RetryStrategy": {"MaximumRetryAttempts": 2}}}`` + job_name: str | None, default = None + Name of the training job that runs the prediction. Auto-generated if not set. + volume_size: int, default = 100 + Size in GB of the storage volume to use for the job. + custom_image_uri: str | None, default = None + Custom container image URI. If set, ``framework_version`` is ignored. + timeout: int, default = 86400 + Maximum job runtime in seconds. Defaults to 24 hours. + SageMaker API ------------- * :sm-api:`CreateTrainingJob`: runs the prediction as a training job (not a batch transform job) on ``instance_type``. Predictions are written to ``predictions_path``. """ + # Checked here too so warnings and errors name `predict()`; the check in `predict_proba()` is then a no-op. + kwargs = check_backend_kwargs(kwargs, TrainingJobKwargs, "predict", IGNORED_FM_JOB_KWARGS) result = self.predict_proba( test_data, train_data, @@ -903,9 +931,9 @@ def predict( hyperparameters=hyperparameters, instance_type=instance_type, framework_version=framework_version, - custom_image_uri=custom_image_uri, wait=wait, - **backend_kwargs, + backend_overrides=backend_overrides, + **kwargs, ) if not wait: return JobPredictionFuture( @@ -927,9 +955,9 @@ def predict_proba( hyperparameters: dict[str, Any] | None = None, instance_type: str | None = None, framework_version: str = DEFAULT_FRAMEWORK_VERSION, - custom_image_uri: str | None = None, wait: bool = True, - **backend_kwargs, + backend_overrides: dict[str, dict[str, Any]] | None = None, + **kwargs: Unpack[TrainingJobKwargs], ) -> tuple[pd.Series, pd.DataFrame | pd.Series] | pd.DataFrame | pd.Series | JobPredictionFuture: """ Predict class probabilities for ``test_data`` in a one-off SageMaker job, using the labeled ``train_data`` as @@ -957,23 +985,10 @@ def predict_proba( Instance type for the prediction job. Defaults to the model registry value. framework_version: str, optional AutoGluon version, e.g. "1.6". Uses the official AutoGluon DLC image for this version. - custom_image_uri: str | None, default = None - Custom Docker image URI for the container. wait: bool, default = True If True, block until the job completes and return the result. If False, return a :class:`~autogluon.cloud.endpoint.prediction_future.JobPredictionFuture` immediately; call its ``.status()`` to check progress and ``.result()`` to get the result. - **backend_kwargs: Any - Additional SageMaker arguments: - - * ``job_name``: Name of the training job that runs the prediction. Auto-generated if not set. - * ``volume_size``: Size in GB of the storage volume to use for the job. Defaults to 100. - * ``backend_overrides``: raw SageMaker request fields for settings without a dedicated argument. - - * Keys: request names from the *SageMaker API* section below. - * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request - built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. - * Example: ``{"CreateTrainingJob": {"RetryStrategy": {"MaximumRetryAttempts": 2}}}`` Returns ------- @@ -981,12 +996,31 @@ def predict_proba( ``(prediction, predict_probability)`` if ``include_predict`` is True, otherwise ``predict_probability``. A :class:`~autogluon.cloud.endpoint.prediction_future.JobPredictionFuture` if ``wait=False``. + Other Parameters + ---------------- + backend_overrides: dict[str, dict[str, Any]] | None, default = None + Raw SageMaker request fields for settings without a dedicated argument. + + * Keys: request names from the *SageMaker API* section below. + * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request + built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. + * Example: ``{"CreateTrainingJob": {"RetryStrategy": {"MaximumRetryAttempts": 2}}}`` + job_name: str | None, default = None + Name of the training job that runs the prediction. Auto-generated if not set. + volume_size: int, default = 100 + Size in GB of the storage volume to use for the job. + custom_image_uri: str | None, default = None + Custom container image URI. If set, ``framework_version`` is ignored. + timeout: int, default = 86400 + Maximum job runtime in seconds. Defaults to 24 hours. + SageMaker API ------------- * :sm-api:`CreateTrainingJob`: runs the prediction as a training job (not a batch transform job) on ``instance_type``. Predictions are written to ``predictions_path``. """ - self._check_framework_version(framework_version, custom_image_uri) + kwargs = check_backend_kwargs(kwargs, TrainingJobKwargs, "predict_proba", IGNORED_FM_JOB_KWARGS) + self._check_framework_version(framework_version, kwargs.get("custom_image_uri")) if instance_type is None: instance_type = self._config.predict_instance_type @@ -999,7 +1033,7 @@ def predict_proba( extra_ag_args: dict[str, Any] = {"predict_after_fit": True, "save_predictor": False} if predictions_path is not None: extra_ag_args["predictions_path"] = predictions_path - backend_kwargs["leaderboard"] = False + kwargs["leaderboard"] = False self._backend.fit( predictor_init_args=self._build_predictor_init_args(label=label), @@ -1007,11 +1041,11 @@ def predict_proba( data_channels={"train_data": train_data, "tuning_data": tuning_data, "test_data": test_data}, framework_version=framework_version, instance_type=instance_type, - custom_image_uri=custom_image_uri, wait=wait, + backend_overrides=backend_overrides, extra_ag_args=extra_ag_args, extra_tags=[{"Key": "autogluon-cloud-model-id", "Value": self.model_id}], - **backend_kwargs, + **kwargs, ) if not wait: diff --git a/src/autogluon/cloud/predictor/cloud_predictor.py b/src/autogluon/cloud/predictor/cloud_predictor.py index 6f3fa9b..1e20028 100644 --- a/src/autogluon/cloud/predictor/cloud_predictor.py +++ b/src/autogluon/cloud/predictor/cloud_predictor.py @@ -13,7 +13,7 @@ import boto3 import pandas as pd -from typing_extensions import deprecated +from typing_extensions import Self, Unpack, deprecated from autogluon.common.loaders import load_pkl from autogluon.common.savers import save_pkl @@ -26,8 +26,15 @@ from ..backend.constant import SAGEMAKER from ..endpoint.endpoint import Endpoint from ..utils.aws_utils import resolve_cloud_output_path -from ..utils.constants import DEFAULT_FRAMEWORK_VERSION, DEFAULT_VOLUME_SIZE -from ..utils.sagemaker_api import reject_legacy_kwargs +from ..utils.constants import DEFAULT_FRAMEWORK_VERSION +from ..utils.sagemaker_api import ( + IGNORED_TRAINING_KWARGS, + BatchTransformKwargs, + DeployKwargs, + TrainingJobKwargs, + check_backend_kwargs, + reject_legacy_kwargs, +) from ..utils.utils import safe_unpack_archive logger = logging.getLogger(__name__) @@ -39,6 +46,15 @@ _DEPRECATED_REAL_TIME = "Deploy an endpoint with `deploy()` and call `predict()` / `predict_proba()` on it instead." +def _reject_image_column(kwargs: dict[str, Any], name: str) -> None: + """Raise a pointer to MultiModalCloudPredictor if ``name`` is passed to a predictor without image support.""" + if name in kwargs: + raise ValueError( + f"`{name}` is no longer supported for tabular predictors: image models in AutoGluon-Tabular require " + "autogluon.multimodal, which is being deprecated. Use `MultiModalCloudPredictor` for image data." + ) + + class CloudPredictor(ABC, Generic[EndpointT]): predictor_file_name = "CloudPredictor.pkl" backend_map = {} @@ -182,19 +198,12 @@ def fit( 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, - leaderboard: bool = True, framework_version: str = DEFAULT_FRAMEWORK_VERSION, - job_name: str | None = None, instance_type: str = "ml.m5.2xlarge", - instance_count: int | str = "auto", - 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, - **kwargs, - ) -> CloudPredictor: + **kwargs: Unpack[TrainingJobKwargs], + ) -> Self: """ Fit the predictor in a SageMaker training job. @@ -209,27 +218,22 @@ def fit( predictor_fit_args: dict | None, 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. - leaderboard: bool, default = True - Whether to save the leaderboard in the output artifact, so that :meth:`leaderboard` can read it. framework_version: str, optional AutoGluon version, e.g. ``"1.6"``. Training uses the official AutoGluon DLC image for this version. Ignored if ``custom_image_uri`` is set. - job_name: str | None, default = None - Name of the training job. If ``None``, a unique name with prefix ``ag-cloud-`` is - generated. instance_type: str, default = "ml.m5.2xlarge" Instance type of the training job. - instance_count: int | str, default = "auto" - Number of training instances. Only single-instance training is supported, so this is always 1. - volume_size: int, default = 100 - Size in GB of the EBS volume that stores the training data and model artifacts. - custom_image_uri: str | None, default = None - Custom training container image URI. If set, ``framework_version`` is ignored. - timeout: int, default = 24*60*60 - Maximum training job runtime in seconds. wait: bool, default = True Whether to block until the job completes. If ``False``, returns once the job is launched; use :meth:`get_fit_job_status` to poll it. + + Returns + ------- + CloudPredictor + The fitted predictor (``self``). + + Other Parameters + ---------------- backend_overrides: dict[str, dict[str, Any]] | None, default = None Raw SageMaker request fields for settings without a dedicated argument. @@ -237,62 +241,65 @@ def fit( * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. * Example: ``{"CreateTrainingJob": {"RetryStrategy": {"MaximumRetryAttempts": 2}}}`` - - Returns - ------- - CloudPredictor - The fitted predictor (``self``). + job_name: str | None, default = None + Name of the training job. Defaults to a unique name with prefix ``ag-cloud-``. + volume_size: int, default = 100 + Size in GB of the EBS volume that stores the training data and model artifacts. + custom_image_uri: str | None, default = None + Custom training container image URI. If set, ``framework_version`` is ignored. + timeout: int, default = 86400 + Maximum training job runtime in seconds. Defaults to 24 hours. SageMaker API ------------- - * :sm-api:`CreateTrainingJob`: trains the predictor on ``instance_count`` x ``instance_type`` and writes the - artifact to ``cloud_output_path``. + * :sm-api:`CreateTrainingJob`: trains the predictor on ``instance_type`` and writes the artifact to + ``cloud_output_path``. """ # noqa: E501 + _reject_image_column(kwargs, "image_column") + 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, + framework_version=framework_version, + instance_type=instance_type, + wait=wait, + backend_overrides=backend_overrides, + **kwargs, + ) + return self + + def _fit( + self, + *, + data_channels: dict[str, str | Path | pd.DataFrame | None], + predictor_init_args: dict[str, Any], + predictor_fit_args: dict[str, Any] | None, + **backend_args, + ) -> None: + """Validate the inputs shared by ``fit()`` and ``fit_predict()``, then launch the training job. + + ``backend_args`` are forwarded to ``backend.fit()``, including internal ones such as ``extra_ag_args``. + """ assert not self.backend.is_fit, ( "Predictor is already fit! To fit additional models, create a new `CloudPredictor`" ) - # `test_data` / `extra_ag_args` are internal channels for the fit_predict path (see - # TabularCloudPredictor.fit_predict_proba); they are intentionally not part of the public `fit()` signature. - test_data = kwargs.pop("test_data", None) - extra_ag_args = kwargs.pop("extra_ag_args", None) - if kwargs: - raise TypeError(f"fit() got unexpected keyword arguments: {sorted(kwargs)}") predictor_fit_args = {} if predictor_fit_args is None else dict(predictor_fit_args) - data_channels = {"train_data": train_data, "tuning_data": tuning_data, "test_data": test_data} - for key in ("train_data", "tuning_data"): + for key in data_channels: if key in predictor_fit_args: raise TypeError( f"`{key}` can no longer be passed via `predictor_fit_args`. " - f"Pass `{key}` as an explicit argument to `fit()` instead." + f"Pass `{key}` as an explicit argument instead." ) 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." - ) + raise TypeError("missing required argument: 'train_data'") self.backend.fit( predictor_init_args=predictor_init_args, predictor_fit_args=predictor_fit_args, data_channels=data_channels, - image_column=image_column, - leaderboard=leaderboard, - framework_version=framework_version, - job_name=job_name, - instance_type=instance_type, - 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, + **backend_args, ) - return self - def attach_job(self, job_name: str) -> None: """ Attach to an existing SageMaker training job, e.g. after the local process that launched it crashed. @@ -393,17 +400,16 @@ def to_local_predictor(self, predictor_path: str | None = None, save_path: str | @reject_legacy_kwargs def deploy( self, + *, predictor_path: str | None = None, endpoint_name: str | None = None, framework_version: str | None = None, instance_type: str | None = None, - initial_instance_count: int = 1, - custom_image_uri: str | None = None, - volume_size: int | None = None, wait: bool = True, inference_mode: Literal["realtime", "serverless"] = "realtime", inference_config: dict[str, Any] | None = None, backend_overrides: dict[str, dict[str, Any]] | None = None, + **kwargs: Unpack[DeployKwargs], ) -> EndpointT: """ Deploy a predictor to a real-time inference endpoint. @@ -424,12 +430,6 @@ def deploy( instance_type: str | None, default = None Instance type of the endpoint. Defaults to ``ml.m5.2xlarge``. Must be ``None`` when ``inference_mode="serverless"``. - initial_instance_count: int, default = 1 - Initial number of endpoint instances. Ignored when ``inference_mode="serverless"``. - custom_image_uri: str | None, default = None - Custom inference container image URI. If set, ``framework_version`` is ignored. - volume_size: int | None, default = None - Size in GB of the endpoint's EBS volume. Ignored for GPU instances (``ml.g*`` / ``ml.p*``). wait: bool, default = True Whether to block until the endpoint is in service. inference_mode: {"realtime", "serverless"}, default = "realtime" @@ -437,6 +437,14 @@ def deploy( (no instance management, scales to zero). inference_config: dict[str, Any] | None, default = None Serverless settings (``memory_size_in_mb``, ``max_concurrency``, ``provisioned_concurrency``). + + Returns + ------- + TabularEndpoint | TimeSeriesEndpoint | MultiModalEndpoint + Handle to the deployed endpoint, matching the predictor type. + + Other Parameters + ---------------- backend_overrides: dict[str, dict[str, Any]] | None, default = None Raw SageMaker request fields for settings without a dedicated argument. @@ -444,11 +452,12 @@ def deploy( * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. * Example: ``{"ProductionVariant": {"ModelDataDownloadTimeoutInSeconds": 1200}}`` - - Returns - ------- - TabularEndpoint | TimeSeriesEndpoint | MultiModalEndpoint - Handle to the deployed endpoint, matching the predictor type. + initial_instance_count: int, default = 1 + Initial number of endpoint instances. Ignored when ``inference_mode="serverless"``. + volume_size: int | None, default = None + Size in GB of the endpoint's EBS volume. Ignored for GPU instances (``ml.g*`` / ``ml.p*``). + custom_image_uri: str | None, default = None + Custom inference container image URI. If set, ``framework_version`` is ignored. SageMaker API ------------- @@ -463,19 +472,18 @@ def deploy( raise ValueError("`instance_type` must not be set when `inference_mode='serverless'`.") if instance_type is None and inference_mode == "realtime": instance_type = "ml.m5.2xlarge" + kwargs = check_backend_kwargs(kwargs, DeployKwargs, "deploy") self._warn_if_endpoint_active() self.backend.deploy( predictor_path=predictor_path, endpoint_name=endpoint_name, framework_version=framework_version, instance_type=instance_type, - initial_instance_count=initial_instance_count, - custom_image_uri=custom_image_uri, - volume_size=volume_size, wait=wait, inference_mode=inference_mode, inference_config=inference_config, backend_overrides=backend_overrides, + **kwargs, ) return self._endpoint_cls( endpoint_name=self.backend.endpoint_name, @@ -661,16 +669,14 @@ def predict_proba_real_time( 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, - job_name: str | None = None, instance_type: str = "ml.m5.2xlarge", - instance_count: int = 1, - custom_image_uri: str | None = None, 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. @@ -681,28 +687,27 @@ def predict( ---------- test_data: str | pd.DataFrame Data to predict on, as a ``pd.DataFrame`` or local path to a data file. - test_data_image_column: str | None, default = None - Name of the column containing absolute paths to images, if the data contains images. predictor_path: str | None, default = None Local or S3 path of the predictor tarball. If ``None``, uses the predictor trained by :meth:`fit`. framework_version: str | None, default = None AutoGluon version, e.g. ``"1.6"``. Inference uses the official AutoGluon DLC image for this version. Defaults to the version used by :meth:`fit`. Ignored if ``custom_image_uri`` is set. - job_name: str | None, default = None - Name of the batch transform job. If ``None``, a unique name with prefix ``ag-cloud-`` is - generated. instance_type: str, default = "ml.m5.2xlarge" Instance type of the batch transform job. - instance_count: int, default = 1 - Number of batch transform instances. - custom_image_uri: str | None, default = None - Custom inference container image URI. If set, ``framework_version`` is ignored. wait: bool, default = True Whether to block until the job completes and return the predictions. If ``False``, returns ``None`` once the job is launched; use :meth:`get_batch_inference_job_status` to poll it. predictions_path: str | None, default = None S3 prefix under which the batch transform job writes its results (``/.out``). Defaults to ``{cloud_output_path}/batch_transform//results``. + + Returns + ------- + pd.Series | None + Predictions, or ``None`` if ``wait=False``. + + Other Parameters + ---------------- backend_overrides: dict[str, dict[str, Any]] | None, default = None Raw SageMaker request fields for settings without a dedicated argument. @@ -710,11 +715,12 @@ def predict( * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. * Example: ``{"CreateTransformJob": {"BatchStrategy": "SingleRecord", "MaxPayloadInMB": 20}}`` - - Returns - ------- - pd.Series | None - Predictions, or ``None`` if ``wait=False``. + job_name: str | None, default = None + Name of the batch transform job. Defaults to a unique name with prefix ``ag-cloud-``. + instance_count: int, default = 1 + Number of batch transform instances. + custom_image_uri: str | None, default = None + Custom inference container image URI. If set, ``framework_version`` is ignored. SageMaker API ------------- @@ -725,35 +731,32 @@ def predict( The model is deleted when the job finishes. With ``wait=False`` it is kept; delete it with :sm-api:`DeleteModel`. """ + _reject_image_column(kwargs, "test_data_image_column") + 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, - job_name=job_name, instance_type=instance_type, - instance_count=instance_count, - custom_image_uri=custom_image_uri, 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, - job_name: str | None = None, instance_type: str = "ml.m5.2xlarge", - instance_count: int = 1, - custom_image_uri: str | None = None, 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. @@ -765,8 +768,6 @@ def predict_proba( ---------- test_data: str | pd.DataFrame Data to predict on, as a ``pd.DataFrame`` or local path to a data file. - test_data_image_column: str | None, default = None - Name of the column containing absolute paths to images, if the data contains images. include_predict: bool, default = True Whether to also return the predictions, saving a second batch transform job. predictor_path: str | None, default = None @@ -774,28 +775,14 @@ def predict_proba( framework_version: str | None, default = None AutoGluon version, e.g. ``"1.6"``. Inference uses the official AutoGluon DLC image for this version. Defaults to the version used by :meth:`fit`. Ignored if ``custom_image_uri`` is set. - job_name: str | None, default = None - Name of the batch transform job. If ``None``, a unique name with prefix ``ag-cloud-`` is - generated. instance_type: str, default = "ml.m5.2xlarge" Instance type of the batch transform job. - instance_count: int, default = 1 - Number of batch transform instances. - custom_image_uri: str | None, default = None - Custom inference container image URI. If set, ``framework_version`` is ignored. wait: bool, default = True Whether to block until the job completes and return the predictions. If ``False``, returns ``None`` once the job is launched; use :meth:`get_batch_inference_job_status` to poll it. predictions_path: str | None, default = None S3 prefix under which the batch transform job writes its results (``/.out``). Defaults to ``{cloud_output_path}/batch_transform//results``. - backend_overrides: dict[str, dict[str, Any]] | None, default = None - Raw SageMaker request fields for settings without a dedicated argument. - - * Keys: request names from the *SageMaker API* section below. - * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request - built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. - * Example: ``{"CreateTransformJob": {"BatchStrategy": "SingleRecord", "MaxPayloadInMB": 20}}`` Returns ------- @@ -804,6 +791,22 @@ def predict_proba( ``predict_probability`` is a ``pd.DataFrame`` with one column per class, or a ``pd.Series`` for regression. If ``wait=False``, returns ``(None, None)`` or ``None``, respectively. + Other Parameters + ---------------- + backend_overrides: dict[str, dict[str, Any]] | None, default = None + Raw SageMaker request fields for settings without a dedicated argument. + + * Keys: request names from the *SageMaker API* section below. + * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request + built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. + * Example: ``{"CreateTransformJob": {"BatchStrategy": "SingleRecord", "MaxPayloadInMB": 20}}`` + job_name: str | None, default = None + Name of the batch transform job. Defaults to a unique name with prefix ``ag-cloud-``. + instance_count: int, default = 1 + Number of batch transform instances. + custom_image_uri: str | None, default = None + Custom inference container image URI. If set, ``framework_version`` is ignored. + SageMaker API ------------- * :sm-api:`CreateModel`: registers the predictor artifact and inference image as a SageMaker model. @@ -813,19 +816,18 @@ def predict_proba( The model is deleted when the job finishes. With ``wait=False`` it is kept; delete it with :sm-api:`DeleteModel`. """ + _reject_image_column(kwargs, "test_data_image_column") + 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, - job_name=job_name, instance_type=instance_type, - instance_count=instance_count, - custom_image_uri=custom_image_uri, wait=wait, predictions_path=predictions_path, backend_overrides=backend_overrides, + **kwargs, ) def get_batch_inference_job_info(self, job_name: str | None = None) -> dict[str, Any]: diff --git a/src/autogluon/cloud/predictor/multimodal_cloud_predictor.py b/src/autogluon/cloud/predictor/multimodal_cloud_predictor.py index 72efc21..bbd4f90 100644 --- a/src/autogluon/cloud/predictor/multimodal_cloud_predictor.py +++ b/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, + ) diff --git a/src/autogluon/cloud/predictor/tabular_cloud_predictor.py b/src/autogluon/cloud/predictor/tabular_cloud_predictor.py index a7fd8da..602a1d1 100644 --- a/src/autogluon/cloud/predictor/tabular_cloud_predictor.py +++ b/src/autogluon/cloud/predictor/tabular_cloud_predictor.py @@ -5,11 +5,17 @@ from typing import Any import pandas as pd +from typing_extensions import Unpack from ..backend.constant import SAGEMAKER, TABULAR_SAGEMAKER from ..endpoint.tabular_endpoint import TabularEndpoint -from ..utils.constants import DEFAULT_FRAMEWORK_VERSION, DEFAULT_VOLUME_SIZE -from ..utils.sagemaker_api import reject_legacy_kwargs +from ..utils.constants import DEFAULT_FRAMEWORK_VERSION +from ..utils.sagemaker_api import ( + IGNORED_TRAINING_KWARGS, + TrainingJobKwargs, + check_backend_kwargs, + reject_legacy_kwargs, +) from ..utils.utils import split_pred_and_pred_proba from .cloud_predictor import CloudPredictor @@ -48,16 +54,12 @@ def fit_predict( *, predictor_init_args: dict[str, Any], predictor_fit_args: dict[str, Any] | None = None, - leaderboard: bool = True, framework_version: str = DEFAULT_FRAMEWORK_VERSION, - job_name: str | None = None, instance_type: str = "ml.m5.2xlarge", - instance_count: int = 1, - volume_size: int = DEFAULT_VOLUME_SIZE, - custom_image_uri: str | None = None, wait: bool = True, predictions_path: str | None = None, backend_overrides: dict[str, dict[str, Any]] | None = None, + **kwargs: Unpack[TrainingJobKwargs], ) -> pd.Series | None: """ Fit and predict in a single SageMaker training job. @@ -78,26 +80,24 @@ def fit_predict( predictor_fit_args: dict | None, default = None Additional fit args forwarded to ``TabularPredictor.fit()``. Must NOT contain ``train_data`` or ``tuning_data``. - leaderboard: bool, default = True - Whether to include the leaderboard in the output artifact. framework_version: str, optional AutoGluon version, e.g. ``"1.6"``. Training uses the official AutoGluon DLC image for this version. Ignored if ``custom_image_uri`` is set. - job_name: str | None, default = None - Name of the training job. If ``None``, a unique name with prefix ``ag-cloud-tabular`` is generated. instance_type: str, default = "ml.m5.2xlarge" Instance type of the training job. - instance_count: int, default = 1 - Number of training instances. Only single-instance training is supported. - volume_size: int, default = 100 - Size in GB of the EBS volume that stores the training data and model artifacts. - custom_image_uri: str | None, default = None - Custom training container image URI. If set, ``framework_version`` is ignored. wait: bool, default = True Whether to block until the job completes. If ``False``, returns ``None`` once the job is launched. predictions_path: str | None, default = None S3 URL of the predictions file, ending in ``.csv`` or ``.parquet``. Defaults to ``{cloud_output_path}/{job_name}/predictions.csv``. + + Returns + ------- + pd.Series | None + Predictions, or ``None`` if ``wait=False``; fetch them later with :meth:`get_fit_predict_results`. + + Other Parameters + ---------------- backend_overrides: dict[str, dict[str, Any]] | None, default = None Raw SageMaker request fields for settings without a dedicated argument. @@ -105,33 +105,34 @@ def fit_predict( * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. * Example: ``{"CreateTrainingJob": {"RetryStrategy": {"MaximumRetryAttempts": 2}}}`` - - Returns - ------- - pd.Series | None - Predictions, or ``None`` if ``wait=False``; fetch them later with :meth:`get_fit_predict_results`. + job_name: str | None, default = None + Name of the training job. Defaults to a unique name with prefix ``ag-cloud-tabular``. + volume_size: int, default = 100 + Size in GB of the EBS volume that stores the training data and model artifacts. + custom_image_uri: str | None, default = None + Custom training container image URI. If set, ``framework_version`` is ignored. + timeout: int, default = 86400 + Maximum training job runtime in seconds. Defaults to 24 hours. SageMaker API ------------- - * :sm-api:`CreateTrainingJob`: trains the predictor and predicts in the same job on ``instance_count`` x - ``instance_type``. Predictions are written to ``predictions_path``. + * :sm-api:`CreateTrainingJob`: trains the predictor and predicts in the same job on ``instance_type``. + Predictions are written to ``predictions_path``. """ + # Checked here too so warnings and errors name `fit_predict()`; the check in `fit_predict_proba()` is then a no-op. + kwargs = check_backend_kwargs(kwargs, TrainingJobKwargs, "fit_predict", IGNORED_TRAINING_KWARGS) result = self.fit_predict_proba( train_data=train_data, test_data=test_data, predictor_init_args=predictor_init_args, predictor_fit_args=predictor_fit_args, include_predict=True, - leaderboard=leaderboard, framework_version=framework_version, - job_name=job_name, instance_type=instance_type, - instance_count=instance_count, - volume_size=volume_size, - custom_image_uri=custom_image_uri, wait=wait, predictions_path=predictions_path, backend_overrides=backend_overrides, + **kwargs, ) if result is None: # wait=False return None @@ -147,16 +148,12 @@ def fit_predict_proba( predictor_init_args: dict[str, Any], predictor_fit_args: dict[str, Any] | None = None, include_predict: bool = True, - leaderboard: bool = True, framework_version: str = DEFAULT_FRAMEWORK_VERSION, - job_name: str | None = None, instance_type: str = "ml.m5.2xlarge", - instance_count: int = 1, - volume_size: int = DEFAULT_VOLUME_SIZE, - custom_image_uri: str | None = None, wait: bool = True, predictions_path: str | None = None, backend_overrides: dict[str, dict[str, Any]] | None = None, + **kwargs: Unpack[TrainingJobKwargs], ) -> tuple[pd.Series, pd.DataFrame | pd.Series] | pd.DataFrame | pd.Series | None: """ Fit and predict probabilities in a single SageMaker training job. @@ -178,26 +175,25 @@ def fit_predict_proba( ``tuning_data``. include_predict: bool, default = True Whether to also return the predictions. The job always computes both, so this adds no cost. - leaderboard: bool, default = True - Whether to include the leaderboard in the output artifact. framework_version: str, optional AutoGluon version, e.g. ``"1.6"``. Training uses the official AutoGluon DLC image for this version. Ignored if ``custom_image_uri`` is set. - job_name: str | None, default = None - Name of the training job. If ``None``, a unique name with prefix ``ag-cloud-tabular`` is generated. instance_type: str, default = "ml.m5.2xlarge" Instance type of the training job. - instance_count: int, default = 1 - Number of training instances. Only single-instance training is supported. - volume_size: int, default = 100 - Size in GB of the EBS volume that stores the training data and model artifacts. - custom_image_uri: str | None, default = None - Custom training container image URI. If set, ``framework_version`` is ignored. wait: bool, default = True Whether to block until the job completes. If ``False``, returns ``None`` once the job is launched. predictions_path: str | None, default = None S3 URL of the predictions file, ending in ``.csv`` or ``.parquet``. Defaults to ``{cloud_output_path}/{job_name}/predictions.csv``. + + Returns + ------- + tuple[pd.Series, pd.DataFrame | pd.Series] | pd.DataFrame | pd.Series | None + ``(prediction, predict_probability)`` if ``include_predict=True``, otherwise ``predict_probability``. + Returns ``None`` if ``wait=False``; fetch them later with :meth:`get_fit_predict_proba_results`. + + Other Parameters + ---------------- backend_overrides: dict[str, dict[str, Any]] | None, default = None Raw SageMaker request fields for settings without a dedicated argument. @@ -205,37 +201,35 @@ def fit_predict_proba( * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. * Example: ``{"CreateTrainingJob": {"RetryStrategy": {"MaximumRetryAttempts": 2}}}`` - - Returns - ------- - tuple[pd.Series, pd.DataFrame | pd.Series] | pd.DataFrame | pd.Series | None - ``(prediction, predict_probability)`` if ``include_predict=True``, otherwise ``predict_probability``. - Returns ``None`` if ``wait=False``; fetch them later with :meth:`get_fit_predict_proba_results`. + job_name: str | None, default = None + Name of the training job. Defaults to a unique name with prefix ``ag-cloud-tabular``. + volume_size: int, default = 100 + Size in GB of the EBS volume that stores the training data and model artifacts. + custom_image_uri: str | None, default = None + Custom training container image URI. If set, ``framework_version`` is ignored. + timeout: int, default = 86400 + Maximum training job runtime in seconds. Defaults to 24 hours. SageMaker API ------------- - * :sm-api:`CreateTrainingJob`: trains the predictor and predicts in the same job on ``instance_count`` x - ``instance_type``. Predictions are written to ``predictions_path``. + * :sm-api:`CreateTrainingJob`: trains the predictor and predicts in the same job on ``instance_type``. + Predictions are written to ``predictions_path``. """ + kwargs = check_backend_kwargs(kwargs, TrainingJobKwargs, "fit_predict_proba", IGNORED_TRAINING_KWARGS) extra_ag_args = {"predict_after_fit": True} if predictions_path is not None: extra_ag_args["predictions_path"] = predictions_path - self.fit( - train_data=train_data, - test_data=test_data, + self._fit( + data_channels={"train_data": train_data, "tuning_data": None, "test_data": test_data}, predictor_init_args=predictor_init_args, predictor_fit_args=predictor_fit_args, - leaderboard=leaderboard, framework_version=framework_version, - job_name=job_name, instance_type=instance_type, - instance_count=instance_count, - volume_size=volume_size, - custom_image_uri=custom_image_uri, wait=wait, backend_overrides=backend_overrides, extra_ag_args=extra_ag_args, + **kwargs, ) if not wait: diff --git a/src/autogluon/cloud/predictor/timeseries_cloud_predictor.py b/src/autogluon/cloud/predictor/timeseries_cloud_predictor.py index 59ed148..2d39bfe 100644 --- a/src/autogluon/cloud/predictor/timeseries_cloud_predictor.py +++ b/src/autogluon/cloud/predictor/timeseries_cloud_predictor.py @@ -5,12 +5,18 @@ from typing import Any import pandas as pd -from typing_extensions import deprecated +from typing_extensions import Self, Unpack, deprecated from ..backend.constant import SAGEMAKER, TIMESERIES_SAGEMAKER from ..endpoint.timeseries_endpoint import TimeSeriesEndpoint -from ..utils.constants import DEFAULT_FRAMEWORK_VERSION, DEFAULT_VOLUME_SIZE -from ..utils.sagemaker_api import reject_legacy_kwargs +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 _DEPRECATED_REAL_TIME, CloudPredictor logger = logging.getLogger(__name__) @@ -47,20 +53,16 @@ def fit( predictor_init_args: dict[str, Any], predictor_fit_args: dict[str, Any] | None = None, tuning_data: str | Path | pd.DataFrame | None = None, + known_covariates: str | Path | pd.DataFrame | None = None, static_features: str | Path | pd.DataFrame | None = None, id_column: str = "item_id", timestamp_column: str = "timestamp", framework_version: str = DEFAULT_FRAMEWORK_VERSION, - job_name: str | None = None, instance_type: str = "ml.m5.2xlarge", - instance_count: int = 1, - volume_size: int = DEFAULT_VOLUME_SIZE, - custom_image_uri: str | None = None, wait: bool = True, backend_overrides: dict[str, dict[str, Any]] | None = None, - known_covariates: str | Path | pd.DataFrame | None = None, - **kwargs, - ) -> TimeSeriesCloudPredictor: + **kwargs: Unpack[TrainingJobKwargs], + ) -> Self: """ Fit the predictor in a SageMaker training job. @@ -94,19 +96,19 @@ def fit( framework_version: str, optional AutoGluon version, e.g. ``"1.6"``. Training uses the official AutoGluon DLC image for this version. Ignored if ``custom_image_uri`` is set. - job_name: str | None, default = None - Name of the training job. If ``None``, a unique name with prefix ``ag-cloud-timeseries`` is generated. instance_type: str, default = "ml.m5.2xlarge" Instance type of the training job. - instance_count: int, default = 1 - Number of training instances. Only single-instance training is supported. - volume_size: int, default = 100 - Size in GB of the EBS volume that stores the training data and model artifacts. - custom_image_uri: str | None, default = None - Custom training container image URI. If set, ``framework_version`` is ignored. wait: bool, default = True Whether to block until the job completes. If ``False``, returns once the job is launched; use :meth:`get_fit_job_status` to poll it. + + Returns + ------- + TimeSeriesCloudPredictor + The fitted predictor (``self``). + + Other Parameters + ---------------- backend_overrides: dict[str, dict[str, Any]] | None, default = None Raw SageMaker request fields for settings without a dedicated argument. @@ -114,57 +116,38 @@ def fit( * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. * Example: ``{"CreateTrainingJob": {"RetryStrategy": {"MaximumRetryAttempts": 2}}}`` - - Returns - ------- - TimeSeriesCloudPredictor - The fitted predictor (``self``). + job_name: str | None, default = None + Name of the training job. Defaults to a unique name with prefix ``ag-cloud-timeseries``. + volume_size: int, default = 100 + Size in GB of the EBS volume that stores the training data and model artifacts. + custom_image_uri: str | None, default = None + Custom training container image URI. If set, ``framework_version`` is ignored. + timeout: int, default = 86400 + Maximum training job runtime in seconds. Defaults to 24 hours. SageMaker API ------------- - * :sm-api:`CreateTrainingJob`: trains the predictor on ``instance_count`` x ``instance_type`` and writes the - artifact to ``cloud_output_path``. + * :sm-api:`CreateTrainingJob`: trains the predictor on ``instance_type`` and writes the artifact to + ``cloud_output_path``. """ - assert not self.backend.is_fit, ( - "Predictor is already fit! To fit additional models, create a new `CloudPredictor`" - ) - # `extra_ag_args` is an internal channel for `fit_predict`; it is intentionally not part of the public signature. - extra_ag_args = kwargs.pop("extra_ag_args", None) - if kwargs: - raise TypeError(f"fit() got unexpected keyword arguments: {sorted(kwargs)}") - predictor_fit_args = {} if predictor_fit_args is None else dict(predictor_fit_args) - data_channels = { - "train_data": train_data, - "tuning_data": tuning_data, - "known_covariates": known_covariates, - "static_features": static_features, - } - for key in ("train_data", "tuning_data", "known_covariates"): - if key in predictor_fit_args: - raise TypeError( - f"`{key}` can no longer be passed via `predictor_fit_args`. " - f"Pass `{key}` as an explicit argument to `fit()` instead." - ) - if data_channels["train_data"] is None: - raise TypeError("fit() missing required argument: 'train_data'") - - self.backend.fit( + kwargs = check_backend_kwargs(kwargs, TrainingJobKwargs, "fit", IGNORED_TRAINING_KWARGS) + self._fit( + data_channels={ + "train_data": train_data, + "tuning_data": tuning_data, + "known_covariates": known_covariates, + "static_features": static_features, + }, predictor_init_args=predictor_init_args, predictor_fit_args=predictor_fit_args, - data_channels=data_channels, id_column=id_column, timestamp_column=timestamp_column, framework_version=framework_version, - job_name=job_name, instance_type=instance_type, - instance_count=instance_count, - volume_size=volume_size, - custom_image_uri=custom_image_uri, wait=wait, backend_overrides=backend_overrides, - extra_ag_args=extra_ag_args, + **kwargs, ) - return self @deprecated(_DEPRECATED_REAL_TIME, category=None) @@ -233,17 +216,16 @@ def predict_proba_real_time(self, **kwargs) -> pd.DataFrame: def predict( self, data: str | pd.DataFrame, + *, static_features: str | pd.DataFrame | None = None, known_covariates: str | pd.DataFrame | None = None, predictor_path: str | None = None, framework_version: str | None = None, - job_name: str | None = None, instance_type: str = "ml.m5.2xlarge", - instance_count: int = 1, - custom_image_uri: str | None = None, wait: bool = True, predictions_path: str | None = None, backend_overrides: dict[str, dict[str, Any]] | None = None, + **kwargs: Unpack[BatchTransformKwargs], ) -> pd.DataFrame | None: """ Forecast with a SageMaker batch transform job. @@ -266,21 +248,23 @@ def predict( framework_version: str | None, default = None AutoGluon version, e.g. ``"1.6"``. Inference uses the official AutoGluon DLC image for this version. Defaults to the version used by :meth:`fit`. Ignored if ``custom_image_uri`` is set. - job_name: str | None, default = None - Name of the batch transform job. If ``None``, a unique name with prefix ``ag-cloud-timeseries`` is - generated. instance_type: str, default = "ml.m5.2xlarge" Instance type of the batch transform job. - instance_count: int, default = 1 - Number of batch transform instances. - custom_image_uri: str | None, default = None - Custom inference container image URI. If set, ``framework_version`` is ignored. wait: bool, default = True Whether to block until the job completes and return the forecast. If ``False``, returns ``None`` once the job is launched; use :meth:`get_batch_inference_job_status` to poll it. predictions_path: str | None, default = None S3 prefix under which the batch transform job writes its results (``/.out``). Defaults to ``{cloud_output_path}/batch_transform//results``. + + Returns + ------- + pd.DataFrame | None + Forecast in long format with ``item_id``, ``timestamp``, ``mean``, and one column per quantile, or + ``None`` if ``wait=False``. + + Other Parameters + ---------------- backend_overrides: dict[str, dict[str, Any]] | None, default = None Raw SageMaker request fields for settings without a dedicated argument. @@ -288,12 +272,12 @@ def predict( * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. * Example: ``{"CreateTransformJob": {"BatchStrategy": "SingleRecord", "MaxPayloadInMB": 20}}`` - - Returns - ------- - pd.DataFrame | None - Forecast in long format with ``item_id``, ``timestamp``, ``mean``, and one column per quantile, or - ``None`` if ``wait=False``. + job_name: str | None, default = None + Name of the batch transform job. Defaults to a unique name with prefix ``ag-cloud-timeseries``. + instance_count: int, default = 1 + Number of batch transform instances. + custom_image_uri: str | None, default = None + Custom inference container image URI. If set, ``framework_version`` is ignored. SageMaker API ------------- @@ -304,19 +288,18 @@ def predict( The model is deleted when the job finishes. With ``wait=False`` it is kept; delete it with :sm-api:`DeleteModel`. """ + kwargs = check_backend_kwargs(kwargs, BatchTransformKwargs, "predict") return self.backend.predict( test_data=data, static_features=static_features, known_covariates=known_covariates, predictor_path=predictor_path, framework_version=framework_version, - job_name=job_name, instance_type=instance_type, - instance_count=instance_count, - custom_image_uri=custom_image_uri, wait=wait, predictions_path=predictions_path, backend_overrides=backend_overrides, + **kwargs, ) def predict_proba( @@ -341,13 +324,10 @@ def fit_predict( timestamp_column: str = "timestamp", predictions_path: str | None = None, framework_version: str = DEFAULT_FRAMEWORK_VERSION, - job_name: str | None = None, instance_type: str = "ml.m5.2xlarge", - instance_count: int = 1, - volume_size: int = DEFAULT_VOLUME_SIZE, - custom_image_uri: str | None = None, wait: bool = True, backend_overrides: dict[str, dict[str, Any]] | None = None, + **kwargs: Unpack[TrainingJobKwargs], ) -> pd.DataFrame | None: """ Fit and predict in a single SageMaker training job. @@ -384,25 +364,10 @@ def fit_predict( framework_version: str, optional AutoGluon version, e.g. ``"1.6"``. Training uses the official AutoGluon DLC image for this version. Ignored if ``custom_image_uri`` is set. - job_name: str | None, default = None - Name of the training job. If ``None``, a unique name with prefix ``ag-cloud-timeseries`` is generated. instance_type: str, default = "ml.m5.2xlarge" Instance type of the training job. - instance_count: int, default = 1 - Number of training instances. Only single-instance training is supported. - volume_size: int, default = 100 - Size in GB of the EBS volume that stores the training data and model artifacts. - custom_image_uri: str | None, default = None - Custom training container image URI. If set, ``framework_version`` is ignored. wait: bool, default = True Whether to block until the job completes. If ``False``, returns ``None`` once the job is launched. - backend_overrides: dict[str, dict[str, Any]] | None, default = None - Raw SageMaker request fields for settings without a dedicated argument. - - * Keys: request names from the *SageMaker API* section below. - * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request - built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. - * Example: ``{"CreateTrainingJob": {"RetryStrategy": {"MaximumRetryAttempts": 2}}}`` Returns ------- @@ -411,32 +376,51 @@ def fit_predict( ``None`` if ``wait=False``; fetch it later with :meth:`get_fit_predict_results`. Columns are named ``item_id`` and ``timestamp`` regardless of ``id_column`` / ``timestamp_column``. + Other Parameters + ---------------- + backend_overrides: dict[str, dict[str, Any]] | None, default = None + Raw SageMaker request fields for settings without a dedicated argument. + + * Keys: request names from the *SageMaker API* section below. + * Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request + built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones. + * Example: ``{"CreateTrainingJob": {"RetryStrategy": {"MaximumRetryAttempts": 2}}}`` + job_name: str | None, default = None + Name of the training job. Defaults to a unique name with prefix ``ag-cloud-timeseries``. + volume_size: int, default = 100 + Size in GB of the EBS volume that stores the training data and model artifacts. + custom_image_uri: str | None, default = None + Custom training container image URI. If set, ``framework_version`` is ignored. + timeout: int, default = 86400 + Maximum training job runtime in seconds. Defaults to 24 hours. + SageMaker API ------------- - * :sm-api:`CreateTrainingJob`: trains the predictor and predicts in the same job on ``instance_count`` x - ``instance_type``. Predictions are written to ``predictions_path``. + * :sm-api:`CreateTrainingJob`: trains the predictor and predicts in the same job on ``instance_type``. + Predictions are written to ``predictions_path``. """ + kwargs = check_backend_kwargs(kwargs, TrainingJobKwargs, "fit_predict", IGNORED_TRAINING_KWARGS) extra_ag_args = {"predict_after_fit": True} if predictions_path is not None: extra_ag_args["predictions_path"] = predictions_path - self.fit( - train_data=train_data, - known_covariates=known_covariates, - static_features=static_features, + self._fit( + data_channels={ + "train_data": train_data, + "tuning_data": None, + "known_covariates": known_covariates, + "static_features": static_features, + }, predictor_init_args=predictor_init_args, predictor_fit_args=predictor_fit_args, id_column=id_column, timestamp_column=timestamp_column, framework_version=framework_version, - job_name=job_name, instance_type=instance_type, - instance_count=instance_count, - volume_size=volume_size, - custom_image_uri=custom_image_uri, wait=wait, backend_overrides=backend_overrides, extra_ag_args=extra_ag_args, + **kwargs, ) if not wait: diff --git a/src/autogluon/cloud/utils/sagemaker_api.py b/src/autogluon/cloud/utils/sagemaker_api.py index b1ec1aa..f27a336 100644 --- a/src/autogluon/cloud/utils/sagemaker_api.py +++ b/src/autogluon/cloud/utils/sagemaker_api.py @@ -4,7 +4,7 @@ import functools import logging from collections.abc import Callable, Iterable, Mapping -from typing import Any +from typing import Any, TypedDict from .aws_utils import AwsSession @@ -75,6 +75,65 @@ def check_override_keys(overrides: Mapping[str, Any] | None, allowed_keys: Itera return overrides +class TrainingJobKwargs(TypedDict, total=False): + """Less common settings of methods that run a SageMaker training job, passed as ``**kwargs``.""" + + job_name: str + volume_size: int + custom_image_uri: str + timeout: int + + +class BatchTransformKwargs(TypedDict, total=False): + """Less common settings of methods that run a SageMaker batch transform job, passed as ``**kwargs``.""" + + job_name: str + instance_count: int + custom_image_uri: str + + +class DeployKwargs(TypedDict, total=False): + """Less common settings of methods that deploy a SageMaker endpoint, passed as ``**kwargs``.""" + + initial_instance_count: int + volume_size: int + custom_image_uri: str + + +# Training kwargs that are no longer supported, but accepted with a warning: name -> why the value is ignored. +IGNORED_TRAINING_KWARGS = { + "instance_count": "only single-instance training is supported", + "leaderboard": "the leaderboard is always saved", +} +# FoundationModel jobs never save a leaderboard, and never accepted `leaderboard`, so only `instance_count` applies. +IGNORED_FM_JOB_KWARGS = {"instance_count": IGNORED_TRAINING_KWARGS["instance_count"]} + + +def check_backend_kwargs( + kwargs: Mapping[str, Any], + kwargs_type: type, + method: str, + ignored: Mapping[str, str] | None = None, +) -> dict[str, Any]: + """Return ``kwargs`` as a dict, raising ``TypeError`` for keys that the ``kwargs_type`` TypedDict doesn't + define. Keys in ``ignored`` (name -> reason) are dropped with a warning instead. + + ``kwargs_type`` is annotated as ``type`` because typing has no way to spell "a TypedDict class". + """ + kwargs = dict(kwargs) + for name, reason in (ignored or {}).items(): + if name in kwargs: + kwargs.pop(name) + logger.warning(f"`{name}` is no longer supported by {method}() and is ignored: {reason}.") + allowed = kwargs_type.__required_keys__ | kwargs_type.__optional_keys__ + unknown = sorted(set(kwargs) - set(allowed)) + if unknown: + raise TypeError( + f"{method}() got unexpected keyword argument(s) {unknown}. Supported keyword arguments: {sorted(allowed)}." + ) + return kwargs + + def delete_quietly(delete: Callable[..., Any], **kwargs) -> None: """Call a ``delete_*`` API during rollback, logging instead of raising so the original error propagates.""" try: diff --git a/tests/unittests/general/test_backend_kwargs.py b/tests/unittests/general/test_backend_kwargs.py new file mode 100644 index 0000000..5ac7eb8 --- /dev/null +++ b/tests/unittests/general/test_backend_kwargs.py @@ -0,0 +1,153 @@ +"""Unit tests for the ``**kwargs`` catch-all of CloudPredictor and FoundationModel methods.""" + +import logging +from unittest import mock + +import pandas as pd +import pytest + +from autogluon.cloud import ( + MultiModalCloudPredictor, + TabularCloudPredictor, + TabularFoundationModel, + TimeSeriesCloudPredictor, + TimeSeriesFoundationModel, +) + +TRAIN_DATA = pd.DataFrame({"x": [1, 2], "y": [0, 1]}) +TS_DATA = pd.DataFrame({"item_id": ["A", "A"], "timestamp": ["2024-01-01", "2024-01-02"], "target": [1.0, 2.0]}) + + +@pytest.fixture(autouse=True) +def _stub_aws(monkeypatch): + """Avoid touching AWS / config files during construction.""" + for module in ("predictor.cloud_predictor", "model.foundation_model"): + monkeypatch.setattr( + f"autogluon.cloud.{module}.resolve_cloud_output_path", + lambda path, backend_name: path or "s3://stub/output", + ) + monkeypatch.setattr( + "autogluon.cloud.backend.backend_factory.BackendFactory.get_backend", + lambda **kwargs: mock.MagicMock(is_fit=False, endpoint_name="ep"), + ) + monkeypatch.setattr("autogluon.cloud.endpoint.endpoint.setup_sagemaker_session", lambda boto_session: boto_session) + # autogluon.common sets propagate=False on the `autogluon` logger, which hides records from caplog. + monkeypatch.setattr(logging.getLogger("autogluon"), "propagate", True) + + +@pytest.fixture +def predictor(tmp_path): + return TabularCloudPredictor(local_output_path=str(tmp_path)) + + +def test_fit_forwards_backend_kwargs(predictor): + predictor.fit(TRAIN_DATA, predictor_init_args={"label": "y"}, job_name="job", custom_image_uri="img", timeout=60) + kwargs = predictor.backend.fit.call_args.kwargs + assert (kwargs["job_name"], kwargs["custom_image_uri"], kwargs["timeout"]) == ("job", "img", 60) + + +@pytest.mark.parametrize("name, value", [("instance_count", 2), ("leaderboard", False)]) +def test_fit_warns_and_drops_ignored_kwargs(predictor, caplog, name, value): + predictor.fit(TRAIN_DATA, predictor_init_args={"label": "y"}, **{name: value}) + assert f"`{name}` is no longer supported by fit() and is ignored" in caplog.text + assert name not in predictor.backend.fit.call_args.kwargs + + +@pytest.mark.parametrize( + "call", + [ + lambda p: p.fit(TRAIN_DATA, predictor_init_args={"label": "y"}, instance_typo="ml.m5.xlarge"), + lambda p: p.fit_predict(TRAIN_DATA, TRAIN_DATA, predictor_init_args={"label": "y"}, instance_typo="x"), + lambda p: p.predict(TRAIN_DATA, volume_size=10), # batch transform has no volume_size + lambda p: p.predict_proba(TRAIN_DATA, job_nme="job"), + lambda p: p.deploy(timeout=10), + ], + ids=["fit", "fit_predict", "predict", "predict_proba", "deploy"], +) +def test_unknown_backend_kwargs_raise(predictor, call): + with pytest.raises(TypeError, match="unexpected keyword argument"): + call(predictor) + + +def test_legacy_backend_kwargs_dict_raises_migration_hint(predictor): + with pytest.raises(TypeError, match="`backend_kwargs` was removed.*backend_overrides"): + predictor.predict(TRAIN_DATA, backend_kwargs={}) + + +@pytest.mark.parametrize( + "name, call", + [ + ("image_column", lambda p: p.fit(TRAIN_DATA, predictor_init_args={"label": "y"}, image_column="image")), + ("test_data_image_column", lambda p: p.predict(TRAIN_DATA, test_data_image_column="image")), + ("test_data_image_column", lambda p: p.predict_proba(TRAIN_DATA, test_data_image_column="image")), + ], + ids=["fit", "predict", "predict_proba"], +) +def test_tabular_rejects_image_columns_with_pointer_to_multimodal(predictor, name, call): + with pytest.raises(ValueError, match=f"`{name}` is no longer supported for tabular predictors.*MultiModal"): + call(predictor) + + +@pytest.mark.filterwarnings("ignore:AutoGluon Multimodal is on a deprecation path") +def test_multimodal_forwards_image_columns(tmp_path): + predictor = MultiModalCloudPredictor(local_output_path=str(tmp_path)) + predictor.fit(TRAIN_DATA, predictor_init_args={"label": "y"}, image_column="image") + assert predictor.backend.fit.call_args.kwargs["image_column"] == "image" + predictor.predict(TRAIN_DATA, test_data_image_column="image", instance_count=2) + kwargs = predictor.backend.predict.call_args.kwargs + assert (kwargs["test_data_image_column"], kwargs["instance_count"]) == ("image", 2) + + +def test_timeseries_fit_forwards_timeout(tmp_path): + predictor = TimeSeriesCloudPredictor(local_output_path=str(tmp_path)) + predictor.fit(TS_DATA, predictor_init_args={"prediction_length": 1}, timeout=60) + assert predictor.backend.fit.call_args.kwargs["timeout"] == 60 + + +@pytest.mark.parametrize( + "call", + [ + lambda p: p.predict(TRAIN_DATA, "s3://bucket/predictor.tar.gz"), + lambda p: p.deploy("s3://bucket/predictor.tar.gz"), + ], + ids=["predict", "deploy"], +) +def test_options_are_keyword_only(predictor, call): + with pytest.raises(TypeError, match="positional argument"): + call(predictor) + + +def test_foundation_model_deploy_forwards_backend_kwargs(): + fm = TimeSeriesFoundationModel("chronos-2", cloud_output_path="s3://b") + fm.deploy(initial_instance_count=2, custom_image_uri="img", backend_overrides={"CreateEndpoint": {}}) + kwargs = fm._backend.deploy.call_args.kwargs + assert (kwargs["initial_instance_count"], kwargs["custom_image_uri"]) == (2, "img") + assert kwargs["backend_overrides"] == {"CreateEndpoint": {}} + + +def test_foundation_model_predict_rejects_unknown_backend_kwargs(): + fm = TabularFoundationModel("mitra-classifier", cloud_output_path="s3://b") + with pytest.raises(TypeError, match="predict\\(\\) got unexpected keyword argument"): + fm.predict(TRAIN_DATA, TRAIN_DATA, label="y", instance_typo="ml.m5.xlarge") + fm._backend.fit.assert_not_called() + + +@pytest.mark.parametrize( + "fm_cls, model_id, args", + [ + (TabularFoundationModel, "mitra-classifier", dict(test_data=TRAIN_DATA, train_data=TRAIN_DATA, label="y")), + (TimeSeriesFoundationModel, "chronos-2", dict(data=TS_DATA)), + ], + ids=["tabular", "timeseries"], +) +def test_foundation_model_predict_warns_and_drops_instance_count(caplog, fm_cls, model_id, args): + fm = fm_cls(model_id, cloud_output_path="s3://b") + fm.predict(**args, wait=False, instance_count=2) + assert "`instance_count` is no longer supported by predict() and is ignored" in caplog.text + assert "instance_count" not in fm._backend.fit.call_args.kwargs + + +def test_timeseries_foundation_model_predict_is_keyword_only(): + fm = TimeSeriesFoundationModel("chronos-2", cloud_output_path="s3://b") + with pytest.raises(TypeError, match="positional argument"): + fm.predict(TS_DATA, "target") diff --git a/tests/unittests/general/test_tabular_foundation_model_predict.py b/tests/unittests/general/test_tabular_foundation_model_predict.py index 7301a29..82fa2bf 100644 --- a/tests/unittests/general/test_tabular_foundation_model_predict.py +++ b/tests/unittests/general/test_tabular_foundation_model_predict.py @@ -128,10 +128,10 @@ def test_wait_false_returns_future_without_fetching(): def _make_tcp(result=CLASSIFICATION_FRAME): - """A TabularCloudPredictor with `fit` and the backend mocked out — no AWS interaction.""" + """A TabularCloudPredictor with `_fit` and the backend mocked out — no AWS interaction.""" with mock.patch.object(TabularCloudPredictor, "__init__", lambda self: None): tcp = TabularCloudPredictor() - tcp.fit = mock.MagicMock() + tcp._fit = mock.MagicMock() tcp.backend = mock.MagicMock() tcp.backend.get_fit_predict_results.return_value = result return tcp diff --git a/tests/unittests/tabular/test_tabular_fit_predict.py b/tests/unittests/tabular/test_tabular_fit_predict.py index 94116a7..845211d 100644 --- a/tests/unittests/tabular/test_tabular_fit_predict.py +++ b/tests/unittests/tabular/test_tabular_fit_predict.py @@ -23,10 +23,10 @@ @pytest.fixture def cloud_predictor(): - """A TabularCloudPredictor with `fit` and the backend mocked out — no AWS interaction.""" + """A TabularCloudPredictor with `_fit` and the backend mocked out — no AWS interaction.""" with mock.patch.object(TabularCloudPredictor, "__init__", lambda self: None): predictor = TabularCloudPredictor() - predictor.fit = mock.MagicMock() + predictor._fit = mock.MagicMock() predictor.backend = mock.MagicMock() predictor.backend.get_fit_predict_results.return_value = CLASSIFICATION_FRAME return predictor @@ -34,7 +34,7 @@ def cloud_predictor(): def test_when_fit_predict_then_launches_predict_job_and_returns_prediction_series(cloud_predictor): pred = cloud_predictor.fit_predict(**DEFAULT_ARGS) - extra_ag_args = cloud_predictor.fit.call_args.kwargs["extra_ag_args"] + extra_ag_args = cloud_predictor._fit.call_args.kwargs["extra_ag_args"] assert extra_ag_args["predict_after_fit"] is True assert "predictions_path" not in extra_ag_args # not passed -> backend fills in a default @@ -44,7 +44,7 @@ def test_when_fit_predict_then_launches_predict_job_and_returns_prediction_serie def test_when_predictions_path_given_then_forwarded_to_backend(cloud_predictor): cloud_predictor.fit_predict(**DEFAULT_ARGS, predictions_path="s3://bucket/key/predictions.csv") - extra_ag_args = cloud_predictor.fit.call_args.kwargs["extra_ag_args"] + extra_ag_args = cloud_predictor._fit.call_args.kwargs["extra_ag_args"] assert extra_ag_args["predictions_path"] == "s3://bucket/key/predictions.csv"