Add Mitra-v2 classifier; enforce minimum AutoGluon version per foundation model

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

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
 

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,33 +60,38 @@ 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",
         ag_model_key="Toto2",
         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",
         ag_model_key="Toto2",
         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",
         ag_model_key="Toto2",
         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,39 +129,45 @@ 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",
         ag_model_key="TABICL",
         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",
         ag_model_key="NORI",
         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",
         ag_model_key="NORI",
         model_source_uri="Synthefy/Nori-30M",
         inference_hyperparameters={"model": "nori-30m"},
         predict_instance_type="ml.g5.xlarge",
+        min_framework_version="1.6",
     ),
 }
 

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"