fix: Updated Error Type to TaskclusterRestFailure to align with the switch to the taskcluster package

Author: abhishekmadan30

src/taskgraph/actions/cancel.py

@@ -7,6 +7,7 @@
 
 import requests
 
+from taskcluster import TaskclusterRestFailure
 from taskgraph.util.taskcluster import cancel_task
 
 from .registry import register_callback_action
@@ -27,8 +28,14 @@ def cancel_action(parameters, graph_config, input, task_group_id, task_id):
     # only cancel tasks with the level-specific schedulerId.
     try:
         cancel_task(task_id)
-    except requests.HTTPError as e:
-        if e.response.status_code == 409:
+    except (requests.HTTPError, TaskclusterRestFailure) as e:
+        status_code = None
+        if isinstance(e, requests.HTTPError):
+            status_code = e.response.status_code if e.response else None
+        elif isinstance(e, TaskclusterRestFailure):
+            status_code = e.status_code
+
+        if status_code == 409:
             # A 409 response indicates that this task is past its deadline.  It
             # cannot be cancelled at this time, but it's also not running
             # anymore, so we can ignore this error.

src/taskgraph/actions/cancel_all.py

@@ -9,6 +9,7 @@
 
 import requests
 
+from taskcluster import TaskclusterRestFailure
 from taskgraph.util.taskcluster import (
     CONCURRENCY,
     cancel_task,
@@ -36,8 +37,14 @@ def do_cancel_task(task_id):
         logger.info(f"Cancelling task {task_id}")
         try:
             cancel_task(task_id)
-        except requests.HTTPError as e:
-            if e.response.status_code == 409:
+        except (requests.HTTPError, TaskclusterRestFailure) as e:
+            status_code = None
+            if isinstance(e, requests.HTTPError):
+                status_code = e.response.status_code if e.response else None
+            elif isinstance(e, TaskclusterRestFailure):
+                status_code = e.status_code
+
+            if status_code == 409:
                 # A 409 response indicates that this task is past its deadline.  It
                 # cannot be cancelled at this time, but it's also not running
                 # anymore, so we can ignore this error.

src/taskgraph/actions/util.py

@@ -11,6 +11,7 @@
 from functools import reduce
 
 from requests.exceptions import HTTPError
+from taskcluster.exceptions import TaskclusterRestFailure
 
 from taskgraph import create
 from taskgraph.decision import read_artifact, rename_artifact, write_artifact
@@ -36,7 +37,7 @@ def fetch_graph_and_labels(parameters, graph_config, task_group_id=None):
     try:
         # Look up the decision_task id in the index
         decision_task_id = find_decision_task(parameters, graph_config)
-    except KeyError:
+    except (KeyError, TaskclusterRestFailure):
         if not task_group_id:
             raise
         # Not found (e.g. from github-pull-request), fall back to the task group id.

src/taskgraph/optimize/strategies.py

@@ -1,6 +1,8 @@
 import logging
 from datetime import datetime
 
+from taskcluster.exceptions import TaskclusterRestFailure
+
 from taskgraph.optimize.base import OptimizationStrategy, register_strategy
 from taskgraph.util.path import match as match_path
 from taskgraph.util.taskcluster import find_task_id, status_task
@@ -64,7 +66,7 @@ def should_replace_task(self, task, params, deadline, arg):
                     continue
 
                 return task_id
-            except KeyError:
+            except (KeyError, TaskclusterRestFailure):
                 # go on to the next index path
                 pass
 

src/taskgraph/util/taskcluster.py

@@ -476,11 +476,15 @@ def get_ancestors(task_ids: Union[list[str], str]) -> dict[str, str]:
     for task_id in task_ids:
         try:
             task_def = get_task_definition(task_id)
-        except requests.HTTPError as e:
+        except (requests.HTTPError, taskcluster.TaskclusterRestFailure) as e:
             # Task has most likely expired, which means it's no longer a
             # dependency for the purposes of this function.
-            if e.response.status_code == 404:
-                continue
+            if isinstance(e, requests.HTTPError):
+                if e.response and e.response.status_code == 404:
+                    continue
+            elif isinstance(e, taskcluster.TaskclusterRestFailure):
+                if e.status_code == 404:
+                    continue
             raise e
 
         dependencies = task_def.get("dependencies", [])