Skip to main content

Module workers

Module workers 

Source
Expand description

Which worker the calling thread is, and how that is discovered.

A sharded pipeline needs to know two things: how many workers are reading it, and which one this thread is. Everything else — splitting shards across workers, deriving per-worker seeds — follows from that.

How the answer is stored depends on what the target offers:

buildstorage
stda thread-local, set for the duration of with_worker
no_std + threadsa lock-free slot table, keyed by a host-supplied thread token
no_std without threadsa single global, which is all a single-threaded target needs

The middle row is the interesting one. Without the standard library there is no portable way to ask “which thread am I on?”, so the host has to say. Call set_thread_id_hook once with something that identifies the current execution context — a task id, a core id, a worker index — and per-worker state starts working:

use webdataset_core::workers::{set_thread_id_hook, with_worker};
use webdataset_core::worker_info;

// On a real RTOS this would return the current task's id.
set_thread_id_hook(|| 0).ok();

let info = with_worker(2, 4, worker_info);
assert_eq!((info.worker, info.num_workers), (2, 4));

Without a hook the table is bypassed and the single global is used, so a single-threaded no_std program needs no setup at all.

Functions§

clear_worker
Forget this context’s worker identity.
currentstd
The worker identity bound to this thread, if any.
replacestd
Bind an identity to this thread, returning the one it replaced.
set_thread_id_hookstd or non-threads
Install the hook that tells the library which execution context is running.
set_worker
Bind a worker identity to this context until it is replaced.
with_worker
Run body with this context presenting as worker worker of num_workers.