Skip to content

[Feature] Unify transport selection for inference services#3958

Open
vmoens wants to merge 3 commits into
mainfrom
service-transport-selection
Open

[Feature] Unify transport selection for inference services#3958
vmoens wants to merge 3 commits into
mainfrom
service-transport-selection

Conversation

@vmoens

@vmoens vmoens commented Jul 8, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Description

Gives inference services one clear, configurable story for moving payloads independently of where the service owner runs:

  • service_backend (and the stable collector backend/policy_backend spellings) chooses ownership and placement;
  • transport chooses how payloads move ("queue", "shared_memory", "distributed", "ray", "monarch", or "auto");
  • .client() hides both choices from consumer code.

Follow-up to #3955 (SharedMemoryTransport), which this PR wires into collector/server-level selection.

What's included

Common configuration envelope (torchrl._comm.transport, re-exported from torchrl._comm and torchrl.services):

  • TransportKind literal, frozen TransportConfig(kind, options) dataclass, normalize_transport_kind / as_transport_config validation helpers, and a per-domain transport-resolver registry (register_transport_resolver / resolve_transport). Deliberately no common transport implementation ABC: queues, shared memory, object stores, and collectives have different semantics; domains register a resolver instead.

Public inference transport resolution (make_inference_transport):

  • Accepts None/"auto", a kind string, a TransportConfig, or a concrete InferenceTransport (returned unchanged), plus topology context (policy_backend, num_slots, request_spec/response_spec, transport_options).
  • "auto" resolves deterministically: thread actors with a known slot count get SlotTransport; process actors with a static TensorDict contract get SharedMemoryTransport; process actors otherwise get MPTransport; Ray/Monarch topologies get their transports. Partial specs and incompatible forced combinations raise with actionable messages instead of silently degrading. The resolution is logged at debug level.

DistributedTransport (new): placement and payload movement compose independently — in particular Ray placement + torch.distributed payloads:

  • Fixed-spec request/reply over torch.distributed point-to-point ops for actors in other processes or on other nodes (Ray actors being the primary case). backend="gloo" moves CPU tensors; backend="nccl" moves CUDA tensors GPU-direct with no host round trip.
  • Every server/client pair gets standalone ProcessGroupGloo/ProcessGroupNCCL objects rendezvoused through a transport-owned TCPStore — the default process group is never touched, so training code keeps init_process_group for its own collectives.
  • Small headers and pickled exceptions ride gloo control pairs; NCCL data receives are posted only after a header announces them (keeps the NCCL watchdog away from idle receives, keeps each communicator single-threaded). Blocking receiver threads feed in-process queues on both sides (gloo Work.is_completed() does not report completion for receives, so irecv-polling designs do not work — verified empirically).
  • Clients are picklable and join lazily inside the actor process; the server discovers clients through the store, so client creation is not ordered with respect to server startup.
  • Resolved by transport="distributed" for any non-thread actor topology; never chosen by "auto" (it needs a static payload contract and a rendezvous).

Wiring:

  • AsyncBatchedCollector, InferenceServer, and ProcessInferenceServer accept string/config transport choices plus transport_options; the collector's private _make_transport factory is replaced by the public resolver (num_slots defaults to the number of envs). ProcessInferenceServer resolves with a multiprocessing topology, so transport="auto" + specs gives the shared-memory fast path with zero extra plumbing.
  • Resolved choices are inspectable: transport_kind (derived from the concrete instance, so injected transports report too; unknown classes report "custom") and transport_config (None when an instance was injected) on all three classes.

Compatibility

  • backend/env_backend/policy_backend spellings unchanged; service_backend keeps placement semantics.
  • Concrete transport injection unchanged (transport=MPTransport(...) etc. still works everywhere).
  • Default resolution is behavior-compatible: threading -> SlotTransport(num_envs), multiprocessing -> MPTransport() unless specs enable shared memory, ray/monarch unchanged.

Tests

