Instrument the Metaflow Client to trace metadata requests

[valayDave]

When an AI agent (or any program) uses the Metaflow Client API, it triggers a chain of HTTP requests to the metadata service. For example, finding all failed tasks in a run requires iterating Run → Steps → Tasks — but it is currently not possible to know, without external tooling, how many requests this generates or how the count scales with run size.

The MetadataProvider base class is the single choke point through which all client queries flow via the get_object() method. Adding optional request tracing here would give a precise, per-operation count of API calls made by any client operation, enabling systematic measurement of client-side efficiency.

Goals

• Add an opt-in tracing mode to MetadataProvider that records: request count, object type requested, and hierarchy depth for each call to get_object()
• Verify the tracer captures all requests for a complete Run → Steps → Tasks traversal
• Document the interface clearly so it can be used by benchmark scripts in a follow-up issue

Instructions

This issue is a venue for questions and design discussion. The implementation should be non-breaking — tracing is disabled by default and has no effect on normal usage.

Spin up the dev stack (see Resources) and run a flow against your local service before starting, so you have real data to trace against.

When you have met the goals, show your work via a fork with an issue / PRs explaining your work. You can add comments with what you achieve on the issue thread. You can reference the Fork/PR/issue in this issue.

Resources

Dev stack:

• Setting Up the Dev Stack — run metaflow-dev up to get a full local environment (metadata service, DB, UI); use metaflow-dev shell to get a pre-configured shell pointing at it
• Once running, execute flows against your local stack from within the metaflow-dev shell — real data is needed to trace requests against

Key code:

• metaflow/metadata_provider/metadata.py#L79 — MetadataProvider base class
• metaflow/metadata_provider/metadata.py#L352 — get_object(), the method all client queries go through
• metaflow/client/core.py#L376 — MetaflowObject.__iter__(), where traversal is initiated
• metaflow/client/core.py#L2149 — Run class

Docs:

• Metaflow Client API reference
• Service architecture overview

[gupta-tilak]

Hi @valayDave ,

After going through metadata_provider/metadata.py and client/core.py, here's how I see the request chain working: every time the client API traverses the object hierarchy, it fires a separate HTTP call per level. Something like iterating over all tasks in a run to find failures plays out roughly like:

• Run.__init__ → get_object('flow', 'run', ...) — resolves the specific run
• Iterating the run → get_object('run', 'step', ...) — lists all steps in it
• Iterating each step → get_object('step', 'task', ...) — tasks per step, one call per step

So even for a modest run with 6 steps and 1 task each, that's already 8 requests before you've touched a single artifact or piece of metadata. Scale to a real production flow and an AI agent doing this in a loop and it adds up fast and the problem right now is you'd need external tooling (like a proxy or network sniffer) to see it happening.

The good news from reading the code is that get_object() at L352 really is the single chokepoint the issue describes. Every provider — whether it's the HTTP service provider or the local one — goes through this classmethod before hitting _get_object_internal(). And the ObjectOrder class is already sitting there mapping each type to its numeric depth (root=0, flow=1, run=2, step=3, task=4, artifact=5), so the "hierarchy depth" field the issue asks for is derivable for free from type_order which is already computed inside get_object(). No extra bookkeeping needed.

There are few ambiguities that I noticed

• get_object is a classmethod. Any tracer state has to live at the class level or be threaded in from outside. Can't be per-instance naturally.
• Two code paths exist inside get_object. The method calls _get_object_internal() and then optionally applies _reconstruct_metadata_for_attempt(). A trace entry at the top of get_object covers the outer call cleanly. The metadata reconstruction is a local in-memory filter (no HTTP), so it probably doesn't affect request counts, but worth confirming.
• Thread safety. Class-level mutable state (_tracer = self) means concurrent client calls from different threads would corrupt each other's trace. This might not matter for single-threaded benchmark scripts, but if the goal is eventually to use this inside AI agent loops that may be async, it's a real footgun.
• The local provider also routes through get_object. Tracing would capture filesystem lookups the same as HTTP calls. Could be useful for comparing providers, but I wasn't sure if that was intentional or not.

Proposed approach

A context manager that sets a class-level _tracer attribute on MetadataProvider when active, None when not. Inside get_object(), add a single if cls._tracer is not None check before calling _get_object_internal(). When tracing is off, that check is essentially free.

# metadata.py — two additions:

class MetadataProvider(object):
    _tracer = None  # opt-in; None means tracing is off

    @classmethod
    def get_object(cls, obj_type, sub_type, filters, attempt, *args):
        type_order = ObjectOrder.type_to_order(obj_type)
        sub_order = ObjectOrder.type_to_order(sub_type)
        # ... existing validation unchanged ...

        if cls._tracer is not None:                          # ← new (2 lines)
            cls._tracer._record(obj_type, sub_type, type_order, attempt, args)

        pre_filter = cls._get_object_internal(...)
        # ... rest unchanged ...

# tracer.py — new file:

