feat(pypi): derive Metaflow @pypi environments from uv.lock

[Abhishek Patil (abhishek-pattern)]

What

Adds @uv_pypi_base / @uv_pypi — Metaflow's @pypi_base / @pypi with the Python version and packages filled in from the repo's own uv.lock, so a flow's environment cannot drift from what uv sync installs.

from metaflow import FlowSpec, step

from ds_platform_utils.metaflow import uv_pypi_base


@uv_pypi_base
class MyFlow(FlowSpec):
    @step
    def start(self):
        self.next(self.end)

    @step
    def end(self):
        pass

No python=, no packages=, nothing to keep in sync. Both decorators work bare or called:

uv_pypi_base(flow=None, *, dependency_groups=None, python=None, project_root=None)
uv_pypi(step=None,      *, dependency_groups=None, python=None, project_root=None)

Docs: docs/metaflow/pypi_packages.md.

Why uv.lock rather than pyproject.toml

pyproject.toml holds constraints, not versions. A >= bound, or a git dependency pinned to rev = "main", means two bakes a week apart can produce different images. uv.lock records what uv actually resolved, including the commit SHA for git dependencies, so a bake is reproducible.

Resolving a universal lockfile

uv.lock is a universal resolution — every Python version and platform in range at once, each tagged with a marker. @pypi takes a flat {name: version} map with nowhere to put a marker, so markers are resolved at decoration time:

Lock state	Emitted
One entry, or several with a marker that selects one	that exact version
Marker excludes this environment (sys_platform == 'darwin' on a Linux bake)	omitted
Several entries, nothing to tell them apart	"", left for @pypi

Concretely, from one unchanged lockfile:

python=3.10  ->  pandas 2.3.3
python=3.11  ->  pandas 3.0.5

The Python version is resolved first and handed down, so the interpreter and the packages cannot disagree. Markers evaluate against Linux, which is what Metaflow builds for a remote task. Only the root project's direct dependencies are emitted — @pypi resolves transitives itself, and per-platform wheel availability is its job.

Where the Python version comes from

1. .python-version — the interpreter uv pinned
2. requires-python in uv.lock, else pyproject.toml — a range, so its floor
3. The running interpreter

An installed package cannot find the flow repo from __file__, so project files are located by walking up from the launch directory; project_root= overrides with an exact directory.

Output, once per run

Each decorator prints what it resolved:

@uv_pypi_base on MyFlow: python 3.10, 10 package(s) from uv.lock
  jinja2                      3.1.6
  outerbounds                 0.12.39
  polars                      (unpinned)
  snowflake-connector-python  4.7.2

Getting this to appear once per run rather than once per step took two attempts, and the reason is worth knowing for anyone touching it. Metaflow re-imports the flow module — re-evaluating every decorator — in each process that touches a task: the local worker it spawns per task, and the container. current.is_running_flow cannot gate that, because current is populated by the task runtime, which starts after the module is imported and the decorators ran. So three checks share one _is_running_task() helper:

• MF_PATHSPEC — exported by Metaflow's mflog helper into every remote task command (Kubernetes, Batch, Argo, Step Functions).
• step / spin-step in sys.argv — covers the local worker, whose environment is only a copy of the client's and so carries no marker variable.
• current.is_running_flow — catches what neither of the others can see: a flow module re-imported mid-run, e.g. one flow triggering another.

Verified on a real local run: one block at the top, nothing under any task.

Failing loudly on a missing lockfile

A missing uv.lock used to yield an empty packages map, so a flow launched from the wrong directory baked an empty environment and failed later somewhere less informative. It now raises FileNotFoundError at import, naming the directory searched and suggesting the fixes.

This cannot be unconditional: Metaflow's code package carries only .py files, so uv.lock never reaches the container, and the task re-imports the flow module — every remote run would die at import. A process already running a task therefore still gets an empty map, which is correct there.

pyproject.toml is deliberately not required; it is only the middle tier of the Python-version fallback, and plenty of projects declare no requires-python.

Reviewer notes

