Skip to content

feat(measurement): add W&B benchmark support#191

Merged
binaryaaron merged 20 commits into
mainfrom
binaryaaron/wandb-support
Jul 8, 2026
Merged

feat(measurement): add W&B benchmark support#191
binaryaaron merged 20 commits into
mainfrom
binaryaaron/wandb-support

Conversation

@binaryaaron

@binaryaaron binaryaaron commented Jun 16, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add opt-in W&B publication, reports, workspaces, and sweep orchestration for benchmark tooling only.
  • Replace loose dictionary and Pandas egress with strict measurement ingress and typed W&B config, history, summary, table, and report models.
  • Add atomic completion seals and a strict, retry-safe importer for externally produced Slurm benchmark cases.
  • Preserve historical v1 report reads while publishing v2 metadata with explicit native-suite, sweep-arm, and imported-config comparison axes.
  • Centralize scalar measurement fields in the product package and derive W&B aggregation, validation, and report metric paths from shared immutable registries.

Security and Privacy Boundary

  • Capture measurement JSONL once through pinned descriptors with owner, regular-file, hard-link, byte, line, record, nesting, mutation, and terminal-state checks.
  • Reject trace types, unknown records, coercions, non-finite values, negative counts, mixed runs, symlinks, FIFOs, sockets, devices, parent redirection, growth, and truncation before SDK import.
  • Build complete outbound payloads through extra="forbid" models. Raw text, prompts, model responses, replacement maps, entity payloads, local paths, provider secrets, and trace data cannot enter W&B payloads.
  • Isolate W&B in a guarded process environment. WANDB_API_KEY and ambient routing variables are removed; online runs use the local W&B credential file under the preserved home directory.
  • Keep native publication best-effort so observability failures do not change benchmark or sweep outcomes. Keep sealed Slurm imports strict across init, config, log, summary, and finish.
  • Keep table exposure as an explicit fail-closed allowlist while deriving operational scalar policy from package-owned field groups.

Slurm Import Contract

  • write_completion_seal.py atomically binds case identity, Slurm provenance, producer commit, terminal status, byte and record counts, and measurement SHA-256.
  • import_wandb_run.py verifies the seal and snapshot before initialization, derives a destination-scoped stable run ID, and publishes imported history at deterministic step 0 with resume="allow".
  • Identical imports become remote no-ops. Changed sealed content receives a different run ID.
  • The companion anonymizer-experiments launcher patch writes seals after report completion, requires a seal for SKIP_COMPLETED, and removes stale seals before forced reruns.

Reports and Sweeps

  • Parse historical v1 and current v2 runs into a closed WandbRunView.
  • Escape Markdown by rendering context, reconstruct run URLs from validated identifiers, and reject mixed or ambiguous comparison groups.
  • Preserve completed sweep-arm results when report or workspace creation fails, while returning a redacted view error and nonzero CLI status.
  • Compose report metric families from namespace helpers and derive complete report coverage from the panel registry.

Validation

  • make test (1025 passed)
  • Measurement tool suite (145 passed, council verifier)
  • Product measurement plus tool suites (189 passed; two sandbox-only integration tests excluded locally)
  • make format-check
  • Scoped ty check over changed measurement modules
  • git diff --check
  • Real sequential offline SDK publication and strict sealed retry
  • Online strict retry using local-file authentication with WANDB_API_KEY removed
  • Remote API and file readback: one history row at step 0, matching publication markers and seal digest, no configured privacy canaries
  • Four-member correctness, architecture, and privacy council review with no remaining actionable findings

Notes

  • W&B remains disabled by default.
  • Product CLI and SDK runs do not initialize W&B.
  • Workspaces and reports remain noncanonical presentation layers over sanitized benchmark data.
  • The package-owned scalar field catalog is a precursor to the broader typed measurement-event migration. Closed event models, shared codecs, typed emitters, and producer migration remain deferred to separate changes.

@binaryaaron binaryaaron marked this pull request as ready for review June 25, 2026 18:28
@binaryaaron binaryaaron requested review from a team as code owners June 25, 2026 18:28
@greptile-apps

greptile-apps Bot commented Jun 25, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR introduces opt-in W&B (Weights & Biases) benchmark observability to the measurement tooling, gated behind WandbMode.disabled as the default so no production path is affected. The implementation is notably defensive: field-level outbound-policy registry with compile-time completeness validation, O_NOFOLLOW/O_NONBLOCK/O_CLOEXEC file-descriptor traversal to prevent TOCTOU attacks, atomic seal writes with os.replace + os.fsync, and a process-environment transaction that strips credentials before any SDK call.

  • New ingress & publication pipeline (wandb_ingress.py, wandb_setup.py, wandb_completion.py): snapshot capture with SHA-256 integrity, two-phase seal verification before any SDK call, and idempotent publication via a resume="allow" / step-0 pattern.
  • Strict outbound privacy boundary (wandb_models.py, wandb_metric_schema.py, wandb_logging.py): all SDK payloads are extra="forbid" Pydantic models; a field-level policy registry (OUTBOUND_FIELD_POLICIES) validated at import time ensures every measurement field has an explicit exposure decision.
  • Sweep & import CLIs (sweep_benchmarks.py, import_wandb_run.py): per-arm failure isolation, deterministic 128-bit stable run ID for idempotent Slurm imports, and a SweepCliOptions dataclass replacing fragile locals() forwarding.

