Fix crashes in the ring and jaccl distributed backends

[jzdziarski]

Proposed changes

Three unrelated stability issues in the CPU-side distributed backends, all of which show up as SIGSEGVs on long-running multi-node jobs.

Use-after-free in the collectives. Encoder::dispatch only enqueues the lambda, and set_input_array/set_output_array are no-ops on the CPU backend, so nothing keeps the arrays' buffers alive until the task runs. cpu::eval queues a keep-alive holding the inputs, but excludes the output, which the collectives allocate themselves. Under memory pressure the allocator can reclaim and reuse the output block before the collective runs, and the stream thread then writes through a stale pointer. Capture data_shared_ptr() alongside the raw pointers in the jaccl and ring collectives so the buffers are pinned for the task's lifetime.

Unvalidated work completions in the jaccl ring. A failed or spurious completion carries an undefined wr_id, and the wire/buff decoded from it was used unchecked to index the per-wire counters, offsets, limits and the buffer pools. Out-of-range values produced a wild pointer in reduce_op/std::copy. Check wc.status and range-check the decoded indices, then drop the completion and keep draining. The in-flight counter is decremented before the check so the loop still terminates.

32-bit offset overflow. read_offset/write_offset in the mesh all_gather and n_steps in the ring all_gather were int while indexing payloads sized by int64_t, wrapping past 2GB and ~4GB per rank respectively. Widen them to int64_t.

Checklist

NOTE: Only reproducible around 420GB, so was not able to write a test for this fix.

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[zcbenz]

Do you happen to still have the stack trace of the crashes? It would help diagnose whether the fix is correct.

[zcbenz]

Since wr_id is undefined on failure, checking lr/lw/buff would not be meaningful?

I think the code should:

1. ensure wc[i].status is not IBV_WC_SUCCESS before polling;
2. only check wc[i].status != IBV_WC_SUCCESS for error.