We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Handle cases where Prefect is deployed to multiple environments with differently-calibrated clocks
It's possible that if a Prefect flow is deployed in an environment with an incorrect clock, then it may not run (if it receives states from an external system like Prefect Cloud). 

For example, say the clock in a cluster is running 5 minutes slow, and a flow is supposed to start at 9:00am UTC. At 9:00am UTC, the agent in that cluster will receive a payload including a `Scheduled` state with a start time of `9:00:00 UTC`. Note that in the cluster, however, the time is 8:55am UTC. The agent will take that `Scheduled` state and submit it for execution, but the `Runner` classes will refuse to run it because it hasn't reached its start time yet -- and won't for 5 minutes. 

We have two possible remedies for this situation:
- A global `time_adjust_ms` config. This could be set by querying a known/trusted UTC time source and comparing the result to the local time. The difference (in milliseconds) could be used by the runner whenever a scheduled state was examined for execution. If this were exposed as a configuration, it could be set by the agent as an env var whenever it deployed a flow (the agent could do the checking, then set the variable appropriately). The downside of this approach is that it might create bad behavior for any state generated locally. For example, a local retry wouldn't want this adjustment to be applied. If the adjustment were positive (meaning it *adds* time to state start times because the the local clock is running fast), then the local retry wouldn't happen at all because it wouldn't have reached its presumed start time! However, if that's the only case that's problematic, we could handle it by only applying negative adjustments.

- An alternative is for the agent to "hold" any runs that haven't reached their start time yet, rather than deploying them for execution immediately. (Note: this dovetails with an idea @joshmeek and I had to give scheduled runs to the agent slightly ahead of schedule but "hold" them so they started exactly on time, with no poll delay). The downside of this approach is that any run which involved a dynamic task, like a mapped task that gets its state directly from Cloud and consequently bypasses the agent, would fail to run. 

Perhaps a combination of these approaches is possible? Perhaps there's a third?

