Add more tests

src/prefect/engine/executors/dask.py

@@ -79,6 +79,29 @@ def start(self) -> Iterator[None]:
             self.client = None
             self.is_started = False
 
+    def _prep_dask_kwargs(self) -> dict:
+        dask_kwargs = {"pure": False}  # type: dict
+
+        ## set a key for the dask scheduler UI
+        if context.get("task_full_name"):
+            key = context.get("task_full_name", "") + "-" + str(uuid.uuid4())
+            dask_kwargs.update(key=key)
+
+        ## 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)})
+            dask_kwargs.update(resources=resources)
+
+        return dask_kwargs
+
     def queue(self, maxsize: int = 0, client: Client = None) -> Queue:
         """
         Creates an executor-compatible Queue object that can share state
@@ -114,30 +137,15 @@ 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)
+        dask_kwargs = self._prep_dask_kwargs()
+        kwargs.update(dask_kwargs)
 
         if self.is_started and hasattr(self, "client"):
-            future = self.client.submit(fn, *args, pure=False, key=key, **kwargs)
+            future = self.client.submit(fn, *args, **kwargs)
         elif self.is_started:
             with worker_client(separate_thread=True) as client:
-                future = client.submit(fn, *args, pure=False, key=key, **kwargs)
+                future = client.submit(fn, *args, **kwargs)
         else:
             raise ValueError("This executor has not been started.")
 
@@ -161,16 +169,14 @@ def map(self, fn: Callable, *args: Any, **kwargs: Any) -> List[Future]:
         if not args:
             return []
 
-        if context.get("task_full_name"):
-            key = context.get("task_full_name", "") + "-" + str(uuid.uuid4())
-        else:
-            key = None
+        dask_kwargs = self._prep_dask_kwargs()
+        kwargs.update(dask_kwargs)
 
         if self.is_started and hasattr(self, "client"):
-            futures = self.client.map(fn, *args, pure=False, key=key, **kwargs)
+            futures = self.client.map(fn, *args, **kwargs)
         elif self.is_started:
             with worker_client(separate_thread=True) as client:
-                futures = client.map(fn, *args, pure=False, key=key, **kwargs)
+                futures = client.map(fn, *args, **kwargs)
                 return client.gather(futures)
         else:
             raise ValueError("This executor has not been started.")

tests/engine/executors/test_executors.py

@@ -237,6 +237,26 @@ def test_init_kwargs_are_passed_to_init(self, monkeypatch):
         assert client.called
         assert client.call_args[-1]["test_kwarg"] == "test_value"
 
+    def test_task_names_are_passed_to_submit(self, monkeypatch):
+        client = MagicMock()
+        monkeypatch.setattr(prefect.engine.executors.dask, "Client", client)
+        executor = DaskExecutor()
+        with executor.start():
+            with prefect.context(task_full_name="FISH!"):
+                executor.submit(lambda: None)
+        kwargs = client.return_value.__enter__.return_value.submit.call_args[1]
+        assert kwargs["key"].startswith("FISH!")
+
+    def test_task_names_are_passed_to_map(self, monkeypatch):
+        client = MagicMock()
+        monkeypatch.setattr(prefect.engine.executors.dask, "Client", client)
+        executor = DaskExecutor()
+        with executor.start():
+            with prefect.context(task_full_name="FISH![0]"):
+                executor.map(lambda: None, [1, 2])
+        kwargs = client.return_value.__enter__.return_value.map.call_args[1]
+        assert kwargs["key"].startswith("FISH![0]")
+
     def test_context_tags_are_passed_to_submit(self, monkeypatch):
         client = MagicMock()
         monkeypatch.setattr(prefect.engine.executors.dask, "Client", client)
@@ -247,6 +267,16 @@ def test_context_tags_are_passed_to_submit(self, monkeypatch):
         kwargs = client.return_value.__enter__.return_value.submit.call_args[1]
         assert kwargs["resources"] == {"GPU": 1.0}
 
+    def test_context_tags_are_passed_to_map(self, monkeypatch):
+        client = MagicMock()
+        monkeypatch.setattr(prefect.engine.executors.dask, "Client", client)
+        executor = DaskExecutor()
+        with executor.start():
+            with prefect.context(task_tags=["dask-resource:GPU=1"]):
+                executor.map(lambda: None, [1, 2])
+        kwargs = client.return_value.__enter__.return_value.map.call_args[1]
+        assert kwargs["resources"] == {"GPU": 1.0}
+
     def test_debug_is_converted_to_silence_logs(self, monkeypatch):
         client = MagicMock()
         monkeypatch.setattr(prefect.engine.executors.dask, "Client", client)