tfx.orchestration.experimental.core.task_scheduler.TaskSchedulerRegistry

A registry for task schedulers.

Methods

create_task_scheduler

View source

Creates a task scheduler for the given task.

Note that this assumes deployment_config packed in the pipeline IR is of type IntermediateDeploymentConfig. This detail may change in the future.

Args
mlmd_handle A handle to the MLMD db.
pipeline The pipeline IR.
task The task that needs to be scheduled.

Returns
An instance of TaskScheduler for the given task.

Raises
NotImplementedError Raised if not an ExecNodeTask.
ValueError Deployment config not present in the IR proto or if executor spec for the node corresponding to task not configured in the IR.

register

View source

Registers a new task scheduler for the given executor spec type url.

Args
executor_spec_type_url The URL of the executor spec type.
scheduler_class The class that will be instantiated for a matching task.

Raises
ValueError If executor_spec_type_url is already in the registry.