get_worker_info()Return a WorkerInfo object in a worker process, else None.
Inside the worker_init_fn or dataset __getitem__ /
__iter__ of a multi-process DataLoader, this returns
the WorkerInfo for the current worker. In the main
process — or when num_workers=0 — it returns None.
Returns
WorkerInfo or NonePer-worker context if called from a worker, else None.
Notes
The dominant use case is sharding an IterableDataset
across workers: each worker pulls its WorkerInfo.id and
advances through a distinct slice of the underlying stream so that
samples are not duplicated across the pool.
Examples
>>> from lucid.utils.data import IterableDataset, get_worker_info
>>> class ShardedStream(IterableDataset):
... def __init__(self, n):
... self.n = n
... def __iter__(self):
... info = get_worker_info()
... if info is None:
... yield from range(self.n)
... else:
... yield from range(info.id, self.n, info.num_workers)
>>> list(ShardedStream(5)) # main process: no worker info, whole stream
[0, 1, 2, 3, 4]