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(*args: Any, **kwargs: Any)[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: airflow.sdk.definitions.context.Context) → str[source]

Adds the work requirement, task groups and tasks.

template_fields: Collection[str] = ('connection_id', 'work_requirement', 'task_groups_and_tasks')
class yellowdog_provider.operators.yellowdog_operators.AddTaskGroupsToWorkRequirement(*args: Any, **kwargs: Any)[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: airflow.sdk.definitions.context.Context) → list[str][source]

Adds the task groups to the work requirement.

template_fields: Collection[str] = ('connection_id', 'work_requirement_id', 'namespace', 'work_requirement_name', 'task_groups')
class yellowdog_provider.operators.yellowdog_operators.AddTasksToTaskGroup(*args: Any, **kwargs: Any)[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: airflow.sdk.definitions.context.Context) → list[str][source]

Adds the tasks to the task group.

template_fields: Collection[str] = ('connection_id', 'namespace', 'task_group_id', 'work_requirement_name', 'task_group_name', 'tasks')
class yellowdog_provider.operators.yellowdog_operators.AddWorkRequirement(*args: Any, **kwargs: Any)[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: airflow.sdk.definitions.context.Context) → str[source]

Adds the work requirement.

template_fields: Collection[str] = ('connection_id', 'work_requirement')
class yellowdog_provider.operators.yellowdog_operators.CancelWorkRequirement(*args: Any, **kwargs: Any)[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: airflow.sdk.definitions.context.Context) → str[source]

Cancels the work requirement, optionally aborting running tasks.

template_fields: Collection[str] = ('connection_id', 'work_requirement_id', 'namespace', 'work_requirement_name')
class yellowdog_provider.operators.yellowdog_operators.ProvisionWorkerPool(*args: Any, **kwargs: Any)[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: airflow.sdk.definitions.context.Context) → str[source]

Provisions the worker pool.

template_fields: Collection[str] = ('connection_id', 'compute_requirement_template_usage', 'provisioned_worker_pool_properties')
class yellowdog_provider.operators.yellowdog_operators.ShutdownProvisionedWorkerPool(*args: Any, **kwargs: Any)[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 (str | Callable | None) – 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: airflow.sdk.definitions.context.Context) → str | None[source]

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

template_fields: Collection[str] = ('connection_id', 'worker_pool_id', 'namespace', 'worker_pool_name')
class yellowdog_provider.operators.yellowdog_operators.YellowDogOperator(*args: Any, **kwargs: Any)[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: Collection[str] = ('connection_id',)

Module contents