Extends test/test_inference_server.py: config parsing/validation/registry, the full auto and forced-kind resolution matrix, incompatible-combination errors, options forwarding, instance passthrough, kind introspection, server construction from strings, collector-level resolution, and an end-to-end collection run over a forced "queue" transport. DistributedTransport is covered in-process, across spawned processes, and from a real Ray task (verified locally with Ray: the pickled client rendezvouses via the TCPStore and moves payloads over gloo p2p while the driver-side server resolves), plus out-of-order futures, exception propagation, nested keys, timeout semantics, and spec validation.

Deferred (tracked follow-ups from the same design note)

  • NCCL end-to-end CI coverage for DistributedTransport (needs a multi-process single-GPU runner); the gloo path and NCCL spec validation are covered.
  • Hybrid inline/out-of-band payload handling for mailbox-backed logger calls (video artifacts).
  • Replay-buffer transport configuration and a process-backed replay data plane.
  • Feeding the envelope into weight-sync factories (placement=ray, transport=distributed(nccl) spelling for RayWeightSyncScheme).
  • The cross-domain support matrix and the common transport benchmark harness.

🤖 Generated with Claude Code

@vmoens @claude
[Feature] Unify transport selection for inference services
4fcdaab 
Give inference services one configurable story for moving payloads
independently of where the service owner runs: service_backend (and the
collector backend/policy_backend spellings) chooses placement,
transport chooses how payloads move, and .client() hides both from
consumer code.

Changes:
- Add TransportKind, TransportConfig, normalization helpers, and a
  per-domain transport-resolver registry in torchrl._comm.transport
  (re-exported from torchrl._comm and torchrl.services). There is
  deliberately no common transport implementation ABC: queues, shared
  memory, object stores, and collectives keep their own semantics.
- Add make_inference_transport, a public resolver for the inference
  request/reply domain that accepts None/"auto", a kind string, a
  TransportConfig, or a concrete InferenceTransport, plus topology
  context (policy_backend, num_slots, request/response specs). "auto"
  deterministically picks the fastest safe transport: per-producer
  slots for threads, shared-memory TensorDict slots for process actors
  with a static payload contract, plain queues otherwise, and the
  Ray/Monarch transports for their topologies. Incompatible forced
  choices raise instead of silently degrading.
- Accept string/config transport choices (plus transport_options) in
  InferenceServer, ProcessInferenceServer, and AsyncBatchedCollector,
  replacing the private collector transport factory. The resolved
  choice is inspectable through the new transport_kind and
  transport_config properties; inference_transport_kind reports the
  kind of injected instances too.
- Document the placement/transport separation and auto resolution
  rules in the inference-server reference.

Existing spellings are preserved: backend/env_backend/policy_backend
keep their meaning, concrete transport injection is unchanged, and
default resolution is behavior-compatible (threading -> SlotTransport,
multiprocessing -> MPTransport unless specs enable shared memory).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@pytorch-bot

pytorch-bot Bot commented Jul 8, 2026 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3958

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit af64a0f with merge base 2a0c554 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 8, 2026
@github-actions

github-actions Bot commented Jul 8, 2026

Copy link
Copy Markdown
Contributor

Benchmark Results: PR 4fcdaab6 vs main d26bf573

Benchmark run: https://github.com/pytorch/rl/actions/runs/28973096443

Higher ops/sec is better. Tables are sorted by largest absolute change.

CPU

Compared 220 benchmarks. Regressions over 5%: 11. Improvements over 5%: 20.