cc @cicdw @joshmeek @zdhughes  


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/prefect/agent/agent.py`
Content:
```
1 import logging
2 from typing import Union
3 
4 import pendulum
5 import time
6 
7 from prefect import config
8 from prefect.client import Client
9 from prefect.serialization import state
10 from prefect.engine.state import Submitted
11 from prefect.utilities.exceptions import AuthorizationError
12 from prefect.utilities.graphql import with_args
13 
14 
15 ascii_name = r"""
16  ____            __           _        _                    _
17 |  _ \ _ __ ___ / _| ___  ___| |_     / \   __ _  ___ _ __ | |_
18 | |_) | '__/ _ \ |_ / _ \/ __| __|   / _ \ / _` |/ _ \ '_ \| __|
19 |  __/| | |  __/  _|  __/ (__| |_   / ___ \ (_| |  __/ | | | |_
20 |_|   |_|  \___|_|  \___|\___|\__| /_/   \_\__, |\___|_| |_|\__|
21                                            |___/
22 """
23 
24 
25 class Agent:
26     """
27     Base class for Agents.
28 
29     This Agent class is a standard point for executing Flows in Prefect Cloud. It is meant
30     to have subclasses which inherit functionality from this class. The only piece that
31     the subclasses should implement is the `deploy_flows` function, which specifies how to run a Flow on the given platform. It is built in this
32     way to keep Prefect Cloud logic standard but allows for platform specific
33     customizability.
34 
35     In order for this to operate `PREFECT__CLOUD__AGENT__AUTH_TOKEN` must be set as an
36     environment variable or in your user configuration file.
37     """
38 
39     def __init__(self) -> None:
40 
41         token = config.cloud.agent.get("auth_token")
42 
43         self.client = Client(api_token=token)
44         self._verify_token(token)
45 
46         logger = logging.getLogger("agent")
47         logger.setLevel(logging.DEBUG)
48         ch = logging.StreamHandler()
49         ch.setLevel(logging.DEBUG)
50         formatter = logging.Formatter(
51             "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
52         )
53         ch.setFormatter(formatter)
54         logger.addHandler(ch)
55 
56         self.logger = logger
57 
58     def _verify_token(self, token: str) -> None:
59         """
60         Checks whether a token with a `RUNNER` scope was provided
61 
62         Args:
63             - token (str): The provided agent token to verify
64 
65         Raises:
66             - AuthorizationError: if token is empty or does not have a RUNNER role
67         """
68         if not token:
69             raise AuthorizationError("No agent API token provided.")
70 
71         # Check if RUNNER role
72         result = self.client.graphql(query="query { authInfo { apiTokenScope } }")
73         if (
74             not result.data  # type: ignore
75             or result.data.authInfo.apiTokenScope != "RUNNER"  # type: ignore
76         ):
77             raise AuthorizationError("Provided token does not have a RUNNER scope.")
78 
79     def start(self) -> None:
80         """
81         The main entrypoint to the agent. This function loops and constantly polls for
82         new flow runs to deploy
83         """
84         tenant_id = self.agent_connect()
85 
86         # Loop intervals for query sleep backoff
87         loop_intervals = {0: 0.25, 1: 0.5, 2: 1.0, 3: 2.0, 4: 4.0, 5: 8.0, 6: 10.0}
88 
89         index = 0
90         while True:
91             self.heartbeat()
92 
93             runs = self.agent_process(tenant_id)
94             if runs:
95                 index = 0
96             elif index < max(loop_intervals.keys()):
97                 index += 1
98 
99             time.sleep(loop_intervals[index])
100 
101     def agent_connect(self) -> str:
102         """
103         Verify agent connection to Prefect Cloud by finding and returning a tenant id
104 
105         Returns:
106             - str: The current tenant id
107         """
108         print(ascii_name)
109         self.logger.info("Starting {}".format(type(self).__name__))
110         self.logger.info(
111             "Agent documentation can be found at https://docs.prefect.io/cloud/"
112         )
113         tenant_id = self.query_tenant_id()
114 
115         if not tenant_id:
116             raise ConnectionError(
117                 "Tenant ID not found. Verify that you are using the proper API token."
118             )
119 
120         self.logger.info("Agent successfully connected to Prefect Cloud")
121         self.logger.info("Waiting for flow runs...")
122 
123         return tenant_id
124 
125     def agent_process(self, tenant_id: str) -> bool:
126         """
127         Full process for finding flow runs, updating states, and deploying.
128 
129         Args:
130             - tenant_id (str): The tenant id to use in the query
131 
132         Returns:
133             - bool: whether or not flow runs were found
134         """
135         flow_runs = None
136         try:
137             flow_runs = self.query_flow_runs(tenant_id=tenant_id)
138 
139             if flow_runs:
140                 self.logger.info(
141                     "Found {} flow run(s) to submit for execution.".format(
142                         len(flow_runs)
143                     )
144                 )
145 
146                 self.update_states(flow_runs)
147                 self.deploy_flows(flow_runs)
148                 self.logger.info(
149                     "Submitted {} flow run(s) for execution.".format(len(flow_runs))
150                 )
151         except Exception as exc:
152             self.logger.error(exc)
153 
154         return bool(flow_runs)
155 
156     def query_tenant_id(self) -> Union[str, None]:
157         """
158         Query Prefect Cloud for the tenant id that corresponds to the agent's auth token
159 
160         Returns:
161             - Union[str, None]: The current tenant id if found, None otherwise
162         """
163         query = {"query": {"tenant": {"id"}}}
164         result = self.client.graphql(query)
165 
166         if result.data.tenant:  # type: ignore
167             return result.data.tenant[0].id  # type: ignore
168 
169         return None
170 
171     def query_flow_runs(self, tenant_id: str) -> list:
172         """
173         Query Prefect Cloud for flow runs which need to be deployed and executed
174 
175         Args:
176             - tenant_id (str): The tenant id to use in the query
177 
178         Returns:
179             - list: A list of GraphQLResult flow run objects
180         """
181 
182         # Get scheduled flow runs from queue
183         mutation = {
184             "mutation($input: getRunsInQueueInput!)": {
185                 "getRunsInQueue(input: $input)": {"flow_run_ids"}
186             }
187         }
188 
189         result = self.client.graphql(
190             mutation, variables={"input": {"tenantId": tenant_id}}
191         )
192         flow_run_ids = result.data.getRunsInQueue.flow_run_ids  # type: ignore
193         now = pendulum.now("UTC")
194 
195         # Query metadata fow flow runs found in queue
196         query = {
197             "query": {
198                 with_args(
199                     "flow_run",
200                     {
201                         # match flow runs in the flow_run_ids list
202                         "where": {
203                             "id": {"_in": flow_run_ids},
204                             "_or": [
205                                 # who are EITHER scheduled...
206                                 {"state": {"_eq": "Scheduled"}},
207                                 # OR running with task runs scheduled to start more than 3 seconds ago
208                                 {
209                                     "state": {"_eq": "Running"},
210                                     "task_runs": {
211                                         "state_start_time": {
212                                             "_lte": str(now.subtract(seconds=3))
213                                         }
214                                     },
215                                 },
216                             ],
217                         }
218                     },
219                 ): {
220                     "id": True,
221                     "version": True,
222                     "tenant_id": True,
223                     "state": True,
224                     "serialized_state": True,
225                     "parameters": True,
226                     "flow": {"id", "name", "environment", "storage"},
227                     with_args(
228                         "task_runs",
229                         {
230                             "where": {
231                                 "state_start_time": {
232                                     "_lte": str(now.subtract(seconds=3))
233                                 }
234                             }
235                         },
236                     ): {"id", "version", "task_id", "serialized_state"},
237                 }
238             }
239         }
240 
241         result = self.client.graphql(query)
242         return result.data.flow_run  # type: ignore
243 
244     def update_states(self, flow_runs: list) -> None:
245         """
246         After a flow run is grabbed this function sets the state to Submitted so it
247         won't be picked up by any other processes
248 
249         Args:
250             - flow_runs (list): A list of GraphQLResult flow run objects
251         """
252         for flow_run in flow_runs:
253 
254             # Set flow run state to `Submitted` if it is currently `Scheduled`
255             if state.StateSchema().load(flow_run.serialized_state).is_scheduled():
256                 self.client.set_flow_run_state(
257                     flow_run_id=flow_run.id,
258                     version=flow_run.version,
259                     state=Submitted(
260                         message="Submitted for execution",
261                         state=state.StateSchema().load(flow_run.serialized_state),
262                     ),
263                 )
264 
265             # Set task run states to `Submitted` if they are currently `Scheduled`
266             for task_run in flow_run.task_runs:
267                 if state.StateSchema().load(task_run.serialized_state).is_scheduled():
268                     self.client.set_task_run_state(
269                         task_run_id=task_run.id,
270                         version=task_run.version,
271                         state=Submitted(
272                             message="Submitted for execution",
273                             state=state.StateSchema().load(task_run.serialized_state),
274                         ),
275                     )
276 
277     def deploy_flows(self, flow_runs: list) -> None:
278         """
279         Meant to be overridden by a platform specific deployment option
280 
281         Args:
282             - flow_runs (list): A list of GraphQLResult flow run objects
283         """
284         pass
285 
286     def heartbeat(self) -> None:
287         """
288         Meant to be overridden by a platform specific heartbeat option
289         """
290         pass
291 
292 
293 if __name__ == "__main__":
294     Agent().start()
295
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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
@@ -186,11 +186,12 @@
             }
         }
 
