yellowdog_provider.operators package

Submodules

yellowdog_provider.operators.yellowdog_operators module

YellowDog Airflow Operators for managing work requirements and worker pools.

class yellowdog_provider.operators.yellowdog_operators.AddPopulatedWorkRequirement(task_id: str, connection_id: Callable[[Context, Environment], str] | str, work_requirement: Callable[[Context, Environment], WorkRequirement] | WorkRequirement, task_groups_and_tasks: Callable[[Context, Environment], list[tuple[TaskGroup, list[Task]]]] | list[tuple[TaskGroup, list[Task]]], **kwargs)[source]

Bases: YellowDogOperator

Add a ‘one-shot’, populated YellowDog work requirement to the platform, including its task groups and tasks.

Parameters:

  • task_id (str) – the Airflow task ID

  • connection_id (str | Callable) – connection to run the operator with (templated) or a Callable that generates the connection ID

  • work_requirement (WorkRequirement | Callable) – a YellowDog WorkRequirement object or a Callable that generates a WorkRequirement object (templated)

  • task_groups_and_tasks (list[tuple[TaskGroup, list[Task]]] | Callable) – a list of YellowDog TaskGroup objects with their constituent Task objects, or a Callable that generates the list (templated)

Returns:

the YellowDog ID of the added work requirement

Return type:

str

execute(context: Context) → str[source]

Adds the work requirement, task groups and tasks.

template_fields: Sequence[str] = ('connection_id', 'work_requirement', 'task_groups_and_tasks')
class yellowdog_provider.operators.yellowdog_operators.AddTaskGroupsToWorkRequirement(task_id: str, connection_id: Callable[[Context, Environment], str] | str, task_groups: Callable[[Context, Environment], list[TaskGroup]] | list[TaskGroup], work_requirement_id: Callable[[Context, Environment], str] | str | None = None, namespace: Callable[[Context, Environment], str] | str | None = None, work_requirement_name: Callable[[Context, Environment], str] | str | None = None, **kwargs)[source]

Bases: YellowDogOperator

Add task groups to a YellowDog work requirement. Either the work_requirement_id or both the namespace and work_requirement_name must be supplied.

Parameters:

  • task_id (str) – the Airflow task ID

  • connection_id (str | Callable) – connection to run the operator with (templated) or a Callable that generates the connection ID

  • task_groups (list[TaskGroup] | Callable) – a list of YellowDog TaskGroup objects or a Callable that generates a list of TaskGroup objects (templated)

  • work_requirement_id (str | Callable | None) – the ID of the work requirement (templated) or a Callable that generates the ID

  • namespace (str | Callable | None) – the namespace of the work requirement (templated) or a Callable that generates the namespace

  • work_requirement_name (str | Callable | None) – the name of a work requirement (templated) or a Callable that generates a work requirement name

Returns:

the YellowDog IDs of the added task groups

Return type:

list[str]

execute(context: Context) → list[str][source]

Adds the task groups to the work requirement.

template_fields: Sequence[str] = ('connection_id', 'work_requirement_id', 'namespace', 'work_requirement_name', 'task_groups')
class yellowdog_provider.operators.yellowdog_operators.AddTasksToTaskGroup(task_id: str, connection_id: Callable[[Context, Environment], str] | str, tasks: Callable[[Context, Environment], list[Task]] | list[Task], task_group_id: Callable[[Context, Environment], str] | str | None = None, namespace: Callable[[Context, Environment], str] | str | None = None, work_requirement_name: Callable[[Context, Environment], str] | str | None = None, task_group_name: Callable[[Context, Environment], str] | str | None = None, **kwargs)[source]

Bases: YellowDogOperator

Add a list of tasks to a task group. Either the task_group_id, or all of: namespace, work_requirement_name and task_group_name, must be supplied.

Parameters:

  • task_id (str) – the Airflow task ID

  • connection_id (str | Callable) – connection to run the operator with (templated) or a Callable that generates the connection ID

  • tasks (list[Task] | Callable) – a list of YellowDog Task objects or a Callable that generates a list of Task objects (templated)

  • task_group_id (str | Callable | None) – the ID of the task_group (templated) or a Callable that generates the ID

  • namespace (str | Callable | None) – the namespace of the work requirement (templated) or a Callable that generates the namespace

  • work_requirement_name (str | Callable | None) – the name of a work requirement (templated) or a Callable that generates a work requirement name

  • task_group_name (str | Callable | None) – the name of a task group (templated) or a Callable that generates a task group name

Returns:

the YellowDog IDs of the added tasks

Return type:

list[str]

execute(context: Context) → list[str][source]

Adds the tasks to the task group.

template_fields: Sequence[str] = ('connection_id', 'namespace', 'task_group_id', 'work_requirement_name', 'task_group_name', 'tasks')
class yellowdog_provider.operators.yellowdog_operators.AddWorkRequirement(task_id: str, connection_id: Callable[[Context, Environment], str] | str, work_requirement: Callable[[Context, Environment], WorkRequirement] | WorkRequirement, **kwargs)[source]

Bases: YellowDogOperator

Add a YellowDog work requirement to the platform.

