Merge pull request #3042 from jameslamb/fix/duplicate-labels

prevent assignment of duplicate labels to agents
PrefectHQ · Jul 28, 2020 · 8e4db16 · 8e4db16
2 parents 38a983f + efaef5b
commit 8e4db16
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 15 deletions.
diff --git a/changes/pr3042.yaml b/changes/pr3042.yaml
@@ -0,0 +1,5 @@
+fix:
+  - 'Prevent duplicated agent labels - [#3029](https://github.com/PrefectHQ/prefect/pull/3042)'
+
+contributor:
+  - "[James Lamb](https://github.com/jameslamb)"
diff --git a/src/prefect/agent/local/agent.py b/src/prefect/agent/local/agent.py
@@ -85,14 +85,16 @@ def __init__(
         if hostname_label and (hostname not in self.labels):
             assert isinstance(self.labels, list)
             self.labels.append(hostname)
-        self.labels.extend(
-            [
-                "azure-flow-storage",
-                "gcs-flow-storage",
-                "s3-flow-storage",
-                "github-flow-storage",
-            ]
-        )
+
+        all_storage_labels = [
+            "azure-flow-storage",
+            "gcs-flow-storage",
+            "s3-flow-storage",
+            "github-flow-storage",
+        ]
+        for label in all_storage_labels:
+            if label not in self.labels:
+                self.labels.append(label)
 
         self.logger.debug(f"Import paths: {self.import_paths}")
         self.logger.debug(f"Show flow logs: {self.show_flow_logs}")

diff --git a/src/prefect/cli/agent.py b/src/prefect/cli/agent.py
@@ -246,10 +246,12 @@ def start(
             k, v = env_var.split("=")
             env_vars[k] = v
 
+        labels = list(set(label))
+
         if agent_option == "local":
             from_qualified_name(retrieved_agent)(
                 name=name,
-                labels=list(label),
+                labels=labels,
                 env_vars=env_vars,
                 max_polls=max_polls,
                 agent_address=agent_address,
@@ -260,7 +262,7 @@ def start(
         elif agent_option == "docker":
             from_qualified_name(retrieved_agent)(
                 name=name,
-                labels=list(label),
+                labels=labels,
                 env_vars=env_vars,
                 max_polls=max_polls,
                 agent_address=agent_address,
@@ -274,7 +276,7 @@ def start(
         elif agent_option == "fargate":
             from_qualified_name(retrieved_agent)(
                 name=name,
-                labels=list(label),
+                labels=labels,
                 env_vars=env_vars,
                 max_polls=max_polls,
                 agent_address=agent_address,
@@ -284,15 +286,15 @@ def start(
             from_qualified_name(retrieved_agent)(
                 namespace=namespace,
                 name=name,
-                labels=list(label),
+                labels=labels,
                 env_vars=env_vars,
                 max_polls=max_polls,
                 agent_address=agent_address,
             ).start()
         else:
             from_qualified_name(retrieved_agent)(
                 name=name,
-                labels=list(label),
+                labels=labels,
                 env_vars=env_vars,
                 max_polls=max_polls,
                 agent_address=agent_address,
@@ -476,6 +478,7 @@ def install(
         k, v = env_var.split("=")
         env_vars[k] = v
 
+    labels = list(set(label))
     if name == "kubernetes":
         deployment = from_qualified_name(retrieved_agent).generate_deployment_yaml(
             token=token,
@@ -491,15 +494,15 @@ def install(
             cpu_limit=cpu_limit,
             image_pull_policy=image_pull_policy,
             service_account_name=service_account_name,
-            labels=list(label),
+            labels=labels,
             env_vars=env_vars,
             backend=backend,
         )
         click.echo(deployment)
     elif name == "local":
         conf = from_qualified_name(retrieved_agent).generate_supervisor_conf(
             token=token,
-            labels=list(label),
+            labels=labels,
             import_paths=list(import_path),
             show_flow_logs=show_flow_logs,
         )

diff --git a/tests/agent/test_local_agent.py b/tests/agent/test_local_agent.py
@@ -25,6 +25,19 @@ def test_local_agent_init(runner_token):
     assert agent.name == "agent"
 
 
+def test_local_agent_deduplicates_labels(runner_token):
+    agent = LocalAgent(labels=["azure-flow-storage"])
+    assert agent
+    assert set(agent.labels) == {
+        socket.gethostname(),
+        "azure-flow-storage",
+        "s3-flow-storage",
+        "gcs-flow-storage",
+        "github-flow-storage",
+    }
+    assert len(agent.labels) == len(set(agent.labels))
+
+
 def test_local_agent_config_options(runner_token):
     with set_temporary_config(
         {"cloud.agent.auth_token": "TEST_TOKEN", "logging.log_to_cloud": True}