Benchmark main ops PR ops Change
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-ListStorage-None-400] 188.36 54.90 -70.85%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-None] 62.06 89.33 +43.93%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-400] 862.98 1,122 +30.00%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400] 48.18 36.91 -23.40%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-32-512] 28.96 34.44 +18.92%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-same] 2,840 2,448 -13.81%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-backward] 894.31 1,018 +13.81%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400] 510.96 448.42 -12.24%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 2,090 2,331 +11.53%
benchmarks/test_collectors_benchmark.py::test_single_with_rb 7.3545 6.5083 -11.51%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-backward] 127.67 139.79 +9.49%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-64] 9.2013 8.3636 -9.10%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 3,299 3,001 -9.04%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 1,971 2,130 +8.05%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400] 553.16 508.93 -8.00%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[reduce-overhead-None] 262.52 282.09 +7.45%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 3,233 3,038 -6.06%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 3,088 3,271 +5.92%
benchmarks/test_objectives_benchmarks.py::test_values[generalized_advantage_estimate-True-True] 96.67 90.98 -5.89%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-parallel-buffers-True] 0.4962 0.5246 +5.73%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-10000] 2,181 2,299 +5.42%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape1-atari] 263.17 277.41 +5.41%
benchmarks/test_objectives_benchmarks.py::test_values[td0_return_estimate-False-False] 7,482 7,879 +5.30%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-None] 261.58 275.39 +5.28%
benchmarks/test_envs_benchmark.py::test_simple 1.8044 1.7105 -5.21%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-None] 113.69 119.57 +5.17%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-backward] 116.40 122.39 +5.15%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[safetensors] 22,800 23,972 +5.14%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400] 524.53 551.35 +5.11%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[True-None] 276.98 291.13 +5.11%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-None] 327.89 344.47 +5.05%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-backward] 122.27 128.30 +4.93%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-backward] 58.58 61.45 +4.89%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 2,917 2,776 -4.85%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-False-False] 46,551 44,311 -4.81%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-64] 5.9519 5.6827 -4.52%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400] 889.48 929.48 +4.50%
benchmarks/test_envs_benchmark.py::test_transformed 0.9056 0.8659 -4.38%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 1,987 2,071 +4.20%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape2-large_img] 174.12 167.15 -4.00%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-backward] 408.79 424.66 +3.88%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[reduce-overhead-None] 258.11 267.99 +3.83%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[generalized_advantage_estimate-False-1-512] 108.23 112.30 +3.76%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-64] 11.19 10.77 -3.72%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400] 1,038 1,074 +3.49%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-False] 78,538 75,806 -3.48%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-None] 1,786 1,724 -3.43%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-None] 210.58 217.71 +3.39%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape1-atari] 699.59 722.05 +3.21%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape1-atari] 642.45 662.86 +3.18%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-64] 4.4789 4.6170 +3.08%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-True-1-512] 643.97 663.62 +3.05%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-True] 25.19 25.95 +3.02%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-1] 489.02 503.80 +3.02%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-backward] 81.84 84.30 +3.01%
benchmarks/test_objectives_benchmarks.py::test_values[td_lambda_return_estimate-True-False] 24.54 25.27 +3.00%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[False-backward] 64.61 62.67 -3.00%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[reduce-overhead-None] 115.75 119.21 +2.99%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-same] 1,854 1,909 +2.97%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[reduce-overhead-None] 332.31 341.69 +2.82%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-16] 47.76 49.10 +2.79%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[reduce-overhead-None] 283.41 291.14 +2.73%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-32-512] 552.25 566.83 +2.64%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[False-None] 49.52 50.83 +2.63%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 3,564 3,471 -2.62%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape2-large_img] 405.09 415.71 +2.62%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-True] 39,127 38,136 -2.53%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-1] 282.49 275.56 -2.46%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[200-img_shape3-large_batch] 329.72 337.70 +2.42%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-backward] 284.47 291.35 +2.42%
benchmarks/test_envs_benchmark.py::test_parallel 0.9569 0.9338 -2.42%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-False] 64,979 63,429 -2.39%
benchmarks/test_replaybuffer_benchmark.py::TestWindowingTransformsBenchmark::test_action_chunk_transform[no_done] 1,582 1,619 +2.38%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-16] 18.01 18.43 +2.36%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-False] 32,223 31,475 -2.32%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-backward] 58.02 59.35 +2.29%
benchmarks/test_replaybuffer_benchmark.py::TestWindowingTransformsBenchmark::test_catframes_offline 529.35 541.33 +2.26%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-backward] 131.75 134.70 +2.24%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-4] 48.49 49.57 +2.24%
benchmarks/test_collectors_benchmark.py::test_sync_preempt 16.14 16.50 +2.24%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[untyped_storage] 8.2890 8.1044 -2.23%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[50-img_shape0-small] 3,515 3,593 +2.22%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 3,414 3,489 +2.22%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-True] 20,029 19,591 -2.19%
benchmarks/test_replaybuffer_benchmark.py::test_replay_buffer_direct_client_identity 8,100,540 7,927,533 -2.14%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[reduce-overhead-None] 83.15 84.92 +2.12%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-None] 259.49 264.85 +2.07%
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sample_mixed_devices[1000000-memmap_cpu_storage_cpu... 75.23 76.76 +2.03%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[50-img_shape0-small] 4,418 4,507 +2.03%
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sampler_sample_scale[10000000-cpu] 50.47 51.49 +2.01%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb[100-img_shape0-atari] 29.32 29.90 +1.99%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-4] 71.80 73.20 +1.95%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[reduce-overhead-None] 1,797 1,832 +1.94%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-1] 620.26 632.30 +1.94%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-16] 42.01 42.82 +1.93%
benchmarks/test_envs_benchmark.py::test_serial 0.5727 0.5837 +1.93%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-None] 176.23 179.60 +1.91%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-16] 12.20 12.43 +1.89%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-None] 83.72 85.29 +1.88%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-True] 23.58 24.01 +1.83%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-True] 33,532 32,921 -1.82%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-False-True] 42,859 42,085 -1.80%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-False] 56,639 57,650 +1.79%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-1] 192.54 195.94 +1.77%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-False] 28,838 29,333 +1.72%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-gru] 2.9658 3.0166 +1.71%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-None] 341.27 346.95 +1.67%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-backward] 114.96 116.88 +1.67%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-backward] 250.54 254.63 +1.64%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[False-backward] 28.29 27.83 -1.62%
benchmarks/test_collectors_benchmark.py::test_single 8.7992 8.9378 +1.58%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-4000] 168.36 165.71 -1.57%
benchmarks/test_objectives_benchmarks.py::test_values[td1_return_estimate-False-False] 36.45 37.02 +1.57%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-None] 559.78 568.51 +1.56%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[False-backward] 77.47 78.68 +1.56%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-True-True-True] 24,204 23,828 -1.55%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[200-img_shape3-large_batch] 309.49 314.26 +1.54%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb[200-img_shape1-large_batch] 14.97 15.19 +1.52%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-backward] 242.16 245.82 +1.51%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-64] 3.0529 3.0983 +1.49%
... ... ... Showing 120 of 220 comparisons, sorted by absolute change.

