Repository: PrefectHQ/prefect

Issue:
Prepare Agents for "registering" themselves with Cloud
Whenever a user spins up an agent against the `"cloud"` backend, we will soon be registering the agent as a first step.  At registration time:

- we will call an agent registration route (name TBD)
- send a payload of information about this agent (type, name, labels, etc.)
- receives back an ID which is stored on the agent instance
- sends this ID within a header for every future query

We should hook up the lifecycle of this into our base agent class so we're ready once the backend is ready.


Discussion:

Resolving patch:
diff --git a/src/prefect/agent/agent.py b/src/prefect/agent/agent.py
--- a/src/prefect/agent/agent.py
+++ b/src/prefect/agent/agent.py
@@ -94,6 +94,7 @@ def __init__(
         self.client = Client(api_token=token)
         if config.backend == "cloud":
             self._verify_token(token)
+            self.client.attach_headers({"X-PREFECT-AGENT-ID": self._register_agent()})
 
         logger = logging.getLogger(self.name)
         logger.setLevel(config.cloud.agent.get("level"))
@@ -126,6 +127,18 @@ def _verify_token(self, token: str) -> None:
         ):
             raise AuthorizationError("Provided token does not have a RUNNER scope.")
 
+    def _register_agent(self) -> str:
+        """
+        Register this agent with Prefect Cloud and retrieve agent ID
+
+        Returns:
+            - The agent ID as a string
+        """
+        agent_id = self.client.register_agent(
+            agent_type=type(self).__name__, name=self.name, labels=self.labels
+        )
+        return agent_id
+
     def start(self) -> None:
         """
         The main entrypoint to the agent. This function loops and constantly polls for
diff --git a/src/prefect/client/client.py b/src/prefect/client/client.py
--- a/src/prefect/client/client.py
+++ b/src/prefect/client/client.py
@@ -81,6 +81,7 @@ def __init__(self, api_server: str = None, api_token: str = None):
         self._refresh_token = None
         self._access_token_expires_at = pendulum.now()
         self._active_tenant_id = None
+        self._attached_headers = {}  # type: Dict[str, str]
         self.logger = create_diagnostic_logger("Diagnostics")
 
         # store api server
@@ -275,6 +276,9 @@ def _request(
             headers["Authorization"] = "Bearer {}".format(token)
         headers["X-PREFECT-CORE-VERSION"] = str(prefect.__version__)
 
+        if self._attached_headers:
+            headers.update(self._attached_headers)
+
         session = requests.Session()
         retries = requests.packages.urllib3.util.retry.Retry(
             total=6,
@@ -315,6 +319,16 @@ def _request(
 
         return response
 
+    def attach_headers(self, headers: dict) -> None:
+        """
+        Set headers to be attached to this Client
+
+        Args:
+            - headers (dict): A dictionary of headers to attach to this client. These headers
+                get added on to the existing dictionary of headers.
+        """
+        self._attached_headers.update(headers)
+
     # -------------------------------------------------------------------------
     # Auth
     # -------------------------------------------------------------------------
@@ -1229,3 +1243,34 @@ def write_run_logs(self, logs: List[Dict]) -> None:
 
         if not result.data.write_run_logs.success:
             raise ValueError("Writing logs failed.")
+
+    def register_agent(
+        self, agent_type: str, name: str = None, labels: List[str] = None
+    ) -> str:
+        """
+        Register an agent with Cloud
+
+        Args:
+            - agent_type (str): The type of agent being registered
+            - name: (str, optional): The name of the agent being registered
+            - labels (List[str], optional): A list of any present labels on the agent
+                being registered
+
+        Returns:
+            - The agent ID as a string
+        """
+        mutation = {
+            "mutation($input: register_agent_input!)": {
+                "register_agent(input: $input)": {"id"}
+            }
+        }
+
+        result = self.graphql(
+            mutation,
+            variables=dict(input=dict(type=agent_type, name=name, labels=labels)),
+        )
+
+        if not result.data.register_agent.id:
+            raise ValueError("Error registering agent")
+
+        return result.data.register_agent.id


Test patch:
diff --git a/tests/agent/test_agent.py b/tests/agent/test_agent.py
--- a/tests/agent/test_agent.py
+++ b/tests/agent/test_agent.py
@@ -10,17 +10,17 @@
 from prefect.utilities.graphql import GraphQLResult
 
 
-def test_agent_init(runner_token):
+def test_agent_init(runner_token, cloud_api):
     agent = Agent()
     assert agent
 
 
-def test_multiple_agent_init_doesnt_duplicate_logs(runner_token):
+def test_multiple_agent_init_doesnt_duplicate_logs(runner_token, cloud_api):
     a, b, c = Agent(), Agent(), Agent()
     assert len(c.logger.handlers) == 1
 
 
-def test_agent_config_options(runner_token):
+def test_agent_config_options(runner_token, cloud_api):
     with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
         agent = Agent()
         assert agent.labels == []
@@ -32,7 +32,7 @@ def test_agent_config_options(runner_token):
         assert agent.logger.name == "agent"
 
 
-def test_agent_name_set_options(monkeypatch, runner_token):
+def test_agent_name_set_options(monkeypatch, runner_token, cloud_api):
     # Default
     agent = Agent()
     assert agent.name == "agent"
@@ -50,13 +50,13 @@ def test_agent_name_set_options(monkeypatch, runner_token):
         assert agent.logger.name == "test2"
 
 
-def test_agent_log_level(runner_token):
+def test_agent_log_level(runner_token, cloud_api):
     with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
         agent = Agent()
         assert agent.logger.level == 20
 
 
-def test_agent_log_level_responds_to_config(runner_token):
+def test_agent_log_level_responds_to_config(runner_token, cloud_api):
     with set_temporary_config(
         {"cloud.agent.auth_token": "TEST_TOKEN", "cloud.agent.level": "DEBUG"}
     ):
@@ -64,19 +64,19 @@ def test_agent_log_level_responds_to_config(runner_token):
         assert agent.logger.level == 10
 
 
-def test_agent_env_vars(runner_token):
+def test_agent_env_vars(runner_token, cloud_api):
     with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
         agent = Agent(env_vars=dict(AUTH_THING="foo"))
         assert agent.env_vars == dict(AUTH_THING="foo")
 
 
-def test_agent_max_polls(runner_token):
+def test_agent_max_polls(runner_token, cloud_api):
     with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
         agent = Agent(max_polls=10)
         assert agent.max_polls == 10
 
 
-def test_agent_labels(runner_token):
+def test_agent_labels(runner_token, cloud_api):
     with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
         agent = Agent(labels=["test", "2"])
         assert agent.labels == ["test", "2"]
@@ -119,7 +119,7 @@ def test_agent_fails_no_runner_token(monkeypatch, cloud_api):
         agent = Agent()
 
 
-def test_query_flow_runs(monkeypatch, runner_token):
+def test_query_flow_runs(monkeypatch, runner_token, cloud_api):
     gql_return = MagicMock(
         return_value=MagicMock(
             data=MagicMock(
@@ -137,7 +137,9 @@ def test_query_flow_runs(monkeypatch, runner_token):
     assert flow_runs == [{"id": "id"}]
 
 
-def test_query_flow_runs_ignores_currently_submitting_runs(monkeypatch, runner_token):
+def test_query_flow_runs_ignores_currently_submitting_runs(
+    monkeypatch, runner_token, cloud_api
+):
     gql_return = MagicMock(
         return_value=MagicMock(
             data=MagicMock(
@@ -162,7 +164,7 @@ def test_query_flow_runs_ignores_currently_submitting_runs(monkeypatch, runner_t
 
 
 def test_query_flow_runs_does_not_use_submitting_flow_runs_directly(
-    monkeypatch, runner_token, caplog
+    monkeypatch, runner_token, caplog, cloud_api
 ):
     gql_return = MagicMock(
         return_value=MagicMock(
@@ -188,7 +190,7 @@ def test_query_flow_runs_does_not_use_submitting_flow_runs_directly(
     copy_mock.assert_called_once_with()
 
 
-def test_update_states_passes_no_task_runs(monkeypatch, runner_token):
+def test_update_states_passes_no_task_runs(monkeypatch, runner_token, cloud_api):
     gql_return = MagicMock(
         return_value=MagicMock(
             data=MagicMock(set_flow_run_state=None, set_task_run_state=None)
@@ -211,7 +213,7 @@ def test_update_states_passes_no_task_runs(monkeypatch, runner_token):
     )
 
 
-def test_update_states_passes_task_runs(monkeypatch, runner_token):
+def test_update_states_passes_task_runs(monkeypatch, runner_token, cloud_api):
     gql_return = MagicMock(
         return_value=MagicMock(
             data=MagicMock(set_flow_run_state=None, set_task_run_state=None)
@@ -242,7 +244,7 @@ def test_update_states_passes_task_runs(monkeypatch, runner_token):
     )
 
 
-def test_mark_failed(monkeypatch, runner_token):
+def test_mark_failed(monkeypatch, runner_token, cloud_api):
     gql_return = MagicMock(
         return_value=MagicMock(
             data=MagicMock(set_flow_run_state=None, set_task_run_state=None)
@@ -266,18 +268,18 @@ def test_mark_failed(monkeypatch, runner_token):
     )
 
 
-def test_deploy_flows_passes_base_agent(runner_token):
+def test_deploy_flows_passes_base_agent(runner_token, cloud_api):
     agent = Agent()
     with pytest.raises(NotImplementedError):
         agent.deploy_flow(None)
 
 
-def test_heartbeat_passes_base_agent(runner_token):
+def test_heartbeat_passes_base_agent(runner_token, cloud_api):
     agent = Agent()
     assert not agent.heartbeat()
 
 
-def test_agent_connect(monkeypatch, runner_token):
+def test_agent_connect(monkeypatch, runner_token, cloud_api):
     post = MagicMock(return_value=MagicMock(json=MagicMock(return_value="hello")))
     session = MagicMock()
     session.return_value.post = post
@@ -287,7 +289,7 @@ def test_agent_connect(monkeypatch, runner_token):
     assert agent.agent_connect() is None
 
 
-def test_agent_connect_handled_error(monkeypatch, runner_token):
+def test_agent_connect_handled_error(monkeypatch, runner_token, cloud_api):
     post = MagicMock(side_effect=Exception)
     session = MagicMock()
     session.return_value.post = post
@@ -297,14 +299,14 @@ def test_agent_connect_handled_error(monkeypatch, runner_token):
     assert agent.agent_connect() is None
 
 
-def test_on_flow_run_deploy_attempt_removes_id(monkeypatch, runner_token):
+def test_on_flow_run_deploy_attempt_removes_id(monkeypatch, runner_token, cloud_api):
     agent = Agent()
     agent.submitting_flow_runs.add("id")
     agent.on_flow_run_deploy_attempt(None, "id")
     assert len(agent.submitting_flow_runs) == 0
 
 
-def test_agent_process(monkeypatch, runner_token):
+def test_agent_process(monkeypatch, runner_token, cloud_api):
     gql_return = MagicMock(
         return_value=MagicMock(
             data=MagicMock(
@@ -346,7 +348,7 @@ def test_agent_process(monkeypatch, runner_token):
     assert future_mock.add_done_callback.called
 
 
-def test_agent_process_no_runs_found(monkeypatch, runner_token):
+def test_agent_process_no_runs_found(monkeypatch, runner_token, cloud_api):
     gql_return = MagicMock(
         return_value=MagicMock(
             data=MagicMock(
@@ -368,7 +370,7 @@ def test_agent_process_no_runs_found(monkeypatch, runner_token):
     assert not executor.submit.called
 
 
-def test_agent_logs_flow_run_exceptions(monkeypatch, runner_token, caplog):
+def test_agent_logs_flow_run_exceptions(monkeypatch, runner_token, caplog, cloud_api):
     gql_return = MagicMock(
         return_value=MagicMock(data=MagicMock(write_run_logs=MagicMock(success=True)))
     )
@@ -404,7 +406,7 @@ def test_agent_logs_flow_run_exceptions(monkeypatch, runner_token, caplog):
     assert "Logging platform error for flow run" in caplog.text
 
 
-def test_agent_process_raises_exception_and_logs(monkeypatch, runner_token):
+def test_agent_process_raises_exception_and_logs(monkeypatch, runner_token, cloud_api):
     client = MagicMock()
     client.return_value.graphql.side_effect = ValueError("Error")
     monkeypatch.setattr("prefect.agent.agent.Client", client)
@@ -417,7 +419,7 @@ def test_agent_process_raises_exception_and_logs(monkeypatch, runner_token):
         assert client.write_run_log.called
 
 
-def test_agent_start_max_polls(monkeypatch, runner_token):
+def test_agent_start_max_polls(monkeypatch, runner_token, cloud_api):
     on_shutdown = MagicMock()
     monkeypatch.setattr("prefect.agent.agent.Agent.on_shutdown", on_shutdown)
 
@@ -438,7 +440,7 @@ def test_agent_start_max_polls(monkeypatch, runner_token):
     assert heartbeat.called
 
 
-def test_agent_start_max_polls_count(monkeypatch, runner_token):
+def test_agent_start_max_polls_count(monkeypatch, runner_token, cloud_api):
     on_shutdown = MagicMock()
     monkeypatch.setattr("prefect.agent.agent.Agent.on_shutdown", on_shutdown)
 
@@ -459,7 +461,7 @@ def test_agent_start_max_polls_count(monkeypatch, runner_token):
     assert heartbeat.call_count == 2
 
 
-def test_agent_start_max_polls_zero(monkeypatch, runner_token):
+def test_agent_start_max_polls_zero(monkeypatch, runner_token, cloud_api):
     on_shutdown = MagicMock()
     monkeypatch.setattr("prefect.agent.agent.Agent.on_shutdown", on_shutdown)
 
@@ -478,3 +480,14 @@ def test_agent_start_max_polls_zero(monkeypatch, runner_token):
     assert on_shutdown.call_count == 1
     assert agent_process.call_count == 0
     assert heartbeat.call_count == 0
+
+
+def test_agent_registration_and_id(monkeypatch, cloud_api):
+    monkeypatch.setattr("prefect.agent.agent.Agent._verify_token", MagicMock())
+    monkeypatch.setattr(
+        "prefect.agent.agent.Client.register_agent", MagicMock(return_value="ID")
+    )
+
+    agent = Agent()
+    assert agent._register_agent() == "ID"
+    assert agent.client._attached_headers == {"X-PREFECT-AGENT-ID": "ID"}
diff --git a/tests/client/test_client.py b/tests/client/test_client.py
--- a/tests/client/test_client.py
+++ b/tests/client/test_client.py
@@ -62,6 +62,22 @@ def test_version_header_cant_be_overridden(monkeypatch):
     )
 
 
+def test_client_attached_headers(monkeypatch, cloud_api):
+    get = MagicMock()
+    session = MagicMock()
+    session.return_value.get = get
+    monkeypatch.setattr("requests.Session", session)
+    with set_temporary_config({"cloud.auth_token": "secret_token"}):
+        client = Client()
+        assert client._attached_headers == {}
+
+        client.attach_headers({"1": "1"})
+        assert client._attached_headers == {"1": "1"}
+
+        client.attach_headers({"2": "2"})
+        assert client._attached_headers == {"1": "1", "2": "2"}
+
+
 def test_client_posts_graphql_to_api_server(patch_post):
     post = patch_post(dict(data=dict(success=True)))
 
@@ -931,3 +947,31 @@ def test_get_cloud_url_different_regex(patch_post, cloud_api):
 
         url = client.get_cloud_url(subdirectory="flow-run", id="id2")
         assert url == "http://hello.prefect.io/tslug/flow-run/id2"
+
+
+def test_register_agent(patch_post, cloud_api):
+    response = {"data": {"register_agent": {"id": "ID"}}}
+
+    patch_post(response)
+
+    with set_temporary_config({"cloud.auth_token": "secret_token"}):
+        client = Client()
+
+        agent_id = client.register_agent(
+            agent_type="type", name="name", labels=["1", "2"]
+        )
+        assert agent_id == "ID"
+
+
+def test_register_agent_raises_error(patch_post, cloud_api):
+    response = {"data": {"register_agent": {"id": None}}}
+
+    patch_post(response)
+
+    with set_temporary_config({"cloud.auth_token": "secret_token"}):
+        client = Client()
+
+        with pytest.raises(ValueError):
+            agent_id = client.register_agent(
+                agent_type="type", name="name", labels=["1", "2"]
+            )
diff --git a/tests/conftest.py b/tests/conftest.py
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -145,6 +145,7 @@ def patch(responses):
 @pytest.fixture()
 def runner_token(monkeypatch):
     monkeypatch.setattr("prefect.agent.agent.Agent._verify_token", MagicMock())
+    monkeypatch.setattr("prefect.agent.agent.Agent._register_agent", MagicMock())
 
 
 @pytest.fixture()