diff --git a/docs/source/reference/modules_inference_server.rst b/docs/source/reference/modules_inference_server.rst index d6a016cd326..96e8aca07b5 100644 --- a/docs/source/reference/modules_inference_server.rst +++ b/docs/source/reference/modules_inference_server.rst @@ -21,6 +21,7 @@ Core API InferenceDeviceConfig ProcessInferenceServer InferenceClient + PolicyClientModule InferenceTransport Transport Backends @@ -109,6 +110,22 @@ drive the collector-side transfers: ), ) +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 ffb96fe5da3..8810d296cb1 100644 --- a/test/test_inference_server.py +++ b/test/test_inference_server.py @@ -7,6 +7,7 @@ import concurrent.futures import importlib.util import multiprocessing as mp +import pickle import threading import time @@ -25,6 +26,7 @@ InferenceServerConfig, InferenceTransport, MPTransport, + PolicyClientModule, ProcessInferenceServer, RayTransport, SlotTransport, @@ -407,6 +409,83 @@ def test_submit_returns_future(self): assert "action" in result.keys() +def _echo_client(td): + """Picklable stand-in client for pickling tests.""" + return td + + +class TestPolicyClientModule: + 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() + + def test_nested_keys(self): + """in_keys/out_keys accept nested keys end to end.""" + transport = ThreadingTransport() + policy = TensorDictModule( + nn.Linear(4, 2), + in_keys=[("agents", "observation")], + out_keys=[("agents", "action")], + ) + with InferenceServer(policy, transport, max_batch_size=4): + remote_policy = PolicyClientModule( + transport, + in_keys=[("agents", "observation")], + out_keys=[("agents", "action")], + ) + td = TensorDict({("agents", "observation"): torch.randn(4)}) + result = remote_policy(td) + assert result["agents", "action"].shape == (2,) + assert remote_policy.in_keys == [("agents", "observation")] + assert remote_policy.out_keys == [("agents", "action")] + + def test_client_contract_picklable_no_lifecycle(self): + """Clients pickle cleanly and expose no lifecycle methods.""" + remote_policy = PolicyClientModule( + _echo_client, in_keys=["observation"], out_keys=["observation"] + ) + restored = pickle.loads(pickle.dumps(remote_policy)) + td = TensorDict({"observation": torch.randn(4)}) + result = restored(td) + assert "observation" in result.keys() + assert restored.in_keys == ["observation"] + # Clients carry no lifecycle rights over the service + for lifecycle in ("start", "shutdown", "close", "flush"): + assert not hasattr(restored, lifecycle) + + def test_plain_callable_client_defers_errors(self): + """A plain-callable client defers exceptions to result().""" + + def failing_client(td): + raise ValueError("local policy failure") + + remote_policy = PolicyClientModule(failing_client) + future = remote_policy.submit(TensorDict({})) + assert future.done() + with pytest.raises(ValueError, match="local policy failure"): + future.result() + + # ============================================================================= # Tests: ThreadingTransport (Commit 2) # ============================================================================= diff --git a/torchrl/collectors/_async_batched.py b/torchrl/collectors/_async_batched.py index 4f73f06fe8f..dd550a79dbd 100644 --- a/torchrl/collectors/_async_batched.py +++ b/torchrl/collectors/_async_batched.py @@ -20,6 +20,7 @@ InferenceDeviceConfig, InferenceServer, InferenceServerConfig, + PolicyClientModule, ProcessInferenceServer, ThreadingTransport, ) @@ -432,7 +433,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..765eea3c59a 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 from torchrl.modules.inference_server._config import ( InferenceDeviceConfig, InferenceServerConfig, @@ -27,6 +28,7 @@ "InferenceTransport", "MonarchTransport", "MPTransport", + "PolicyClientModule", "ProcessInferenceServer", "RayTransport", "SlotTransport", diff --git a/torchrl/modules/inference_server/_client.py b/torchrl/modules/inference_server/_client.py new file mode 100644 index 00000000000..ba17eba1fa8 --- /dev/null +++ b/torchrl/modules/inference_server/_client.py @@ -0,0 +1,126 @@ +# 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 concurrent.futures import Future + +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. + + This class is the reference implementation of TorchRL's service *client* + contract: it duck-types the domain interface (a policy client IS a + TensorDict policy, so consumer code cannot tell local from remote), it is + cheap and picklable (it can be handed to spawned workers), and it carries + no lifecycle rights -- clients can call the service but never start or + shut it down; only the owner that constructed the server can. + + .. note:: + Unlike a local :class:`~tensordict.nn.TensorDictModule`, the result + crosses a transport boundary, so :meth:`forward` returns a *new* + TensorDict rather than writing the ``out_keys`` into the input + TensorDict. Use the return value; do not rely on in-place updates of + the input. + + 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) -> Future | _ImmediateFuture: + """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. + When the wrapped client exposes ``submit`` this is the transport's + :class:`~concurrent.futures.Future` and submission errors raise + synchronously; for a plain callable client the call runs eagerly + and errors are deferred to ``result()`` on a reduced future that + only implements ``done()`` and ``result()``. + """ + submit = getattr(self.client, "submit", None) + if submit is None: + try: + result = self.client(tensordict) + return _ImmediateFuture(result) + except Exception as exc: + return _ImmediateFuture(exc) + return submit(tensordict) + + def forward(self, tensordict: TensorDictBase) -> TensorDictBase: + return self.submit(tensordict).result()