[Refactor] Add structured inference server config objects

ghstack-source-id: 9f2dd45
Pull-Request: #3893

Author: vmoens

benchmarks/bench_collectors.py

@@ -39,6 +39,10 @@
     TransformedEnv,
 )
 from torchrl.modules import ConvNet, MLP
+from torchrl.modules.inference_server import (
+    InferenceDeviceConfig,
+    InferenceServerConfig,
+)
 
 OBS_SHAPE = (3, 84, 84)
 ACTION_DIM = 6
@@ -577,13 +581,17 @@ def main() -> None:
                                 policy=policy_factory(),
                                 frames_per_batch=args.frames_per_batch,
                                 total_frames=-1,
-                                max_batch_size=max_batch_size,
-                                min_batch_size=min_batch_size,
-                                server_timeout=timeout,
                                 env_backend="threading",
-                                server_backend="thread",
-                                policy_device=policy_device,
-                                output_device=output_device,
+                                server_config=InferenceServerConfig(
+                                    backend="thread",
+                                    max_batch_size=max_batch_size,
+                                    min_batch_size=min_batch_size,
+                                    timeout=timeout,
+                                ),
+                                device_config=InferenceDeviceConfig(
+                                    policy_device=policy_device,
+                                    output_device=output_device,
+                                ),
                             ),
                             env_name=args.env,
                             num_envs=num_envs,
@@ -614,13 +622,17 @@ def main() -> None:
                                 policy=policy_factory(),
                                 frames_per_batch=args.frames_per_batch,
                                 total_frames=-1,
-                                max_batch_size=max_batch_size,
-                                min_batch_size=min_batch_size,
-                                server_timeout=timeout,
                                 env_backend="multiprocessing",
-                                server_backend="thread",
-                                policy_device=policy_device,
-                                output_device=output_device,
+                                server_config=InferenceServerConfig(
+                                    backend="thread",
+                                    max_batch_size=max_batch_size,
+                                    min_batch_size=min_batch_size,
+                                    timeout=timeout,
+                                ),
+                                device_config=InferenceDeviceConfig(
+                                    policy_device=policy_device,
+                                    output_device=output_device,
+                                ),
                             ),
                             env_name=args.env,
                             num_envs=num_envs,
@@ -651,13 +663,17 @@ def main() -> None:
                                 policy_factory=policy_factory,
                                 frames_per_batch=args.frames_per_batch,
                                 total_frames=-1,
-                                max_batch_size=max_batch_size,
-                                min_batch_size=min_batch_size,
-                                server_timeout=timeout,
                                 env_backend="multiprocessing",
-                                server_backend="process",
-                                policy_device=policy_device,
-                                output_device=output_device,
+                                server_config=InferenceServerConfig(
+                                    backend="process",
+                                    max_batch_size=max_batch_size,
+                                    min_batch_size=min_batch_size,
+                                    timeout=timeout,
+                                ),
+                                device_config=InferenceDeviceConfig(
+                                    policy_device=policy_device,
+                                    output_device=output_device,
+                                ),
                             ),
                             env_name=args.env,
                             num_envs=num_envs,

docs/source/reference/modules_inference_server.rst

@@ -17,6 +17,8 @@ Core API
     :template: rl_template_noinherit.rst
 
     InferenceServer
+    InferenceServerConfig
+    InferenceDeviceConfig
     ProcessInferenceServer
     InferenceClient
     InferenceTransport
@@ -67,6 +69,46 @@ threads in the same process:
 
     server.shutdown()
 
+Structured Configuration
+^^^^^^^^^^^^^^^^^^^^^^^^
+
+Server execution, batching, and device placement are grouped into two
+dataclasses instead of loose keyword arguments: :class:`InferenceServerConfig`
+collects the execution ``backend`` (``"thread"`` or ``"process"``) and the
+batching/instrumentation knobs (``max_batch_size``, ``min_batch_size``,
+``timeout``, ``collect_stats``, ``stats_window_size``), and
+:class:`InferenceDeviceConfig` describes device placement across the
+collection pipeline (``policy_device``, ``output_device``, ``env_device``,
+``storing_device``). Both :class:`InferenceServer` and
+:class:`~torchrl.collectors.AsyncBatchedCollector` accept them through the
+``server_config`` and ``device_config`` keyword arguments; a config object is
+mutually exclusive with the individual keyword arguments it replaces, and the
+config objects are the only way to set the per-role devices and the server
+backend on the collector. Servers consume only the
+``policy_device``/``output_device`` fields (``env_device`` doubles as an
+``output_device`` fallback), while ``env_device`` and ``storing_device``
+drive the collector-side transfers:
+
+.. code-block:: python
+
+    from torchrl.collectors import AsyncBatchedCollector
+    from torchrl.modules.inference_server import (
+        InferenceDeviceConfig,
+        InferenceServerConfig,
+    )
+
+    collector = AsyncBatchedCollector(
+        create_env_fn=[make_env] * 8,
+        policy=my_policy,
+        frames_per_batch=200,
+        server_config=InferenceServerConfig(max_batch_size=8, timeout=0.005),
+        device_config=InferenceDeviceConfig(
+            policy_device="cuda:0",
+            env_device="cpu",
+            storing_device="cpu",
+        ),
+    )
+
 Weight Synchronisation
 ^^^^^^^^^^^^^^^^^^^^^^
 

test/test_inference_server.py

