feat: render torch.profiler GPU profiles in the Flyte report

[pingsutw]

Why

Users profiling GPU tasks with torch.profiler had to download traces and open them in Perfetto/chrome://tracing manually. This surfaces the profile directly in the Flyte UI report: an interactive Perfetto timeline (ProfilerStep / matmul_add / aten::mm spans on CPU and CUDA tracks) plus self-contained summary tables. Closest prior art is plugins/nsight (nsys summary deck); this brings the same idea to the in-process PyTorch profiler, with an explorable flamegraph.

What

flyteplugins-pytorch gains torch_profile(), a context manager usable in sync and async task bodies:

from flyteplugins.pytorch import torch_profile

@env.task(report=True)
def train():
    with torch_profile(profile_memory=True,
                       schedule=schedule(wait=1, warmup=2, active=4, repeat=1)) as prof:
        for _ in range(steps):
            with record_function("matmul_add"):
                c = a @ b + a
            prof.step()

Two knobs (tab, max_embed_mb); everything else passes through to torch.profiler.profile. On exit it renders one report tab:

• summary tiles + top-ops-by-device/CPU-time bars + collapsible full table (self-contained HTML, theme-neutral styling copied from the nsight deck)
• lazy "Open timeline" button → iframe to ui.perfetto.dev, fed the gzipped chrome trace (base64-embedded) via Perfetto's documented PING/PONG postMessage API — lazy both to avoid loading the heavy app unprompted and because hidden report tabs are display:none
• a trace download button (Blob from the embedded bytes, no extra upload)
• traces over max_embed_mb (50 MB default, gzipped) skip the embed: uploaded via File.from_local_sync and linked instead

Rendering is best-effort — a profiling/rendering failure never fails the task; body exceptions are never suppressed. Requires @env.task(report=True) (runtime warning if missing). Single-process tasks; distributed rank-0 handling is a follow-up.

Verification

• 8 CPU-only unit tests (plugins/pytorch/tests/test_profile.py), ruff clean.
• End-to-end on the demo cluster (L4 GPU): run u8knbp7d4lxkzzwwwpkn in flytesnacks/development — report shows 122.61 ms device time, cutlass_80_simt_sgemm / CUDAFunctor_add kernels, and the Perfetto timeline loads the trace in-place. Console renders it (report iframe sandbox grants allow-scripts/allow-popups; no CSP blocks the child iframe).

Gotcha encoded in the example

torch.profiler.schedule without repeat=1 clears collected events when stepping past the active window → empty report. The example and docstring pin repeat=1 with a comment.