maxray 0.7.0

No project description provided

maxray

the problem with doing weird metaprogramming shit is having to deal with other people's weird metaprogramming shit

Trace and modify the result of (almost) every single expression executed in a Python program in realtime, requiring zero changes to your source code. Very WIP.

[!NOTE] Will not work in a REPL -- save and run the code as a script or package.

from maxray import transform, xray, maxray

from torch import tensor, Tensor, device


def move_tensor(x, ctx):
    if isinstance(x, Tensor) and x.device != device("cuda"):
        return x.to("cuda")
    return x


@transform(move_tensor)
def show_multiply(a, b):
    print(a @ b)

# Source code is rewritten to be equivalent to:

def _show_multiply(a, b):
    move_tensor(move_tensor(print, ...)(move_tensor(a, ...) @ move_tensor(b, ...)), ...)

# ---

show_multiply(
    tensor([[0.0, 1.0], [1.0, 1.0]], device="cpu"),
    tensor([[1.0], [1.0]], device="cuda"),
)  # Without the decorator, you'd expect `RuntimeError: Expected all tensors to be on the same device`

# tensor([[1.],
#         [2.]], device='cuda:0')

The ctx argument contains context information about the location of the original source code, which may be useful to build editor/LSP integrations.

The *xray decorators will recursively trace and patch every single callable they encounter until reaching either builtins, native, or generated code.

Usage

This package comes with 2 CLI commands (see --help for limited docs):

• xpy: TUI for debugging and observability. Replace python your_script.py with xpy your_script.py to apply transformations (via -W args) globally
  • -W takes file:symbol args... (be careful with shell quoting) generated by maxray template <new_script.py> and hot-reloads it on every filesystem change. Can be used multiple times to compose transformations (applied sequentially)
  • --restrict to only transform "your" (same source file or module) source code (i.e. does not descend into external packages). Significantly faster, more robust, and probably more useful - should probably be the default...
  • -m is equivalent to python -m some_module
• maxray
  • template: see --help for options

Examples

• callgraph: draws the runtime function call graph
  • compared to tools relying on static analysis, can exactly resolve calls like [fn_a, fn_b][rand() < 0.5]()

pip install maxray[all]
xpy -W 'callgraph:Draw --labels' <your_script.py>

• capture: dump all runtime captured information to an Arrow IPC file for analysis
• rerun: replace all uses of print with logging to a Rerun viewer, attaching source file and line information for filtering
• tensorcheck: global nan and inf checking for tensors
• strings: grepping of runtime values to find and rewrite URLs and filesystem paths
• progress: tqdm, but like, everywhere. Automatically shows a progress bar and the last yielded value for every for _ in iterable expected to take longer than 1s to complete
• ...: Write your own! Argument to -W can be any source file or installed package/module

Issues

If running a script or module with xpy deterministically results in differing output or exceptions compared to when run with python (CPython >= 3.11), that's a bug. A stopgap solution is to restrict the offending modules to prevent them from being transformed.

[!TIP] export MAXRAY_LOG_LEVEL=1 to show internal logging