Add static and runtime dag info, API to fetch ancestor and successor tasks #2124

talsperre · 2024-10-31T09:07:57Z

Add runtime DAG info so that we can query the ancestor and successor tasks for a given task easily.

Usage

from metaflow import Task, namespace
namespace(None)
task = Task('RuntimeDAGFlow/18/step_c/32076012', attempt=0)

To get ancestors, progenies, and siblings, use the following API:

ancestors = task.immediate_ancestors()
successors = task.immediate_successors()
siblings = task.closest_siblings()

saikonen · 2024-11-11T09:45:07Z

metaflow/metadata/metadata.py

+    @classmethod
+    def _filter_tasks_by_metadata(
+        cls, flow_id, run_id, query_step, field_name, field_value
+    ):
+        raise NotImplementedError()
+
+    @classmethod
+    def filter_tasks_by_metadata(
+        cls, flow_id, run_id, query_step, field_name, field_value
+    ):
+        # TODO: Do we need to do anything wrt to task attempt?
+        task_ids = cls._filter_tasks_by_metadata(
+            flow_id, run_id, query_step, field_name, field_value
+        )
+        return task_ids
+


is there a need for the private method, or could this simply be contained in the public-facing one? right now its not doing anything before calling the private one.

also, did you have an implementation of this for service.py yet?

saikonen · 2024-11-11T09:55:09Z

metaflow/metadata/metadata.py

+    def filter_tasks_by_metadata(
+        cls, flow_id, run_id, query_step, field_name, field_value
+    ):
+        # TODO: Do we need to do anything wrt to task attempt?


probably not, as the ancestors for task attempts should be identical, right? What about the immediate_siblings though, will they include or exclude attempts of the same task?

romain-intel

A few comments. I think it's pretty close though. I haven't looked at hte metadata service changes. We may also want to raise a better error message if the service is not new enough?

romain-intel · 2025-01-13T19:59:19Z

metaflow/client/core.py

+        run_id: str,
+        cur_foreach_stack_len: int,
+        steps: List[str],
+        query_type: str,


nit: I would just use a boolean. Something like is_ancestor. These are internal functions anyways and slightly more efficient to use bools :)

romain-intel · 2025-01-13T20:01:09Z

metaflow/client/core.py

@@ -1123,6 +1124,202 @@ def _iter_filter(self, x):
        # exclude private data artifacts
        return x.id[0] != "_"

+    def _get_task_for_queried_step(self, flow_id, run_id, query_step):


Could this function simply be replaced by Step(flow_id/run_id/step_name).task? It doesn't seem to bring much more to it.

Wasn't aware that this returns the first available task.

romain-intel · 2025-01-13T20:02:59Z

metaflow/client/core.py

+        if query_foreach_stack_len == cur_foreach_stack_len:
+            # The successor or ancestor tasks belong to the same foreach stack level
+            field_name = "foreach-indices"
+            field_value = self.metadata_dict.get(field_name)


We don't currently cache metadata_dict so either we could fix that or cache it here to avoid making multiple calls to the metadata service and then sorts. It would need to be cached across _get_related_tasks and this function.

romain-intel · 2025-01-13T20:06:36Z

metaflow/client/core.py

+                # Current Task: foreach-indices = [0, 1, 2], foreach-indices-truncated = [0, 1]
+                # Ancestor Task: foreach-indices = [0, 1], foreach-indices-truncated = [0]
+                # We will compare the foreach-indices value of ancestor task with the
+                # foreach-indices value of current task


nit: foreach-indices-truncated value of the current task

romain-intel · 2025-01-13T20:08:27Z

metaflow/client/core.py

+        return field_name, field_value
+
+    def _get_related_tasks(self, relation_type: str) -> Dict[str, List[str]]:
+        start_time = time.time()


not used -- can strip or use.

romain-intel · 2025-01-13T20:14:31Z

metaflow/metaflow_config.py

@@ -248,8 +248,7 @@
 # Default container registry
 DEFAULT_CONTAINER_REGISTRY = from_conf("DEFAULT_CONTAINER_REGISTRY")
 # Controls whether to include foreach stack information in metadata.