class MetadataTracer:
    def __init__(self):
        self.calls = []

    def _record(self, obj_type, sub_type, depth, attempt, path_args):
        self.calls.append({
            "obj_type": obj_type,
            "sub_type": sub_type,
            "depth": depth,
            "attempt": attempt,
            "path": "/".join(str(a) for a in path_args if a),
            "ts": time.time(),
        })

    def __enter__(self):
        MetadataProvider._tracer = self
        return self

    def __exit__(self, *_):
        MetadataProvider._tracer = None

    def summary(self):
        counts = Counter(c["obj_type"] for c in self.calls)
        return {"total": len(self.calls), "by_type": dict(counts)}

    def report(self):
        print(f"Total get_object() calls : {len(self.calls)}")
        print(f"Calls by type            : {self.summary()['by_type']}")
        for i, c in enumerate(self.calls, 1):
            print(f"  {i:3}.  {c['obj_type']:<10} {c['sub_type']:<10} depth={c['depth']}  {c['path']}")

Why not subclassing? I thought about a TracingMetadataProvider that wraps _get_object_internal, but providers are resolved through the METADATA_PROVIDERS registry — in that case we will have to swap out the registered class or complicate the init path.

Why not ContextVar? It would make the tracer thread/context-local, which is the "correct" approach for async code. My hesitation is that it's meaningfully more complex to expose for what sounds like a benchmark scripting use case. But this is being used inside async agent loops, I'd prefer to go the ContextVar route from the start rather than bolt it on later.

I will soon share the PR and other implementation details for your review.

[gupta-tilak]

Hi @valayDave , I went ahead with the local dev setup and implementation, pls checkout the approach and let me know if I can go ahead with the PR

Dev stack is running

The full local stack came up clean via metaflow-dev up.

The flow I ran

Wrote a small Issue4DemoFlow with a branching structure — start → process_a / process_b / process_c → join → end — so there'd be meaningful step and task counts to trace against. Ran it from within metaflow-dev shell pointing at the local service.

$ python issue4_demo_flow.py run
Metaflow 2.x.x executing Issue4DemoFlow for user:tilakgupta
...
2026-03-06 07:56:22 [1/start/2] Task is starting.
2026-03-06 07:56:26 [1/process_a/3] Task is starting.
2026-03-06 07:56:30 [1/join/6] Task is starting.
2026-03-06 07:56:36 [1/end/7] Task is starting.
Done! The flow Issue4DemoFlow/1 completed successfully.

You can see the run in the Metaflow UI (run ID 1, completed, 6 steps, 7 tasks, duration 14s):

What I changed (3 files in the installed metaflow package)

metaflow/metadata_provider/tracer.py (new file) — The MetadataTracer context manager. __enter__ sets MetadataProvider._tracer = self, __exit__ clears it. _record() appends a dict per call. summary() returns totals; report() prints a formatted table.

metaflow/metadata_provider/metadata.py (two additions) — Added _tracer = None as a class attribute on MetadataProvider, and a 2-line hook at the top of get_object(), placed just after the type_order/sub_order are computed (before the validation block, so we can capture the call even if it would have raised), or optionally after validation — I put it after so we only count well-formed calls. Happy to discuss which is more useful.

metaflow/metadata_provider/__init__.py (one line) — Added from .tracer import MetadataTracer so it's importable cleanly.

The actual tracer output

Wrote a trace_demo.py script that runs this against the live local service:

from metaflow import Run
from metaflow.metadata_provider import MetadataTracer

with MetadataTracer() as tracer:
    run = Run("Issue4DemoFlow/1")
    failed_tasks = [t for s in run for t in s if not t.successful]

tracer.report()

Output:

Total get_object() calls : 14
Calls by type            : {'run': 2, 'step': 6, 'artifact': 6}

   1.  run        self      depth=2  Issue4DemoFlow/1              ← fetch the run itself
   2.  run        step      depth=2  Issue4DemoFlow/1              ← list all steps
   3.  step       task      depth=3  Issue4DemoFlow/1/start        ← tasks per step (×6)
   4.  step       task      depth=3  Issue4DemoFlow/1/process_a
   5.  step       task      depth=3  Issue4DemoFlow/1/process_b
   6.  step       task      depth=3  Issue4DemoFlow/1/process_c
   7.  step       task      depth=3  Issue4DemoFlow/1/join
   8.  step       task      depth=3  Issue4DemoFlow/1/end
   9.  artifact   self      depth=5  Issue4DemoFlow/1/start/2/_success    ← t.successful (×6)
  10.  artifact   self      depth=5  Issue4DemoFlow/1/process_a/3/_success
  11.  artifact   self      depth=5  Issue4DemoFlow/1/process_b/4/_success
  12.  artifact   self      depth=5  Issue4DemoFlow/1/process_c/5/_success
  13.  artifact   self      depth=5  Issue4DemoFlow/1/join/6/_success
  14.  artifact   self      depth=5  Issue4DemoFlow/1/end/7/_success

I found that the minimum hierarchy traversal (run → steps → tasks) for this 6-step, 7-task flow would be 1 + 1 + 6 = 8 calls. But the actual count is 14, because the if not t.successful guard on each task triggers a separate artifact → self fetch for the _success artifact, one per task.

