From c633f453a805338c02c747ef12b0f68d5260569f Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Sun, 21 Jun 2026 07:35:20 +0100 Subject: [PATCH] Update [ghstack-poisoned] --- .../reference/modules_inference_server.rst | 18 +++ test/test_inference_server.py | 31 +++++ torchrl/collectors/_async_batched.py | 6 +- torchrl/modules/inference_server/__init__.py | 3 + torchrl/modules/inference_server/_client.py | 108 ++++++++++++++++++ 5 files changed, 165 insertions(+), 1 deletion(-) create mode 100644 torchrl/modules/inference_server/_client.py diff --git a/docs/source/reference/modules_inference_server.rst b/docs/source/reference/modules_inference_server.rst index 60cbd0383da..f04285b9743 100644 --- a/docs/source/reference/modules_inference_server.rst +++ b/docs/source/reference/modules_inference_server.rst @@ -21,6 +21,8 @@ Core API InferenceDeviceConfig ProcessInferenceServer InferenceClient + PolicyClientModule + RemotePolicy InferenceTransport Transport Backends @@ -69,6 +71,22 @@ threads in the same process: server.shutdown() +Remote policy module +^^^^^^^^^^^^^^^^^^^^ + +Use :class:`PolicyClientModule` when an actor or collector expects a regular +TensorDict policy but inference should be served by the policy server: + +.. code-block:: python + + remote_policy = PolicyClientModule( + transport, + in_keys=["observation"], + out_keys=["action"], + ) + + data = remote_policy(data) + Weight Synchronisation ^^^^^^^^^^^^^^^^^^^^^^ diff --git a/test/test_inference_server.py b/test/test_inference_server.py index c7e2160b574..a80de416ca4 100644 --- a/test/test_inference_server.py +++ b/test/test_inference_server.py @@ -24,8 +24,10 @@ InferenceServerConfig, InferenceTransport, MPTransport, + PolicyClientModule, ProcessInferenceServer, RayTransport, + RemotePolicy, SlotTransport, ThreadingTransport, ) @@ -316,6 +318,35 @@ def test_submit_returns_future(self): assert "action" in result.keys() +class TestPolicyClientModule: + def test_remote_policy_alias(self): + assert RemotePolicy is PolicyClientModule + + def test_forward_as_tensordict_module(self): + transport = ThreadingTransport() + policy = _make_policy() + with InferenceServer(policy, transport, max_batch_size=4): + remote_policy = PolicyClientModule( + transport, + in_keys=["observation"], + out_keys=["action"], + ) + td = TensorDict({"observation": torch.randn(4)}) + result = remote_policy(td) + assert result["action"].shape == (2,) + assert remote_policy.in_keys == ["observation"] + assert remote_policy.out_keys == ["action"] + + def test_submit(self): + transport = ThreadingTransport() + policy = _make_policy() + with InferenceServer(policy, transport, max_batch_size=4): + remote_policy = PolicyClientModule(transport) + future = remote_policy.submit(TensorDict({"observation": torch.randn(4)})) + result = future.result(timeout=5.0) + assert "action" in result.keys() + + # ============================================================================= # Tests: ThreadingTransport (Commit 2) # ============================================================================= diff --git a/torchrl/collectors/_async_batched.py b/torchrl/collectors/_async_batched.py index 5fe42693215..bcab3350e70 100644 --- a/torchrl/collectors/_async_batched.py +++ b/torchrl/collectors/_async_batched.py @@ -20,6 +20,7 @@ InferenceDeviceConfig, InferenceServer, InferenceServerConfig, + PolicyClientModule, ProcessInferenceServer, ThreadingTransport, ) @@ -453,7 +454,10 @@ def _ensure_started(self) -> None: # Create clients before a process server starts so response queues are # inherited by the child process. if self._clients is None: - self._clients = [self._transport.client() for _ in range(self._num_envs)] + self._clients = [ + PolicyClientModule(self._transport.client()) + for _ in range(self._num_envs) + ] # Start inference server if not self._server.is_alive: diff --git a/torchrl/modules/inference_server/__init__.py b/torchrl/modules/inference_server/__init__.py index f543ceaaffd..2ccdd70baf9 100644 --- a/torchrl/modules/inference_server/__init__.py +++ b/torchrl/modules/inference_server/__init__.py @@ -3,6 +3,7 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. +from torchrl.modules.inference_server._client import PolicyClientModule, RemotePolicy from torchrl.modules.inference_server._config import ( InferenceDeviceConfig, InferenceServerConfig, @@ -27,8 +28,10 @@ "InferenceTransport", "MonarchTransport", "MPTransport", + "PolicyClientModule", "ProcessInferenceServer", "RayTransport", + "RemotePolicy", "SlotTransport", "ThreadingTransport", ] diff --git a/torchrl/modules/inference_server/_client.py b/torchrl/modules/inference_server/_client.py new file mode 100644 index 00000000000..f0190f07149 --- /dev/null +++ b/torchrl/modules/inference_server/_client.py @@ -0,0 +1,108 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +from __future__ import annotations + +from collections.abc import Callable, Sequence +from tensordict.base import TensorDictBase +from tensordict.nn import TensorDictModuleBase +from tensordict.utils import NestedKey + +from torchrl.modules.inference_server._transport import InferenceTransport + + +class _ImmediateFuture: + def __init__(self, result: TensorDictBase | BaseException): + self._result = result + + def done(self) -> bool: + return True + + def result(self, timeout: float | None = None) -> TensorDictBase: + if isinstance(self._result, BaseException): + raise self._result + return self._result + + +class PolicyClientModule(TensorDictModuleBase): + """TensorDict policy wrapper for remote inference-server clients. + + ``PolicyClientModule`` makes a transport client look like a TorchRL policy: + it accepts a :class:`~tensordict.TensorDictBase`, submits it to an + :class:`~torchrl.modules.inference_server.InferenceServer`, and returns the + TensorDict produced by the remote policy. It can be passed anywhere a + TensorDict policy module is expected. + + Args: + client (Callable or InferenceTransport): actor-side inference client. + If a transport is provided, ``transport.client()`` is called. + + Keyword Args: + in_keys (sequence of NestedKey, optional): input keys advertised by the + module. The full input TensorDict is still sent to the server. + out_keys (sequence of NestedKey, optional): output keys advertised by + the module. + + Examples: + >>> import torch + >>> import torch.nn as nn + >>> from tensordict import TensorDict + >>> from tensordict.nn import TensorDictModule + >>> from torchrl.modules.inference_server import ( + ... InferenceServer, + ... PolicyClientModule, + ... ThreadingTransport, + ... ) + >>> policy = TensorDictModule( + ... nn.Linear(4, 2), in_keys=["observation"], out_keys=["action"] + ... ) + >>> transport = ThreadingTransport() + >>> server = InferenceServer(policy, transport).start() + >>> remote_policy = PolicyClientModule( + ... transport, in_keys=["observation"], out_keys=["action"] + ... ) + >>> td = remote_policy(TensorDict({"observation": torch.randn(4)})) + >>> "action" in td.keys() + True + >>> server.shutdown() + """ + + def __init__( + self, + client: Callable[[TensorDictBase], TensorDictBase] | InferenceTransport, + *, + in_keys: Sequence[NestedKey] | None = None, + out_keys: Sequence[NestedKey] | None = None, + ) -> None: + super().__init__() + if isinstance(client, InferenceTransport): + client = client.client() + self.client = client + self.in_keys = list(in_keys or []) + self.out_keys = list(out_keys or []) + + def submit(self, tensordict: TensorDictBase): + """Submit a TensorDict request and return a future-like object. + + Args: + tensordict (TensorDictBase): observation TensorDict to send to the + remote policy. + + Returns: + Future-like object whose ``result()`` method returns a TensorDict. + """ + submit = getattr(self.client, "submit", None) + if submit is None: + try: + result = self.client(tensordict) + return _ImmediateFuture(result) + except BaseException as exc: + return _ImmediateFuture(exc) + return submit(tensordict) + + def forward(self, tensordict: TensorDictBase) -> TensorDictBase: + return self.submit(tensordict).result() + + +RemotePolicy = PolicyClientModule