Repository: PrefectHQ/prefect

Issue:
Option to include item string in mapped task name
## Use Case
Minor, non-urgent feature request: mapped tasks are currently referenced in logs and the Cloud UI using their numeric index, e.g. my_task[0], my_task[1]. These numeric values aren't particularly meaningful and, in some cases, the string value of the mapped item would be more useful during log examination or when viewing Flow runs in the Cloud UI.

## Solution
Having an opt-in way to enable use of the string version of a mapped item in the task name would be helpful, e.g.

```
    @task(mapped_item_in_task_name=True)
    def my_task(item):
        ...

    my_task.map(["apple", "banana", "orange"])
```

Which would yield task names of `my_task[apple]`, `my_task[banana]`, and `my_task[orange]` rather than `my_task[0]`, `my_task[1]`. and `my_task[2]`

For tasks mapping over multiple parameters, this option could simply choose to use the first parameter. Elements with long string representations could be truncated to a reasonable size for display purposes. This should *definitely* be opt-in to avoid logging sensitive data.

A better implementation might be to pass an iterable/list to be used for the task names (as a specific kwarg) which might or might not match the parameter to map over (but would need to have the same length):

```
    fruit = {
        0: "apple",
        1: "banana",
        15: "orange"
    }
    my_task.map(fruit.keys(), mapped_task_names=fruit.values())`