So for an AI agent doing this pattern naively across many runs:

total_calls = 1 (run) + 1 (steps) + N_steps (tasks) + N_tasks (artifact fetch per t.successful)
            = 2 + N_steps + N_tasks

For a real flow with 20 steps and 100 tasks that's 122 calls — not 22.

What I'm not sure about yet

Where to hook the record call exactly. I placed it after the type/order validation (so only well-formed calls get counted) but before _get_object_internal. Alternative is placing it before validation — that would catch malformed calls too, which might be useful if you're debugging incorrect client usage. Currently leaning toward post-validation since the benchmark use case is measuring successful traversals.

Whether the local provider output (filesystem reads) should be visually distinguished from HTTP calls in the report. Right now they look identical in the trace. Could add a provider_type field derived from cls.TYPE, which is already set on each provider subclass. Not sure if that's in scope here or better as a follow-up.

Test strategy. I've been testing manually against the dev stack. For a proper PR I'd want to add either a unit test that mocks _get_object_internal and asserts the trace records are correct, or an integration test that spins up the local provider. Which would you prefer?

The code is sitting in my fork. I'll clean it up and open a PR once I hear back on these points - particularly the hook placement and whether provider_type is worth including.

[mandeepsingh2007]

hii @valayDave

i reviewed the @gupta-tilak 's pr and found that modifying the core MetadataProvider or ServiceMetadataProvider files directly introduces unnecessary tech-debt and thread-safety risks (like using class-level variables for tracers, which will break in multi-threaded/async agent loops).

I'd like to propose a Zero-Footprint Runtime Tracing approach. Instead of changing Metaflow's internal source code, we can build a standalone context manager that uses dynamic runtime patching (via unittest.mock.patch.object or a transparent Proxy class) to intercept _get_object_internal or _request calls purely from the client script.

Why this approach is better:

Zero Impact on Core: No modifications to metaflow/plugins/ or metaflow/client/. The tracer lives entirely in a separate benchmark utilities folder.

Zero Review Cost for Core Maintainers: It doesn't bloat the production codebase with observability logic.

Thread-Safe by Design: We can bind the trace context to the specific execution thread/async-task using contextvars.

If this aligns with your vision for keeping the core library clean, I already have the architecture mapped out and can share a quick script demonstrating this non-invasive tracer. Let me know your thoughts!

[gupta-tilak]

Hi @mandeepsingh2007,

Thanks for the review - I’d like to clarify a few technical points.

The issue explicitly asks for an opt-in tracer inside MetadataProvider, not external instrumentation. This matters because get_object() computes type_order and sub_order as local variables, the exact hierarchy data we need. External patching would have to re-derive this logic, which is brittle.

Patching _request also won’t work: it operates at the HTTP level and lacks semantic context (obj_type, hierarchy depth). The tracer is meant to capture logical operations, not raw requests.

Regarding fragility, _get_object_internal is private and can change anytime. In-tree instrumentation is safer since it evolves with the code.

On tech debt: the change is minimal (one attribute + a guarded call). When disabled, it’s just a None check with negligible overhead.

Thread safety is already acknowledged with a possible ContextVar approach under consideration.

Finally, downstream issues require this tracer to be part of the package, not an external utility.

If you want to contribute, a thread-safe (ContextVar) version or scaling benchmarks would be especially valuable.

[mandeepsingh2007]

Hii @gupta-tilak

Thanks for the detailed breakdown! You make a very fair point about type_order and sub_order being computed locally inside get_object(). Re-deriving that logic externally would indeed be brittle, and you're completely right that hooking at the _request level loses the semantic context (hierarchy depth) that the issue explicitly asks for. I agree that an opt-in, in-tree approach is the cleanest way to capture the exact logical operations here.

I'd be glad to take you up on the collaboration offer regarding the thread-safety aspect!

Implementing this with contextvars.ContextVar instead of a class-level attribute will ensure that the tracer's state is strictly isolated per concurrent thread/async-task. This is especially crucial for the downstream agentic loops where multiple parallel metadata queries might be firing simultaneously.

Would you prefer I open a quick PR targeting your fork to add the ContextVar implementation for the MetadataTracer, or should I submit it directly against the main repo referencing this discussion? Happy to help push this forward!

[Youssef-SH]

Hi @valayDave,
I added an opt-in tracer at MetadataProvider.get_object() using a ContextVar instead of class-level state, so tracing stays isolated per execution context.

Each call records (obj_type, sub_type, depth, attempt, path), with request count derived from len(records). The hook is placed after validation so only valid semantic queries are counted. When disabled, the overhead is just a single None check.

Ran this against a local dev stack flow (DemoFlow, 3 steps / 3 tasks):

Run -> Steps -> Tasks
observed: 4 calls (run=1, step=3)
expected: 1 + N_steps = 4

Same traversal with t.successful:
observed: 7 calls (run=1, step=3, artifact=3)
expected: 1 + N_steps + N_tasks = 7

Counts match the expected expansion.

Implementation:
Youssef-SH/metaflow@02d0d32