reference/generated/torchrl.modules.inference_server.InferenceServerConfig

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

InferenceServerConfig

class torchrl.modules.inference_server.InferenceServerConfig(backend: Literal['thread', 'process'] = 'thread', max_batch_size: int = 64, min_batch_size: int = 1, timeout: float = 0.01, collect_stats: bool = True, stats_window_size: int = 1024, max_inflight_per_env: int | None = None)

Server-side execution, batching, timeout, and instrumentation settings.

Parameters:

• backend (str, optional) – execution backend for the policy server. "thread" runs the serve loop in a background thread of the constructing process; "process" runs a dedicated server process (which requires a picklable policy_factory and a multiprocessing-capable transport). Defaults to "thread".

• max_batch_size (int, optional) – maximum number of requests per forward pass. Defaults to 64.

• min_batch_size (int, optional) – minimum number of requests to accumulate after the first request arrives. Defaults to 1.

• timeout (float, optional) – seconds to wait for more requests before flushing a partial batch. Defaults to 0.01.

• collect_stats (bool, optional) – whether to collect lightweight throughput and latency stats. Defaults to True.

• stats_window_size (int, optional) – number of recent timing samples kept for percentile stats. Defaults to 1024.

• max_inflight_per_env (int, optional) – maximum unresolved remote-policy requests each environment coordinator may have inflight (consumed by AsyncBatchedCollector when building its clients). Defaults to None (unbounded), so the guard never throttles by surprise; set an explicit bound when backpressure is wanted.

Examples

>>> import torch
>>> import torch.nn as nn
>>> from tensordict import TensorDict
>>> from tensordict.nn import TensorDictModule
>>> from torchrl.modules.inference_server import (
...     InferenceServer,
...     InferenceServerConfig,
...     ThreadingTransport,
... )
>>> policy = TensorDictModule(
...     nn.Linear(4, 2), in_keys=["observation"], out_keys=["action"]
... )
>>> transport = ThreadingTransport()
>>> config = InferenceServerConfig(max_batch_size=8, timeout=0.001)
>>> with InferenceServer(policy, transport, server_config=config) as server:
...     client = transport.client()
...     result = client(TensorDict({"observation": torch.randn(4)}))
>>> result["action"].shape
torch.Size([2])
>>> server.max_batch_size
8