Parameters:

  • task_id (str) – the Airflow task ID

  • connection_id (str | Callable) – connection to run the operator with (templated) or a Callable that generates the connection ID

  • work_requirement (WorkRequirement or Callable) – a YellowDog WorkRequirement object or a Callable that generates a WorkRequirement object (templated)

Returns:

the YellowDog ID of the added work requirement

Return type:

str

execute(context: Context) → str[source]

Adds the work requirement.

template_fields: Sequence[str] = ('connection_id', 'work_requirement')
class yellowdog_provider.operators.yellowdog_operators.CancelWorkRequirement(task_id: str, connection_id: Callable[[Context, Environment], str] | str, work_requirement_id: Callable[[Context, Environment], str] | str | None = None, namespace: Callable[[Context, Environment], str] | str | None = None, work_requirement_name: Callable[[Context, Environment], str] | str | None = None, abort_running_tasks: bool = False, **kwargs)[source]

Bases: YellowDogOperator

Cancel a YellowDog work requirement. Either the work_requirement_id or both namespace and work_requirement_name must be supplied.

Parameters:

  • task_id (str) – the Airflow task ID

  • connection_id (str | Callable) – connection to run the operator with (templated) or a Callable that generates the connection ID

  • work_requirement_id (str | Callable | None) – the ID of the work requirement (templated) or a Callable that generates the ID

  • namespace (str | Callable | None) – the namespace of the work requirement (templated) or a Callable that generates the namespace

  • work_requirement_name (str | Callable | None) – the name of a work requirement (templated) or a Callable that generates a work requirement name

  • abort_running_tasks (bool) – abort currently running tasks; defaults to ‘False’

Returns:

the YellowDog ID of the cancelled work requirement

Return type:

str

execute(context: Context) → str[source]

Cancels the work requirement, optionally aborting running tasks.

template_fields: Sequence[str] = ('connection_id', 'work_requirement_id', 'namespace', 'work_requirement_name')
class yellowdog_provider.operators.yellowdog_operators.ProvisionWorkerPool(task_id: str, connection_id: Callable[[Context, Environment], str] | str, compute_requirement_template_usage: Callable[[Context, Environment], ComputeRequirementTemplateUsage] | ComputeRequirementTemplateUsage, provisioned_worker_pool_properties: Callable[[Context, Environment], ProvisionedWorkerPoolProperties] | ProvisionedWorkerPoolProperties | None, **kwargs)[source]

Bases: YellowDogOperator

Provision a YellowDog worker pool.

Parameters:

  • task_id (str) – the Airflow task ID

  • connection_id (str | Callable) – connection to run the operator with (templated) or a Callable that generates the connection ID

  • compute_requirement_template_usage (ComputeRequirementTemplateUsage | Callable) – a ComputeRequirementTemplateUsage object or a Callable that generates the object (templated)

  • provisioned_worker_pool_properties (ProvisionedWorkerPoolProperties | Callable | None) – a ProvisionedWorkerPoolProperties object or a Callable that generates the object (templated)

Returns:

the YellowDog ID of the provisioned worker pool

Return type:

str

execute(context: Context) → str[source]

Provisions the worker pool.

template_fields: Sequence[str] = ('connection_id', 'compute_requirement_template_usage', 'provisioned_worker_pool_properties')
class yellowdog_provider.operators.yellowdog_operators.ShutdownProvisionedWorkerPool(task_id: str, connection_id: Callable[[Context, Environment], str] | str, worker_pool_id: Callable[[Context, Environment], str] | str | None = None, namespace: Callable[[Context, Environment], str] | str | None = None, worker_pool_name: Callable[[Context, Environment], str] | str | None = None, terminate_immediately: bool = False, **kwargs)[source]

Bases: YellowDogOperator

Shuts down a YellowDog provisioned worker pool, and optionally immediately terminates its associated compute requirement. Either the worker_pool_id or both the namespace and worker_pool_name must be supplied.

Parameters:

  • task_id (str) – the Airflow task ID

  • connection_id (str | Callable) – connection to run the operator with (templated) or a Callable that generates the connection ID

  • worker_pool_id – the ID of the worker pool (templated) or a Callable that generates the ID

  • namespace (str | Callable | None) – the namespace of the worker pool (templated) or a Callable that generates the namespace

  • worker_pool_name (str | Callable | None) – the name of the worker pool (templated) or a Callable that generates a worker pool name

  • terminate_immediately (bool) – immediately terminate the associated compute requirement (default: False)

Returns:

the YellowDog ID of the worker pool that was shut down

Return type:

str

execute(context: Context) → str | None[source]

Shuts down a worker pool and optionally terminates its compute requirement.

template_fields: Sequence[str] = ('connection_id', 'worker_pool_id', 'namespace', 'worker_pool_name')
class yellowdog_provider.operators.yellowdog_operators.YellowDogOperator(connection_id: Callable[[Context, Environment], str] | str, **kwargs)[source]

Bases: BaseOperator

Base class for YellowDog operators.

Parameters:

connection_id (str | Callable) – connection to run the operator with (templated) or a Callable that generates the connection ID

template_fields: Sequence[str] = ('connection_id',)

Module contents