From 0bae9d3de507826dedf1a045f0b15a89e8f800c4 Mon Sep 17 00:00:00 2001
From: Taylor McKay <taylorm@zillowgroup.com>
Date: Thu, 4 Jan 2024 11:12:46 -0800
Subject: [PATCH] Move product label to _get_flow_labels

---
 metaflow/plugins/aip/aip.py                         | 13 +++++++------
 metaflow/plugins/aip/argo_client.py                 |  3 ---
 metaflow/plugins/aip/tests/run_integration_tests.py |  1 +
 3 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/metaflow/plugins/aip/aip.py b/metaflow/plugins/aip/aip.py
index dce4516daa5..cb7739a1a81 100644
--- a/metaflow/plugins/aip/aip.py
+++ b/metaflow/plugins/aip/aip.py
@@ -1041,21 +1041,25 @@ def _get_flow_labels(
 
             ret_flow_labels[annotation_name] = annotation_value
 
+        zodiac_prefix = "zodiac.zillowgroup.net"
+
         # - In context of Zillow CICD username == "cicd_compile"
         # - In the context of a Zillow NB username == METAFLOW_USER (user_alias)
         # - In the context of Metaflow integration tests username == USER=$GITLAB_USER_EMAIL
         owner = username
         if "@" in owner:
             owner = owner.split("@")[0]
-        ret_flow_labels["zodiac.zillowgroup.net/owner"] = owner
+        ret_flow_labels[f"{zodiac_prefix}/owner"] = owner
 
         # If the Zodiac environment variable is present in the notebook (individual profile notebooks only),
         # the Zodiac service and team labels are added to the AIP pods and set. These labels are not added
         # by the AIP webhook to support user-supplied Zodiac service per AIP Notebook. Workflows launched
         # in project CICD profiles will still have these labels added via the AIP webhook.
         if ZILLOW_ZODIAC_SERVICE and ZILLOW_ZODIAC_TEAM:
-            ret_flow_labels["zodiac.zillowgroup.net/service"] = ZILLOW_ZODIAC_SERVICE
-            ret_flow_labels["zodiac.zillowgroup.net/team"] = ZILLOW_ZODIAC_TEAM
+            ret_flow_labels[f"{zodiac_prefix}/service"] = ZILLOW_ZODIAC_SERVICE
+            ret_flow_labels[f"{zodiac_prefix}/team"] = ZILLOW_ZODIAC_TEAM
+
+        ret_flow_labels[f"{zodiac_prefix}/product"] = "batch"
 
         return ret_flow_labels
 
@@ -1082,9 +1086,6 @@ def _set_container_labels(self, container_op: ContainerOp):
         container_op.add_pod_label(
             "tags.ledger.zgtools.net/ai-step-name", container_op.name
         )
-        container_op.add_pod_label(
-            "zodiac.zillowgroup.net/product", "batch"
-        )
         if self.experiment:
             container_op.add_pod_label(
                 "tags.ledger.zgtools.net/ai-experiment-name", self.experiment
diff --git a/metaflow/plugins/aip/argo_client.py b/metaflow/plugins/aip/argo_client.py
index b12e702bb8b..076a6bec3d4 100644
--- a/metaflow/plugins/aip/argo_client.py
+++ b/metaflow/plugins/aip/argo_client.py
@@ -90,9 +90,6 @@ def trigger_workflow_template(self, name: str, parameters: Optional[Dict] = None
             "kind": "Workflow",
             "metadata": {
                 "generateName": name + "-",
-                "labels": {
-                    "zodiac.zillowgroup.net/product": "batch"
-                }
             },
             "spec": {
                 "workflowTemplateRef": {"name": name},
diff --git a/metaflow/plugins/aip/tests/run_integration_tests.py b/metaflow/plugins/aip/tests/run_integration_tests.py
index d7d6346174b..d53e54e6839 100644
--- a/metaflow/plugins/aip/tests/run_integration_tests.py
+++ b/metaflow/plugins/aip/tests/run_integration_tests.py
@@ -495,6 +495,7 @@ def test_flow_labels():
     assert tags["metaflow.org/tag_t1"] == "true"
     assert tags["metaflow.org/tag_s1"] == "true"
     assert tags["zodiac.zillowgroup.net/owner"] == "foo"
+    assert tags["zodiac.zillowgroup.net/product"] == "batch"
 
     with pytest.raises(ValueError) as e:
         KubeflowPipelines._get_flow_labels(