diff --git a/sagemaker-serve/pyproject.toml b/sagemaker-serve/pyproject.toml index 58c3c5b97e..4b1d6c475f 100644 --- a/sagemaker-serve/pyproject.toml +++ b/sagemaker-serve/pyproject.toml @@ -27,7 +27,12 @@ dependencies = [ "boto3>=1.42.2,<2.0", "botocore>=1.35.75,<2.0", "deepdiff", - "mlflow", + "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", @@ -39,6 +44,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-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() 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",