From 710652dd5b7567f22445ae389704f702452ffe34 Mon Sep 17 00:00:00 2001
From: Mathieu Larose <mathieu@retain.ai>
Date: Tue, 17 Sep 2024 16:11:31 -0400
Subject: [PATCH] feat: future proofing the tag used to trigger run metrics
 collection

---
 python_modules/dagster/dagster/_cli/api.py           | 11 ++++++++---
 python_modules/dagster/dagster/_core/storage/tags.py |  3 +++
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/python_modules/dagster/dagster/_cli/api.py b/python_modules/dagster/dagster/_cli/api.py
index 82be6a3205652..a3edb09ab0980 100644
--- a/python_modules/dagster/dagster/_cli/api.py
+++ b/python_modules/dagster/dagster/_cli/api.py
@@ -36,6 +36,11 @@
     get_python_environment_entry_point,
 )
 from dagster._core.storage.dagster_run import DagsterRun, DagsterRunStatus
+from dagster._core.storage.tags import (
+    RUN_METRIC_TAGS,
+    RUN_METRICS_POLLING_INTERVAL_TAG,
+    RUN_METRICS_PYTHON_RUNTIME_TAG,
+)
 from dagster._core.types.loadable_target_origin import LoadableTargetOrigin
 from dagster._core.utils import FuturesAwareThreadPoolExecutor
 from dagster._grpc import DagsterGrpcClient, DagsterGrpcServer
@@ -87,11 +92,11 @@ def send_to_buffer(event):
 
 
 def _should_start_metrics_thread(dagster_run: DagsterRun) -> bool:
-    return get_boolean_tag_value(dagster_run.tags.get("dagster/run_metrics"))
+    return any(get_boolean_tag_value(dagster_run.tags.get(tag)) for tag in RUN_METRIC_TAGS)
 
 
 def _enable_python_runtime_metrics(dagster_run: DagsterRun) -> bool:
-    return get_boolean_tag_value(dagster_run.tags.get("dagster/python_runtime_metrics"))
+    return get_boolean_tag_value(dagster_run.tags.get(RUN_METRICS_PYTHON_RUNTIME_TAG))
 
 
 def _metrics_polling_interval(
@@ -100,7 +105,7 @@ def _metrics_polling_interval(
     try:
         return float(
             dagster_run.tags.get(
-                "dagster/run_metrics_polling_interval_seconds",
+                RUN_METRICS_POLLING_INTERVAL_TAG,
                 DEFAULT_RUN_METRICS_POLL_INTERVAL_SECONDS,
             )
         )
diff --git a/python_modules/dagster/dagster/_core/storage/tags.py b/python_modules/dagster/dagster/_core/storage/tags.py
index 7dc7031479b67..2f767ad0d2987 100644
--- a/python_modules/dagster/dagster/_core/storage/tags.py
+++ b/python_modules/dagster/dagster/_core/storage/tags.py
@@ -96,6 +96,9 @@
     f"{SYSTEM_TAG_PREFIX}run_metrics",
 ]
 
+RUN_METRICS_POLLING_INTERVAL_TAG = f"{HIDDEN_TAG_PREFIX}run_metrics_polling_interval"
+RUN_METRICS_PYTHON_RUNTIME_TAG = f"{HIDDEN_TAG_PREFIX}python_runtime_metrics"
+
 
 class TagType(Enum):
     # Custom tag provided by a user