Repository: apache/airflow

Issue:
[Graph view] After clearing the task (and its downstream tasks) in a task group the task group becomes disconnected from the dag
### Apache Airflow version

2.3.4

### What happened

n the graph view of the dag, after clearing the task (and its downstream tasks) in a task group and refreshing the page the browser the task group becomes disconnected from the dag. See attached gif.
![airflow_2_3_4_task_group_bug](https://user-images.githubusercontent.com/6542519/187409008-767e13e6-ab91-4875-9f3e-bd261b346d0f.gif)
The issue is not persistent and consistent. The graph view becomes disconnected from time to time as you can see on the attached video.

### What you think should happen instead

The graph should be rendered properly and consistently.

### How to reproduce

1. Add the following dag to the dag folder:
```
import logging
import time
from typing import List

import pendulum
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.utils.task_group import TaskGroup


def log_function(message: str, **kwargs):
    logging.info(message)
    time.sleep(3)


def create_file_handling_task_group(supplier):
    with TaskGroup(group_id=f"file_handlig_task_group_{supplier}", ui_color='#666666') as file_handlig_task_group:
        entry = PythonOperator(
            task_id='entry',
            python_callable=log_function,
            op_kwargs={'message': 'create_file_handlig_task_group-Entry-task'}
        )
        with TaskGroup(group_id=f"file_handling_task_sub_group-{supplier}",
                       ui_color='#666666') as file_handlig_task_sub_group:
            sub_group_submit = PythonOperator(
                task_id='sub_group_submit',
                python_callable=log_function,
                op_kwargs={'message': 'create_file_handlig_sub_group_submit'}
            )
            sub_group_monitor = PythonOperator(
                task_id='sub_group_monitor',
                python_callable=log_function,
                op_kwargs={'message': 'create_file_handlig_sub_group_monitor'}
            )
            sub_group_submit >> sub_group_monitor
        entry >> file_handlig_task_sub_group
    return file_handlig_task_group


def get_stage_1_taskgroups(supplierlist: List) -> List[TaskGroup]:
    return [create_file_handling_task_group(supplier) for supplier in supplierlist]


def connect_stage1_to_stage2(self, stage1_tasks: List[TaskGroup], stage2_tasks: List[TaskGroup]) -> None:
    if stage2_tasks:
        for stage1_task in stage1_tasks:
            supplier_code: str = self.get_supplier_code(stage1_task)
            stage2_task = self.get_suppliers_tasks(supplier_code, stage2_tasks)
            stage1_task >> stage2_task


def get_stage_2_taskgroup(taskgroup_id: str):
    with TaskGroup(group_id=taskgroup_id, ui_color='#666666') as stage_2_taskgroup:
        sub_group_submit = PythonOperator(
            task_id='sub_group_submit',
            python_callable=log_function,
            op_kwargs={'message': 'create_file_handlig_sub_group_submit'}
        )
        sub_group_monitor = PythonOperator(
            task_id='sub_group_monitor',
            python_callable=log_function,
            op_kwargs={'message': 'create_file_handlig_sub_group_monitor'}
        )
        sub_group_submit >> sub_group_monitor
    return stage_2_taskgroup


def create_dag():
    with DAG(
            dag_id="horizon-task-group-bug",
            start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
            catchup=False,
            description="description"
    ) as dag:
        start = PythonOperator(
            task_id='start_main',
            python_callable=log_function,
            op_kwargs={'message': 'Entry-task'}
        )
        end = PythonOperator(
            task_id='end_main',
            python_callable=log_function,
            op_kwargs={'message': 'End-task'}
        )
        with TaskGroup(group_id=f"main_file_task_group", ui_color='#666666') as main_file_task_group:
            end_main_file_task_stage_1 = PythonOperator(
                task_id='end_main_file_task_stage_1',
                python_callable=log_function,
                op_kwargs={'message': 'end_main_file_task_stage_1'}
            )
            first_stage = get_stage_1_taskgroups(['9001', '9002'])
            first_stage >> get_stage_2_taskgroup("stage_2_1_taskgroup")
            first_stage >> get_stage_2_taskgroup("stage_2_2_taskgroup")
            first_stage >> end_main_file_task_stage_1
        start >> main_file_task_group >> end
        return dag


dag = create_dag()

```
2. Go to de graph view of the dag.
3. Run the dag.
4. After the dag run has finished. Clear the "sub_group_submit" task within the "stage_2_1_taskgroup" with downstream tasks.
5. Refresh the page multiple times and notice how from time to time the "stage_2_1_taskgroup" becomes disconnected from the dag.
6.  Clear the "sub_group_submit" task within the "stage_2_2_taskgroup" with downstream tasks.
7.  Refresh the page multiple times and notice how from time to time the "stage_2_2_taskgroup" becomes disconnected from the dag.

### Operating System

Mac OS, Linux

### Versions of Apache Airflow Providers

_No response_

### Deployment

Other Docker-based deployment

### Deployment details

Custom docker image based on apache/airflow:2.3.4-python3.10

### Anything else

_No response_

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:
Thanks for opening your first issue here! Be sure to follow the issue template!

The issue is  more of area:ui but I could not add that label. I haven't experienced any issues with dag processing.
Having similar issues with Taskgroups that lose their dependency on other taskgroups. Would be great to have this fixed!
I can replicate this so now let me poke around to see what's the root cause.
The webserver isn't calculating upstream and downstream join ids correctly, leading it to lose some nodes and edges. Restarting the webserver does fix the issue until you clear tasks again. Since the task group is programatically generated I wonder if `task_group.downstream_group_ids` becomes None?
@bbovenzi Is your new graph view (https://github.com/apache/airflow/pull/29413) also going to hit this problem or not?
It does because both graphs use the same data from the webserver. I'll look into this again though.
Did a bunch of testing. 
To replicate: clear a any task in a task group that contains a join node. In our example_dags, `example_task_group.section_2.task_1` shows this well. The join node will disappear after clearing the task and since the node is missing we cannot connect the two task groups anymore and the graph breaks.

This happens because clear calls `dag.partial_subset()` and that function is not properly copying the `dag.task_group` in its memo [here](https://github.com/apache/airflow/blob/main/airflow/models/dag.py#L2183). Removing the task_group memo and calling `filter_task_group` [here](https://github.com/apache/airflow/blob/main/airflow/models/dag.py#L2240) with the copied `dag.task_group` instead of `self.task_group`. It works, but is significantly slower for large dags (2000+ tasks)
I'm not quite sure the best way to fix our deep copy memo.

We do use `partial_subset` for filtering upstream/downstream too and you can replicate that way too.
Got it, thanks!

Fix incoming.

Resolving patch:
diff --git a/airflow/models/dag.py b/airflow/models/dag.py
--- a/airflow/models/dag.py
+++ b/airflow/models/dag.py
@@ -2215,21 +2215,32 @@ def _deepcopy_task(t) -> Operator:
 
         def filter_task_group(group, parent_group):
             """Exclude tasks not included in the subdag from the given TaskGroup."""
+            # We want to deepcopy _most but not all_ attributes of the task group, so we create a shallow copy
+            # and then manually deep copy the instances. (memo argument to deepcopy only works for instances
+            # of classes, not "native" properties of an instance)
             copied = copy.copy(group)
-            copied.used_group_ids = set(copied.used_group_ids)
-            copied._parent_group = parent_group
 
-            copied.children = {}
+            memo[id(group.children)] = {}
+            if parent_group:
+                memo[id(group.parent_group)] = parent_group
+            for attr, value in copied.__dict__.items():
+                if id(value) in memo:
+                    value = memo[id(value)]
+                else:
+                    value = copy.deepcopy(value, memo)
+                copied.__dict__[attr] = value
+
+            proxy = weakref.proxy(copied)
 
             for child in group.children.values():
                 if isinstance(child, AbstractOperator):
                     if child.task_id in dag.task_dict:
                         task = copied.children[child.task_id] = dag.task_dict[child.task_id]
-                        task.task_group = weakref.proxy(copied)
+                        task.task_group = proxy
                     else:
                         copied.used_group_ids.discard(child.task_id)
                 else:
-                    filtered_child = filter_task_group(child, copied)
+                    filtered_child = filter_task_group(child, proxy)
 
                     # Only include this child TaskGroup if it is non-empty.
                     if filtered_child.children:
@@ -2237,7 +2248,7 @@ def filter_task_group(group, parent_group):
 
             return copied
 
-        dag._task_group = filter_task_group(self._task_group, None)
+        dag._task_group = filter_task_group(self.task_group, None)
 
         # Removing upstream/downstream references to tasks and TaskGroups that did not make
         # the cut.


Test patch:
diff --git a/tests/models/test_dag.py b/tests/models/test_dag.py
--- a/tests/models/test_dag.py
+++ b/tests/models/test_dag.py
@@ -24,6 +24,7 @@
 import pickle
 import re
 import sys
+import weakref
 from contextlib import redirect_stdout
 from datetime import timedelta
 from pathlib import Path
@@ -1381,7 +1382,7 @@ def test_duplicate_task_ids_for_same_task_is_allowed(self):
         assert dag.task_dict == {op1.task_id: op1, op3.task_id: op3}
         assert dag.task_dict == {op2.task_id: op2, op3.task_id: op3}
 
-    def test_sub_dag_updates_all_references_while_deepcopy(self):
+    def test_partial_subset_updates_all_references_while_deepcopy(self):
         with DAG("test_dag", start_date=DEFAULT_DATE) as dag:
             op1 = EmptyOperator(task_id="t1")
             op2 = EmptyOperator(task_id="t2")
@@ -1389,11 +1390,37 @@ def test_sub_dag_updates_all_references_while_deepcopy(self):
             op1 >> op2
             op2 >> op3
 
-        sub_dag = dag.partial_subset("t2", include_upstream=True, include_downstream=False)
-        assert id(sub_dag.task_dict["t1"].downstream_list[0].dag) == id(sub_dag)
+        partial = dag.partial_subset("t2", include_upstream=True, include_downstream=False)
+        assert id(partial.task_dict["t1"].downstream_list[0].dag) == id(partial)
 
         # Copied DAG should not include unused task IDs in used_group_ids
-        assert "t3" not in sub_dag._task_group.used_group_ids
+        assert "t3" not in partial.task_group.used_group_ids
+
+    def test_partial_subset_taskgroup_join_ids(self):
+        with DAG("test_dag", start_date=DEFAULT_DATE) as dag:
+            start = EmptyOperator(task_id="start")
+            with TaskGroup(group_id="outer", prefix_group_id=False) as outer_group:
+                with TaskGroup(group_id="tg1", prefix_group_id=False) as tg1:
+                    EmptyOperator(task_id="t1")
+                with TaskGroup(group_id="tg2", prefix_group_id=False) as tg2:
+                    EmptyOperator(task_id="t2")
+
+                start >> tg1 >> tg2
+
+        # Pre-condition checks
+        task = dag.get_task("t2")
+        assert task.task_group.upstream_group_ids == {"tg1"}
+        assert isinstance(task.task_group.parent_group, weakref.ProxyType)
+        assert task.task_group.parent_group == outer_group
+
+        partial = dag.partial_subset(["t2"], include_upstream=True, include_downstream=False)
+        copied_task = partial.get_task("t2")
+        assert copied_task.task_group.upstream_group_ids == {"tg1"}
+        assert isinstance(copied_task.task_group.parent_group, weakref.ProxyType)
+        assert copied_task.task_group.parent_group
+
+        # Make sure we don't affect the original!
+        assert task.task_group.upstream_group_ids is not copied_task.task_group.upstream_group_ids
 
     def test_schedule_dag_no_previous_runs(self):
         """