[RFC] Unify device configuration across collectors, evaluators, and the inference server
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 3.6k
- Forks
- 487
- Avg merge
- 1d 1h
- Merged PRs (30d)
- 207
Description
Context
#3893 introduced InferenceDeviceConfig, a dataclass grouping the four devices
involved in asynchronous policy-server collection (policy_device,
output_device, env_device, storing_device), consumed by
InferenceServer, ProcessInferenceServer, and AsyncBatchedCollector.
Review of that PR raised design questions that were deliberately deferred so
the stack could land (the PR itself only received mechanical fixes:
sentinel-based mutual-exclusivity checks, documented/rejected unused fields,
stronger device tests):
- Scope. Regular collectors (
Collector,MultiCollector, batched
variants) and evaluators expose the same concepts as loose kwargs
(device,policy_device,env_device,storing_device, plus
no_cuda_syncetc.). If a structured device config is the right API, it
should not be exclusive to the inference-server path -- otherwise we grow
two parallel conventions. - Deprecation story. If collectors adopt
device_config, do we softly
deprecate the explicit device kwargs? That touches a very wide public
surface and needs the standard two-minor-release cycle with
FutureWarnings naming the target version. - Multi-collector semantics. Collectors with sub-collectors accept
per-worker device lists today. Should they take a
list[InferenceDeviceConfig](one per worker), a single config broadcast
to all workers, or both? - API bloat. The concern is that config objects only pay off if they
replace boilerplate rather than adding a second way to say the same
thing. A decision is needed on whetherInferenceServerConfig/
InferenceDeviceConfigbecome the canonical convention (with kwargs as
thin sugar), or stay an inference-server-local convenience. - Interaction with Hydra configs.
torchrl/trainers/algorithms/configs/
already defines Hydra dataclasses for collectors. If runtime config
objects and Hydra configs coexist, the mapping between them should be
mechanical (field names and defaults aligned) to keep rule-14 parity
cheap.
Proposal sketch (to be debated)
- Promote
InferenceDeviceConfigtotorchrl.collectors(re-exported from
torchrl.modules.inference_serverfor BC) and rename or alias it
DevicePlacementConfig. - Accept
device_configonCollector/MultiCollector/ evaluator
constructors; single config broadcasts,Sequence[DevicePlacementConfig]
maps per worker. - Keep explicit kwargs working for at least two minor releases; emit
FutureWarningonly once a decision on final removal is made. - One resolution function shared by all consumers so precedence rules
(devicealias,env_devicefallback foroutput_device, etc.) are
defined in exactly one place.
Non-goals
- Changing runtime semantics of any device transfer.
- Touching the Hydra config tree beyond keeping field parity.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with InferenceDeviceConfig, InferenceServer, ProcessInferenceServer, AsyncBatchedCollector, Collector, MultiCollector, evaluator constructors, and the Hydra dataclasses under torchrl/trainers/algorithms/configs/. Trace how current device kwargs and per-worker lists are resolved, then review the deferred design questions and existing device tests. Done means a decided, consistent configuration convention with aligned field behavior and an explicit deprecation plan, without changing device-transfer semantics.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend-api-design, machine-learning
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 30/100