GPU

Compared 230 benchmarks. Regressions over 5%: 13. Improvements over 5%: 13.

Benchmark main ops PR ops Change
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400] 42.09 195.18 +363.78%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[False-None] 83.16 115.44 +38.82%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 3,463 2,512 -27.48%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 3,695 2,816 -23.77%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[reduce-overhead-None] 104.12 81.47 -21.75%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-400] 816.06 989.53 +21.26%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 2,543 3,045 +19.75%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-10000] 3,275 2,665 -18.62%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 1,831 2,169 +18.47%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[reduce-overhead-None] 120.25 100.32 -16.57%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 3,391 2,898 -14.55%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-RandomSampler-400] 474.20 526.16 +10.96%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-backward] 1,014 915.24 -9.73%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-LazyTensorStorage-SamplerWithoutReplacement-10000] 2,721 2,961 +8.79%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyTensorStorage-RandomSampler-10000] 2,716 2,907 +7.01%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-400] 489.82 455.93 -6.92%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-sampler6-10000] 791.93 738.86 -6.70%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[generalized_advantage_estimate-False-1-512] 50.40 47.14 -6.47%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[200-img_shape3-large_batch] 134.69 142.86 +6.07%
benchmarks/test_envs_benchmark.py::test_simple 1.2808 1.2076 -5.72%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-400] 501.15 473.24 -5.57%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[safetensors] 24,851 23,472 -5.55%
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sample_mixed_devices[1000000-cuda_storage_cpu_sampler] 83.79 88.37 +5.46%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[True-None] 1,919 2,023 +5.41%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-same] 3,561 3,743 +5.12%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[16-constant] 4,727 4,964 +5.03%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-None] 734.64 771.25 +4.98%
benchmarks/test_replaybuffer_benchmark.py::TestWindowingTransformsBenchmark::test_action_chunk_transform[done_aware] 1,519 1,595 +4.96%
benchmarks/test_non_tensor_env_benchmark.py::test_non_tensor_env_rollout_speed[1000-single-True] 1.3273 1.3928 +4.93%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-True] 38,296 40,152 +4.85%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape1-atari] 698.35 664.59 -4.83%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-1-512] 1,399 1,333 -4.74%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[100-img_shape2-large_img] 409.09 389.95 -4.68%
benchmarks/test_objectives_benchmarks.py::test_sac_speed[True-backward] 349.13 332.94 -4.64%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-False] 55,714 58,276 +4.60%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-False-True] 33,588 35,048 +4.35%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-1] 581.38 606.08 +4.25%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-4] 69.77 72.73 +4.24%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-False] 27,135 28,280 +4.22%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-True] 20,963 21,840 +4.18%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-False] 32,003 33,341 +4.18%
benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed[torch.save] 7,055 7,345 +4.11%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictReplayBuffer-LazyMemmapStorage-SamplerWithoutReplacement-10000] 3,471 3,331 -4.03%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-False] 64,674 67,248 +3.98%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-True-False] 34,932 36,318 +3.97%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[50-img_shape0-small] 4,355 4,183 -3.95%
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sample_mixed_devices[1000000-memmap_cpu_storage_cud... 995.93 956.69 -3.94%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-False-True-True] 19,789 20,569 +3.94%
benchmarks/test_collectors_benchmark.py::test_single 6.9125 6.6448 -3.87%
benchmarks/test_objectives_benchmarks.py::test_dqn_speed[reduce-overhead-None] 1,885 1,958 +3.84%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[100-img_shape2-large_img] 565.49 543.77 -3.84%
benchmarks/test_objectives_benchmarks.py::test_td3_speed[True-backward] 405.44 390.08 -3.79%
benchmarks/test_replaybuffer_benchmark.py::TestWindowingTransformsBenchmark::test_action_chunk_transform[no_done] 1,581 1,640 +3.72%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-False] 32,249 33,447 +3.72%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-True] 31,495 30,325 -3.71%
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sampler_sample_scale[1000000-cuda] 2,330 2,246 -3.62%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictPrioritizedReplayBuffer-LazyTensorStorage-None-400] 809.22 838.38 +3.60%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-lstm] 21.62 22.35 +3.36%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[True-backward] 357.52 369.47 +3.34%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-True-False] 34,566 35,705 +3.29%
benchmarks/test_envs_benchmark.py::test_cat_frames_functional[4-constant] 4,858 5,017 +3.28%
benchmarks/test_replaybuffer_benchmark.py::TestPrioritizedReplayBufferBenchmark::test_sample_mixed_devices[1000000-cuda_storage_cuda_samp... 1,553 1,502 -3.27%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-256-256-1] 188.28 194.40 +3.25%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-False-True-True] 17,927 18,494 +3.16%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_lazystack[100-img_shape2-large_img] 420.71 408.05 -3.01%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-224-224-16] 17.62 18.14 +2.95%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-False-True-False-True] 37,779 38,882 +2.92%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-True-True] 22,086 22,719 +2.86%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[False-backward] 249.15 242.05 -2.85%
benchmarks/test_replaybuffer_benchmark.py::test_rb_sample[TensorDictPrioritizedReplayBuffer-LazyMemmapStorage-None-10000] 1,956 2,011 +2.83%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-224-224-4] 178.72 183.67 +2.77%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-False] 29,851 30,678 +2.77%
benchmarks/test_objectives_benchmarks.py::test_a2c_speed[False-backward] 155.79 151.69 -2.63%
benchmarks/test_replaybuffer_benchmark.py::test_replay_buffer_direct_client_identity 7,962,244 8,171,414 +2.63%
benchmarks/test_objectives_benchmarks.py::test_values[vec_generalized_advantage_estimate-True-True] 462.16 474.12 +2.59%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[False-backward] 289.08 281.63 -2.58%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-1] 465.22 453.29 -2.56%
benchmarks/test_objectives_benchmarks.py::test_gae_speed[vec_generalized_advantage_estimate-False-32-512] 1,366 1,331 -2.56%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-True-True] 19,775 20,282 +2.56%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[200-img_shape3-large_batch] 721.58 740.02 +2.56%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-True-0-gru] 52.96 51.61 -2.54%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-False-False] 50,698 51,973 +2.52%
benchmarks/test_objectives_benchmarks.py::test_iql_speed[True-backward] 254.56 248.23 -2.49%
benchmarks/test_replaybuffer_benchmark.py::TestWindowingTransformsBenchmark::test_catframes_offline 568.02 554.30 -2.42%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-16] 35.33 36.17 +2.40%
benchmarks/test_envs_benchmark.py::test_parallel 0.5460 0.5585 +2.30%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[100-img_shape2-large_img] 171.51 175.31 +2.22%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb_cuda[100-img_shape0-atari] 13.76 13.47 -2.15%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[True-True-False-False-False] 64,707 66,098 +2.15%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-True-False-False] 57,988 59,208 +2.10%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_stack_then_write[50-img_shape0-small] 903.96 885.19 -2.08%
benchmarks/test_envs_benchmark.py::test_transformed 0.7302 0.7155 -2.02%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-True-False-False-False] 50,898 51,924 +2.02%
benchmarks/test_objectives_benchmarks.py::test_values[vec_td1_return_estimate-False-False] 866.94 849.86 -1.97%
benchmarks/test_objectives_benchmarks.py::test_cql_speed[True-backward] 229.01 233.40 +1.92%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-256-256-16] 41.50 42.28 +1.88%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[50-img_shape0-small] 3,553 3,489 -1.81%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb[100-img_shape0-atari] 20.93 20.55 -1.80%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_storage_write_contiguous[50-img_shape0-small] 6,142 6,034 -1.76%
benchmarks/test_objectives_benchmarks.py::test_redq_deprec_speed[True-None] 441.01 433.38 -1.73%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-True-0-lstm] 82.53 81.10 -1.72%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[torchvision-480-640-64] 6.9127 7.0309 +1.71%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-1000000-10000-100-False] 47.64 48.45 +1.69%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[reduce-overhead-None] 839.53 853.63 +1.68%
benchmarks/test_objectives_benchmarks.py::test_values[td_lambda_return_estimate-True-False] 11.83 12.03 +1.67%
benchmarks/test_objectives_benchmarks.py::test_ppo_speed[True-None] 716.88 728.76 +1.66%
benchmarks/test_replaybuffer_benchmark.py::test_rb_extend_sample[ReplayBuffer-LazyTensorStorage-RandomSampler-10000-10000-100-False] 53.20 52.33 -1.65%
benchmarks/test_objectives_benchmarks.py::test_ddpg_speed[True-backward] 462.46 470.05 +1.64%
benchmarks/test_collectors_benchmark.py::test_single_with_rb_pixels 5.3740 5.2860 -1.64%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_without_rb[200-img_shape1-large_batch] 15.09 14.85 -1.61%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-480-640-1] 79.70 80.96 +1.59%
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-RandomSampler-400] 194.93 197.96 +1.55%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-480-640-4] 19.86 20.17 +1.55%
benchmarks/test_vla_preprocessing_benchmark.py::test_openvla_preprocessing_throughput[pil-480-640-64] 1.2489 1.2680 +1.53%
benchmarks/test_envs_benchmark.py::test_step_mdp_speed[False-False-True-True-True] 18,694 18,978 +1.52%
benchmarks/test_objectives_benchmarks.py::test_reinforce_speed[reduce-overhead-None] 122.42 120.57 -1.51%
benchmarks/test_storage_write_benchmark.py::TestCollectorIntegrationBenchmark::test_collector_with_rb_cuda[200-img_shape1-large_batch] 6.8968 6.7935 -1.50%
benchmarks/test_replaybuffer_benchmark.py::test_rb_iterate[TensorDictReplayBuffer-ListStorage-RandomSampler-4000] 170.97 173.52 +1.49%
benchmarks/test_storage_write_benchmark.py::TestStorageWriteBenchmark::test_collector_lazystack_then_write[200-img_shape3-large_batch] 299.56 303.96 +1.47%
benchmarks/test_rnn_reset_backends_benchmark.py::test_rnn_rollout_with_intermediate_resets[b256-t128-i32-h512-scan-False-0-gru] 23.20 23.54 +1.46%
... ... ... Showing 120 of 230 comparisons, sorted by absolute change.

