Add pattern matching features on allow and block lists (GoogleCloudPlatform#956)

* Add DAG item rules to manage allow and block lists. Update readme. Migrate allow_list.txt

* refactore

* refactor for review

Author: severus-ho

plugins/README.md

@@ -10,22 +10,36 @@ By default, this plugin is **disabled** for all DAGs.
 
 ##### Enabling the Plugin for a DAG
 
-If you want your DAG to **trigger this plugin**, add your DAG ID as a new line in **plugins/allow_list.txt**.
+If you want your DAG to **trigger this plugin**, add your DAG item as a new line in **plugins/allow_list.txt**.
 
 ##### Disabling the Plugin for a DAG
 
-To prevent your DAG from ever triggering this plugin, add your DAG ID as a new line in **plugins/block_list.txt**. This list is intended to avoid repeated issue postings for DAGs that are not currently resolvable.
-
-> **Note:** Each DAG ID should be placed on its own line, without extra spaces or quotes.
+To prevent your DAG from ever triggering this plugin, add your DAG item as a new line in **plugins/block_list.txt**. This list is intended to avoid repeated issue postings for DAGs that are not currently resolvable.
 
+> **Note:** Each item should be placed on its own line, without extra spaces or quotes.
+##### DAG item rules
+##### 1. id
+```
+id:<DAG_ID>
+```
+<DAG_ID> should be the same as the DAG ID you want to track.
+##### 2. tag
+```
+tag:<DAG_TAG> 
+```
+The DAG will be tracked if <DAG_TAG> matches one of its tags. An empty tag ('') is invalid and will be ignored.
+##### 3. pattern
+```
+pattern:<DAG_ID_PATTERN>
+```
+<DAG_ID_PATTERN> follows the same rules as UNIX shell globbing, not regular expression.
 ##### Example
-
 ```
-<DAG_ID 1>
-<DAG_ID 2>
+id:on_failure_action_trigger
+tag:xlml
+pattern:on_failure_action_*
 ...
 ```
-
 ##### Plugin Default Toggle
 The `enable_plugin_by_default` attribute in `config.json` should be set to true to have all completed DAGs trigger this plugin, unless those DAGs are explicitly listed in the `allow_list.txt` or `block_list.txt`.
 If you don't want this plugin to be triggered by default, please set `enable_plugin_by_default` to false.

plugins/allow_list.txt

@@ -1,26 +1,11 @@
-on_failure_actions_trigger
-maxtext_sft_trainer
-maxtext_configs_aot
-maxtext_convergence
-maxtext_end_to_end
-maxtext_profiling
-gke_node_pool_status
-multi-host-availability-rollback
-validate_interruption_count_gce_bare_metal_preemption
-validate_interruption_count_gce_defragmentation
-validate_interruption_count_gce_eviction
-validate_interruption_count_gce_host_error
-validate_interruption_count_gce_hwsw_maintenance
-validate_interruption_count_gce_migrate_on_hwsw_maintenance
-validate_interruption_count_gce_other
-validate_interruption_count_gke_bare_metal_preemption
-validate_interruption_count_gke_defragmentation
-validate_interruption_count_gke_eviction
-validate_interruption_count_gke_host_error
-validate_interruption_count_gke_hwsw_maintenance
-validate_interruption_count_gke_migrate_on_hwsw_maintenance
-validate_interruption_count_gke_other
-maxtext_emc_and_mtc_orbax_save_local
-maxtext_mtc_orbax_save_gcs
-maxtext_emc_orbax_res_local
-maxtext_mtc_orbax_res_local
+id:on_failure_actions_trigger
+id:maxtext_sft_trainer
+id:maxtext_configs_aot
+id:maxtext_convergence
+id:maxtext_end_to_end
+id:maxtext_profiling
+id:gke_node_pool_status
+id:multi-host-availability-rollback
+pattern:validate_interruption_*
+pattern:maxtext_mtc_*
+pattern:maxtext_emc_*

plugins/on_failure_actions.py

@@ -1,8 +1,9 @@
+import fnmatch
 import json
 import logging
 import os
 import time
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Set
 
 import google.auth.transport.requests
 import jwt
@@ -135,7 +136,7 @@ def create_comment(issue, body):
   issue.create_comment(body=body)
 
 
-def read_list_from_gcs(bucket_name: str, blob_name: str) -> set[str]:
+def read_items_from_gcs(bucket_name: str, blob_name: str) -> set[str]:
   storage_client = storage.Client()
   bucket = storage_client.bucket(bucket_name)
   blob = bucket.blob(blob_name)
@@ -190,18 +191,19 @@ def on_dag_finished(self, dag_run: DagRun, msg: str):
     try:
       # A DAG would trigger the following logic if the DAG ID is in allow list and not in block list.
       enable_plugin = None
-      if DagRunListener.is_in_allow_list(dag_run.dag_id):
+      if DagRunListener.is_in_allow_list(dag_run):
         enable_plugin = True
         logging.info(
             f"[{self.log_prefix}] DAG {dag_run.dag_id} is in allow_list.txt"
         )
-      if DagRunListener.is_in_block_list(dag_run.dag_id):
+
+      if DagRunListener.is_in_block_list(dag_run):
         enable_plugin = False
         logging.info(
             f"[{self.log_prefix}] DAG {dag_run.dag_id} is in block_list.txt"
         )
 
-      default_enabled = DagRunListener.enable_plugin_by_default(dag_run.dag_id)
+      default_enabled = DagRunListener.enable_plugin_by_default()
 
       if enable_plugin is False or (
           enable_plugin is None and not default_enabled
@@ -331,27 +333,71 @@ def get_airflow_url(project: str, region: str, env: str) -> str:
     return configs["config"]["airflowUri"]
 
   @staticmethod
-  def is_in_allow_list(dag_id: str) -> bool:
-    allow_list = read_list_from_gcs(
+  def is_in_allow_list(dag_run: DagRun) -> bool:
+    allow_items = read_items_from_gcs(
         os.environ.get("GCS_BUCKET"), ALLOW_LIST_PATH
     )
-    return True if dag_id in allow_list else False
+    return DagRunListener.is_in_list(allow_items, dag_run)
 
   @staticmethod
-  def is_in_block_list(dag_id: str) -> bool:
-    block_list = read_list_from_gcs(
+  def is_in_block_list(dag_run: DagRun) -> bool:
+    block_items = read_items_from_gcs(
         os.environ.get("GCS_BUCKET"), BLOCK_LIST_PATH
     )
-    return True if dag_id in block_list else False
+    return DagRunListener.is_in_list(block_items, dag_run)
 
   @staticmethod
-  def enable_plugin_by_default(dag_id: str) -> bool:
+  def is_in_list(items: Set[str], dag_run: DagRun):
+    return (
+        DagRunListener.contains_id(items, dag_run.dag_id)
+        or DagRunListener.has_any_tag(items, dag_run.dag.tags)
+        or DagRunListener.matches_pattern(items, dag_run.dag_id)
+    )
+
+  @staticmethod
+  def enable_plugin_by_default() -> bool:
     config = read_json_from_gcs(os.environ.get("GCS_BUCKET"), CONFIG_PATH)
     logging.info(
         f"[DagRunListener] Enable plugin by default: {config['enable_plugin_by_default']}"
     )
     return config["enable_plugin_by_default"]
 
+  @staticmethod
+  def contains_id(items: Set[str], target_id: str):
+    id_set = set()
+    for item in items:
+      parts = item.split(":", 1)
+      if len(parts) == 2:
+        key = parts[0].lower()
+        value = parts[1].strip()
+        if key == "id":
+          id_set.add(value)
+    return target_id in id_set
+
+  @staticmethod
+  def has_any_tag(items: Set[str], target_tags: List[str]):
+    tag_set = set()
+    for item in items:
+      parts = item.split(":", 1)
+      if len(parts) == 2:
+        key = parts[0].lower()
+        value = parts[1].strip()
+        if key == "tag" and value:
+          tag_set.add(value)
+    return not tag_set.isdisjoint(target_tags)
+
+  @staticmethod
+  def matches_pattern(items: Set[str], target_id: str):
+    pattern_set = set()
+    for item in items:
+      parts = item.split(":", 1)
+      if len(parts) == 2:
+        key = parts[0].lower()
+        value = parts[1].strip()
+        if key == "pattern" and value:
+          pattern_set.add(value)
+    return any(fnmatch.fnmatch(target_id, pattern) for pattern in pattern_set)
+
 
 class ListenerPlugins(AirflowPlugin):
   name = "listener_plugins"