|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# |
| 3 | +# This source code is licensed under the MIT license found in the |
| 4 | +# LICENSE file in the root directory of this source tree. |
| 5 | +from __future__ import annotations |
| 6 | + |
| 7 | +from collections.abc import Callable, Sequence |
| 8 | +from concurrent.futures import Future |
| 9 | + |
| 10 | +from tensordict.base import TensorDictBase |
| 11 | +from tensordict.nn import TensorDictModuleBase |
| 12 | +from tensordict.utils import NestedKey |
| 13 | + |
| 14 | +from torchrl.modules.inference_server._transport import InferenceTransport |
| 15 | + |
| 16 | + |
| 17 | +class _ImmediateFuture: |
| 18 | + def __init__(self, result: TensorDictBase | BaseException): |
| 19 | + self._result = result |
| 20 | + |
| 21 | + def done(self) -> bool: |
| 22 | + return True |
| 23 | + |
| 24 | + def result(self, timeout: float | None = None) -> TensorDictBase: |
| 25 | + if isinstance(self._result, BaseException): |
| 26 | + raise self._result |
| 27 | + return self._result |
| 28 | + |
| 29 | + |
| 30 | +class PolicyClientModule(TensorDictModuleBase): |
| 31 | + """TensorDict policy wrapper for remote inference-server clients. |
| 32 | +
|
| 33 | + ``PolicyClientModule`` makes a transport client look like a TorchRL policy: |
| 34 | + it accepts a :class:`~tensordict.TensorDictBase`, submits it to an |
| 35 | + :class:`~torchrl.modules.inference_server.InferenceServer`, and returns the |
| 36 | + TensorDict produced by the remote policy. It can be passed anywhere a |
| 37 | + TensorDict policy module is expected. |
| 38 | +
|
| 39 | + This class is the reference implementation of TorchRL's service *client* |
| 40 | + contract: it duck-types the domain interface (a policy client IS a |
| 41 | + TensorDict policy, so consumer code cannot tell local from remote), it is |
| 42 | + cheap and picklable (it can be handed to spawned workers), and it carries |
| 43 | + no lifecycle rights -- clients can call the service but never start or |
| 44 | + shut it down; only the owner that constructed the server can. |
| 45 | +
|
| 46 | + .. note:: |
| 47 | + Unlike a local :class:`~tensordict.nn.TensorDictModule`, the result |
| 48 | + crosses a transport boundary, so :meth:`forward` returns a *new* |
| 49 | + TensorDict rather than writing the ``out_keys`` into the input |
| 50 | + TensorDict. Use the return value; do not rely on in-place updates of |
| 51 | + the input. |
| 52 | +
|
| 53 | + Args: |
| 54 | + client (Callable or InferenceTransport): actor-side inference client. |
| 55 | + If a transport is provided, ``transport.client()`` is called. |
| 56 | +
|
| 57 | + Keyword Args: |
| 58 | + in_keys (sequence of NestedKey, optional): input keys advertised by the |
| 59 | + module. The full input TensorDict is still sent to the server. |
| 60 | + out_keys (sequence of NestedKey, optional): output keys advertised by |
| 61 | + the module. |
| 62 | +
|
| 63 | + Examples: |
| 64 | + >>> import torch |
| 65 | + >>> import torch.nn as nn |
| 66 | + >>> from tensordict import TensorDict |
| 67 | + >>> from tensordict.nn import TensorDictModule |
| 68 | + >>> from torchrl.modules.inference_server import ( |
| 69 | + ... InferenceServer, |
| 70 | + ... PolicyClientModule, |
| 71 | + ... ThreadingTransport, |
| 72 | + ... ) |
| 73 | + >>> policy = TensorDictModule( |
| 74 | + ... nn.Linear(4, 2), in_keys=["observation"], out_keys=["action"] |
| 75 | + ... ) |
| 76 | + >>> transport = ThreadingTransport() |
| 77 | + >>> server = InferenceServer(policy, transport).start() |
| 78 | + >>> remote_policy = PolicyClientModule( |
| 79 | + ... transport, in_keys=["observation"], out_keys=["action"] |
| 80 | + ... ) |
| 81 | + >>> td = remote_policy(TensorDict({"observation": torch.randn(4)})) |
| 82 | + >>> "action" in td.keys() |
| 83 | + True |
| 84 | + >>> server.shutdown() |
| 85 | + """ |
| 86 | + |
| 87 | + def __init__( |
| 88 | + self, |
| 89 | + client: Callable[[TensorDictBase], TensorDictBase] | InferenceTransport, |
| 90 | + *, |
| 91 | + in_keys: Sequence[NestedKey] | None = None, |
| 92 | + out_keys: Sequence[NestedKey] | None = None, |
| 93 | + ) -> None: |
| 94 | + super().__init__() |
| 95 | + if isinstance(client, InferenceTransport): |
| 96 | + client = client.client() |
| 97 | + self.client = client |
| 98 | + self.in_keys = list(in_keys or []) |
| 99 | + self.out_keys = list(out_keys or []) |
| 100 | + |
| 101 | + def submit(self, tensordict: TensorDictBase) -> Future | _ImmediateFuture: |
| 102 | + """Submit a TensorDict request and return a future-like object. |
| 103 | +
|
| 104 | + Args: |
| 105 | + tensordict (TensorDictBase): observation TensorDict to send to the |
| 106 | + remote policy. |
| 107 | +
|
| 108 | + Returns: |
| 109 | + Future-like object whose ``result()`` method returns a TensorDict. |
| 110 | + When the wrapped client exposes ``submit`` this is the transport's |
| 111 | + :class:`~concurrent.futures.Future` and submission errors raise |
| 112 | + synchronously; for a plain callable client the call runs eagerly |
| 113 | + and errors are deferred to ``result()`` on a reduced future that |
| 114 | + only implements ``done()`` and ``result()``. |
| 115 | + """ |
| 116 | + submit = getattr(self.client, "submit", None) |
| 117 | + if submit is None: |
| 118 | + try: |
| 119 | + result = self.client(tensordict) |
| 120 | + return _ImmediateFuture(result) |
| 121 | + except Exception as exc: |
| 122 | + return _ImmediateFuture(exc) |
| 123 | + return submit(tensordict) |
| 124 | + |
| 125 | + def forward(self, tensordict: TensorDictBase) -> TensorDictBase: |
| 126 | + return self.submit(tensordict).result() |
0 commit comments