Join the SIG TFX-Addons community and help make TFX even better!

tfx.orchestration.experimental.core.task_scheduler.TaskSchedulerRegistry

A registry for task schedulers.

Methods

clear

View source

create_task_scheduler

View source

Creates a task scheduler for the given task.

The task is matched as follows:

  1. The node type name of the node associated with the task is looked up in the registry and a scheduler is instantiated if present.
  2. Next, the executor spec url of the node (if one exists) is looked up in the registry and a scheduler is instantiated if present. This assumes deployment_config packed in the pipeline IR is of type IntermediateDeploymentConfig.
  3. Lastly, a ValueError is raised if no match can be found.

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 If a scheduler could not be found in the registry for the given task.

register

View source

Registers a new task scheduler for the given url.

Args
url The URL associated with the task scheduler. It should either be the node type url or executor spec url.
scheduler_class The class that will be instantiated for a matching task.

Raises
ValueError If url is already in the registry.