@vmoens @claude
[Feature] Add torch.distributed inference transport
8309d6d 
Support the ray + distributed combination: Ray (or any remote-process)
actors moving inference payloads through torch.distributed
point-to-point ops instead of object-store serialization.

DistributedTransport gives every server/client pair standalone
ProcessGroupGloo (and, for CUDA payloads, ProcessGroupNCCL) groups
rendezvoused through a transport-owned TCPStore, so the default
process group stays free for training collectives. Small headers and
pickled exceptions ride gloo control pairs; fixed-spec tensor payloads
ride gloo (CPU) or NCCL (GPU-direct) data pairs, with NCCL receives
posted only after a header announces them so idle receives never trip
the NCCL watchdog and each communicator stays single-threaded.
Clients are picklable and connect lazily inside the actor process; the
server discovers clients through the store, so client creation is not
ordered with respect to server startup.

Blocking receiver threads feed in-process queues on both sides (gloo
Work.is_completed() does not report completion for receives, so
irecv-polling designs do not work).

The resolver maps transport="distributed" to this transport for any
non-thread actor topology (requires request/response specs; never
chosen by "auto"), and inference_transport_kind reports it. Tests
cover in-process, spawned cross-process, and Ray-actor clients,
out-of-order futures, exception propagation, nested keys, timeout
semantics, and spec validation.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
github-advanced-security[bot]
github-advanced-security AI found potential problems Jul 9, 2026
View reviewed changes
Comment thread torchrl/modules/inference_server/_distributed.py Fixed
@vmoens @claude
[Quality] Harden DistributedTransport rendezvous binding
af64a0f 
Address the CodeQL finding on the free-port probe: bind the probe
socket to the advertised host instead of the wildcard address, which
also validates early that the host is a local interface.

Document the actual security posture in the class docstring: like all
of torch.distributed, the rendezvous TCPStore accepts unauthenticated
connections on all interfaces regardless of the host argument, control
headers and exception payloads are pickled, and the payload channels
are unencrypted, so the transport must only be deployed on trusted,
isolated networks.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@vmoens vmoens mentioned this pull request Jul 9, 2026
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

1 more reviewer
@github-advanced-security github-advanced-security[bot] github-advanced-security[bot] left review comments
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Collectors Documentation Improvements or additions to documentation Feature New feature Integrations/torch_geometric Integrations Modules Services

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@vmoens @github-advanced-security