+        now = pendulum.now("UTC")
         result = self.client.graphql(
-            mutation, variables={"input": {"tenantId": tenant_id}}
+            mutation,
+            variables={"input": {"tenantId": tenant_id, "before": now.isoformat()}},
         )
         flow_run_ids = result.data.getRunsInQueue.flow_run_ids  # type: ignore
-        now = pendulum.now("UTC")
 
         # Query metadata fow flow runs found in queue
         query = {

{"golden_diff": "diff --git a/src/prefect/agent/agent.py b/src/prefect/agent/agent.py\n--- a/src/prefect/agent/agent.py\n+++ b/src/prefect/agent/agent.py\n@@ -186,11 +186,12 @@\n             }\n         }\n \n+        now = pendulum.now(\"UTC\")\n         result = self.client.graphql(\n-            mutation, variables={\"input\": {\"tenantId\": tenant_id}}\n+            mutation,\n+            variables={\"input\": {\"tenantId\": tenant_id, \"before\": now.isoformat()}},\n         )\n         flow_run_ids = result.data.getRunsInQueue.flow_run_ids  # type: ignore\n-        now = pendulum.now(\"UTC\")\n \n         # Query metadata fow flow runs found in queue\n         query = {\n", "issue": "Handle cases where Prefect is deployed to multiple environments with differently-calibrated clocks\nIt's possible that if a Prefect flow is deployed in an environment with an incorrect clock, then it may not run (if it receives states from an external system like Prefect Cloud). \r\n\r\nFor example, say the clock in a cluster is running 5 minutes slow, and a flow is supposed to start at 9:00am UTC. At 9:00am UTC, the agent in that cluster will receive a payload including a `Scheduled` state with a start time of `9:00:00 UTC`. Note that in the cluster, however, the time is 8:55am UTC. The agent will take that `Scheduled` state and submit it for execution, but the `Runner` classes will refuse to run it because it hasn't reached its start time yet -- and won't for 5 minutes. \r\n\r\nWe have two possible remedies for this situation:\r\n- A global `time_adjust_ms` config. This could be set by querying a known/trusted UTC time source and comparing the result to the local time. The difference (in milliseconds) could be used by the runner whenever a scheduled state was examined for execution. If this were exposed as a configuration, it could be set by the agent as an env var whenever it deployed a flow (the agent could do the checking, then set the variable appropriately). The downside of this approach is that it might create bad behavior for any state generated locally. For example, a local retry wouldn't want this adjustment to be applied. If the adjustment were positive (meaning it *adds* time to state start times because the the local clock is running fast), then the local retry wouldn't happen at all because it wouldn't have reached its presumed start time! However, if that's the only case that's problematic, we could handle it by only applying negative adjustments.\r\n\r\n- An alternative is for the agent to \"hold\" any runs that haven't reached their start time yet, rather than deploying them for execution immediately. (Note: this dovetails with an idea @joshmeek and I had to give scheduled runs to the agent slightly ahead of schedule but \"hold\" them so they started exactly on time, with no poll delay). The downside of this approach is that any run which involved a dynamic task, like a mapped task that gets its state directly from Cloud and consequently bypasses the agent, would fail to run. \r\n\r\nPerhaps a combination of these approaches is possible? Perhaps there's a third?\r\n\r\ncc @cicdw @joshmeek @zdhughes  \r\n\n", "before_files": [{"content": "import logging\nfrom typing import Union\n\nimport pendulum\nimport time\n\nfrom prefect import config\nfrom prefect.client import Client\nfrom prefect.serialization import state\nfrom prefect.engine.state import Submitted\nfrom prefect.utilities.exceptions import AuthorizationError\nfrom prefect.utilities.graphql import with_args\n\n\nascii_name = r\"\"\"\n ____            __           _        _                    _\n|  _ \\ _ __ ___ / _| ___  ___| |_     / \\   __ _  ___ _ __ | |_\n| |_) | '__/ _ \\ |_ / _ \\/ __| __|   / _ \\ / _` |/ _ \\ '_ \\| __|\n|  __/| | |  __/  _|  __/ (__| |_   / ___ \\ (_| |  __/ | | | |_\n|_|   |_|  \\___|_|  \\___|\\___|\\__| /_/   \\_\\__, |\\___|_| |_|\\__|\n                                           |___/\n\"\"\"\n\n\nclass Agent:\n    \"\"\"\n    Base class for Agents.\n\n    This Agent class is a standard point for executing Flows in Prefect Cloud. It is meant\n    to have subclasses which inherit functionality from this class. The only piece that\n    the subclasses should implement is the `deploy_flows` function, which specifies how to run a Flow on the given platform. It is built in this\n    way to keep Prefect Cloud logic standard but allows for platform specific\n    customizability.\n\n    In order for this to operate `PREFECT__CLOUD__AGENT__AUTH_TOKEN` must be set as an\n    environment variable or in your user configuration file.\n    \"\"\"\n\n    def __init__(self) -> None:\n\n        token = config.cloud.agent.get(\"auth_token\")\n\n        self.client = Client(api_token=token)\n        self._verify_token(token)\n\n        logger = logging.getLogger(\"agent\")\n        logger.setLevel(logging.DEBUG)\n        ch = logging.StreamHandler()\n        ch.setLevel(logging.DEBUG)\n        formatter = logging.Formatter(\n            \"%(asctime)s - %(name)s - %(levelname)s - %(message)s\"\n        )\n        ch.setFormatter(formatter)\n        logger.addHandler(ch)\n\n        self.logger = logger\n\n    def _verify_token(self, token: str) -> None:\n        \"\"\"\n        Checks whether a token with a `RUNNER` scope was provided\n\n        Args:\n            - token (str): The provided agent token to verify\n\n        Raises:\n            - AuthorizationError: if token is empty or does not have a RUNNER role\n        \"\"\"\n        if not token:\n            raise AuthorizationError(\"No agent API token provided.\")\n\n        # Check if RUNNER role\n        result = self.client.graphql(query=\"query { authInfo { apiTokenScope } }\")\n        if (\n            not result.data  # type: ignore\n            or result.data.authInfo.apiTokenScope != \"RUNNER\"  # type: ignore\n        ):\n            raise AuthorizationError(\"Provided token does not have a RUNNER scope.\")\n\n    def start(self) -> None:\n        \"\"\"\n        The main entrypoint to the agent. This function loops and constantly polls for\n        new flow runs to deploy\n        \"\"\"\n        tenant_id = self.agent_connect()\n\n        # Loop intervals for query sleep backoff\n        loop_intervals = {0: 0.25, 1: 0.5, 2: 1.0, 3: 2.0, 4: 4.0, 5: 8.0, 6: 10.0}\n\n        index = 0\n        while True:\n            self.heartbeat()\n\n            runs = self.agent_process(tenant_id)\n            if runs:\n                index = 0\n            elif index < max(loop_intervals.keys()):\n                index += 1\n\n            time.sleep(loop_intervals[index])\n\n    def agent_connect(self) -> str:\n        \"\"\"\n        Verify agent connection to Prefect Cloud by finding and returning a tenant id\n\n        Returns:\n            - str: The current tenant id\n        \"\"\"\n        print(ascii_name)\n        self.logger.info(\"Starting {}\".format(type(self).__name__))\n        self.logger.info(\n            \"Agent documentation can be found at https://docs.prefect.io/cloud/\"\n        )\n        tenant_id = self.query_tenant_id()\n\n        if not tenant_id:\n            raise ConnectionError(\n                \"Tenant ID not found. Verify that you are using the proper API token.\"\n            )\n\n        self.logger.info(\"Agent successfully connected to Prefect Cloud\")\n        self.logger.info(\"Waiting for flow runs...\")\n\n        return tenant_id\n\n    def agent_process(self, tenant_id: str) -> bool:\n        \"\"\"\n        Full process for finding flow runs, updating states, and deploying.\n\n        Args:\n            - tenant_id (str): The tenant id to use in the query\n\n        Returns:\n            - bool: whether or not flow runs were found\n        \"\"\"\n        flow_runs = None\n        try:\n            flow_runs = self.query_flow_runs(tenant_id=tenant_id)\n\n            if flow_runs:\n                self.logger.info(\n                    \"Found {} flow run(s) to submit for execution.\".format(\n                        len(flow_runs)\n                    )\n                )\n\n                self.update_states(flow_runs)\n                self.deploy_flows(flow_runs)\n                self.logger.info(\n                    \"Submitted {} flow run(s) for execution.\".format(len(flow_runs))\n                )\n        except Exception as exc:\n            self.logger.error(exc)\n\n        return bool(flow_runs)\n\n    def query_tenant_id(self) -> Union[str, None]:\n        \"\"\"\n        Query Prefect Cloud for the tenant id that corresponds to the agent's auth token\n\n        Returns:\n            - Union[str, None]: The current tenant id if found, None otherwise\n        \"\"\"\n        query = {\"query\": {\"tenant\": {\"id\"}}}\n        result = self.client.graphql(query)\n\n        if result.data.tenant:  # type: ignore\n            return result.data.tenant[0].id  # type: ignore\n\n        return None\n\n    def query_flow_runs(self, tenant_id: str) -> list:\n        \"\"\"\n        Query Prefect Cloud for flow runs which need to be deployed and executed\n\n        Args:\n            - tenant_id (str): The tenant id to use in the query\n\n        Returns:\n            - list: A list of GraphQLResult flow run objects\n        \"\"\"\n\n        # Get scheduled flow runs from queue\n        mutation = {\n            \"mutation($input: getRunsInQueueInput!)\": {\n                \"getRunsInQueue(input: $input)\": {\"flow_run_ids\"}\n            }\n        }\n\n        result = self.client.graphql(\n            mutation, variables={\"input\": {\"tenantId\": tenant_id}}\n        )\n        flow_run_ids = result.data.getRunsInQueue.flow_run_ids  # type: ignore\n        now = pendulum.now(\"UTC\")\n\n        # Query metadata fow flow runs found in queue\n        query = {\n            \"query\": {\n                with_args(\n                    \"flow_run\",\n                    {\n                        # match flow runs in the flow_run_ids list\n                        \"where\": {\n                            \"id\": {\"_in\": flow_run_ids},\n                            \"_or\": [\n                                # who are EITHER scheduled...\n                                {\"state\": {\"_eq\": \"Scheduled\"}},\n                                # OR running with task runs scheduled to start more than 3 seconds ago\n                                {\n                                    \"state\": {\"_eq\": \"Running\"},\n                                    \"task_runs\": {\n                                        \"state_start_time\": {\n                                            \"_lte\": str(now.subtract(seconds=3))\n                                        }\n                                    },\n                                },\n                            ],\n                        }\n                    },\n                ): {\n                    \"id\": True,\n                    \"version\": True,\n                    \"tenant_id\": True,\n                    \"state\": True,\n                    \"serialized_state\": True,\n                    \"parameters\": True,\n                    \"flow\": {\"id\", \"name\", \"environment\", \"storage\"},\n                    with_args(\n                        \"task_runs\",\n                        {\n                            \"where\": {\n                                \"state_start_time\": {\n                                    \"_lte\": str(now.subtract(seconds=3))\n                                }\n                            }\n                        },\n                    ): {\"id\", \"version\", \"task_id\", \"serialized_state\"},\n                }\n            }\n        }\n\n        result = self.client.graphql(query)\n        return result.data.flow_run  # type: ignore\n\n    def update_states(self, flow_runs: list) -> None:\n        \"\"\"\n        After a flow run is grabbed this function sets the state to Submitted so it\n        won't be picked up by any other processes\n\n        Args:\n            - flow_runs (list): A list of GraphQLResult flow run objects\n        \"\"\"\n        for flow_run in flow_runs:\n\n            # Set flow run state to `Submitted` if it is currently `Scheduled`\n            if state.StateSchema().load(flow_run.serialized_state).is_scheduled():\n                self.client.set_flow_run_state(\n                    flow_run_id=flow_run.id,\n                    version=flow_run.version,\n                    state=Submitted(\n                        message=\"Submitted for execution\",\n                        state=state.StateSchema().load(flow_run.serialized_state),\n                    ),\n                )\n\n            # Set task run states to `Submitted` if they are currently `Scheduled`\n            for task_run in flow_run.task_runs:\n                if state.StateSchema().load(task_run.serialized_state).is_scheduled():\n                    self.client.set_task_run_state(\n                        task_run_id=task_run.id,\n                        version=task_run.version,\n                        state=Submitted(\n                            message=\"Submitted for execution\",\n                            state=state.StateSchema().load(task_run.serialized_state),\n                        ),\n                    )\n\n    def deploy_flows(self, flow_runs: list) -> None:\n        \"\"\"\n        Meant to be overridden by a platform specific deployment option\n\n        Args:\n            - flow_runs (list): A list of GraphQLResult flow run objects\n        \"\"\"\n        pass\n\n    def heartbeat(self) -> None:\n        \"\"\"\n        Meant to be overridden by a platform specific heartbeat option\n        \"\"\"\n        pass\n\n\nif __name__ == \"__main__\":\n    Agent().start()\n", "path": "src/prefect/agent/agent.py"}], "after_files": [{"content": "import logging\nfrom typing import Union\n\nimport pendulum\nimport time\n\nfrom prefect import config\nfrom prefect.client import Client\nfrom prefect.serialization import state\nfrom prefect.engine.state import Submitted\nfrom prefect.utilities.exceptions import AuthorizationError\nfrom prefect.utilities.graphql import with_args\n\n\nascii_name = r\"\"\"\n ____            __           _        _                    _\n|  _ \\ _ __ ___ / _| ___  ___| |_     / \\   __ _  ___ _ __ | |_\n| |_) | '__/ _ \\ |_ / _ \\/ __| __|   / _ \\ / _` |/ _ \\ '_ \\| __|\n|  __/| | |  __/  _|  __/ (__| |_   / ___ \\ (_| |  __/ | | | |_\n|_|   |_|  \\___|_|  \\___|\\___|\\__| /_/   \\_\\__, |\\___|_| |_|\\__|\n                                           |___/\n\"\"\"\n\n\nclass Agent:\n    \"\"\"\n    Base class for Agents.\n\n    This Agent class is a standard point for executing Flows in Prefect Cloud. It is meant\n    to have subclasses which inherit functionality from this class. The only piece that\n    the subclasses should implement is the `deploy_flows` function, which specifies how to run a Flow on the given platform. It is built in this\n    way to keep Prefect Cloud logic standard but allows for platform specific\n    customizability.\n\n    In order for this to operate `PREFECT__CLOUD__AGENT__AUTH_TOKEN` must be set as an\n    environment variable or in your user configuration file.\n    \"\"\"\n\n    def __init__(self) -> None:\n\n        token = config.cloud.agent.get(\"auth_token\")\n\n        self.client = Client(api_token=token)\n        self._verify_token(token)\n\n        logger = logging.getLogger(\"agent\")\n        logger.setLevel(logging.DEBUG)\n        ch = logging.StreamHandler()\n        ch.setLevel(logging.DEBUG)\n        formatter = logging.Formatter(\n            \"%(asctime)s - %(name)s - %(levelname)s - %(message)s\"\n        )\n        ch.setFormatter(formatter)\n        logger.addHandler(ch)\n\n        self.logger = logger\n\n    def _verify_token(self, token: str) -> None:\n        \"\"\"\n        Checks whether a token with a `RUNNER` scope was provided\n\n        Args:\n            - token (str): The provided agent token to verify\n\n        Raises:\n            - AuthorizationError: if token is empty or does not have a RUNNER role\n        \"\"\"\n        if not token:\n            raise AuthorizationError(\"No agent API token provided.\")\n\n        # Check if RUNNER role\n        result = self.client.graphql(query=\"query { authInfo { apiTokenScope } }\")\n        if (\n            not result.data  # type: ignore\n            or result.data.authInfo.apiTokenScope != \"RUNNER\"  # type: ignore\n        ):\n            raise AuthorizationError(\"Provided token does not have a RUNNER scope.\")\n\n    def start(self) -> None:\n        \"\"\"\n        The main entrypoint to the agent. This function loops and constantly polls for\n        new flow runs to deploy\n        \"\"\"\n        tenant_id = self.agent_connect()\n\n        # Loop intervals for query sleep backoff\n        loop_intervals = {0: 0.25, 1: 0.5, 2: 1.0, 3: 2.0, 4: 4.0, 5: 8.0, 6: 10.0}\n\n        index = 0\n        while True:\n            self.heartbeat()\n\n            runs = self.agent_process(tenant_id)\n            if runs:\n                index = 0\n            elif index < max(loop_intervals.keys()):\n                index += 1\n\n            time.sleep(loop_intervals[index])\n\n    def agent_connect(self) -> str:\n        \"\"\"\n        Verify agent connection to Prefect Cloud by finding and returning a tenant id\n\n        Returns:\n            - str: The current tenant id\n        \"\"\"\n        print(ascii_name)\n        self.logger.info(\"Starting {}\".format(type(self).__name__))\n        self.logger.info(\n            \"Agent documentation can be found at https://docs.prefect.io/cloud/\"\n        )\n        tenant_id = self.query_tenant_id()\n\n        if not tenant_id:\n            raise ConnectionError(\n                \"Tenant ID not found. Verify that you are using the proper API token.\"\n            )\n\n        self.logger.info(\"Agent successfully connected to Prefect Cloud\")\n        self.logger.info(\"Waiting for flow runs...\")\n\n        return tenant_id\n\n    def agent_process(self, tenant_id: str) -> bool:\n        \"\"\"\n        Full process for finding flow runs, updating states, and deploying.\n\n        Args:\n            - tenant_id (str): The tenant id to use in the query\n\n        Returns:\n            - bool: whether or not flow runs were found\n        \"\"\"\n        flow_runs = None\n        try:\n            flow_runs = self.query_flow_runs(tenant_id=tenant_id)\n\n            if flow_runs:\n                self.logger.info(\n                    \"Found {} flow run(s) to submit for execution.\".format(\n                        len(flow_runs)\n                    )\n                )\n\n                self.update_states(flow_runs)\n                self.deploy_flows(flow_runs)\n                self.logger.info(\n                    \"Submitted {} flow run(s) for execution.\".format(len(flow_runs))\n                )\n        except Exception as exc:\n            self.logger.error(exc)\n\n        return bool(flow_runs)\n\n    def query_tenant_id(self) -> Union[str, None]:\n        \"\"\"\n        Query Prefect Cloud for the tenant id that corresponds to the agent's auth token\n\n        Returns:\n            - Union[str, None]: The current tenant id if found, None otherwise\n        \"\"\"\n        query = {\"query\": {\"tenant\": {\"id\"}}}\n        result = self.client.graphql(query)\n\n        if result.data.tenant:  # type: ignore\n            return result.data.tenant[0].id  # type: ignore\n\n        return None\n\n    def query_flow_runs(self, tenant_id: str) -> list:\n        \"\"\"\n        Query Prefect Cloud for flow runs which need to be deployed and executed\n\n        Args:\n            - tenant_id (str): The tenant id to use in the query\n\n        Returns:\n            - list: A list of GraphQLResult flow run objects\n        \"\"\"\n\n        # Get scheduled flow runs from queue\n        mutation = {\n            \"mutation($input: getRunsInQueueInput!)\": {\n                \"getRunsInQueue(input: $input)\": {\"flow_run_ids\"}\n            }\n        }\n\n        now = pendulum.now(\"UTC\")\n        result = self.client.graphql(\n            mutation,\n            variables={\"input\": {\"tenantId\": tenant_id, \"before\": now.isoformat()}},\n        )\n        flow_run_ids = result.data.getRunsInQueue.flow_run_ids  # type: ignore\n\n        # Query metadata fow flow runs found in queue\n        query = {\n            \"query\": {\n                with_args(\n                    \"flow_run\",\n                    {\n                        # match flow runs in the flow_run_ids list\n                        \"where\": {\n                            \"id\": {\"_in\": flow_run_ids},\n                            \"_or\": [\n                                # who are EITHER scheduled...\n                                {\"state\": {\"_eq\": \"Scheduled\"}},\n                                # OR running with task runs scheduled to start more than 3 seconds ago\n                                {\n                                    \"state\": {\"_eq\": \"Running\"},\n                                    \"task_runs\": {\n                                        \"state_start_time\": {\n                                            \"_lte\": str(now.subtract(seconds=3))\n                                        }\n                                    },\n                                },\n                            ],\n                        }\n                    },\n                ): {\n                    \"id\": True,\n                    \"version\": True,\n                    \"tenant_id\": True,\n                    \"state\": True,\n                    \"serialized_state\": True,\n                    \"parameters\": True,\n                    \"flow\": {\"id\", \"name\", \"environment\", \"storage\"},\n                    with_args(\n                        \"task_runs\",\n                        {\n                            \"where\": {\n                                \"state_start_time\": {\n                                    \"_lte\": str(now.subtract(seconds=3))\n                                }\n                            }\n                        },\n                    ): {\"id\", \"version\", \"task_id\", \"serialized_state\"},\n                }\n            }\n        }\n\n        result = self.client.graphql(query)\n        return result.data.flow_run  # type: ignore\n\n    def update_states(self, flow_runs: list) -> None:\n        \"\"\"\n        After a flow run is grabbed this function sets the state to Submitted so it\n        won't be picked up by any other processes\n\n        Args:\n            - flow_runs (list): A list of GraphQLResult flow run objects\n        \"\"\"\n        for flow_run in flow_runs:\n\n            # Set flow run state to `Submitted` if it is currently `Scheduled`\n            if state.StateSchema().load(flow_run.serialized_state).is_scheduled():\n                self.client.set_flow_run_state(\n                    flow_run_id=flow_run.id,\n                    version=flow_run.version,\n                    state=Submitted(\n                        message=\"Submitted for execution\",\n                        state=state.StateSchema().load(flow_run.serialized_state),\n                    ),\n                )\n\n            # Set task run states to `Submitted` if they are currently `Scheduled`\n            for task_run in flow_run.task_runs:\n                if state.StateSchema().load(task_run.serialized_state).is_scheduled():\n                    self.client.set_task_run_state(\n                        task_run_id=task_run.id,\n                        version=task_run.version,\n                        state=Submitted(\n                            message=\"Submitted for execution\",\n                            state=state.StateSchema().load(task_run.serialized_state),\n                        ),\n                    )\n\n    def deploy_flows(self, flow_runs: list) -> None:\n        \"\"\"\n        Meant to be overridden by a platform specific deployment option\n\n        Args:\n            - flow_runs (list): A list of GraphQLResult flow run objects\n        \"\"\"\n        pass\n\n    def heartbeat(self) -> None:\n        \"\"\"\n        Meant to be overridden by a platform specific heartbeat option\n        \"\"\"\n        pass\n\n\nif __name__ == \"__main__\":\n    Agent().start()\n", "path": "src/prefect/agent/agent.py"}]}