First pass implementation at dask task affinity tags

src/prefect/engine/executors/dask.py

@@ -110,10 +110,25 @@ def submit(self, fn: Callable, *args: Any, **kwargs: Any) -> Future:
         Returns:
             - Future: a Future-like object that represents the computation of `fn(*args, **kwargs)`
         """
+        ## set a key for the dask scheduler UI
         if context.get("task_full_name"):
             key = context.get("task_full_name", "") + "-" + str(uuid.uuid4())
         else:
             key = None
+
+        ## infer from context if dask resources are being utilized
+        dask_resource_tags = [
+            tag
+            for tag in context.get("task_tags", [])
+            if tag.lower().startswith("dask-resource")
+        ]
+        if dask_resource_tags:
+            resources = {}
+            for tag in dask_resource_tags:
+                prefix, val = tag.split("=")
+                resources.update({prefix.split(":")[1]: float(val)})
+            kwargs.update(resources=resources)
+
         if self.is_started and hasattr(self, "client"):
             future = self.client.submit(fn, *args, pure=False, key=key, **kwargs)
         elif self.is_started:

src/prefect/engine/task_runner.py

@@ -146,7 +146,9 @@ def initialize_run(  # type: ignore
         if isinstance(state, Resume):
             context.update(resume=True)
 
-        context.update(task_run_count=run_count, task_name=self.task.name)
+        context.update(
+            task_run_count=run_count, task_name=self.task.name, task_tags=self.task.tags
+        )
 
         return TaskRunnerInitializeResult(state=state, context=context)
 

tests/engine/test_task_runner.py

@@ -363,6 +363,19 @@ def test_task_runner_puts_resume_in_context_if_state_is_resume(self):
             result = TaskRunner(Task()).initialize_run(state=Resume(), context=ctx)
             assert result.context.resume is True
 
+    def test_task_runner_puts_tags_in_context(self):
+        with prefect.context() as ctx:
+            assert "task_tags" not in ctx
+            result = TaskRunner(Task()).initialize_run(state=None, context=ctx)
+            assert result.context.task_tags == set()
+
+        with prefect.context() as ctx:
+            assert "task_tags" not in ctx
+            result = TaskRunner(Task(tags=["foo", "bar"])).initialize_run(
+                state=None, context=ctx
+            )
+            assert result.context.task_tags == {"foo", "bar"}
+
     @pytest.mark.parametrize(
         "state", [Success(), Failed(), Pending(), Scheduled(), Skipped(), Cached()]
     )