@@ -20,7 +20,9 @@
 
 from torchrl.modules.inference_server import (
     InferenceClient,
+    InferenceDeviceConfig,
     InferenceServer,
+    InferenceServerConfig,
     InferenceTransport,
     MPTransport,
     ProcessInferenceServer,
@@ -283,6 +285,88 @@ def test_stats_accounting(self):
         assert stats["avg_batch_size"] > 0
         assert stats["p95_forward_ms"] >= 0
 
+    def test_structured_config(self):
+        transport = ThreadingTransport()
+        policy = _make_policy()
+        server_config = InferenceServerConfig(max_batch_size=2, timeout=0.001)
+        device_config = InferenceDeviceConfig(policy_device="cpu", output_device="cpu")
+        with InferenceServer(
+            policy,
+            transport,
+            server_config=server_config,
+            device_config=device_config,
+        ) as server:
+            # The config values must actually land on the server
+            assert server.max_batch_size == 2
+            assert server.timeout == 0.001
+            assert server.policy_device == torch.device("cpu")
+            assert server.output_device == torch.device("cpu")
+            client = transport.client()
+            result = client(TensorDict({"observation": torch.randn(4)}))
+            stats = server.stats()
+        assert result["action"].device.type == "cpu"
+        assert stats["requests"] == 1
+
+    def test_server_config_exclusive_even_at_default_values(self):
+        """Passing an explicit kwarg equal to the default still raises."""
+        transport = ThreadingTransport()
+        policy = _make_policy()
+        with pytest.raises(ValueError, match="mutually exclusive"):
+            InferenceServer(
+                policy,
+                transport,
+                max_batch_size=64,
+                server_config=InferenceServerConfig(max_batch_size=8),
+            )
+        with pytest.raises(ValueError, match="mutually exclusive"):
+            InferenceServer(
+                policy,
+                transport,
+                device="cpu",
+                device_config=InferenceDeviceConfig(policy_device="cpu"),
+            )
+
+    def test_device_config_env_device_fallback_and_storing_device_rejected(self):
+        config = InferenceDeviceConfig(env_device="cpu")
+        assert config.server_output_device() == torch.device("cpu")
+        config = InferenceDeviceConfig(env_device="cpu", output_device="meta")
+        assert config.server_output_device() == torch.device("meta")
+        transport = ThreadingTransport()
+        policy = _make_policy()
+        with pytest.raises(ValueError, match="storing_device is a collector-level"):
+            InferenceServer(
+                policy,
+                transport,
+                device_config=InferenceDeviceConfig(storing_device="cpu"),
+            )
+
+    @pytest.mark.gpu
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+    def test_device_config_cuda_roundtrip(self):
+        """device_config must actually drive the policy and output moves."""
+        transport = ThreadingTransport()
+        policy = _make_policy()
+        seen_devices = []
+        inner_forward = policy.module.forward
+
+        def recording_forward(x):
+            seen_devices.append(x.device)
+            return inner_forward(x)
+
+        policy.module.forward = recording_forward
+        with InferenceServer(
+            policy,
+            transport,
+            device_config=InferenceDeviceConfig(
+                policy_device="cuda:0", output_device="cpu"
+            ),
+        ):
+            client = transport.client()
+            result = client(TensorDict({"observation": torch.randn(4)}, device="cpu"))
+        assert result["action"].device.type == "cpu"
+        assert all(device.type == "cuda" for device in seen_devices)
+        assert next(policy.parameters()).device.type == "cuda"
+
     @pytest.mark.gpu
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
     def test_cuda_policy_cpu_output(self):
@@ -1041,9 +1125,8 @@ def test_process_server_backend_smoke(self):
             policy_factory=_make_counting_policy,
             frames_per_batch=10,
             total_frames=20,
-            max_batch_size=2,
             env_backend="threading",
-            server_backend="process",
+            server_config=InferenceServerConfig(backend="process", max_batch_size=2),
         )
         total = 0
         for batch in collector:
@@ -1052,13 +1135,8 @@ def test_process_server_backend_smoke(self):
         assert total >= 20
 
     def test_invalid_server_backend_raises(self):
-        with pytest.raises(ValueError, match="server_backend"):
-            AsyncBatchedCollector(
-                create_env_fn=[_counting_env_factory] * 2,
-                policy_factory=_make_counting_policy,
-                frames_per_batch=10,
-                server_backend="not-a-backend",
-            )
+        with pytest.raises(ValueError, match="backend"):
+            InferenceServerConfig(backend="not-a-backend")
 
     def test_server_death_raises_instead_of_hanging(self):
         """Killing the server process surfaces an error in the iterator."""
@@ -1067,9 +1145,8 @@ def test_server_death_raises_instead_of_hanging(self):
             policy_factory=_make_counting_policy,
             frames_per_batch=10,
             total_frames=-1,
-            max_batch_size=2,
             env_backend="threading",
-            server_backend="process",
+            server_config=InferenceServerConfig(backend="process", max_batch_size=2),
         )
         try:
             iterator = iter(collector)
@@ -1083,6 +1160,30 @@ def test_server_death_raises_instead_of_hanging(self):
         finally:
             collector.shutdown(timeout=0.5)
 
+    def test_device_config_and_server_config(self):
+        """Collector accepts structured device and server config objects."""
+        collector = AsyncBatchedCollector(
+            create_env_fn=[_counting_env_factory] * 2,
+            policy=_make_counting_policy(),
+            frames_per_batch=10,
+            total_frames=20,
+            server_config=InferenceServerConfig(max_batch_size=2),
+            device_config=InferenceDeviceConfig(
+                policy_device="cpu",
+                output_device="cpu",
+                env_device="cpu",
+                storing_device="cpu",
+            ),
+        )
+        total = 0
+        for batch in collector:
+            assert batch.device is None or batch.device.type == "cpu"
+            total += batch.numel()
+        stats = collector.server_stats()
+        collector.shutdown()
+        assert total >= 20
+        assert stats["requests"] > 0
+
 
 # =============================================================================
 # Tests: SlotTransport