• d105d2f changes what flows bake. Moving the Python floor 3.9 → 3.10 is what selects pandas 2.3.3 over 3.0.5. requires-python is a published constraint, so this affects consumers of the library, not just this repo. Worth a deliberate look.
• fb42cc9 bumps the version to 0.6.0.
• packaging is now a declared runtime dependency — it was only present transitively, and is used for PEP 508 marker and specifier parsing.
• Public surface is just the two decorators. The helpers that build the map are private and tested through the module.
• docs/metaflow/private_repo_access.md (be9ffc5) is unrelated to the code: it documents the Fast Bakery Git credentials needed for private patterninc git dependencies, and records that the prod perimeter has no GIT_PYPI_REPOSITORY integration — so a flow deployed there with a private dependency will fail to bake until one is created.
• Pre-existing, not from this PR: the 3.10 floor moves ruff's inferred target, so ruff 0.16 now wants PEP 604 annotations across ~158 sites repo-wide. This PR keeps the existing Optional/Union style rather than mixing. Either declare requires-python and modernize in one dedicated pass, or pin target-version in [tool.ruff]. Expect CI lint to flag it.

Testing

34 unit tests in tests/unit_tests/metaflow/test__pypi_packages.py. Coverage: version pinning, git SHA references, marker resolution by Python version, platform-gated dependencies being dropped, the unpinned fallback, all three Python-version tiers, dependency groups, the upward project-root search, the raise and its task-process exemption, and the print suppression — parametrized over each signal in isolation, including a remote task whose argv looks like the client's.

Decorator tests assert against a spy on Metaflow's pypi_base, since the contract is "call Metaflow with this environment" and Metaflow 2.19 already moved where decorators are recorded (_flow_decorators → _flow_state). One test exercises the real decorator end to end.

Not covered: a real Kubernetes run with an actual Fast Bakery image build. Suppression happens at decoration time and is independent of baking, but that path has only been exercised locally.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

Introduces Metaflow decorators that derive @pypi_base / @pypi configuration (Python version + direct dependency pins) from the repository’s uv.lock, aiming to prevent environment drift between local uv sync and Metaflow-baked images.

Changes:

• Added uv_pypi_base and uv_pypi decorators that compute python= and packages= from .python-version / requires-python and uv.lock.
• Implemented uv.lock parsing + PEP 508 marker evaluation to emit a flat name -> version/direct-reference map for Metaflow.
• Added docs and unit tests covering version selection, marker handling, git SHA pinning, dependency groups, and project-root discovery; updated runtime deps (packaging) and bumped default Python pin to 3.10.

Reviewed changes

Copilot reviewed 7 out of 8 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/unit_tests/metaflow/test__pypi_packages.py	New unit tests exercising uv.lock parsing and decorator behavior.
src/ds_platform_utils/metaflow/pypi_packages.py	New implementation of uv.lock → Metaflow @pypi(_base) environment derivation and wrappers.
src/ds_platform_utils/metaflow/__init__.py	Exports uv_pypi and uv_pypi_base as the public API.
README.md	Adds docs link for the new decorators.
pyproject.toml	Adds runtime dependency on packaging and bumps the polars lower bound.
docs/metaflow/pypi_packages.md	New documentation page for uv_pypi_base / uv_pypi.
.python-version	Updates the repo’s pinned Python version to 3.10.

Suppressed comments (1)

src/ds_platform_utils/metaflow/pypi_packages.py:295

• _split_lock_packages overwrites root every time it sees any local-source package (virtual/editable/directory). If a lock contains additional local directory/editable dependencies, the last one wins and the derived dependency list can come from the wrong project entry. Avoid overwriting root once it has been found.

        if any(key in source for key in _LOCAL_SOURCE_KEYS):
            # the repo itself -- it is the thing depending on everything else, not a dep.
            root = package
            continue

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 9 changed files in this pull request and generated no new comments.

Suppressed comments (2)

src/ds_platform_utils/metaflow/pypi_packages.py:453

• The docstring example imports _get_pypi_kwargs from ds_platform_utils.metaflow, but it is not exported there (only uv_pypi/uv_pypi_base are). As written, the example will raise ImportError.

    from ds_platform_utils.metaflow import _get_pypi_kwargs

src/ds_platform_utils/metaflow/pypi_packages.py:570

• Docstring references _should_log_environment, but that helper doesn't exist; the actual gating logic uses _is_running_task(). This makes the doc misleading for future maintainers.

    The resolved environment is printed, so a run records what it actually built rather than
    leaving you to re-derive it -- but only from a process that is not itself executing a task,
    per `_should_log_environment`. Nothing is printed when the map came back empty either, that
    being the remote task re-importing the flow module inside an already-baked image, where
    there is no lockfile to read and nothing worth reporting.

[Vinay Shende (vinay79n)]

Looks great, will add real value for ds-platfprm-utils users