Confidence Score: 5/5

Safe to merge. W&B integration is opt-in and disabled by default; no existing production paths are altered.

The change is purely additive benchmark tooling with strong defensive coding patterns throughout. Three P2 style/consistency findings were identified, none of which affect correctness or security of the default-off feature. No P0 or P1 issues were found.

No files require special attention — all three findings are non-blocking style issues.

Important Files Changed

Filename Overview
tools/measurement/measurement_tools/wandb_ingress.py New snapshot-capture module with O_NOFOLLOW traversal, before/after fstat integrity check, bounded reads, and JSON pre-scan; no issues found.
tools/measurement/measurement_tools/wandb_setup.py WandbSdkEnvironment + WandbPublisher; minor P2: _open_directory_no_follow uses if 'descriptor' in locals() instead of the sentinel-based pattern used by its companion function.
tools/measurement/measurement_tools/wandb_completion.py CompletionSeal atomic write and verification helpers; minor P2: suppress(FileNotFoundError) in finally block should be suppress(OSError) to avoid masking the original error.
tools/measurement/measurement_tools/wandb_models.py Strict outbound-payload models with compile-time completeness check via validate_outbound_field_policies(); design is sound.
tools/measurement/measurement_tools/wandb_metric_schema.py ScalarMetricAccumulator and aggregation definitions; straightforward and correct.
tools/measurement/measurement_tools/wandb_logging.py build_outbound_measurements and typed-table helpers; closed WANDB_TABLE_ROW_MODELS allowlist prevents unintended exposure.
tools/measurement/sweep_benchmarks.py Sweep orchestration with per-arm failure isolation and SweepCliOptions dataclass; entity guard on view creation is correct.
tools/measurement/import_wandb_run.py Two-phase import with seal verification before SDK call and deterministic stable_import_run_id; design is correct.
tests/conftest.py load_tool fixture mutates sys.path without cleanup, risking cross-test module resolution pollution (P2).
tools/measurement/measurement_tools/wandb_report_models.py Strict report-view models with URL validation and ComparisonAxis discriminated union; no issues found.
src/anonymizer/measurement/fields.py Canonical scalar field group definitions used by both the package and W&B metric schema.
src/anonymizer/measurement/recorders.py extra_fields routing fix: now only injected into local_fields for model_workflow records, preventing extra='forbid' rejection.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant CLI as CLI / Slurm Job
    participant Ingress as wandb_ingress
    participant Seal as wandb_completion
    participant Env as WandbSdkEnvironment
    participant Publisher as WandbPublisher
    participant SDK as W&B SDK

    CLI->>Ingress: capture_ingress_bytes(path)
    Ingress->>Ingress: O_NOFOLLOW open + fstat validation
    Ingress->>Ingress: bounded read + before/after stat
    Ingress->>Ingress: JSON nesting pre-scan
    Ingress-->>CLI: MeasurementRecord

    CLI->>Seal: build_completion_seal(record)
    Seal->>Seal: SHA-256 digest
    Seal->>Seal: atomic write (tmp → os.replace + fsync)
    Seal-->>CLI: CompletionSeal

    CLI->>Publisher: publish_payload(payload, staging_dir)
    Publisher->>Seal: verify_completion_seal(staging_dir)
    Seal-->>Publisher: verified
    Publisher->>Env: __enter__ (strip env, set WANDB_API_KEY)
    Env-->>Publisher: clean environment
    Publisher->>SDK: "wandb.init(resume="allow")"
    Publisher->>SDK: log config / history / summary
    Publisher->>SDK: wandb.finish()
    Publisher->>Env: __exit__ (restore env)
    Publisher-->>CLI: WandbPublishResult
Loading
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant CLI as CLI / Slurm Job
    participant Ingress as wandb_ingress
    participant Seal as wandb_completion
    participant Env as WandbSdkEnvironment
    participant Publisher as WandbPublisher
    participant SDK as W&B SDK

    CLI->>Ingress: capture_ingress_bytes(path)
    Ingress->>Ingress: O_NOFOLLOW open + fstat validation
    Ingress->>Ingress: bounded read + before/after stat
    Ingress->>Ingress: JSON nesting pre-scan
    Ingress-->>CLI: MeasurementRecord

    CLI->>Seal: build_completion_seal(record)
    Seal->>Seal: SHA-256 digest
    Seal->>Seal: atomic write (tmp → os.replace + fsync)
    Seal-->>CLI: CompletionSeal

    CLI->>Publisher: publish_payload(payload, staging_dir)
    Publisher->>Seal: verify_completion_seal(staging_dir)
    Seal-->>Publisher: verified
    Publisher->>Env: __enter__ (strip env, set WANDB_API_KEY)
    Env-->>Publisher: clean environment
    Publisher->>SDK: "wandb.init(resume="allow")"
    Publisher->>SDK: log config / history / summary
    Publisher->>SDK: wandb.finish()
    Publisher->>Env: __exit__ (restore env)
    Publisher-->>CLI: WandbPublishResult