```
This would also yield task names of `my_task[apple]`, `my_task[banana]`, and `my_task[orange]`.

## Alternatives
None available currently.


Discussion:

Patch:
diff --git a/changes/issue2100.yaml b/changes/issue2100.yaml
new file mode 100644
index 0000000000..48566252b3
--- /dev/null
+++ b/changes/issue2100.yaml
@@ -0,0 +1,6 @@
+enhancement:
+  - "Add option to template task run name at runtime when using backend API - [#2100](https://github.com/PrefectHQ/prefect/issues/2100)"
+  - "Add `set_task_run_name` Client function - [#2100](https://github.com/PrefectHQ/prefect/issues/2100)"
+
+fix:
+  - "Fix mypy typing for `target` kwarg on base Task class - [#2100](https://github.com/PrefectHQ/prefect/issues/2100)"
diff --git a/changes/pr3351.yaml b/changes/pr3351.yaml
deleted file mode 100644
index 81df93dfe4..0000000000
--- a/changes/pr3351.yaml
+++ /dev/null
@@ -1,4 +0,0 @@
-enhancement:
-  - "Allow python path in Local storage - [#3351](https://github.com/PrefectHQ/prefect/pull/3351)"
-contributor:
-  - "[Jack D. Sundberg](https://github.com/jacksund)"
diff --git a/src/prefect/client/client.py b/src/prefect/client/client.py
index 7604b7bd02..790b8a83c0 100644
--- a/src/prefect/client/client.py
+++ b/src/prefect/client/client.py
@@ -1254,6 +1254,31 @@ class Client:
             state=state,
         )
 
+    def set_task_run_name(self, task_run_id: str, name: str) -> bool:
+        """
+        Set the name of a task run
+
+        Args:
+            - task_run_id (str): the id of a task run
+            - name (str): a name for this task run
+
+        Returns:
+            - bool: whether or not the task run name was updated
+        """
+        mutation = {
+            "mutation($input: set_task_run_name_input!)": {
+                "set_task_run_name(input: $input)": {
+                    "success": True,
+                }
+            }
+        }
+
+        result = self.graphql(
+            mutation, variables=dict(input=dict(task_run_id=task_run_id, name=name))
+        )
+
+        return result.data.set_task_run_name.success
+
     def get_task_run_state(self, task_run_id: str) -> "prefect.engine.state.State":
         """
         Retrieves the current state for a task run.
diff --git a/src/prefect/core/task.py b/src/prefect/core/task.py
index 446779f670..a57be202c0 100644
--- a/src/prefect/core/task.py
+++ b/src/prefect/core/task.py
@@ -12,6 +12,7 @@ from typing import (
     List,
     Mapping,
     Optional,
+    Union,
 )
 
 import prefect
@@ -173,8 +174,8 @@ class Task(metaclass=SignatureValidator):
             formatted at runtime with values from `prefect.context`. If a callable function
             is provided, it should have signature `callable(**kwargs) -> str` and at write
             time all formatting kwargs will be passed and a fully formatted location is
-            expected as the return value.  Can be used for string formatting logic that
-            `.format(**kwargs)` doesn't support
+            expected as the return value. The callable can be used for string formatting logic that
+            `.format(**kwargs)` doesn't support.
         - state_handlers (Iterable[Callable], optional): A list of state change handlers
             that will be called whenever the task changes state, providing an
             opportunity to inspect or modify the new state. The handler
@@ -188,6 +189,14 @@ class Task(metaclass=SignatureValidator):
             Task enters a failure state
         - log_stdout (bool, optional): Toggle whether or not to send stdout messages to
             the Prefect logger. Defaults to `False`.
+        - task_run_name (Union[str, Callable], optional): a name to set for this task at runtime.
+            `task_run_name` strings can be templated formatting strings which will be
+            formatted at runtime with values from `prefect.context`, task inputs, and parameters.
+            If a callable function is provided, it should have signature `callable(**kwargs) -> str`
+            and at write time all formatting kwargs will be passed and a fully formatted location is
+            expected as the return value. The callable can be used for string formatting logic that
+            `.format(**kwargs)` doesn't support. **Note**: this only works for tasks running against a
+            backend API.
 
     Raises:
         - TypeError: if `tags` is of type `str`
@@ -216,7 +225,8 @@ class Task(metaclass=SignatureValidator):
         on_failure: Callable = None,
         log_stdout: bool = False,
         result: "Result" = None,
-        target: str = None,
+        target: Union[str, Callable] = None,
+        task_run_name: Union[str, Callable] = None,
     ):
         self.name = name or type(self).__name__
         self.slug = slug
@@ -309,7 +319,9 @@ class Task(metaclass=SignatureValidator):
                     stacklevel=2,
                 )
             self.result = self.result.copy()
-            self.result.location = self.target
+            self.result.location = self.target  # type: ignore
+
+        self.task_run_name = task_run_name  # type: ignore
 
         if state_handlers and not isinstance(state_handlers, collections.abc.Sequence):
             raise TypeError("state_handlers should be iterable.")
@@ -422,7 +434,7 @@ class Task(metaclass=SignatureValidator):
                     stacklevel=2,
                 )
             new.result = new.result.copy()
-            new.result.location = new.target
+            new.result.location = new.target  # type: ignore
 
         new.tags = copy.deepcopy(self.tags).union(set(new.tags))
         tags = set(prefect.context.get("tags", set()))
diff --git a/src/prefect/engine/cloud/task_runner.py b/src/prefect/engine/cloud/task_runner.py
index 348be4c5cb..517556e641 100644
--- a/src/prefect/engine/cloud/task_runner.py
+++ b/src/prefect/engine/cloud/task_runner.py
@@ -322,6 +322,33 @@ class CloudTaskRunner(TaskRunner):
             final_state = self.handle_state_change(old_state=state, new_state=new_state)
             raise ENDRUN(final_state) from exc
 
+    def set_task_run_name(self, task_inputs: Dict[str, Result]) -> None:
+        """
+        Sets the name for this task run by calling the `set_task_run_name` mutation.
+
+        Args:
+            - task_inputs (Dict[str, Result]): a dictionary of inputs whose keys correspond
+                to the task's `run()` arguments.
+        """
+        task_run_name = self.task.task_run_name
+
+        if task_run_name:
+            raw_inputs = {k: r.value for k, r in task_inputs.items()}
+            formatting_kwargs = {
+                **prefect.context.get("parameters", {}).copy(),
+                **raw_inputs,
+                **prefect.context,
+            }
+
+            if not isinstance(task_run_name, str):
+                task_run_name = task_run_name(**formatting_kwargs)
+            else:
+                task_run_name = task_run_name.format(**formatting_kwargs)
+
+            self.client.set_task_run_name(
+                task_run_id=self.task_run_id, name=task_run_name  # type: ignore
+            )
+
     @tail_recursive
     def run(
         self,
diff --git a/src/prefect/engine/task_runner.py b/src/prefect/engine/task_runner.py
index 15c486bac5..f557066568 100644
--- a/src/prefect/engine/task_runner.py
+++ b/src/prefect/engine/task_runner.py
@@ -179,7 +179,7 @@ class TaskRunner(Runner):
         # If provided, use task's target as result location
         if self.task.target:
             if not isinstance(self.task.target, str):
-                self.result._formatter = self.task.target
+                self.result._formatter = self.task.target  # type: ignore
                 self.result.location = None
             else:
                 self.result.location = self.task.target
@@ -267,6 +267,9 @@ class TaskRunner(Runner):
                         state, upstream_states=upstream_states
                     )
 
+                # dynamically set task run name
+                self.set_task_run_name(task_inputs=task_inputs)
+
                 if self.task.target:
                     # check to see if there is a Result at the task's target
                     state = self.check_target(state, inputs=task_inputs)
@@ -656,6 +659,16 @@ class TaskRunner(Runner):
         """
         return state, upstream_states
 
+    def set_task_run_name(self, task_inputs: Dict[str, Result]) -> None:
+        """
+        Sets the name for this task run.
+
+        Args:
+            - task_inputs (Dict[str, Result]): a dictionary of inputs whose keys correspond
+                to the task's `run()` arguments.
+        """
+        pass
+
     @call_state_handlers
     def check_target(self, state: State, inputs: Dict[str, Result]) -> State:
         """
@@ -683,8 +696,8 @@ class TaskRunner(Runner):
             if not isinstance(target, str):
                 target = target(**formatting_kwargs)
 
-            if result.exists(target, **formatting_kwargs):
-                known_location = target.format(**formatting_kwargs)
+            if result.exists(target, **formatting_kwargs):  # type: ignore
+                known_location = target.format(**formatting_kwargs)  # type: ignore
                 new_res = result.read(known_location)
                 cached_state = Cached(
                     result=new_res,
@@ -726,7 +739,7 @@ class TaskRunner(Runner):
                 state = Pending("Cache was invalid; ready to run.")
 
         if self.task.cache_for is not None:
-            candidate_states = []
+            candidate_states = []  # type: ignore
             if prefect.context.get("caches"):
                 candidate_states = prefect.context.caches.get(
                     self.task.cache_key or self.task.name, []
diff --git a/src/prefect/environments/storage/local.py b/src/prefect/environments/storage/local.py
index fa6e9988b4..5617834dcb 100644
--- a/src/prefect/environments/storage/local.py
+++ b/src/prefect/environments/storage/local.py
@@ -7,7 +7,7 @@ from slugify import slugify
 import prefect
 from prefect.engine.results import LocalResult
 from prefect.environments.storage import Storage
-from prefect.utilities.storage import extract_flow_from_file, extract_flow_from_module
+from prefect.utilities.storage import extract_flow_from_file
 
 if TYPE_CHECKING:
     from prefect.core.flow import Flow
@@ -32,8 +32,7 @@ class Local(Storage):
             absolute path and created.  Defaults to `True`
         - path (str, optional): a direct path to the location of the flow file if
             `stored_as_script=True`, otherwise this path will be used when storing the serialized,
-            pickled flow. If `stored_as_script=True`, the direct path may be a file path
-            (such as 'path/to/myflow.py') or a direct python path (such as 'myrepo.mymodule.myflow')
+            pickled flow.
         - stored_as_script (bool, optional): boolean for specifying if the flow has been stored
             as a `.py` file. Defaults to `False`
         - **kwargs (Any, optional): any additional `Storage` initialization options
@@ -77,8 +76,7 @@ class Local(Storage):
 
         Args:
             - flow_location (str, optional): the location of a flow within this Storage; in this case,
-                a file path or python path where a Flow has been serialized to. Will use `path`
-                if not provided.
+                a file path where a Flow has been serialized to. Will use `path` if not provided.
 
         Returns:
             - Flow: the requested flow
@@ -94,15 +92,10 @@ class Local(Storage):
         else:
             raise ValueError("No flow location provided")
 
-        # check if the path given is a file path
-        if os.path.isfile(flow_location):
-            if self.stored_as_script:
-                return extract_flow_from_file(file_path=flow_location)
-            else:
-                return prefect.core.flow.Flow.load(flow_location)
-        # otherwise the path is given in the module format
+        if self.stored_as_script:
+            return extract_flow_from_file(file_path=flow_location)
         else:
-            return extract_flow_from_module(module_str=flow_location)
+            return prefect.core.flow.Flow.load(flow_location)
 
     def add_flow(self, flow: "Flow") -> str:
         """
diff --git a/src/prefect/utilities/storage.py b/src/prefect/utilities/storage.py
index 87a98d6119..96c834998a 100644
--- a/src/prefect/utilities/storage.py
+++ b/src/prefect/utilities/storage.py
@@ -1,4 +1,3 @@
-import importlib
 from typing import TYPE_CHECKING
 
 import prefect
@@ -83,34 +82,3 @@ def extract_flow_from_file(
                 return exec_vals[var]
 
     raise ValueError("No flow found in file.")
-
-
-def extract_flow_from_module(module_str: str, flow_name: str = None) -> "Flow":
-    """
-    Extract a flow object from a python module.
-
-    Args:
-        - module_str (str): A module path pointing to a .py file containing a flow.
-            For example, 'myrepo.mymodule.myflow' where myflow.py contains the flow.
-        - flow_name (str, optional): A specific name of a flow to extract from a file.
-            If not set then the first flow object retrieved from file will be returned.
-            The "first" flow object will be based on the dir(module) ordering, which
-            is alphabetical and capitalized letters come first.
-
-    Returns:
-        - Flow: A flow object extracted from a file
-    """
-
-    # load the module
-    module = importlib.import_module(module_str)
-
-    # if flow_name is specified, grab it from the module
-    if flow_name:
-        return getattr(module, flow_name)
-    # otherwise loop until we get a Flow object
-    else:
-        for var in dir(module):
-            if isinstance(getattr(module, var), prefect.Flow):
-                return getattr(module, var)
-
-    raise ValueError("No flow found in module.")


Test patch:
diff --git a/tests/client/test_client.py b/tests/client/test_client.py
index e69bbf6c82..583f5cfd95 100644
--- a/tests/client/test_client.py
+++ b/tests/client/test_client.py
@@ -1037,6 +1037,17 @@ def test_get_task_run_info_with_error(patch_post):
         )
 
 
+def test_set_task_run_name(patch_posts, cloud_api):
+    mutation_resp = {"data": {"set_task_run_name": {"success": True}}}
+
+    post = patch_posts(mutation_resp)
+
+    client = Client()
+    result = client.set_task_run_name(task_run_id="76-salt", name="name")
+
+    assert result == True
+
+
 def test_get_task_run_state(patch_posts, cloud_api, runner_token):
     query_resp = {
         "task_run_by_pk": {
diff --git a/tests/core/test_task.py b/tests/core/test_task.py
index 006f578304..43889c8bc8 100644
--- a/tests/core/test_task.py
+++ b/tests/core/test_task.py
@@ -299,6 +299,16 @@ class TestCreateTask:
         s = Task(log_stdout=True)
         assert s.log_stdout is True
 
+    def test_create_task_with_task_run_name(self):
+        t1 = Task()
+        assert t1.task_run_name is None
+
+        t2 = Task(task_run_name="test")
+        assert t2.task_run_name == "test"
+
+        t2 = Task(task_run_name=lambda: 42)
+        assert t2.task_run_name() == 42
+
 
 def test_task_has_logger():
     t = Task()
diff --git a/tests/engine/cloud/test_cloud_task_runner.py b/tests/engine/cloud/test_cloud_task_runner.py
index b271f29ef1..08e05545a2 100644
--- a/tests/engine/cloud/test_cloud_task_runner.py
+++ b/tests/engine/cloud/test_cloud_task_runner.py
@@ -1050,3 +1050,45 @@ def test_task_runner_handles_version_lock_error(monkeypatch):
     client.get_task_run_state.return_value = s
     with pytest.raises(ENDRUN):
         res = runner.call_runner_target_handlers(Pending(), Running())
+
+
+def test_task_runner_sets_task_name(monkeypatch, cloud_settings):
+
+    client = MagicMock()
+    monkeypatch.setattr(
+        "prefect.engine.cloud.task_runner.Client", MagicMock(return_value=client)
+    )
+    client.set_task_run_name = MagicMock()
+
+    task = Task(name="test", task_run_name="asdf")
+    runner = CloudTaskRunner(task=task)
+    runner.task_run_id = "id"
+
+    runner.set_task_run_name(task_inputs={})
+
+    assert client.set_task_run_name.called
+    assert client.set_task_run_name.call_args[1]["name"] == "asdf"
+    assert client.set_task_run_name.call_args[1]["task_run_id"] == "id"
+
+    task = Task(name="test", task_run_name="{map_index}")
+    runner = CloudTaskRunner(task=task)
+    runner.task_run_id = "id"
+
+    class Temp:
+        value = 100
+
+    runner.set_task_run_name(task_inputs={"map_index": Temp()})
+
+    assert client.set_task_run_name.called
+    assert client.set_task_run_name.call_args[1]["name"] == "100"
+    assert client.set_task_run_name.call_args[1]["task_run_id"] == "id"
+
+    task = Task(name="test", task_run_name=lambda **kwargs: "name")
+    runner = CloudTaskRunner(task=task)
+    runner.task_run_id = "id"
+
+    runner.set_task_run_name(task_inputs={})
+
+    assert client.set_task_run_name.called
+    assert client.set_task_run_name.call_args[1]["name"] == "name"
+    assert client.set_task_run_name.call_args[1]["task_run_id"] == "id"