Metadata-Version: 2.4
Name: metaflow_pyspy
Version: 0.1.2
Summary: Periodic Python process diagnostics for Metaflow tasks
Author: Outerbounds
Author-email: help@outerbounds.com
Requires-Python: >=3.8
Description-Content-Type: text/markdown
Requires-Dist: py-spy<0.5,>=0.4.2
Dynamic: author
Dynamic: author-email
Dynamic: description
Dynamic: description-content-type
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

# metaflow-pyspy

`metaflow-pyspy` adds a step-level `@pyspy` decorator to Metaflow. While a
task is running, it periodically discovers the task's live processes, captures
structured Python stacks with `py-spy`, records Linux blocking and resource
state, and uploads each compressed snapshot through the task's configured
Metaflow datastore.

The collector is best-effort by design. It never changes whether user code
succeeds, it never creates failure artifacts or failure metadata, and every
collector message is written to stderr with an `[@pyspy]` prefix.

## Install

```bash
pip install metaflow-pyspy
```

The package depends on `py-spy` but deliberately does not depend on a specific
Metaflow distribution, so it can be installed alongside Metaflow or
Outerbounds' Metaflow package.

Installing the plugin on the launcher does not necessarily install `py-spy`
inside a remote task image: Metaflow packages extension source with the flow,
but container dependencies are controlled separately. Make the binary
available in one of these ways:

- install `py-spy==0.4.2` in the task with `@pypi`;
- bake it into the task image; or
- set `PYSPY_EXECUTABLE` to an executable path inside the task container.

The collector checks `PYSPY_EXECUTABLE`, the task interpreter's sibling `bin`
directory, and then `PATH`. If it cannot find the binary, it reports the issue
on stderr and continues uploading the remaining process diagnostics.

## Use

```python
from metaflow import FlowSpec, step, pyspy


class TrainingFlow(FlowSpec):
    @pyspy(
        interval=120,
        timeout=25,
        max_parallel=8,
        max_local_mb=512,
        focus_files=["train.py", "my_package/**"],
    )
    @step
    def train(self):
        # User code, library calls, and Python subprocesses are monitored.
        ...
```

Options:

- `interval`: seconds between the start of capture cycles. The first capture
  starts immediately.
- `timeout`: maximum seconds allowed for one `py-spy` attach.
- `nonblocking`: pass `--nonblocking` to `py-spy` (default: `True`).
- `max_parallel`: maximum simultaneous `py-spy` attaches.
- `max_local_mb`: cap for snapshots retained locally after upload failures.
- `process_pattern`: optional regex recorded against scoped processes for
  viewer-side grouping.
- `focus_files`: source path patterns a viewer can use to emphasize relevant
  frames. These do not filter capture data; complete raw stacks are retained.

There is intentionally no native-stack mode in this plugin.

## What is monitored

Every cycle re-discovers live processes. In a remote task, scope is the task
container/cgroup; locally, scope is the task interpreter and its descendants.
The collector excludes itself and its own `py-spy` children.

For every scoped process it records process identity (`pid` plus Linux start
time), ancestry, command line, executable, selected distributed-training
environment variables, status, wait channel, current syscall, kernel stack,
thread blocking state, IO, limits, scheduling counters, and file descriptors.
Pipe and socket ownership is correlated across the process set. Cgroup v2
memory/CPU/IO/PID counters and pressure data plus `/dev/shm` capacity are also
captured.

`py-spy` is attached independently to every live Python interpreter detected
by executable, command line, or loaded `libpython`. This includes different
Python versions and virtual environments and Python children created by
`subprocess.run`, `multiprocessing`, launchers, and `@torchrun`. A failure to
attach to one interpreter does not prevent other captures.

Only an allowlist of distributed runtime variables is read: rank/master/CUDA
variables and `NCCL_`, `TORCH_`, `GLOO_`, and `OMP_` prefixes. Arbitrary task
environment variables are never uploaded.

## Storage contract

The task attempt comes before the diagnostics kind so other collectors can use
the same root later:

```text
<task datastore path>/attempt-<attempt>/diagnostics/
  node-<node identity>/
    index.json
    snapshots/
      <UTC timestamp>-<sequence>.json.gz
```

Snapshots and indexes are versioned structured JSON. A snapshot contains:

```text
schema_version, collector_version, capture, task, focus_files,
processes, ipc, cgroup, filesystem
```

Successful captures are uploaded after every cycle and the small per-node
index is updated last. Failed uploads remain in a bounded local queue and are
retried on the next cycle. No py-spy error text is stored in snapshots; it is
only reported on stderr.

The task also receives normal metadata named `pyspy-diagnostics-root`, which
points to its attempt's diagnostics URI.

The decorator also creates a portable, decorator-owned artifact named
`self.pyspy_diagnostics`. Once the task is persisted, it is available through
the Metaflow Client API as `task.data.pyspy_diagnostics`:

```python
from metaflow import Task

task = Task("TrainingFlow/123/train/abc")
reference = task.data.pyspy_diagnostics

print(reference["root"])
print(reference["attempt"])
```

The artifact is a plain dictionary containing `type`, `schema_version`,
`root`, `pathspec`, `attempt`, and `datastore_type`.

## Reading captures

The top-level `PySpy` API is intended for notebooks, scripts, and an
Outerbounds Deployment viewer:

```python
from metaflow import PySpy
from metaflow import Task

task = Task("TrainingFlow/123/train/abc", attempt=0)
diagnostics = PySpy(task)

print(diagnostics.root)
print(task.data.pyspy_diagnostics)
print(diagnostics.nodes())
entries = diagnostics.snapshots()
latest = diagnostics.latest()
snapshot = diagnostics.load(entries[0])
```

The storage contract and reader API make the data directly suitable for a
separate authenticated Outerbounds App/Deployment showing a capture timeline,
rank/process matrix, repeated stack signatures, IPC relationships, cgroup
pressure, and `focus_files` frame emphasis. The web viewer is intentionally a
separate deliverable; installing this decorator does not deploy an app.

## Lifecycle and limitations

The collector starts in Metaflow's `task_pre_step` hook and stops in its
`task_post_step` or `task_exception` hook, before task finalization. It covers
subsequent decorator setup/wrappers, user step code, library calls, and child
processes. Hook ordering means it cannot guarantee capture of decorator hooks
that run before its `task_pre_step` or after its stop hook. It also cannot
observe environment bootstrap, dependency installation, or task launch. A
future supervisor/bootstrap mode could cover that wider window.

On Linux, `py-spy` attach permissions depend on the runtime's ptrace policy.
The decorator does not mutate Kubernetes `securityContext`. If attachment is
denied, enabling `SYS_PTRACE` or adjusting the runtime ptrace policy is a
deployment-specific troubleshooting step and a possible future integration.

## Integration test

[`integration_tests/pyspy_flow.py`](integration_tests/pyspy_flow.py) launches a
real monitored task and validates the uploaded index, snapshot schema, process
data, and at least one successful py-spy capture. See
[`integration_tests/README.md`](integration_tests/README.md) for local,
Kubernetes, and AWS Batch commands.
