pytorch / pytorch/rl

[RFC] Unify device configuration across collectors, evaluators, and the inference server

Open
#3,943 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
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):

  1. 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_sync etc.). 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.
  2. 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.
  3. 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?
  4. 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 whether InferenceServerConfig /
    InferenceDeviceConfig become the canonical convention (with kwargs as
    thin sugar), or stay an inference-server-local convenience.
  5. 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 InferenceDeviceConfig to torchrl.collectors (re-exported from
    torchrl.modules.inference_server for BC) and rename or alias it
    DevicePlacementConfig.
  • Accept device_config on Collector / MultiCollector / evaluator
    constructors; single config broadcasts, Sequence[DevicePlacementConfig]
    maps per worker.
  • Keep explicit kwargs working for at least two minor releases; emit
    FutureWarning only once a decision on final removal is made.
  • One resolution function shared by all consumers so precedence rules
    (device alias, env_device fallback for output_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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.