Move product label to _get_flow_labels

zillow · Jan 4, 2024 · 0bae9d3 · 0bae9d3
1 parent 1c63366
commit 0bae9d3
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 9 deletions.
diff --git 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
@@ -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
@@ -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(