-# TODO(Darin, 05/01/24): Remove this flag once we are confident with this feature.
-INCLUDE_FOREACH_STACK = from_conf("INCLUDE_FOREACH_STACK", False)
+INCLUDE_FOREACH_STACK = from_conf("INCLUDE_FOREACH_STACK", True)


We should probably change this at some point and remove it to not make it optional anymore.

romain-intel · 2025-01-13T20:19:07Z

metaflow/plugins/metadata_providers/local.py

+            # Filter tasks based on metadata
+            for task in tasks:
+                task_id = task.get("task_id")
+                if not task_id:


when does this happen? Also, task_id of zero is valid iirc.

romain-intel · 2025-01-13T20:21:04Z

metaflow/plugins/metadata_providers/local.py

+            # and the artifact files are saved as: <attempt>_artifact__<artifact_name>.json
+            # We loop over all the JSON files in the directory and find the latest one
+            # that matches the field prefix.
+            json_files = glob.glob(os.path.join(path, "*.json"))


we should be able to do a more efficient globbing so we don't have to filter by field_prefix later on. SOmething like f"{field_prefix}*.json".

romain-intel · 2025-01-13T20:23:33Z

metaflow/task.py

+                        type="foreach-indices-truncated",
+                        tags=metadata_tags,
+                    ),
+                    MetaDatum(


I believe this is only used in the siblings thing. If that's the case, we may be able to get rid of this when we refactor the siblings thing (if we do that). I am also a little confused as to why this is needed.

Will refactor siblings function mostly to return siblings irrespective of whether it is in a for each or not.

Cool -- I think we can now get rid of this metadatum then right?

romain-intel · 2025-01-13T20:24:17Z

metaflow/task.py

+                                tags=metadata_tags,
+                            ),
+                            MetaDatum(
+                                field="previous_steps",


nit: consistency here between previous_steps and foreach-indices for example.

saikonen · 2025-01-14T16:12:20Z

metaflow/plugins/metadata_providers/service.py

+        }
+        url = ServiceMetadataProvider._obj_path(flow_id, run_id, query_step)
+        url = f"{url}/tasks?{urlencode(query_params)}"
+        return cls._request(cls._monitor, url, "GET")


getting an error with this that cls does not have _monitor. All other calls to _request pass in None

Yes, I will pass in None simply.

saikonen · 2025-01-14T16:12:30Z

metaflow/plugins/metadata_providers/service.py

+            "query_step": query_step,
+        }
+        url = ServiceMetadataProvider._obj_path(flow_id, run_id, query_step)
+        url = f"{url}/tasks?{urlencode(query_params)}"


missing import for urlencode. f-strings are probably fine by 2025, as we've gotten rid of the older tests that break with them.

Ya, I was wondering about fstrings but did check and our official minimum version is 3.6 which supports it -- and yes, let's move at least a tad into the future :). I'm going to start using them too and the code will slowly migrate to it (and become infinitesimally faster :) )

romain-intel · 2025-01-16T08:53:17Z

metaflow/client/core.py

+                m.name: m.value
+                for m in sorted(self.metadata, key=lambda m: m.created_at)
+            }
+        return self._metadata_dict


Note: this slightly changes the syntax since now if there is new metadata, the user won't get it. Should check if this impacts other operations. Or scope the caching to just the functions that need it.

saikonen · 2025-01-16T11:39:07Z

metaflow/client/core.py

