From 7a06b76cce17485182cc4792ff57e53a27993490 Mon Sep 17 00:00:00 2001
From: Gabriel Bustamante <gbustamante@mozilla.com>
Date: Wed, 19 Jun 2024 11:37:45 -0500
Subject: [PATCH] Override the existing_tasks explicitly provided in the
 action's input (#683)

* Override the existing_tasks explicitly provided in the action's input

* Log the labels and task ids of any existing tasks that are overridden

---------

Co-authored-by: Evgeny Pavlov <epavlov@mozilla.com>
---
 .../translations_taskgraph/actions/train.py   | 33 +++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/taskcluster/translations_taskgraph/actions/train.py b/taskcluster/translations_taskgraph/actions/train.py
index 28b801bc4..8779c8309 100644
--- a/taskcluster/translations_taskgraph/actions/train.py
+++ b/taskcluster/translations_taskgraph/actions/train.py
@@ -1,6 +1,8 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
+import json
+import logging
 
 from taskgraph.actions.registry import register_callback_action
 from taskgraph.decision import taskgraph_decision
@@ -10,6 +12,8 @@
 
 from translations_taskgraph.parameters import get_defaults
 
+logger = logging.getLogger(__name__)
+
 TRAIN_ON_PROJECTS = (
     "https://github.com/mozilla/firefox-translations-training",
     "https://github.com/mozilla-releng/staging-firefox-translations-training",
@@ -398,6 +402,35 @@ def train_action(parameters, graph_config, input, task_group_id, task_id):
         # an identical name.
         parameters["existing_tasks"] = get_ancestors(start_task_ids)
 
+    # Override the `existing_tasks` explicitly provided in the action's input
+    existing_tasks = input.pop("existing_tasks", {})
+
+    # Find and log `overridden_existing_tasks`
+    overridden_existing_tasks = {
+        existing_task: parameters["existing_tasks"][existing_task]
+        for existing_task in existing_tasks.keys()
+        if existing_task in parameters["existing_tasks"]
+    }
+
+    if overridden_existing_tasks:
+        logger.info(
+            f"Old values for `overridden_existing_tasks`: {json.dumps(overridden_existing_tasks, indent=2)}"
+        )
+
+    # Do the override!
+    parameters["existing_tasks"].update(existing_tasks)
+
+    # Log the new values for the `overridden_existing_tasks`
+    new_values_for_overridden = {
+        existing_task: parameters["existing_tasks"][existing_task]
+        for existing_task in overridden_existing_tasks.keys()
+    }
+
+    if new_values_for_overridden:
+        logger.info(
+            f"New values for `overridden_existing_tasks`: {json.dumps(new_values_for_overridden, indent=2)}"
+        )
+
     parameters["target_tasks_method"] = "train-target-tasks"
     parameters["optimize_target_tasks"] = True
     parameters["tasks_for"] = "action"