Add Mitra-v2 regressor via optional mitra-finetune dependency

docs/tutorials/foundation-model-tabular.md

@@ -40,6 +40,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) |
+| `mitra-v2-regressor` | Regression | [MitraModel](https://auto.gluon.ai/stable/api/autogluon.tabular.models.html#autogluon.tabular.models.MitraModel) | [autogluon/mitra-regressor-2](https://huggingface.co/autogluon/mitra-regressor-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/backend/sagemaker_backend.py

@@ -17,6 +17,7 @@
 from ..scripts import ScriptManager
 from ..utils.ag_sagemaker import (
     SOURCE_DIR_TARBALL_NAME,
+    add_requirements_to_tar,
     repack_model_with_serving_code,
     script_mode_environment,
     staged_serving_code,
@@ -208,6 +209,7 @@ def fit(
         backend_overrides: dict[str, dict[str, Any]] | None = None,
         extra_ag_args: dict[str, Any] | None = None,
         extra_tags: list[dict[str, str]] | None = None,
+        requirements: list[str] | None = None,
     ) -> None:
         """
         Fit the predictor with SageMaker.
@@ -327,6 +329,7 @@ def fit(
             entry_point=entry_point,
             sagemaker_session=self.sagemaker_session,
             s3_uri_prefix=f"{self.cloud_output_path}/code/{job_name}/source",
+            requirements=requirements,
         )
 
         request: dict[str, Any] = {
@@ -438,6 +441,7 @@ def deploy(
         inference_config: dict[str, Any] | None = None,
         repack: bool = True,
         extra_tags: list[dict[str, str]] | None = None,
+        requirements: list[str] | None = None,
     ) -> None:
         """
         Deploy a predictor as a SageMaker endpoint, which can be used to do real-time inference later.
@@ -544,7 +548,7 @@ def deploy(
         # Decide whether the tarball already contains the entry_point script — if yes, use it as-is;
         # if no, repack the script into it.
         if predictor_path is None:
-            model_data = self._create_serve_script_tarball(entry_point, endpoint_name)
+            model_data = self._create_serve_script_tarball(entry_point, endpoint_name, requirements)
         else:
             is_default_fit_output = (
                 self._fit_job is not None
@@ -627,14 +631,17 @@ def deploy(
         if wait:
             client.get_waiter("endpoint_in_service").wait(EndpointName=self.endpoint_name)
 
-    def _create_serve_script_tarball(self, serve_script_path: str, endpoint_name: str) -> str:
-        """Create a minimal model.tar.gz containing the serve script + serving_utils/ under code/."""
+    def _create_serve_script_tarball(
+        self, serve_script_path: str, endpoint_name: str, requirements: list[str] | None = None
+    ) -> str:
+        """Create a minimal model.tar.gz containing the serve script + serving_utils/ (+ requirements.txt) under code/."""
 
         tarball_dir = tempfile.mkdtemp(prefix="ag_serve_")
         tarball_path = os.path.join(tarball_dir, "model.tar.gz")
         with tarfile.open(tarball_path, "w:gz") as tar:
             tar.add(serve_script_path, arcname=f"code/{os.path.basename(serve_script_path)}")
             tar.add(ScriptManager.SAGEMAKER_SERVING_UTILS_DIR, arcname="code/serving_utils")
+            add_requirements_to_tar(tar, requirements, "code/requirements.txt")
         s3_key = f"endpoints/{endpoint_name}/model/model.tar.gz"
         s3_path = self._upload_predictor(tarball_path, s3_key)
         return s3_path

src/autogluon/cloud/backend/timeseries_sagemaker_backend.py

@@ -36,6 +36,7 @@ def fit(
         backend_overrides: dict[str, dict[str, Any]] | None = None,
         extra_ag_args: dict[str, Any] | None = None,
         extra_tags: list[dict[str, str]] | None = None,
+        requirements: list[str] | None = None,
     ) -> None:
         """Fit a TimeSeriesPredictor in SageMaker.
 
@@ -68,6 +69,7 @@ def fit(
             backend_overrides=backend_overrides,
             extra_ag_args=extra_ag_args,
             extra_tags=extra_tags,
+            requirements=requirements,
         )
 
     def predict_real_time(

src/autogluon/cloud/model/foundation_model.py

@@ -23,6 +23,7 @@
 from ..endpoint.tabular_endpoint import TabularEndpoint
 from ..endpoint.timeseries_endpoint import TimeSeriesEndpoint
 from ..scripts.script_manager import ScriptManager
+from ..utils.ag_sagemaker import add_requirements_to_tar
 from ..utils.aws_utils import resolve_cloud_output_path
 from ..utils.constants import DEFAULT_FRAMEWORK_VERSION
 from ..utils.sagemaker_api import (
@@ -283,6 +284,7 @@ def _deploy_backend(
             inference_config=inference_config,
             repack=False,
             extra_tags=[{"Key": "autogluon-cloud-model-id", "Value": self.model_id}],
+            requirements=self._config.requirements,
             **kwargs,
         )
         assert self._backend.endpoint_name is not None
@@ -392,6 +394,7 @@ def cache_model_artifact(self, cache_path: str, *, overwrite: bool = False) -> S
                     tar.add(weights_dir, arcname="weights")
                     tar.add(serve_script, arcname=f"code/{serve_script.name}")
                     tar.add(ScriptManager.SAGEMAKER_SERVING_UTILS_DIR, arcname="code/serving_utils")
+                    add_requirements_to_tar(tar, self._config.requirements, "code/requirements.txt")
                 logger.info(f"Uploading to {cache_key}")
                 s3.upload_file(
                     str(tarball),
@@ -701,6 +704,7 @@ def predict(
             backend_overrides=backend_overrides,
             extra_ag_args=extra_ag_args,
             extra_tags=[{"Key": "autogluon-cloud-model-id", "Value": self.model_id}],
+            requirements=self._config.requirements,
             **kwargs,
         )
 
@@ -1045,6 +1049,7 @@ def predict_proba(
             backend_overrides=backend_overrides,
             extra_ag_args=extra_ag_args,
             extra_tags=[{"Key": "autogluon-cloud-model-id", "Value": self.model_id}],
+            requirements=self._config.requirements,
             **kwargs,
         )
 

src/autogluon/cloud/model/registry.py

@@ -22,6 +22,8 @@ class FoundationModelConfig:
     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"
+    # Lines of a requirements.txt that the container pip-installs before the model is used.
+    requirements: list[str] = field(default_factory=list)
 
 
 FOUNDATION_MODEL_REGISTRY: dict[str, FoundationModelConfig] = {
@@ -121,6 +123,18 @@ class FoundationModelConfig:
         deploy_instance_type="ml.m5.4xlarge",
         min_framework_version="1.6",
     ),
+    "mitra-v2-regressor": FoundationModelConfig(
+        problem_type="regression",
+        ag_model_key="MITRA",
+        model_source_uri="autogluon/mitra-regressor-2",
+        model_source_hyperparameter="hf_reg_model",
+        inference_hyperparameters={"fine_tune": False},
+        predict_instance_type="ml.m5.4xlarge",
+        deploy_instance_type="ml.m5.4xlarge",
+        min_framework_version="1.6",
+        # Patches AutoGluon's Mitra to run the v2 regression (cross-entropy bin head) checkpoint.
+        requirements=["mitra-finetune==0.3.0"],
+    ),
     # 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(

src/autogluon/cloud/scripts/sagemaker_scripts/serving_utils/mitra.py

@@ -0,0 +1,24 @@
+"""Patch AutoGluon's Mitra for the v2 regressor checkpoint, using the optional ``mitra-finetune`` package."""
+
+import json
+import os
+
+
+def patch_mitra_regressor(hyperparameters):
+    """Patch AutoGluon's Mitra to run the v2 regressor, if ``mitra_finetune`` is installed and this is a Mitra regressor."""
+    # TODO(ag-1.7): remove once AutoGluon supports the v2 regressor natively.
+    from huggingface_hub import snapshot_download
+
+    source = hyperparameters.get("hf_reg_model")
+    if source is None:
+        return
+    try:
+        from mitra_finetune.patches import install_d2h_sync_patch, install_reg_ce_patches, install_use_hf_patch
+    except ImportError:
+        return
+    config_dir = source if os.path.isdir(source) else snapshot_download(repo_id=source, allow_patterns=["config.json"])
+    with open(os.path.join(config_dir, "config.json")) as f:
+        n_bins = int(json.load(f)["dim_output"])
+    install_d2h_sync_patch()
+    install_reg_ce_patches(n_bins)
+    install_use_hf_patch()

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

@@ -13,6 +13,7 @@
 
 import pandas as pd
 from huggingface_hub import snapshot_download
+from serving_utils.mitra import patch_mitra_regressor
 
 from autogluon.tabular import TabularPredictor
 
@@ -32,6 +33,7 @@ def model_fn(model_dir):
                 allow_patterns=["config.json", "model.safetensors"],
             )
             break
+    patch_mitra_regressor(hyperparameters)
     return model_config
 
 

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

@@ -17,6 +17,7 @@
 from autogluon.common.utils.s3_utils import s3_path_to_bucket_prefix
 from autogluon.tabular import TabularPredictor, TabularDataset
 from autogluon.timeseries import TimeSeriesDataFrame
+from serving_utils.mitra import patch_mitra_regressor
 
 
 def get_input_path(path):
@@ -164,6 +165,12 @@ def prepare_data(data_file, predictor_type, ag_args, static_features_df=None):
     if predict_after_fit and predictor_type == "tabular":
         assert args.test_dir is not None, "`test_data` channel is required for tabular fit_predict."
 
+    hyperparameters = predictor_fit_args.get("hyperparameters")
+    if isinstance(hyperparameters, dict):  # may also be a preset name or None
+        for model_hyperparameters in hyperparameters.values():
+            if isinstance(model_hyperparameters, dict):
+                patch_mitra_regressor(model_hyperparameters)
+
     predictor = predictor_cls(**predictor_init_args).fit(training_data, tuning_data=tuning_data, **predictor_fit_args)
 
     # When use automm backend, predictor needs to be saved with standalone flag to avoid need of internet access when loading

src/autogluon/cloud/utils/ag_sagemaker.py

@@ -6,6 +6,7 @@
 ``SAGEMAKER_SUBMIT_DIRECTORY`` environment variables (inference).
 """
 
+import io
 import json
 import os
 import shutil
@@ -21,15 +22,33 @@
 SOURCE_DIR_TARBALL_NAME = "sourcedir.tar.gz"
 
 
-def upload_training_code(entry_point: str, sagemaker_session, s3_uri_prefix: str) -> str:
-    """Bundle the training entry point as ``sourcedir.tar.gz`` and upload it.
+def add_requirements_to_tar(tar: tarfile.TarFile, requirements: list[str] | None, arcname: str) -> None:
+    """Write ``requirements`` as a ``requirements.txt`` at ``arcname``, which the DLC pip-installs before running the
+    script (the training toolkit for the submit directory, the serving entrypoint for ``code/``)."""
+    if not requirements:
+        return
+    data = ("\n".join(requirements) + "\n").encode()
+    info = tarfile.TarInfo(arcname)
+    info.size = len(data)
+    tar.addfile(info, io.BytesIO(data))
+
+
+def upload_training_code(
+    entry_point: str, sagemaker_session, s3_uri_prefix: str, requirements: list[str] | None = None
+) -> str:
+    """Bundle the training entry point, ``serving_utils/`` (shared helpers) and an optional ``requirements.txt`` as
+    ``sourcedir.tar.gz`` and upload it.
 
     Returns the S3 URI of the uploaded tarball.
     """
+    from ..scripts import ScriptManager  # deferred: importing scripts pulls in the backend package
+
     with tempfile.TemporaryDirectory(prefix="ag_train_code_") as tmpdir:
         tarball_path = os.path.join(tmpdir, SOURCE_DIR_TARBALL_NAME)
         with tarfile.open(tarball_path, "w:gz") as tar:
             tar.add(entry_point, arcname=os.path.basename(entry_point))
+            tar.add(ScriptManager.SAGEMAKER_SERVING_UTILS_DIR, arcname="serving_utils")
+            add_requirements_to_tar(tar, requirements, "requirements.txt")
         bucket, key_prefix = s3_path_to_bucket_prefix(s3_uri_prefix)
         return sagemaker_session.upload_data(path=tarball_path, bucket=bucket, key_prefix=key_prefix)
 

tests/unittests/general/test_ag_sagemaker.py

@@ -0,0 +1,40 @@
+"""Unit tests for the training sourcedir built by upload_training_code."""
+
+import tarfile
+from unittest import mock
+
+import pytest
+
+from autogluon.cloud.utils.ag_sagemaker import upload_training_code
+
+
+def _sourcedir_contents(tmp_path, requirements):
+    entry_point = tmp_path / "train.py"
+    entry_point.write_text("")
+    contents = {}
+
+    def read_tarball(path, bucket, key_prefix):
+        with tarfile.open(path) as tar:
+            for member in tar.getmembers():
+                f = tar.extractfile(member)
+                contents[member.name] = f.read().decode() if f is not None else None
+        return f"s3://{bucket}/{key_prefix}/sourcedir.tar.gz"
+
+    session = mock.MagicMock()
+    session.upload_data.side_effect = read_tarball
+    uri = upload_training_code(str(entry_point), session, "s3://b/code/job/source", requirements=requirements)
+    assert uri == "s3://b/code/job/source/sourcedir.tar.gz"
+    return contents
+
+
+@pytest.mark.parametrize("requirements", [None, []])
+def test_sourcedir_without_requirements(tmp_path, requirements):
+    contents = _sourcedir_contents(tmp_path, requirements)
+    assert "train.py" in contents
+    assert "serving_utils/mitra.py" in contents
+    assert "requirements.txt" not in contents
+
+
+def test_sourcedir_writes_requirements(tmp_path):
+    contents = _sourcedir_contents(tmp_path, ["mitra-finetune==0.3.0", "foo>=1"])
+    assert contents["requirements.txt"] == "mitra-finetune==0.3.0\nfoo>=1\n"

tests/unittests/general/test_foundation_model.py

@@ -224,6 +224,7 @@ def test_cache_model_artifact_rejects_models_without_weights_source_hyperparamet
         ("tabdpt-turbo-classifier", {}),
         ("nori-30m-regressor", {"model": "nori-30m"}),
         ("mitra-v2-classifier", {"fine_tune": False, "hf_cls_model": "autogluon/mitra-classifier-2"}),
+        ("mitra-v2-regressor", {"fine_tune": False, "hf_reg_model": "autogluon/mitra-regressor-2"}),
     ],
 )
 def test_inference_hyperparameters_for_new_models(model_id, expected_hp):