From f7c0dc9e9ba328eb53bd19912db60c7806ec7174 Mon Sep 17 00:00:00 2001 From: Oleksandr Shchur Date: Wed, 7 Oct 2026 12:27:21 +0000 Subject: [PATCH] Add Mitra-v2 classifier; enforce minimum AutoGluon version per foundation model --- docs/tutorials/foundation-model-tabular.md | 1 + src/autogluon/cloud/model/foundation_model.py | 22 +++++++++++++ src/autogluon/cloud/model/registry.py | 23 ++++++++++++++ .../general/test_foundation_model.py | 31 +++++++++++++++++++ 4 files changed, 77 insertions(+) diff --git a/docs/tutorials/foundation-model-tabular.md b/docs/tutorials/foundation-model-tabular.md index d30476c..1b51cac 100644 --- a/docs/tutorials/foundation-model-tabular.md +++ b/docs/tutorials/foundation-model-tabular.md @@ -39,6 +39,7 @@ The following `model_id` values are currently supported (also available via `Tab |----------|------|---------------|---------| | `mitra-classifier` | Classification | [MitraModel](https://auto.gluon.ai/stable/api/autogluon.tabular.models.html#autogluon.tabular.models.MitraModel) | [autogluon/mitra-classifier](https://huggingface.co/autogluon/mitra-classifier) | | `mitra-regressor` | Regression | [MitraModel](https://auto.gluon.ai/stable/api/autogluon.tabular.models.html#autogluon.tabular.models.MitraModel) | [autogluon/mitra-regressor](https://huggingface.co/autogluon/mitra-regressor) | +| `mitra-v2-classifier` | Classification | [MitraModel](https://auto.gluon.ai/stable/api/autogluon.tabular.models.html#autogluon.tabular.models.MitraModel) | [autogluon/mitra-classifier-2](https://huggingface.co/autogluon/mitra-classifier-2) | | `tabicl-v2-classifier` | Classification | [TabICLModel](https://auto.gluon.ai/stable/api/autogluon.tabular.models.html#autogluon.tabular.models.TabICLModel) | [jingang/TabICL](https://huggingface.co/jingang/TabICL) | | `tabicl-v2-regressor` | Regression | [TabICLModel](https://auto.gluon.ai/stable/api/autogluon.tabular.models.html#autogluon.tabular.models.TabICLModel) | [jingang/TabICL](https://huggingface.co/jingang/TabICL) | | `tabdpt-turbo-classifier` | Classification | [TabDPTTurboModel](https://auto.gluon.ai/stable/api/autogluon.tabular.models.html#autogluon.tabular.models.TabDPTTurboModel) | [Layer6/TabDPT](https://huggingface.co/Layer6/TabDPT) | diff --git a/src/autogluon/cloud/model/foundation_model.py b/src/autogluon/cloud/model/foundation_model.py index c10ea2b..0fa3828 100644 --- a/src/autogluon/cloud/model/foundation_model.py +++ b/src/autogluon/cloud/model/foundation_model.py @@ -11,6 +11,7 @@ from typing import Any, Literal import pandas as pd +from packaging.version import InvalidVersion, Version from typing_extensions import Self from autogluon.common.loaders import load_pd @@ -172,6 +173,24 @@ def _get_hyperparameters( merged.setdefault(self._config.model_source_hyperparameter, self._config.model_source_uri) return merged + def _check_framework_version(self, framework_version: str, custom_image_uri: str | None) -> None: + """Raise if ``framework_version`` is older than the model's ``min_framework_version``. + + Skipped for custom images, whose AutoGluon version is unknown. + """ + if custom_image_uri is not None or framework_version == "latest": + return + min_version = self._config.min_framework_version + try: + requested = Version(framework_version).release[:2] + except InvalidVersion: + return # the backend raises a descriptive error for invalid versions + if requested < Version(min_version).release[:2]: + raise ValueError( + f"Model '{self.model_id}' requires AutoGluon {min_version} or newer, got " + f"framework_version={framework_version!r}." + ) + @abstractmethod def _build_predictor_init_args(self, **user_kwargs) -> dict[str, Any]: """Build predictor_init_args dict from user-provided kwargs. @@ -219,6 +238,7 @@ def _deploy_backend( **backend_kwargs, ) -> None: """Shared deploy logic. Subclasses call this then wrap the endpoint.""" + self._check_framework_version(framework_version, 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": @@ -631,6 +651,7 @@ def predict( * :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) if instance_type is None: instance_type = self._config.predict_instance_type @@ -965,6 +986,7 @@ def predict_proba( * :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) if instance_type is None: instance_type = self._config.predict_instance_type diff --git a/src/autogluon/cloud/model/registry.py b/src/autogluon/cloud/model/registry.py index dbaa611..2f170ad 100644 --- a/src/autogluon/cloud/model/registry.py +++ b/src/autogluon/cloud/model/registry.py @@ -20,6 +20,8 @@ class FoundationModelConfig: inference_hyperparameters: dict[str, Any] = field(default_factory=dict) # defaults for deploy() and predict() training_hyperparameters: dict[str, Any] = field(default_factory=dict) # defaults for fit() fine_tunable: bool = False # whether .fit() is supported + # Oldest AutoGluon version that supports the model. Not checked when a custom image is used. + min_framework_version: str = "1.5" FOUNDATION_MODEL_REGISTRY: dict[str, FoundationModelConfig] = { @@ -58,12 +60,14 @@ class FoundationModelConfig: ag_model_key="Toto2", model_source_uri="Datadog/Toto-2.0-4m", model_source_hyperparameter="model_path", + min_framework_version="1.6", ), "toto-2.0-22m": FoundationModelConfig( problem_type="forecasting", ag_model_key="Toto2", model_source_uri="Datadog/Toto-2.0-22m", model_source_hyperparameter="model_path", + min_framework_version="1.6", ), "toto-2.0-313m": FoundationModelConfig( problem_type="forecasting", @@ -71,6 +75,7 @@ class FoundationModelConfig: model_source_uri="Datadog/Toto-2.0-313m", model_source_hyperparameter="model_path", predict_instance_type="ml.g5.xlarge", + min_framework_version="1.6", ), "toto-2.0-1b": FoundationModelConfig( problem_type="forecasting", @@ -78,6 +83,7 @@ class FoundationModelConfig: model_source_uri="Datadog/Toto-2.0-1B", model_source_hyperparameter="model_path", predict_instance_type="ml.g5.xlarge", + min_framework_version="1.6", ), "toto-2.0-2.5b": FoundationModelConfig( problem_type="forecasting", @@ -85,6 +91,7 @@ class FoundationModelConfig: model_source_uri="Datadog/Toto-2.0-2.5B", model_source_hyperparameter="model_path", predict_instance_type="ml.g5.xlarge", + min_framework_version="1.6", ), "mitra-classifier": FoundationModelConfig( problem_type="multiclass", @@ -104,6 +111,16 @@ class FoundationModelConfig: predict_instance_type="ml.m5.4xlarge", deploy_instance_type="ml.m5.4xlarge", ), + "mitra-v2-classifier": FoundationModelConfig( + problem_type="multiclass", + ag_model_key="MITRA", + model_source_uri="autogluon/mitra-classifier-2", + model_source_hyperparameter="hf_cls_model", + inference_hyperparameters={"fine_tune": False}, + predict_instance_type="ml.m5.4xlarge", + deploy_instance_type="ml.m5.4xlarge", + min_framework_version="1.6", + ), # TabICL, TabDPT-Turbo and Nori download weights pinned by AutoGluon / the model library at fit time, so they # have no weights-source hyperparameter and don't support `cache_model_artifact`. "tabicl-v2-classifier": FoundationModelConfig( @@ -112,6 +129,7 @@ class FoundationModelConfig: model_source_uri="jingang/TabICL", inference_hyperparameters={"checkpoint_version": "tabicl-classifier-v2-20260212.ckpt"}, predict_instance_type="ml.g5.xlarge", + min_framework_version="1.6", ), "tabicl-v2-regressor": FoundationModelConfig( problem_type="regression", @@ -119,18 +137,21 @@ class FoundationModelConfig: model_source_uri="jingang/TabICL", inference_hyperparameters={"checkpoint_version": "tabicl-regressor-v2-20260212.ckpt"}, predict_instance_type="ml.g5.xlarge", + min_framework_version="1.6", ), "tabdpt-turbo-classifier": FoundationModelConfig( problem_type="multiclass", ag_model_key="TABDPT-TURBO", model_source_uri="Layer6/TabDPT", predict_instance_type="ml.g5.xlarge", + min_framework_version="1.6", ), "tabdpt-turbo-regressor": FoundationModelConfig( problem_type="regression", ag_model_key="TABDPT-TURBO", model_source_uri="Layer6/TabDPT", predict_instance_type="ml.g5.xlarge", + min_framework_version="1.6", ), "nori-regressor": FoundationModelConfig( problem_type="regression", @@ -138,6 +159,7 @@ class FoundationModelConfig: model_source_uri="Synthefy/Nori", inference_hyperparameters={"model": "nori"}, predict_instance_type="ml.g5.xlarge", + min_framework_version="1.6", ), "nori-30m-regressor": FoundationModelConfig( problem_type="regression", @@ -145,6 +167,7 @@ class FoundationModelConfig: model_source_uri="Synthefy/Nori-30M", inference_hyperparameters={"model": "nori-30m"}, predict_instance_type="ml.g5.xlarge", + min_framework_version="1.6", ), } diff --git a/tests/unittests/general/test_foundation_model.py b/tests/unittests/general/test_foundation_model.py index 2b8a2be..2abeb6b 100644 --- a/tests/unittests/general/test_foundation_model.py +++ b/tests/unittests/general/test_foundation_model.py @@ -223,6 +223,7 @@ def test_cache_model_artifact_rejects_models_without_weights_source_hyperparamet ("tabicl-v2-regressor", {"checkpoint_version": "tabicl-regressor-v2-20260212.ckpt"}), ("tabdpt-turbo-classifier", {}), ("nori-30m-regressor", {"model": "nori-30m"}), + ("mitra-v2-classifier", {"fine_tune": False, "hf_cls_model": "autogluon/mitra-classifier-2"}), ], ) def test_inference_hyperparameters_for_new_models(model_id, expected_hp): @@ -304,3 +305,33 @@ def test_list_models_filters_by_task(): def test_subclass_rejects_model_id_for_other_task(): with pytest.raises(ValueError, match="Unknown model_id 'mitra-classifier' for TimeSeriesFoundationModel"): TimeSeriesFoundationModel("mitra-classifier", cloud_output_path="s3://b") + + +@pytest.mark.parametrize("framework_version", ["1.5", "1.5.0", "1.4"]) +def test_rejects_framework_version_below_model_minimum(framework_version): + fm = TabularFoundationModel("mitra-v2-classifier", cloud_output_path="s3://b") + with pytest.raises(ValueError, match="requires AutoGluon 1.6 or newer"): + fm.deploy(framework_version=framework_version) + with pytest.raises(ValueError, match="requires AutoGluon 1.6 or newer"): + fm.predict( + pd.DataFrame({"x": [1]}), + pd.DataFrame({"x": [1], "y": [0]}), + label="y", + framework_version=framework_version, + ) + fm._backend.deploy.assert_not_called() + fm._backend.fit.assert_not_called() + + +@pytest.mark.parametrize("framework_version", ["1.6", "1.6.3", "latest"]) +def test_accepts_framework_version_at_or_above_model_minimum(framework_version): + fm = TimeSeriesFoundationModel("toto-2.0-4m", cloud_output_path="s3://b") + fm.predict(pd.DataFrame({"x": [1]}), framework_version=framework_version) + assert fm._backend.fit.call_args.kwargs["framework_version"] == framework_version + + +def test_skips_framework_version_check_for_custom_image(): + fm = TimeSeriesFoundationModel("toto-2.0-4m", cloud_output_path="s3://b") + fm._backend.endpoint_name = "ep" + fm.deploy(framework_version="1.5", custom_image_uri="my-image") + assert fm._backend.deploy.call_args.kwargs["custom_image_uri"] == "my-image"