+
+    def _get_related_tasks(self, is_ancestor: bool) -> Dict[str, List[str]]:
+        flow_id, run_id, _, _ = self.path_components
+        steps = (


data type problem here which leads to the queries not working correctly:
the steps ends up being of type str on OSS metadata-service, so using these you end up iterating over characters instead of step names, e.g.:

/flows/SplitFlow/runs/63/steps/{/filtered_tasks?metadata_field_name=foreach-indices&metadata_field_value=%7B%7D&query_step=%7B /flows/SplitFlow/runs/63/steps/e/filtered_tasks?metadata_field_name=foreach-indices&metadata_field_value=%7B%7D&query_step=e /flows/SplitFlow/runs/63/steps/n/filtered_tasks?metadata_field_name=foreach-indices&metadata_field_value=%7B%7D&query_step=n /flows/SplitFlow/runs/63/steps/d/filtered_tasks?metadata_field_name=foreach-indices&metadata_field_value=%7B%7D&query_step=d /flows/SplitFlow/runs/63/steps/}/filtered_tasks?metadata_field_name=foreach-indices&metadata_field_value=%7B%7D&query_step=%7D

talsperre force-pushed the dev/add-runtime-dag-info branch from 48c771d to ec43f14 Compare November 1, 2024 18:34

saikonen reviewed Nov 11, 2024

View reviewed changes

saikonen mentioned this pull request Nov 22, 2024

feature: add task filtering based on metadata Netflix/metaflow-service#449

Open

talsperre force-pushed the dev/add-runtime-dag-info branch from ffbf68a to c6fb9ac Compare January 2, 2025 23:25

talsperre changed the title ~~Add static and runtime dag info, API to fetch ancestor tasks~~ Add static and runtime dag info, API to fetch ancestor and successor tasks Jan 7, 2025

talsperre force-pushed the dev/add-runtime-dag-info branch 2 times, most recently from d66d32b to 7644058 Compare January 12, 2025 03:12

talsperre requested review from savingoyal and romain-intel January 12, 2025 05:53

romain-intel reviewed Jan 13, 2025

View reviewed changes

saikonen reviewed Jan 14, 2025

View reviewed changes

talsperre added 19 commits January 14, 2025 16:53

Add static and runtime dag info, API to fetch ancestor tasks

f747175

Add API to get immediate successors

aff0355

Add API for getting closest siblings

7f4887f

Update metadata API params

7f7dfaa

Refactor ancestor and successor client code

4801253

Remove unneccessary prints

fcb15e9

Support querying ancestors and successors in local metadata provider

cd988cf

Refactor and simplify client code

06e5340

Make query logic more descriptive

8441852

Add core tests for ancestor task API

0fc3458

Add core test for immediate successor API

883fcbd

Add endpoint in OSS metadata service

29a8407

Add logs to tests

3fd202e

Log for each stack to metadata, update query logic

90560c5

Add more comments to code

f469734

Run black formatting

f5146d2

Set monitor to None in filter tasks API

1b74f34

import urlencode

81c9e4e

Address comments

96c7560

Update logic for siblings, make it work for static splits as well

7cdfb41

talsperre force-pushed the dev/add-runtime-dag-info branch from 17a4489 to 7cdfb41 Compare January 15, 2025 00:53

update service url for filter task requests. update query param names.

b895b95

romain-intel reviewed Jan 16, 2025

View reviewed changes

saikonen requested changes Jan 16, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add static and runtime dag info, API to fetch ancestor and successor tasks #2124

Add static and runtime dag info, API to fetch ancestor and successor tasks #2124

talsperre commented Oct 31, 2024 •

edited

Loading

saikonen Nov 11, 2024

saikonen Nov 11, 2024

saikonen Nov 11, 2024

romain-intel left a comment

romain-intel Jan 13, 2025

romain-intel Jan 13, 2025

talsperre Jan 14, 2025

romain-intel Jan 13, 2025

romain-intel Jan 13, 2025

romain-intel Jan 13, 2025

romain-intel Jan 13, 2025

romain-intel Jan 13, 2025

romain-intel Jan 13, 2025

romain-intel Jan 13, 2025

talsperre Jan 14, 2025

romain-intel Jan 16, 2025

romain-intel Jan 13, 2025

saikonen Jan 14, 2025

talsperre Jan 14, 2025

saikonen Jan 14, 2025

romain-intel Jan 14, 2025

romain-intel Jan 16, 2025

saikonen Jan 16, 2025 •

edited

Loading

Add static and runtime dag info, API to fetch ancestor and successor tasks #2124

Are you sure you want to change the base?

Add static and runtime dag info, API to fetch ancestor and successor tasks #2124

Conversation

talsperre commented Oct 31, 2024 • edited Loading

Usage

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

romain-intel left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

saikonen Jan 16, 2025 • edited Loading

Choose a reason for hiding this comment

talsperre commented Oct 31, 2024 •

edited

Loading

saikonen Jan 16, 2025 •

edited

Loading