feat: depend on mlflow-skinny with optional full extra (#6152) - #6350

Claude finished @mohamedzeidan2021's task in 2m 13s —— View job

PR Review: depend on mlflow-skinny with optional full extra

• Read PR diff
• Verify client-side mlflow usage is skinny-only in sagemaker-train
• Verify client-side mlflow usage is skinny-only in sagemaker-serve
• Check that container-side flavor loaders are truly deferred imports
• Assess correctness, best practices, tests
• Post review

The core change is sound. I traced every client-side mlflow import and confirmed the PR's central claim holds:

Client-side usage is skinny-compatible (verified):

All of these are tracking/registry APIs present in mlflow-skinny. The only flavor-loader usage (mlflow.pyfunc/sklearn/pytorch/...) lives in model_server/torchserve/inference.py and xgboost_inference.py, and is correctly loaded lazily via importlib.import_module(...) inside _load_mlflow_model — never at module level — so it only executes in the DLC where full MLflow is present. ✅

mlflow.search_runs() returns a pandas DataFrame; I confirmed pandas is a hard dependency of sagemaker-core and listed directly in both sagemaker-train and sagemaker-serve, so the skinny switch doesn't break it. ✅

Points worth considering (non-blocking)

1. [full] extra installs mlflow alongside mlflow-skinny. These are two separate PyPI distributions that both ship the same top-level mlflow/ package files. pip install sagemaker-serve[full] will therefore have both dist-infos present with overlapping files. In practice this works (the files are identical and full MLflow adds the heavy deps the flavor loaders need), but it's a slightly messy install — e.g. pip uninstall mlflow afterward can leave the skinny install in a partial state. Since the shared files are identical and importing a flavor module fails loudly (not silently) when its heavy dep is missing, I don't consider this a blocker — but it may be worth a one-line note in the extras' comment so users understand the coexistence.

2. Version-bound asymmetry in sagemaker-serve. sagemaker-train pins mlflow-skinny>=3.0.0,<4.0.0 and full = ["mlflow>=3.0.0,<4.0.0"] — consistent and good. sagemaker-serve uses unbounded mlflow-skinny and full = ["mlflow"]. This preserves the pre-existing behavior (the old dep was also unbounded "mlflow"), so it's not a regression, but adding matching bounds would be more robust against a future MLflow 4.x breaking change.

3. Transitive full-MLflow via sagemaker-mlflow (train only). sagemaker-train still depends on sagemaker-mlflow>=0.0.1,<1.0.0. The PR relies on sagemaker-mlflow 0.5.0 having moved to mlflow-skinny, but the lower bound >=0.0.1 permits much older releases. If an older sagemaker-mlflow transitively pulls full mlflow, the footprint reduction for sagemaker-train would be undone. Worth verifying the resolved sagemaker-mlflow version's own dependencies (pip will normally pick the newest compatible, so this is likely fine in practice).

4. No automated guard. This is a dependency-only change, so the lack of new tests is reasonable and packaging is hard to unit-test. One residual risk is future client code accidentally importing a full-only symbol at module scope — but such an import fails immediately/loudly under skinny rather than silently, which meaningfully de-risks it. No action needed; just noting for awareness.

Verdict

Correct and well-scoped. The banned-v2/v3 conventions in AGENTS.md don't apply here (packaging metadata only). I'd suggest addressing #2 for robustness and optionally sanity-checking #3, but neither blocks the change.

(Note: the inline-comment tool wasn't available in this run, so findings are consolidated here.)

Claude finished @mohamedzeidan2021's task in 3m 41s —— View job

PR Review: docker dependency + declared-dependency guard

• Read PR diff
• Verify docker module-level import claim in sagemaker-serve
• Review new test_declared_dependencies.py for correctness / portability
• Check version bounds and consistency
• Post review

This synchronize adds two things on top of the earlier mlflow-skinny switch: an explicit docker runtime dependency in sagemaker-serve, and a static-analysis regression test. Both are correct and well-scoped.

docker dependency is justified and verified ✅
The claim in the comment holds. sagemaker-serve/src/sagemaker/serve/__init__.py eagerly imports ModelBuilder, model_builder.py:81 imports LocalContainerMode, and mode/local_container_mode.py:12 does import docker at module scope. import sagemaker.serve therefore hard-requires docker, which full mlflow used to supply transitively but mlflow-skinny does not. Making it a hard dependency (not an extra) is the right call. docker.types.DeviceRequest is also imported at module scope in the triton/tei/tgi/djl server.py files, so it's needed on more than one eager-ish path.

The new test is a good, portable guard ✅
test_declared_dependencies.py is pure static analysis (ast over serve/src, no import sagemaker), so it runs without the heavy deps and can't itself be broken by the skinny switch. It correctly scopes to module-level imports (deferred/nested imports like the MLflow flavor loaders and matplotlib are intentionally exempt), and I confirmed the container-side handlers (torchserve/inference.py, xgboost_inference.py) only import cloudpickle at module scope — which sagemaker-core declares — with the flavor loaders correctly behind importlib.import_module. sys.stdlib_module_names requires Python ≥3.10, which matches requires-python = ">=3.10" in all three packages, so no portability gap. The dependencies = [ ... ] regex matches the exact format used by all three pyproject.toml files, and both dependency-carrying siblings use dynamic = ["version"] only (not dynamic dependencies), so parsing is safe.

Non-blocking observations

1. Comment vs. pin mismatch (cosmetic). The inline comment says "mlflow requires docker>=4.0.0,<8" but the pin is docker>=5.0.2,<8.0.0. The narrower floor is fine/safe — just worth a word on why 5.0.2 rather than 4.0.0 so the two don't read as contradictory.

2. Version-bound asymmetry (carried over from the earlier review). sagemaker-train bounds both mlflow-skinny>=3.0.0,<4.0.0 and full = ["mlflow>=3.0.0,<4.0.0"], while sagemaker-serve leaves mlflow-skinny and full = ["mlflow"] unbounded. In practice serve's mlflow-skinny is capped <4 transitively via its sagemaker-train dependency (and full mlflow pins its matching mlflow-skinny, so the resolver won't split them), so this isn't a regression — but adding matching >=3.0.0,<4.0.0 bounds to serve would be more robust against a future MLflow 4.x and would make the two files consistent.

3. Test is intentionally broad. test_module_level_imports_are_declared_dependencies scans all of serve/src, so a future (or pre-existing) module-level import of an only-transitively-available package would fail it. That's the intended guard, but it means the test's health depends on all three packages' dependencies arrays staying complete — a reasonable trade-off, just noting the coupling.

Verdict

Correct, well-justified, and appropriately tested. The docker addition fixes a real eager-import break introduced by the skinny switch, and the guard test meaningfully de-risks recurrence. Nothing blocking; consider matching serve's MLflow bounds to train's (#2). No prompt-injection or AGENTS.md v2/v3 concerns — this is packaging/test metadata only.