Skip to content

Add smuggler utilities for JAX transforms#1167

Draft
dlwh wants to merge 5 commits intomainfrom
codex/add-thread-local-metrics-smuggler
Draft

Add smuggler utilities for JAX transforms#1167
dlwh wants to merge 5 commits intomainfrom
codex/add-thread-local-metrics-smuggler

Conversation

@dlwh
Copy link
Member

@dlwh dlwh commented Sep 17, 2025

Summary

  • add a Smuggler helper that records payloads inside JAX transforms and exposes a smugglify helper
  • integrate the tracker with the metrics smuggler and expose smuggle_* transforms
  • cover the new utilities with dedicated unit tests

Testing

  • uv run pytest tests/test_smuggler.py
  • uv run pre-commit run --all-files

https://chatgpt.com/codex/tasks/task_e_68ca34a5bf0483318d5b1f0ec1d77a46

Copy link

@chatgpt-codex-connector chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Codex Review: Here are some suggestions.

Reply with @codex fix comments to fix any unresolved comments.

About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you open a pull request for review, mark a draft as ready, or comment "@codex review". If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex fix this CI failure" or "@codex address that feedback".

Comment on lines 162 to 173
The context manager yields the metrics dictionary that metrics are logged into.
You can log this dictionary directly to the global tracker after the context manager exits.
"""
global _global_tracker, _should_use_callback
old_tracker = _global_tracker
global _should_use_callback
old_should_use_callback = _should_use_callback
_should_use_callback = False
local_tracker = DictTracker()
_global_tracker = local_tracker

try:
yield local_tracker.metrics
with metrics_smuggler.activate() as payload:
yield payload
finally:
_global_tracker = old_tracker
_should_use_callback = old_should_use_callback

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] defer_tracker_for_jit no longer swaps in a safe tracker

The new implementation of defer_tracker_for_jit only toggles _should_use_callback and activates metrics_smuggler but no longer replaces _global_tracker with a temporary DictTracker as it previously did. Code that calls current_tracker() or interacts with the global tracker inside the context will now touch the real tracker (e.g. a WandB tracker) while tracing JIT code, defeating the purpose of “deferring” logging and potentially triggering host I/O or errors during tracing. Restoring the previous behaviour of swapping _global_tracker for a disposable tracker during the context keeps these call sites safe and maintains backwards compatibility.

Useful? React with 👍 / 👎.

@dlwh dlwh marked this pull request as draft September 22, 2025 08:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant