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/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/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( diff --git a/src/autogluon/cloud/model/foundation_model.py b/src/autogluon/cloud/model/foundation_model.py index fa86372..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 ( @@ -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, ) diff --git a/src/autogluon/cloud/model/registry.py b/src/autogluon/cloud/model/registry.py index 2f170ad..897665a 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" + # 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( 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..66ef3f2 --- /dev/null +++ b/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() 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..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,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 diff --git a/src/autogluon/cloud/scripts/sagemaker_scripts/train.py b/src/autogluon/cloud/scripts/sagemaker_scripts/train.py index f4efaac..6b66ae2 100644 --- a/src/autogluon/cloud/scripts/sagemaker_scripts/train.py +++ b/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 diff --git a/src/autogluon/cloud/utils/ag_sagemaker.py b/src/autogluon/cloud/utils/ag_sagemaker.py index 1c2eab3..16470f6 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,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) 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):