Loading

Reviews (15): Last reviewed commit: "test(measurement): make W&B staging regr..." | Re-trigger Greptile

Comment thread tools/measurement/measurement_tools/wandb_setup.py Outdated
Comment thread tools/measurement/sweep_benchmarks.py Outdated
Comment thread tools/measurement/sweep_benchmarks.py Outdated
Comment thread tools/measurement/sweep_benchmarks.py Outdated
Comment thread tools/measurement/measurement_tools/wandb_setup.py Outdated
Comment thread tools/measurement/measurement_tools/wandb_logging.py Outdated
Comment thread tools/measurement/measurement_tools/wandb_logging.py Outdated
Comment thread tools/measurement/create_wandb_report.py Outdated
Comment thread tools/measurement/create_wandb_report.py Outdated
Comment thread tools/measurement/sweep_benchmarks.py
Comment thread tools/measurement/sweep_benchmarks.py Outdated
Comment thread tools/measurement/sweep_benchmarks.py Outdated
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
@binaryaaron binaryaaron force-pushed the binaryaaron/wandb-support branch from 7831624 to 7e360e5 Compare June 30, 2026 20:30
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Comment thread tools/measurement/sweep_benchmarks.py Outdated

@andreatgretel andreatgretel left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Thanks Aaron. I re-reviewed the current head, and all six of my previous comments are addressed.

I also looked at how the sealed SLURM import would fit RAT Bench and similar distributed runs. I do not think the full integration needs to be part of this PR. The exporter and launcher wiring, stable case identities, seal lifecycle, batch import, sacct collection, and RAT Bench-specific metrics can evolve separately in our internal experiments repo.

I left two focused comments: one correctness issue with stage aggregation and one request to support multiple SLURM jobs per logical imported case. Richer scheduler timing and resource metadata can be added later.

The existing thread about invalid W&B identifiers escaping the post-run error guard also still looks actionable.

Comment thread tools/measurement/measurement_tools/wandb_logging.py
Comment thread tools/measurement/measurement_tools/wandb_completion.py Outdated
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Comment thread src/anonymizer/measurement/recorders.py Outdated
Comment thread tools/measurement/import_wandb_run.py
Comment thread tools/measurement/README.md
@andreatgretel

Copy link
Copy Markdown
Collaborator

Validated the sealed external-import flow end to end with a multi-stage SLURM case. The importer accepted seven strict measurement records, preserved multiple role-tagged SLURM job IDs, and an identical re-import was idempotent. Destination-scoped re-import also produced a distinct completed run as intended.

The only correctness issue found was the native NDD local_fields mismatch described inline.

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
@andreatgretel

Copy link
Copy Markdown
Collaborator

This seems to work well overall. I validated the sealed import flow with three 100-record benchmark runs, including multi-job SLURM provenance and W&B publication.

The main follow-up I noticed is that some of the generic metrics are difficult to interpret. For example, per-record mean F1 can use a different denominator for each workload and differ substantially from micro F1.

I think we could address this in a focused follow-up PR without blocking this one, unless you would prefer to include it now. Some things I'm thinking of:

  • Publish strict and relaxed micro precision/recall/F1 with their denominators.
  • Separate pipeline completion from record failures and coverage.
  • Distinguish unavailable telemetry from measured zero.
  • Reduce duplicate or low-signal summary fields.
  • Leave benchmark-specific dashboards and headline metrics to the experiment tooling.

The underlying typed records already preserve the required counts, so this should be additive and should not require rerunning benchmarks.

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Comment thread tools/measurement/measurement_tools/wandb_setup.py
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

@andreatgretel andreatgretel left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Thanks Aaron. I re-reviewed the current head and resolved my previous threads. The fixes for strict NDD ingress, importer destination/publication state, and credential preflight look good. The integration remains opt-in and CI is green.

One issue before approval: test_wandb_staging_allows_repeated_basename_in_foreign_intermediate_directory resolves descriptors through /proc/self/fd, which is unavailable on macOS and causes make test to fail. Could we make this portable by comparing st_dev and st_ino with temp_root.stat()?

Non-blocking: the PR is still quite large. Reports, workspaces, and sweeps look like the cleanest candidates for a follow-up if reducing the current diff is still practical. The metric-semantics improvements discussed earlier can also remain a focused follow-up.

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
@binaryaaron binaryaaron merged commit 30a621f into main Jul 8, 2026
12 checks passed
@binaryaaron binaryaaron deleted the binaryaaron/wandb-support branch July 8, 2026 13:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants