Fix TaskGroupWithTimeout teardown initialization and remove redundant .as_teardown() calls

[yuna-tzeng]

Description

This PR addresses an issue where tasks (specifically the injected provision_taskgroup_session root node) within a TaskGroupWithTimeout were not
correctly inheriting the is_teardown=True property, and fixes an AttributeError caused by premature initialization.

Tests

• Dag Name : jobset_rollback_ttr
• Airflow test results: https://39b01025d8504fa5afa7226c9e54d883-dot-us-east1.composer.googleusercontent.com/dags/jobset_rollback_ttr/grid

Airflow/Composer

• GCP Composer name: yuna-ml-auto (under GCP project: cloud-ml-auto-solutions)
• GCP Composer version: 2.13.1

Changes Made:

1. **Fix `TaskGroupWithTimeout` teardown assignment**:
   - Moved the `is_teardown = True` assignment from the `add()` method to the `__exit__()` method.
   - **Reason**: When `add()` is called during a `BaseOperator`'s `__init__`, the task is not yet fully initialized (e.g., `_is_setup` attribute doesn't exist yet), causing an `AttributeError` when setting `is_teardown`. Moving it to `__exit__()` ensures all tasks are fully instantiated before properties are assigned.
   - This also guarantees that the implicitly created root node (`provision_taskgroup_session`) correctly receives the teardown status.

2. **Clean up DAGs**:
   - Removed redundant `.as_teardown()` calls from `post_test` teardown tasks across 21 DAGs in `dags/tpu_observability/`.
   - **Reason**: Since `TaskGroupWithTimeout(is_teardown=True)` now robustly handles applying the teardown property to all its child tasks, explicitly linking them with `.as_teardown()` is no longer necessary. This simplifies the DAG code significantly while maintaining the expected cleanup behavior and `TriggerRule.ALL_DONE` guarantee.

## Impact
- DAGs will now correctly report a `Failed` status if the main test fails, while still successfully executing the teardown group.
- The `provision_taskgroup_session` task will properly show up as a teardown task in the Airflow UI.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run one-shot tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[alfredyu-cienet]

not used?

[alfredyu-cienet]

under this block, not in __exit__

as you can see in the doc string of those 2 context manager methods, those were overridden specifically for forcing a root node

in addition, we have different handling for different types of "node", check the instance there is redundant

on the other hand, if your thought was wanted to inject is_teardown = True for the root node, do that under the block L134

[yuna-tzeng]

we cannot set node.is_teardown = True inside the add() method, this causes an AttributeError: PythonOperator object has no attribute _is_setup.

When an Operator is instantiated, its __init__ method runs. Midway through this initialization, if Airflow detects an active TaskGroup context, it immediately calls task_group.add(self).
This means that when our custom add() method receives the node, the object is only partially initialized. Crucial internal attributes, such as _is_setup and _is_teardown, haven't been created yet because the init method hasn't finished executing.
When we attempt to set node.is_teardown = True inside add(), Airflow's property setter triggers a validation check to ensure a task isn't both a setup and a teardown task (if self.is_setup:). Since the underlying _is_setup attribute doesn't exist yet, it throws an AttributeError

Alternatively, we could bypass the AttributeError by manually injecting the missing attribute (e.g., node._is_setup = False) inside the add() method before setting node.is_teardown = True.
However, I feel this might not be a good choice.
Should I go with this approach?

[yuna-tzeng]

setting task.is_teardown = True inside wrapped_execute doesn't work because wrapped_execute only happens at runtime.

Airflow's Scheduler determines the DAG structure and task roles (like is_teardown) at parse-time. Therefore, any attribute changes made dynamically during execution remain invisible to the Scheduler when it calculates the final DAG state.

As a result, the tasks we intended to be teardowns will still be treated as regular tasks, which ultimately masks the upstream failures.

Test result: https://39b01025d8504fa5afa7226c9e54d883-dot-us-east1.composer.googleusercontent.com/dags/jobset_rollback_ttr/grid?dag_run_id=manual__2026-07-15T05%3A43%3A40.884284%2B00%3A00

[alfredyu-cienet]

I see, let me think about it

[alfredyu-cienet]

how about this

dags/common/task_group_with_timeout/task_group_with_timeout.py:

class TaskGroupWithTimeout(TaskGroup):
  ...

  def __init__(
      self,
      group_id,
      timeout: timedelta,
      as_teardown_of: BaseOperator | None = None,
      **kwargs,
  ):
    super().__init__(group_id=group_id, **kwargs)
    self.group_name = f"{self.__class__.__name__}: '{group_id}'"
    self.timeout = timeout
    self.setup_op = as_teardown_of
    self._root_node = None

  # remove __enter__

  def __exit__(self, *args):
    """Wire `_root_node` as upstream of in-group root children on context exit.

    Overridden to guarantee `_root_node` runs first. Only children with no
    in-group sibling upstream get a direct edge; others inherit transitively.
    """
    self.initialize_task_group_session()
    self.mark_teardown()
    return super().__exit__(*args)

  def initialize_task_group_session(self):
    """
    TBD
    """
    self._root_node = PythonOperator(
        task_id=self.ROOT_TASK_ID,
        python_callable=lambda: datetime.now(timezone.utc).isoformat(),
    ).as_teardown(
        setups=self.setup_op,
        on_failure_fail_dagrun=True,
    )

    children_ids = set(self.children.keys())
    for child in self.children.values():
      if child is self._root_node:
        continue
      # If a sibling already chains into this child, the dependency on
      # _root_node is satisfied transitively — no need to add a direct edge.
      if child.upstream_task_ids & children_ids:
        continue
      child.set_upstream(self._root_node)

  def mark_teardown(self):
    """
    TBD
    """
    for child in self.children.values():
      match child:
        case TaskGroup():
          pass # A TaskGroup does not have a teardown attribute.

        case AbstractOperator():
          child.as_teardown(
              setups=self.setup_op,
              on_failure_fail_dagrun=True,
          )