Merge pull request #939 from Yelp/u/wilmerrafael/COMPINFRA-3601_addin…

…g_tron_run_number_label

Adding tron run id to pod labels for k8s

tests/bin/action_runner_test.py

@@ -146,3 +146,44 @@ def test_build_environment_too_long_run_id(self):
             TRON_RUN_NUM="10",
             TRON_ACTION="bar.baz",
         )
+
+
+class TestBuildLabels:
+    def test_build_labels(self):
+        labels = action_runner.build_labels("MASTER.foo.10.bar")
+
+        assert labels == {
+            "tron.yelp.com/run_num": "10",
+        }
+
+    def test_build_labels_with_merging(self):
+        current_labels = {"LABEL1": "value_1"}
+        labels = action_runner.build_labels("MASTER.foo.10.bar", current_labels)
+
+        assert labels == {
+            "tron.yelp.com/run_num": "10",
+            "LABEL1": "value_1",
+        }
+
+    def test_build_labels_with_merging_on_unknown(self):
+        current_labels = {"LABEL1": "value_1"}
+        labels = action_runner.build_labels("asdf", current_labels)
+
+        assert labels == {
+            "tron.yelp.com/run_num": "UNKNOWN",
+            "LABEL1": "value_1",
+        }
+
+    def test_build_labels_invalid_run_id(self):
+        labels = action_runner.build_labels("asdf")
+
+        assert labels == {
+            "tron.yelp.com/run_num": "UNKNOWN",
+        }
+
+    def test_build_labels_too_long_run_id(self):
+        labels = action_runner.build_labels("MASTER.foo.10.bar.baz")
+
+        assert labels == {
+            "tron.yelp.com/run_num": "10",
+        }

tests/core/actionrun_test.py

@@ -1802,6 +1802,9 @@ def mock_k8s_action_run(self):
                 "TRON_RUN_NUM": "42",
                 "TRON_ACTION": "mock_action_name",
             },
+            labels={
+                "tron.yelp.com/run_num": "42",
+            },
         )
 
         return KubernetesActionRun(

tron/bin/action_runner.py

@@ -10,6 +10,8 @@
 import sys
 import threading
 import time
+from typing import Dict
+from typing import Optional
 
 from tron import yaml
 
@@ -93,6 +95,24 @@ def build_environment(run_id, original_env=None):
     return new_env
 
 
+def build_labels(run_id: str, original_labels: Optional[Dict[str, str]] = None) -> Dict[str, str]:
+    if original_labels is None:
+        original_labels = dict()
+
+    try:
+        # reminder: the format here is "namespace.job.run_num.action"
+        _, _, run_num, _ = run_id.split(".", maxsplit=3)
+    except ValueError:
+        # if we can't parse the run_id, we don't want to abort, so just
+        # set these semi-arbitrarily
+        run_num = "UNKNOWN"
+
+    new_labels = dict(original_labels)
+    new_labels["tron.yelp.com/run_num"] = run_num
+
+    return new_labels
+
+
 def run_proc(output_path, command, run_id, proc):
     logging.warning(f"{run_id} running as pid {proc.pid}")
     status_file = StatusFile(os.path.join(output_path, STATUS_FILE))

tron/core/actionrun.py

@@ -20,6 +20,7 @@
 from tron.actioncommand import NoActionRunnerFactory
 from tron.actioncommand import SubprocessActionRunnerFactory
 from tron.bin.action_runner import build_environment  # type: ignore # mypy can't find library stub
+from tron.bin.action_runner import build_labels
 from tron.config.config_utils import StringFormatter
 from tron.config.schema import ExecutorTypes
 from tron.core import action
@@ -1170,7 +1171,7 @@ def submit_command(self, attempt: ActionRunAttempt) -> Optional[KubernetesTask]:
                 cap_drop=attempt.command_config.cap_drop,
                 node_selectors=attempt.command_config.node_selectors,
                 node_affinities=attempt.command_config.node_affinities,
-                pod_labels=attempt.command_config.labels,
+                pod_labels=build_labels(run_id=self.id, original_labels=attempt.command_config.labels),
                 pod_annotations=attempt.command_config.annotations,
                 service_account_name=attempt.command_config.service_account_name,
                 ports=attempt.command_config.ports,
@@ -1244,7 +1245,7 @@ def recover(self) -> Optional[KubernetesTask]:
             task_id=last_attempt.kubernetes_task_id,
             node_selectors=last_attempt.command_config.node_selectors,
             node_affinities=last_attempt.command_config.node_affinities,
-            pod_labels=last_attempt.command_config.labels,
+            pod_labels=build_labels(run_id=self.id, original_labels=last_attempt.command_config.labels),
             pod_annotations=last_attempt.command_config.annotations,
             service_account_name=last_attempt.command_config.service_account_name,
             ports=last_attempt.command_config.ports,