Template Class DataShuttle

Page Contents

Class Documentation

template<typename Job, typename Result>
class torch::data::detail::DataShuttle

Encapsulates the full life cycle of DataLoader jobs.

When a new job is enqueued to the DataShuttle, a counter for in-flight jobs is bumped. This job is said to be “in-flight” until its result is popped. Worker threads dequeue jobs as soon as they are available. When a worker finishes a job, it enqueues the result. Only when the main thread dequeues a result is the count of in-flight jobs decremented. When the main thread attempts to dequeue a job but no jobs are in-flight, that means the epoch is complete and pop_result returns an empty optional.

Public Functions

void push_job(Job job)

Pushes a new job. Called by the main thread.

void push_result(Result result)

Pushes the result of a job. Called by worker threads.

Job pop_job()

Returns the next job, blocking until there is one available.

Called by worker threads.

optional<Result> pop_result(optional<std::chrono::milliseconds> timeout = nullopt)

Returns the result of a job, or nullopt if all jobs were exhausted.

Called by the main thread.

void drain()

Discards any jobs that are not yet in flight, and waits for all in-flight jobs to finish, discarding their result.

size_t in_flight_jobs() const noexcept

Returns the number of jobs that are still in progress.

When this number is zero, an epoch is finished.


Access comprehensive developer documentation for PyTorch

View Docs


Get in-depth tutorials for beginners and advanced developers

View Tutorials


Find development resources and get your questions answered

View Resources