TaskManager acts on the tasks fetched from the task queues.

TaskManager instance can be used as a context manager:

mlmd_handle ML metadata db connection.
pipeline A pipeline IR proto.
task_queue Task queue.
max_active_task_schedulers Maximum number of task schedulers that can be active at once.
max_dequeue_wait_secs Maximum time to wait when dequeuing if the queue is empty.
process_all_queued_tasks_before_exit All existing items in the queues are processed before exiting the context manager. This is useful for deterministic behavior in tests.



View source


View source