Pass "real" TaskInstance object with xcom methods in context to task execute functions

The K8s system tests are running https://github.com/apache/airflow/blob/cd5ccf0abf5027cf457b36a8a2bee397dcb538b4/airflow/example_dags/example_xcom.py and this is currently failing as we don't pass ti to the operator execute() function.

(We currently pass task_isntance, not ti, but even passing it as ti wouldn't work as the simple class we have for TI in TaskSDK doesn't have xcom_pull or xcom_push methods.

Sub-tasks:

The current values in the Context dict:

airflow/airflow/models/taskinstance.py

Lines 1009 to 1048 in b3f3649

    
           # NOTE: If you add to this dict, make sure to also update the following: 
        
           # * Context in airflow/utils/context.pyi 
        
           # * KNOWN_CONTEXT_KEYS in airflow/utils/context.py 
        
           # * Table in docs/apache-airflow/templates-ref.rst 
        
           context: dict[str, Any] = { 
        
               "conf": conf, 
        
               "dag": dag, 
        
               "dag_run": dag_run, 
        
               "data_interval_end": timezone.coerce_datetime(data_interval.end), 
        
               "data_interval_start": timezone.coerce_datetime(data_interval.start), 
        
               "outlet_events": OutletEventAccessors(), 
        
               "ds": ds, 
        
               "ds_nodash": ds_nodash, 
        
               "expanded_ti_count": expanded_ti_count, 
        
               "inlets": task.inlets, 
        
               "inlet_events": InletEventsAccessors(task.inlets, session=session), 
        
               "logical_date": logical_date, 
        
               "macros": macros, 
        
               "map_index_template": task.map_index_template, 
        
               "outlets": task.outlets, 
        
               "params": validated_params, 
        
               "prev_data_interval_start_success": get_prev_data_interval_start_success(), 
        
               "prev_data_interval_end_success": get_prev_data_interval_end_success(), 
        
               "prev_start_date_success": get_prev_start_date_success(), 
        
               "prev_end_date_success": get_prev_end_date_success(), 
        
               "run_id": task_instance.run_id, 
        
               "task": task, 
        
               "task_instance": task_instance, 
        
               "task_instance_key_str": f"{task.dag_id}__{task.task_id}__{ds_nodash}", 
        
               "test_mode": task_instance.test_mode, 
        
               "ti": task_instance, 
        
               "triggering_asset_events": lazy_object_proxy.Proxy(get_triggering_events), 
        
               "ts": ts, 
        
               "ts_nodash": ts_nodash, 
        
               "ts_nodash_with_tz": ts_nodash_with_tz, 
        
               "var": { 
        
                   "json": VariableAccessor(deserialize_json=True), 
        
                   "value": VariableAccessor(deserialize_json=False), 
        
               }, 
        
               "conn": ConnectionAccessor(),

Template ref: https://airflow.apache.org/docs/apache-airflow/stable/templates-ref.html

The plan is to:

Get certain key:value pairs from the API Server and pass it as a response to start the request to the API server (with TIEnterRunningPayload)
Augment it with key:value pairs from the Task Execution Interface (execution side of the Task SDK)

As a goal: I am planning to get as much info as possible on the Execution side itself as opposed to getting it from the API server

WIP branch: https://github.com/apache/airflow/compare/main...astronomer:airflow:task-context?expand=1

AIP-72: Add some basic Task Context keys (#44894)

AIP-72: Pass context keys from API Server to Workers

AIP-72: Add some basic Task Context keys (apache#44894)

AIP-72: Fix recursion bug with `XComArg` (#45112)

AIP-72: Allow pushing and pulling XCom from Task Context (#45075)

AIP-72: Add Taskflow API support & template rendering in Task SDK (#4…

AIP-72: Add some basic Task Context keys (#44894)

AIP-72: Pass context keys from API Server to Workers (#44899)

AIP-72: Allow retrieving Connection from Task Context (#45043)

AIP-72: Fix recursion bug with `XComArg` (#45112)

AIP-72: Allow pushing and pulling XCom from Task Context (#45075)

AIP-72: Add Taskflow API support & template rendering in Task SDK (#4…

AIP-72: Add some basic Task Context keys (#44894)

AIP-72: Pass context keys from API Server to Workers (#44899)

AIP-72: Allow retrieving Connection from Task Context (#45043)

AIP-72: Fix recursion bug with `XComArg` (#45112)

AIP-72: Allow pushing and pulling XCom from Task Context (#45075)

AIP-72: Add Taskflow API support & template rendering in Task SDK (#4…

Navigation Menu

Repository navigation

Pass "real" TaskInstance object with xcom methods in context to task execute functions #44481

70 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	# NOTE: If you add to this dict, make sure to also update the following:
	# * Context in airflow/utils/context.pyi
	# * KNOWN_CONTEXT_KEYS in airflow/utils/context.py
	# * Table in docs/apache-airflow/templates-ref.rst
	context: dict[str, Any] = {
	"conf": conf,
	"dag": dag,
	"dag_run": dag_run,
	"data_interval_end": timezone.coerce_datetime(data_interval.end),
	"data_interval_start": timezone.coerce_datetime(data_interval.start),
	"outlet_events": OutletEventAccessors(),
	"ds": ds,
	"ds_nodash": ds_nodash,
	"expanded_ti_count": expanded_ti_count,
	"inlets": task.inlets,
	"inlet_events": InletEventsAccessors(task.inlets, session=session),
	"logical_date": logical_date,
	"macros": macros,
	"map_index_template": task.map_index_template,
	"outlets": task.outlets,
	"params": validated_params,
	"prev_data_interval_start_success": get_prev_data_interval_start_success(),
	"prev_data_interval_end_success": get_prev_data_interval_end_success(),
	"prev_start_date_success": get_prev_start_date_success(),
	"prev_end_date_success": get_prev_end_date_success(),
	"run_id": task_instance.run_id,
	"task": task,
	"task_instance": task_instance,
	"task_instance_key_str": f"{task.dag_id}__{task.task_id}__{ds_nodash}",
	"test_mode": task_instance.test_mode,
	"ti": task_instance,
	"triggering_asset_events": lazy_object_proxy.Proxy(get_triggering_events),
	"ts": ts,
	"ts_nodash": ts_nodash,
	"ts_nodash_with_tz": ts_nodash_with_tz,
	"var": {
	"json": VariableAccessor(deserialize_json=True),
	"value": VariableAccessor(deserialize_json=False),
	},
	"conn": ConnectionAccessor(),

Pass "real" TaskInstance object with xcom methods in context to task execute functions #44481

Description

Activity

kaxil commented on Dec 10, 2024

kaxil commented on Dec 11, 2024

70 remaining items

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions