<p>I have a DAG that uses <code>KubernetesPodOperator</code> and the task <code>get_train_test_model_task_count</code> in the DAG pushes an xcom variable and I want to use it in the following tasks.</p>
<pre><code>run_this = BashOperator(
    task_id=&quot;also_run_this&quot;,
    bash_command='echo &quot;ti_key={{ ti.xcom_pull(task_ids=\&quot;get_train_test_model_task_count\&quot;, key=\&quot;return_value\&quot;)[\&quot;models_count\&quot;] }}&quot;',
)
</code></pre>
<p>The above DAG task works and it prints the value as <code>ti_key=24</code>.</p>
<p>I want the same value to be used as a variable,</p>
<pre><code>with TaskGroup(&quot;train_test_model_config&quot;) as train_test_model_config:
    models_count = &quot;{{ ti.xcom_pull(task_ids=\&quot;get_train_test_model_task_count\&quot;, key=\&quot;return_value\&quot;)[\&quot;models_count\&quot;] }}&quot;
    print(models_count)
    for task_num in range(0, int(models_count)):
        generate_train_test_model_config_task(task_num)
</code></pre>
<p><code>int(models_count)</code> doesnot work, by throwing the error -</p>
<blockquote>
<p>ValueError: invalid literal for int() with base 10: '{{ ti.xcom_pull(task_ids=&quot;get_train_test_model_task_count&quot;, key=&quot;return_value&quot;)[&quot;models_count&quot;] }}'</p>
</blockquote>
<p>And the <code>generate_train_test_model_config_task</code> looks as below:</p>
<pre><code>def generate_train_test_model_config_task(task_num):
    task = KubernetesPodOperator(
        name=f&quot;train_test_model_config_{task_num}&quot;,
        image=build_model_image,
        labels=labels,
        cmds=[
            &quot;python3&quot;,
            &quot;-m&quot;,
            &quot;src.models.train_test_model_config&quot;,
            &quot;--tenant=neu&quot;,
            f&quot;--model_tag_id={task_num}&quot;,
            &quot;--line_plan={{ ti.xcom_pull(key=\&quot;file_name\&quot;, task_ids=\&quot;extract_file_name\&quot;) }}&quot;,
            &quot;--staging_bucket=cs-us-ds&quot;
        ],
        task_id=f&quot;train_test_model_config_{task_num}&quot;,
        do_xcom_push=False,
        namespace=&quot;airflow&quot;,
        service_account_name=&quot;airflow-worker&quot;,
        get_logs=True,
        startup_timeout_seconds=300,
        container_resources={&quot;request_memory&quot;: &quot;29G&quot;, &quot;request_cpu&quot;: &quot;7000m&quot;},
        node_selector={&quot;cloud.google.com/gke-nodepool&quot;: NODE_POOL},
        tolerations=[
            {
                &quot;key&quot;: NODE_POOL,
                &quot;operator&quot;: &quot;Equal&quot;,
                &quot;value&quot;: &quot;true&quot;,
                &quot;effect&quot;: &quot;NoSchedule&quot;,
            }
        ],
    )

    return task
</code></pre>

Tom J Muthirenthi
<p>The Jinja template pulls from the Airflow context which you only can do within a task, not in top level code.</p>
<p>Also as a commenter said you will need to use dynamic task mapping to change the DAG structure dynamically, even if you hardcode the model_num or use another way to template it in, those code changes are only picked up every 30s by the scheduler on default and you have no backwards visibility into previous tasks, for example if one day there are only 2 models you can't see model 3 through 8 in the logs from the day before so it gets a bit messy when using a loop like that even if you can get it to work.</p>
<p>The code below shows the structure that I think will achieve what you want, one model config generated for each <code>task_num</code>. This should work in Airflow 2.3+</p>
<pre class="lang-py prettyprint-override"><code>@task
def generate_list_of_model_nums(**context):
    model_count = context[&quot;ti&quot;].xcom_pull(task_ids=&quot;get_train_test_model_task_count&quot;, key=&quot;return_value&quot;)[&quot;models_count&quot;]
    return list(range(model_count + 1))


@task
def generate_train_test_model_config_task(task_num):
    # code that generates the model config
    return model_config

model_nums=generate_list_of_model_nums()
generate_train_test_model_config_task.expand(task_num=model_nums)
</code></pre>
<p>Notes: I did not test the code above so there might be typos, but this is the general idea, create a list of all the task nums, then use dynamic task mapping to expand over the list.</p>
<p>If you pull the XCom from the <code>generate_train_test_model_config_task</code> you should get a list of all the model configs :)</p>
<p>Some resources that might help to adapt this to traditional operators:</p>
<ul>
<li><a href="https://docs.astronomer.io/learn/dynamic-tasks" rel="nofollow noreferrer">Dynamic task mapping guide</a></li>
<li><a href="https://docs.astronomer.io/learn/airflow-context" rel="nofollow noreferrer">Airflow context guide</a></li>
</ul>
<p>Disclaimer: I work at Astronomer the org who created the guides above :)</p>
<p>EDIT: thanks for sharing the KPO code! I see you are using the task_num in two parameters, this means you can try to use <code>.expand_kwargs</code> over a list of sets of inputs in form of a dictionaries and then map the KPO directly. Note that this is an Airflow 2.4+ feature.</p>
<p>Note on the code: I tested the dict generation function but don't have a K8s cluster running rn so I did not test the latter part, I think <code>name</code> and <code>cmd</code> should be expandable 🤞</p>
<pre class="lang-py prettyprint-override"><code>@task
def generate_list_of_param_dicts(**context):
    model_count = context[&quot;ti&quot;].xcom_pull(
        task_ids=&quot;get_train_test_model_task_count&quot;, key=&quot;return_value&quot;
    )[&quot;models_count&quot;]

    param_dicts = []
    for i in range(model_count):
        param_dict = {
            &quot;name&quot;: f&quot;train_test_model_config_{i}&quot;,
            &quot;cmds&quot;: [
                &quot;python3&quot;,
                &quot;-m&quot;,
                &quot;src.models.train_test_model_config&quot;,
                &quot;--tenant=neu&quot;,
                f&quot;--model_tag_id={i}&quot;,
                '--line_plan={{ ti.xcom_pull(key=&quot;file_name&quot;, task_ids=&quot;extract_file_name&quot;) }}',
                &quot;--staging_bucket=cs-us-ds&quot;,
            ],
        }
        param_dicts.append(param_dict)

    return param_dicts


task = KubernetesPodOperator.partial(
    image=build_model_image,
    labels=labels,
    task_id=f&quot;train_test_model_config&quot;,
    do_xcom_push=False,
    namespace=&quot;airflow&quot;,
    service_account_name=&quot;airflow-worker&quot;,
    get_logs=True,
    startup_timeout_seconds=300,
    container_resources={&quot;request_memory&quot;: &quot;29G&quot;, &quot;request_cpu&quot;: &quot;7000m&quot;},
    node_selector={&quot;cloud.google.com/gke-nodepool&quot;: NODE_POOL},
    tolerations=[
        {
            &quot;key&quot;: NODE_POOL,
            &quot;operator&quot;: &quot;Equal&quot;,
            &quot;value&quot;: &quot;true&quot;,
            &quot;effect&quot;: &quot;NoSchedule&quot;,
        }
    ],
).expand_kwargs(generate_list_of_param_dicts())
</code></pre>
