From 03f4f1a177b0f7ca6b321142e06ecd6208d896a7 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Mon, 28 Sep 2026 14:28:02 -0700 Subject: [PATCH 1/2] feat: depend on mlflow-skinny with optional full extra (#6152) --- sagemaker-serve/pyproject.toml | 7 ++++++- sagemaker-train/pyproject.toml | 6 +++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/sagemaker-serve/pyproject.toml b/sagemaker-serve/pyproject.toml index 58c3c5b97e..3f5a845138 100644 --- a/sagemaker-serve/pyproject.toml +++ b/sagemaker-serve/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ "boto3>=1.42.2,<2.0", "botocore>=1.35.75,<2.0", "deepdiff", - "mlflow", + "mlflow-skinny", "sagemaker_schema_inference_artifacts", "pytest", "tqdm", @@ -39,6 +39,11 @@ dependencies = [ ] [project.optional-dependencies] +# Full MLflow. The client only needs the tracking/registry APIs from mlflow-skinny; +# the full MLflow model-flavor loaders (mlflow.pyfunc/sklearn/pytorch/...) are used +# only by the container-side inference handlers that run inside the DLC. Install this +# extra to pull in the full MLflow package locally. +full = ["mlflow"] test = [ "pytest", "pytest-cov", diff --git a/sagemaker-train/pyproject.toml b/sagemaker-train/pyproject.toml index a78dadfac1..51bf2d60f9 100644 --- a/sagemaker-train/pyproject.toml +++ b/sagemaker-train/pyproject.toml @@ -40,7 +40,7 @@ dependencies = [ "paramiko>=2.11.0", "jinja2>=3.0,<4.0", "sagemaker-mlflow>=0.0.1,<1.0.0", - "mlflow>=3.0.0,<4.0.0", + "mlflow-skinny>=3.0.0,<4.0.0", "nest_asyncio>=1.5.0", "requests>=2.28.0", "pydantic>=2.0.0", @@ -52,6 +52,10 @@ Homepage = "https://github.com/aws/sagemaker-python-sdk" [project.optional-dependencies] rft-strands = ["strands-agents"] rft-langchain = ["httpx>=0.24.0"] +# Full MLflow (model flavors, pyfunc/sklearn/pytorch/... loaders). The client only +# needs the tracking/registry APIs from mlflow-skinny; install this extra when you +# want the full MLflow package (e.g. to load MLflow model flavors locally). +full = ["mlflow>=3.0.0,<4.0.0"] test = [ "pytest", "pytest-cov", From 590c5c29ac060e024d2106b9e9dbfb5aa5088541 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Wed, 30 Sep 2026 08:41:22 -0700 Subject: [PATCH 2/2] fix: declare docker explicitly after the mlflow-skinny switch Switching sagemaker-serve from `mlflow` to `mlflow-skinny` silently dropped `docker`. Full mlflow declares `docker>=4.0.0,<8`; mlflow-skinny declares no such requirement, and `docker` was never declared by sagemaker-serve itself. `docker` is imported at module scope and sits on the eager import path out of the package `__init__`: sagemaker/serve/__init__.py:29 -> sagemaker.serve.model_builder sagemaker/serve/model_builder.py:81 -> sagemaker.serve.mode.local_container_mode sagemaker/serve/mode/local_container_mode.py:12 -> import docker so losing it breaks `import sagemaker.serve` outright rather than degrading a niche local-container path. It also affects model_server/{triton,tgi,djl_serving,tei}/server.py. This surfaced as 18 collection errors in integ-tests-us-east-1 with `ModuleNotFoundError: No module named 'docker'`; other jobs stayed green only because their runners happened to have docker installed already, which is what made the wrong dependency set easy to miss. The bound matches the v2 precedent in requirements/extras/local_requirements.txt (`docker>=5.0.2,<8.0.0`) and is compatible with mlflow's own `>=4.0.0,<8`. Audited the rest of the full-mlflow -> mlflow-skinny delta (aiohttp, alembic, cryptography, flask, flask-cors, graphene, gunicorn, huey, matplotlib, numpy, pandas, pyarrow, scikit-learn, scipy, skops, sqlalchemy, waitress): docker is the only one that was both dropped and imported at module scope. numpy, pandas, scipy and cryptography remain satisfied via sagemaker-core, graphene via sagemaker-train's own declaration, and matplotlib/sklearn are imported lazily inside functions on opt-in paths. Adds tests/unit/test_declared_dependencies.py, which fails when any module-level third-party import in sagemaker-serve has no declared distribution across sagemaker-{serve,core,train}. Verified it catches this exact regression: removing the new line makes it fail and name all six importing files. --- sagemaker-serve/pyproject.toml | 5 + .../tests/unit/test_declared_dependencies.py | 138 ++++++++++++++++++ 2 files changed, 143 insertions(+) create mode 100644 sagemaker-serve/tests/unit/test_declared_dependencies.py diff --git a/sagemaker-serve/pyproject.toml b/sagemaker-serve/pyproject.toml index 3f5a845138..4b1d6c475f 100644 --- a/sagemaker-serve/pyproject.toml +++ b/sagemaker-serve/pyproject.toml @@ -28,6 +28,11 @@ dependencies = [ "botocore>=1.35.75,<2.0", "deepdiff", "mlflow-skinny", + # Declared explicitly because full mlflow used to supply it transitively + # (mlflow requires docker>=4.0.0,<8; mlflow-skinny does not). It is imported at + # module scope by sagemaker.serve.mode.local_container_mode, which sagemaker.serve + # imports eagerly, so a missing docker breaks `import sagemaker.serve` outright. + "docker>=5.0.2,<8.0.0", "sagemaker_schema_inference_artifacts", "pytest", "tqdm", diff --git a/sagemaker-serve/tests/unit/test_declared_dependencies.py b/sagemaker-serve/tests/unit/test_declared_dependencies.py new file mode 100644 index 0000000000..8a5b1aea14 --- /dev/null +++ b/sagemaker-serve/tests/unit/test_declared_dependencies.py @@ -0,0 +1,138 @@ +"""Every module-level third-party import must have a declared distribution. + +Regression guard for #6152. Switching ``mlflow`` to ``mlflow-skinny`` silently dropped +``docker``: full mlflow requires ``docker>=4.0.0,<8`` while mlflow-skinny requires +nothing of the sort. ``docker`` is imported at module scope by +``sagemaker.serve.mode.local_container_mode``, which ``sagemaker.serve.__init__`` +imports eagerly, so losing it breaks ``import sagemaker.serve`` outright. Nothing went +red until a CI runner happened not to have ``docker`` already installed, which is +exactly the kind of gap a declared-dependency check closes. + +The check is deliberately limited to *module-level* imports. Imports nested inside a +function are opt-in code paths (matplotlib for plotting, sklearn for framework +detection) and are allowed to rely on an undeclared package. +""" + +from __future__ import absolute_import + +import ast +import os +import re +import sys +from pathlib import Path + +import pytest + +SERVE_ROOT = Path(__file__).resolve().parents[2] +REPO_ROOT = SERVE_ROOT.parent +SERVE_SRC = SERVE_ROOT / "src" + +# sagemaker-serve depends on sagemaker-core and sagemaker-train, so anything those two +# declare is also importable from serve. +PYPROJECTS = ( + SERVE_ROOT / "pyproject.toml", + REPO_ROOT / "sagemaker-core" / "pyproject.toml", + REPO_ROOT / "sagemaker-train" / "pyproject.toml", +) + +# Import name -> distribution name, for the cases where they differ. +IMPORT_TO_DISTRIBUTION = { + "yaml": "pyyaml", +} + +# Modules provided by the runtime environment rather than by our wheels. +EXEMPT_IMPORTS = { + # Injected by the Triton Python backend inside the serving container; it is not + # importable from the SDK and must never be a declared dependency. + "triton_python_backend_utils", +} + +FIRST_PARTY = {"sagemaker"} + +# Matches the top-level `dependencies = [...]` array. The optional-dependency arrays are +# keyed by extra name (`full`, `test`, `dev`), so they cannot match this pattern. +_DEPENDENCIES_ARRAY = re.compile(r"^dependencies = \[(.*?)^\]", re.MULTILINE | re.DOTALL) + + +def _normalize(name): + """Normalize a distribution name per PEP 503.""" + return re.sub(r"[-_.]+", "-", name).lower() + + +def _declared_distributions(): + """Return the normalized distribution names declared across the three packages.""" + declared = set() + for pyproject in PYPROJECTS: + match = _DEPENDENCIES_ARRAY.search(pyproject.read_text(encoding="utf-8")) + assert match, "no top-level dependencies array found in {}".format(pyproject) + for line in match.group(1).splitlines(): + line = line.split("#", 1)[0].strip().rstrip(",").strip() + if not line: + continue + requirement = ast.literal_eval(line) + # Strip extras, version specifiers and environment markers. + name = re.split(r"[\[<>=!~;\s]", requirement, maxsplit=1)[0] + declared.add(_normalize(name)) + return declared + + +def _module_level_imports(): + """Map each module-level third-party import to the files that import it.""" + imports = {} + for dirpath, _, filenames in os.walk(SERVE_SRC): + for filename in filenames: + if not filename.endswith(".py"): + continue + path = Path(dirpath) / filename + try: + tree = ast.parse(path.read_text(encoding="utf-8")) + except (SyntaxError, UnicodeDecodeError): + continue + for node in tree.body: # module level only, by construction + if isinstance(node, ast.Import): + names = [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom) and not node.level and node.module: + names = [node.module] + else: + continue + for name in names: + top = name.split(".")[0] + if top in sys.stdlib_module_names or top in FIRST_PARTY: + continue + if top in EXEMPT_IMPORTS: + continue + imports.setdefault(top, set()).add(str(path.relative_to(REPO_ROOT))) + return imports + + +def test_module_level_imports_are_declared_dependencies(): + """Fail if sagemaker-serve imports a package none of the three packages declare.""" + declared = _declared_distributions() + undeclared = {} + for module, files in _module_level_imports().items(): + distribution = IMPORT_TO_DISTRIBUTION.get(module, module) + if _normalize(distribution) not in declared: + undeclared[module] = sorted(files) + + assert not undeclared, "module-level imports with no declared distribution: " + "; ".join( + "{} (imported by {})".format(module, ", ".join(files)) + for module, files in sorted(undeclared.items()) + ) + + +def test_docker_is_declared(): + """Pin the specific regression from #6152 so it cannot silently come back. + + ``docker`` has to be a hard dependency, not an extra: the import sits at module + scope on the eager path out of ``sagemaker.serve.__init__``. + """ + assert "docker" in _declared_distributions() + + +@pytest.mark.parametrize( + "module", + ["docker", "numpy", "pandas", "yaml"], +) +def test_sanity_known_module_level_imports_are_detected(module): + """Guard the detector itself: these are known module-level imports in serve.""" + assert module in _module_level_imports()