From de13c23f99a41941a5bb3d0d70576a6b1b8dbda8 Mon Sep 17 00:00:00 2001 From: Oleksandr Shchur Date: Thu, 8 Oct 2026 15:20:44 +0000 Subject: [PATCH 1/5] Add Mitra-v2 regressor via optional mitra-finetune dependency (MVP) Add FoundationModelConfig.optional_dependencies, pip-installed in the training / serving container before the model is used. mitra-v2-regressor ships mitra-finetune==0.3.0 (TestPyPI), whose patches switch AutoGluon 1.6 Mitra to the v2 cross-entropy regression head. --- src/autogluon/cloud/model/foundation_model.py | 4 ++ src/autogluon/cloud/model/registry.py | 14 ++++++ .../sagemaker_scripts/serving_utils/mitra.py | 45 +++++++++++++++++ .../sagemaker_scripts/tabular_fm_serve.py | 3 ++ .../cloud/scripts/sagemaker_scripts/train.py | 48 +++++++++++++++++++ 5 files changed, 114 insertions(+) create mode 100644 src/autogluon/cloud/scripts/sagemaker_scripts/serving_utils/mitra.py diff --git a/src/autogluon/cloud/model/foundation_model.py b/src/autogluon/cloud/model/foundation_model.py index fa86372..122a4f3 100644 --- a/src/autogluon/cloud/model/foundation_model.py +++ b/src/autogluon/cloud/model/foundation_model.py @@ -267,6 +267,8 @@ def _deploy_backend( "hyperparameters": merged_hp, "problem_type": self._config.problem_type, } + if self._config.optional_dependencies: + fm_serve_config["optional_dependencies"] = self._config.optional_dependencies # FM deploys never repack: predictor_path is either None (script-only tarball is built locally) or a # pre-bundled cache artifact that already contains the serve script. @@ -1031,6 +1033,8 @@ def predict_proba( tuning_data = train_data.iloc[:2].copy() extra_ag_args: dict[str, Any] = {"predict_after_fit": True, "save_predictor": False} + if self._config.optional_dependencies: + extra_ag_args["optional_dependencies"] = self._config.optional_dependencies if predictions_path is not None: extra_ag_args["predictions_path"] = predictions_path kwargs["leaderboard"] = False diff --git a/src/autogluon/cloud/model/registry.py b/src/autogluon/cloud/model/registry.py index 2f170ad..4cdfdda 100644 --- a/src/autogluon/cloud/model/registry.py +++ b/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" + # pip requirement specs installed in the training / serving container before the model is used. + optional_dependencies: 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. + optional_dependencies=["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( diff --git a/src/autogluon/cloud/scripts/sagemaker_scripts/serving_utils/mitra.py b/src/autogluon/cloud/scripts/sagemaker_scripts/serving_utils/mitra.py new file mode 100644 index 0000000..58ada6c --- /dev/null +++ b/src/autogluon/cloud/scripts/sagemaker_scripts/serving_utils/mitra.py @@ -0,0 +1,45 @@ +"""Optional-dependency install and Mitra regressor patching, shared by the tabular FM serve script.""" + +import json +import os +import subprocess +import sys + +from huggingface_hub import snapshot_download + +# TODO: drop the index override once mitra-finetune is released on PyPI. +_OPTIONAL_DEPENDENCIES_INDEX_URL = "https://test.pypi.org/simple/" + + +def install_optional_dependencies(specs): + """pip-install the registry's ``optional_dependencies`` (without their dependencies) into the container.""" + if specs: + subprocess.check_call( + [ + sys.executable, + "-m", + "pip", + "install", + "--no-deps", + "--index-url", + _OPTIONAL_DEPENDENCIES_INDEX_URL, + *specs, + ] + ) + + +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.""" + 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() diff --git a/src/autogluon/cloud/scripts/sagemaker_scripts/tabular_fm_serve.py b/src/autogluon/cloud/scripts/sagemaker_scripts/tabular_fm_serve.py index fdf5211..7c7beb7 100644 --- a/src/autogluon/cloud/scripts/sagemaker_scripts/tabular_fm_serve.py +++ b/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 install_optional_dependencies, patch_mitra_regressor from autogluon.tabular import TabularPredictor @@ -32,6 +33,8 @@ def model_fn(model_dir): allow_patterns=["config.json", "model.safetensors"], ) break + install_optional_dependencies(model_config.get("optional_dependencies")) + patch_mitra_regressor(hyperparameters) return model_config diff --git a/src/autogluon/cloud/scripts/sagemaker_scripts/train.py b/src/autogluon/cloud/scripts/sagemaker_scripts/train.py index f4efaac..288eb63 100644 --- a/src/autogluon/cloud/scripts/sagemaker_scripts/train.py +++ b/src/autogluon/cloud/scripts/sagemaker_scripts/train.py @@ -6,6 +6,8 @@ import json import os import shutil +import subprocess +import sys import tempfile from pprint import pprint @@ -19,6 +21,47 @@ from autogluon.timeseries import TimeSeriesDataFrame +# Duplicated from serving_utils/mitra.py: the training entry point can't import serving_utils. +# TODO: drop the index override once mitra-finetune is released on PyPI. +_OPTIONAL_DEPENDENCIES_INDEX_URL = "https://test.pypi.org/simple/" + + +def install_optional_dependencies(specs): + """pip-install the registry's ``optional_dependencies`` (without their dependencies) into the container.""" + if specs: + subprocess.check_call( + [ + sys.executable, + "-m", + "pip", + "install", + "--no-deps", + "--index-url", + _OPTIONAL_DEPENDENCIES_INDEX_URL, + *specs, + ] + ) + + +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.""" + 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() + + def get_input_path(path): file = os.listdir(path)[0] if len(os.listdir(path)) > 1: @@ -164,6 +207,11 @@ 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." + install_optional_dependencies(ag_args.get("optional_dependencies")) + for model_hyperparameters in predictor_fit_args.get("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 From d411bad7684f7f806fbb940b079cf5ce5afedb6f Mon Sep 17 00:00:00 2001 From: Oleksandr Shchur Date: Thu, 8 Oct 2026 15:33:53 +0000 Subject: [PATCH 2/5] Install model requirements via requirements.txt in the DLC instead of pip in the scripts Rename FoundationModelConfig.optional_dependencies to requirements. The backend writes them as requirements.txt into the training sourcedir and the model's code/ dir, which the DLC installs before running the script. --- .../cloud/backend/sagemaker_backend.py | 13 +++++++-- src/autogluon/cloud/model/foundation_model.py | 9 +++--- src/autogluon/cloud/model/registry.py | 7 +++-- .../sagemaker_scripts/serving_utils/mitra.py | 28 ++----------------- .../sagemaker_scripts/tabular_fm_serve.py | 3 +- .../cloud/scripts/sagemaker_scripts/train.py | 22 --------------- src/autogluon/cloud/utils/ag_sagemaker.py | 19 +++++++++++-- 7 files changed, 40 insertions(+), 61 deletions(-) diff --git a/src/autogluon/cloud/backend/sagemaker_backend.py b/src/autogluon/cloud/backend/sagemaker_backend.py index 31878c0..2654437 100644 --- a/src/autogluon/cloud/backend/sagemaker_backend.py +++ b/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 diff --git a/src/autogluon/cloud/model/foundation_model.py b/src/autogluon/cloud/model/foundation_model.py index 122a4f3..b85ada6 100644 --- a/src/autogluon/cloud/model/foundation_model.py +++ b/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 ( @@ -267,8 +268,6 @@ def _deploy_backend( "hyperparameters": merged_hp, "problem_type": self._config.problem_type, } - if self._config.optional_dependencies: - fm_serve_config["optional_dependencies"] = self._config.optional_dependencies # FM deploys never repack: predictor_path is either None (script-only tarball is built locally) or a # pre-bundled cache artifact that already contains the serve script. @@ -285,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 @@ -394,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), @@ -703,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, ) @@ -1033,8 +1035,6 @@ def predict_proba( tuning_data = train_data.iloc[:2].copy() extra_ag_args: dict[str, Any] = {"predict_after_fit": True, "save_predictor": False} - if self._config.optional_dependencies: - extra_ag_args["optional_dependencies"] = self._config.optional_dependencies if predictions_path is not None: extra_ag_args["predictions_path"] = predictions_path kwargs["leaderboard"] = False @@ -1049,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, ) diff --git a/src/autogluon/cloud/model/registry.py b/src/autogluon/cloud/model/registry.py index 4cdfdda..bb97372 100644 --- a/src/autogluon/cloud/model/registry.py +++ b/src/autogluon/cloud/model/registry.py @@ -22,8 +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" - # pip requirement specs installed in the training / serving container before the model is used. - optional_dependencies: list[str] = field(default_factory=list) + # 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] = { @@ -133,7 +133,8 @@ class FoundationModelConfig: 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. - optional_dependencies=["mitra-finetune==0.3.0"], + # TODO: drop the index once mitra-finetune is released on PyPI. + requirements=["--extra-index-url https://test.pypi.org/simple/", "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`. diff --git a/src/autogluon/cloud/scripts/sagemaker_scripts/serving_utils/mitra.py b/src/autogluon/cloud/scripts/sagemaker_scripts/serving_utils/mitra.py index 58ada6c..8c8f901 100644 --- a/src/autogluon/cloud/scripts/sagemaker_scripts/serving_utils/mitra.py +++ b/src/autogluon/cloud/scripts/sagemaker_scripts/serving_utils/mitra.py @@ -1,35 +1,13 @@ -"""Optional-dependency install and Mitra regressor patching, shared by the tabular FM serve script.""" +"""Patch AutoGluon's Mitra for the v2 regressor checkpoint, using the optional ``mitra-finetune`` package.""" import json import os -import subprocess -import sys - -from huggingface_hub import snapshot_download - -# TODO: drop the index override once mitra-finetune is released on PyPI. -_OPTIONAL_DEPENDENCIES_INDEX_URL = "https://test.pypi.org/simple/" - - -def install_optional_dependencies(specs): - """pip-install the registry's ``optional_dependencies`` (without their dependencies) into the container.""" - if specs: - subprocess.check_call( - [ - sys.executable, - "-m", - "pip", - "install", - "--no-deps", - "--index-url", - _OPTIONAL_DEPENDENCIES_INDEX_URL, - *specs, - ] - ) 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.""" + from huggingface_hub import snapshot_download + source = hyperparameters.get("hf_reg_model") if source is None: return diff --git a/src/autogluon/cloud/scripts/sagemaker_scripts/tabular_fm_serve.py b/src/autogluon/cloud/scripts/sagemaker_scripts/tabular_fm_serve.py index 7c7beb7..bd36b73 100644 --- a/src/autogluon/cloud/scripts/sagemaker_scripts/tabular_fm_serve.py +++ b/src/autogluon/cloud/scripts/sagemaker_scripts/tabular_fm_serve.py @@ -13,7 +13,7 @@ import pandas as pd from huggingface_hub import snapshot_download -from serving_utils.mitra import install_optional_dependencies, patch_mitra_regressor +from serving_utils.mitra import patch_mitra_regressor from autogluon.tabular import TabularPredictor @@ -33,7 +33,6 @@ def model_fn(model_dir): allow_patterns=["config.json", "model.safetensors"], ) break - install_optional_dependencies(model_config.get("optional_dependencies")) patch_mitra_regressor(hyperparameters) return model_config diff --git a/src/autogluon/cloud/scripts/sagemaker_scripts/train.py b/src/autogluon/cloud/scripts/sagemaker_scripts/train.py index 288eb63..00e9678 100644 --- a/src/autogluon/cloud/scripts/sagemaker_scripts/train.py +++ b/src/autogluon/cloud/scripts/sagemaker_scripts/train.py @@ -6,8 +6,6 @@ import json import os import shutil -import subprocess -import sys import tempfile from pprint import pprint @@ -22,25 +20,6 @@ # Duplicated from serving_utils/mitra.py: the training entry point can't import serving_utils. -# TODO: drop the index override once mitra-finetune is released on PyPI. -_OPTIONAL_DEPENDENCIES_INDEX_URL = "https://test.pypi.org/simple/" - - -def install_optional_dependencies(specs): - """pip-install the registry's ``optional_dependencies`` (without their dependencies) into the container.""" - if specs: - subprocess.check_call( - [ - sys.executable, - "-m", - "pip", - "install", - "--no-deps", - "--index-url", - _OPTIONAL_DEPENDENCIES_INDEX_URL, - *specs, - ] - ) def patch_mitra_regressor(hyperparameters): @@ -207,7 +186,6 @@ 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." - install_optional_dependencies(ag_args.get("optional_dependencies")) for model_hyperparameters in predictor_fit_args.get("hyperparameters", {}).values(): if isinstance(model_hyperparameters, dict): patch_mitra_regressor(model_hyperparameters) diff --git a/src/autogluon/cloud/utils/ag_sagemaker.py b/src/autogluon/cloud/utils/ag_sagemaker.py index 1c2eab3..2afcd73 100644 --- a/src/autogluon/cloud/utils/ag_sagemaker.py +++ b/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,8 +22,21 @@ 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 (plus an optional ``requirements.txt``) as ``sourcedir.tar.gz`` and upload it. Returns the S3 URI of the uploaded tarball. """ @@ -30,6 +44,7 @@ def upload_training_code(entry_point: str, sagemaker_session, s3_uri_prefix: str 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)) + 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) From e1e4ea0bc626b130639446088ad6729993fd72e7 Mon Sep 17 00:00:00 2001 From: Oleksandr Shchur Date: Thu, 8 Oct 2026 16:08:09 +0000 Subject: [PATCH 3/5] Ship serving_utils in the training sourcedir so train.py imports the Mitra patch helper instead of duplicating it --- .../cloud/scripts/sagemaker_scripts/train.py | 23 +------------------ src/autogluon/cloud/utils/ag_sagemaker.py | 6 ++++- 2 files changed, 6 insertions(+), 23 deletions(-) diff --git a/src/autogluon/cloud/scripts/sagemaker_scripts/train.py b/src/autogluon/cloud/scripts/sagemaker_scripts/train.py index 00e9678..620907a 100644 --- a/src/autogluon/cloud/scripts/sagemaker_scripts/train.py +++ b/src/autogluon/cloud/scripts/sagemaker_scripts/train.py @@ -17,28 +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 - - -# Duplicated from serving_utils/mitra.py: the training entry point can't import serving_utils. - - -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.""" - 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() +from serving_utils.mitra import patch_mitra_regressor def get_input_path(path): diff --git a/src/autogluon/cloud/utils/ag_sagemaker.py b/src/autogluon/cloud/utils/ag_sagemaker.py index 2afcd73..16470f6 100644 --- a/src/autogluon/cloud/utils/ag_sagemaker.py +++ b/src/autogluon/cloud/utils/ag_sagemaker.py @@ -36,14 +36,18 @@ def add_requirements_to_tar(tar: tarfile.TarFile, requirements: list[str] | None def upload_training_code( entry_point: str, sagemaker_session, s3_uri_prefix: str, requirements: list[str] | None = None ) -> str: - """Bundle the training entry point (plus an optional ``requirements.txt``) as ``sourcedir.tar.gz`` and upload it. + """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) From b95faa380ee1e0b257458721feadabb6815dd203 Mon Sep 17 00:00:00 2001 From: Oleksandr Shchur Date: Mon, 12 Oct 2026 14:12:09 +0000 Subject: [PATCH 4/5] Install mitra-finetune from PyPI, fix train.py crash on non-dict hyperparameters, add Mitra-v2 regressor docs and tests --- docs/tutorials/foundation-model-tabular.md | 1 + src/autogluon/cloud/model/registry.py | 3 +- .../sagemaker_scripts/serving_utils/mitra.py | 1 + .../cloud/scripts/sagemaker_scripts/train.py | 8 ++-- tests/unittests/general/test_ag_sagemaker.py | 40 +++++++++++++++++++ .../general/test_foundation_model.py | 1 + 6 files changed, 49 insertions(+), 5 deletions(-) create mode 100644 tests/unittests/general/test_ag_sagemaker.py diff --git a/docs/tutorials/foundation-model-tabular.md b/docs/tutorials/foundation-model-tabular.md index 1b51cac..92c9dd4 100644 --- a/docs/tutorials/foundation-model-tabular.md +++ b/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) | diff --git a/src/autogluon/cloud/model/registry.py b/src/autogluon/cloud/model/registry.py index bb97372..897665a 100644 --- a/src/autogluon/cloud/model/registry.py +++ b/src/autogluon/cloud/model/registry.py @@ -133,8 +133,7 @@ class FoundationModelConfig: 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. - # TODO: drop the index once mitra-finetune is released on PyPI. - requirements=["--extra-index-url https://test.pypi.org/simple/", "mitra-finetune==0.3.0"], + 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`. diff --git a/src/autogluon/cloud/scripts/sagemaker_scripts/serving_utils/mitra.py b/src/autogluon/cloud/scripts/sagemaker_scripts/serving_utils/mitra.py index 8c8f901..66ef3f2 100644 --- a/src/autogluon/cloud/scripts/sagemaker_scripts/serving_utils/mitra.py +++ b/src/autogluon/cloud/scripts/sagemaker_scripts/serving_utils/mitra.py @@ -6,6 +6,7 @@ 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") diff --git a/src/autogluon/cloud/scripts/sagemaker_scripts/train.py b/src/autogluon/cloud/scripts/sagemaker_scripts/train.py index 620907a..6b66ae2 100644 --- a/src/autogluon/cloud/scripts/sagemaker_scripts/train.py +++ b/src/autogluon/cloud/scripts/sagemaker_scripts/train.py @@ -165,9 +165,11 @@ 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." - for model_hyperparameters in predictor_fit_args.get("hyperparameters", {}).values(): - if isinstance(model_hyperparameters, dict): - patch_mitra_regressor(model_hyperparameters) + 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) diff --git a/tests/unittests/general/test_ag_sagemaker.py b/tests/unittests/general/test_ag_sagemaker.py new file mode 100644 index 0000000..eae177d --- /dev/null +++ b/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" diff --git a/tests/unittests/general/test_foundation_model.py b/tests/unittests/general/test_foundation_model.py index 2abeb6b..78c882a 100644 --- a/tests/unittests/general/test_foundation_model.py +++ b/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): From daccfdd9b34c6fced7760782cc162b5f110e81e6 Mon Sep 17 00:00:00 2001 From: Oleksandr Shchur Date: Mon, 12 Oct 2026 14:14:25 +0000 Subject: [PATCH 5/5] Forward requirements through TimeSeriesSagemakerBackend.fit --- src/autogluon/cloud/backend/timeseries_sagemaker_backend.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/autogluon/cloud/backend/timeseries_sagemaker_backend.py b/src/autogluon/cloud/backend/timeseries_sagemaker_backend.py index b5f8b94..d127234 100644 --- a/src/autogluon/cloud/backend/timeseries_sagemaker_backend.py +++ b/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(