From f5f9c491fc2e8f35b354b1a1dffb7b772e3248c4 Mon Sep 17 00:00:00 2001 From: Romain Cledat Date: Fri, 5 May 2023 11:03:42 -0700 Subject: [PATCH 1/3] Fix an issue introduced in 2.8.4 that prevented pickled MetaflowObject from being accessed after pickling in a different namespace The core issue is that we now call the constructor when un-pickling which will recheck the namespace. This broke previous behavior where the constructor was not called and no namespace check was done --- metaflow/client/core.py | 10 ++++++++-- test/core/tests/current_singleton.py | 11 ++++++++++- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/metaflow/client/core.py b/metaflow/client/core.py index 3900b60a7f1..75f0cc4a928 100644 --- a/metaflow/client/core.py +++ b/metaflow/client/core.py @@ -516,10 +516,12 @@ def __setstate__(self, state): self._UNPICKLE_FUNC[version](self, state["data"]) else: # For backward compatibility: handles pickled objects that were serialized without a __getstate__ override + # We set namespace_check to False if it doesn't exist for the same + # reason as the one listed in __getstate__ self.__init__( pathspec=state.get("_pathspec", None), attempt=state.get("_attempt", None), - _namespace_check=state.get("_namespace_check", True), + _namespace_check=state.get("_namespace_check", False), ) def __getstate__(self): @@ -531,12 +533,16 @@ def __getstate__(self): from this object) are pickled (serialized) in a later version of Metaflow, it may not be possible to unpickle (deserialize) them in a previous version of Metaflow. """ + # Note that we set _namespace_check to False because we want the user to + # be able to access this object even after unpickling it. If we set it to + # True, it would check the namespace again at the time of unpickling even + # if the user properly got the object in the first place and pickled it. return { "version": "2.8.4", "data": [ self.pathspec, self._attempt, - self._namespace_check, + False, ], } diff --git a/test/core/tests/current_singleton.py b/test/core/tests/current_singleton.py index da1edcbcad4..04e4c92fd86 100644 --- a/test/core/tests/current_singleton.py +++ b/test/core/tests/current_singleton.py @@ -30,6 +30,7 @@ def step_start(self): self.task_data = {current.pathspec: self.uuid} self.tags = current.tags self.task_obj = current.task + self.run_obj = current.run @steps(1, ["join"]) def step_join(self): @@ -72,6 +73,7 @@ def step_join(self): self.task_data[current.pathspec] = self.uuid self.tags.update(current.tags) self.task_obj = current.task + self.run_obj = current.run @steps(2, ["all"]) def step_all(self): @@ -105,17 +107,24 @@ def check_results(self, flow, checker): step.name, "project_names", {"current_singleton"} ) else: - from metaflow import Task + from metaflow import Task, namespace task_data = run.data.task_data for pathspec, uuid in task_data.items(): assert_equals(Task(pathspec).data.uuid, uuid) + namespace("non-existent-namespace-to-test-namespacecheck") for step in run: for task in step: assert_equals(task.data.step_name, step.id) pathspec = "/".join(task.pathspec.split("/")[-4:]) assert_equals(task.data.uuid, task_data[pathspec]) assert_equals(task.data.task_obj.pathspec, task.pathspec) + # Check we can go up and down pickled objects even in a different + # namespace + assert_equals(task.data.parent.parent.id, task.data.run_obj.id) + assert_equals( + task.data.run_obj[task.data.step_name].id, task.data.step_name + ) assert_equals(run.data.run_obj.pathspec, run.pathspec) assert_equals(run.data.project_names, {"current_singleton"}) assert_equals(run.data.branch_names, {"user.tester"}) From d6e4f0cd4f3973329e87d47e9f3b4eb9039b6258 Mon Sep 17 00:00:00 2001 From: Preetam Joshi Date: Fri, 5 May 2023 13:39:46 -0700 Subject: [PATCH 2/3] Restoring the namespace of the checker to fix tests --- test/core/tests/current_singleton.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/test/core/tests/current_singleton.py b/test/core/tests/current_singleton.py index 04e4c92fd86..d52c577d418 100644 --- a/test/core/tests/current_singleton.py +++ b/test/core/tests/current_singleton.py @@ -99,6 +99,8 @@ def step_all(self): def check_results(self, flow, checker): run = checker.get_run() + from metaflow import get_namespace + checker_namespace = get_namespace() if run is None: # very basic sanity check for CLI for step in flow: @@ -112,6 +114,8 @@ def check_results(self, flow, checker): task_data = run.data.task_data for pathspec, uuid in task_data.items(): assert_equals(Task(pathspec).data.uuid, uuid) + + # Override the namespace for the pickling/unpickling checks namespace("non-existent-namespace-to-test-namespacecheck") for step in run: for task in step: @@ -125,6 +129,8 @@ def check_results(self, flow, checker): assert_equals( task.data.run_obj[task.data.step_name].id, task.data.step_name ) + # Restore the original namespace back + namespace(checker_namespace) assert_equals(run.data.run_obj.pathspec, run.pathspec) assert_equals(run.data.project_names, {"current_singleton"}) assert_equals(run.data.branch_names, {"user.tester"}) From 8e15ed31fc7858a245cc3c0eea3d873f2396e635 Mon Sep 17 00:00:00 2001 From: Preetam Joshi Date: Fri, 5 May 2023 13:51:00 -0700 Subject: [PATCH 3/3] Fixing formatting --- test/core/tests/current_singleton.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/core/tests/current_singleton.py b/test/core/tests/current_singleton.py index d52c577d418..1d7be1fac24 100644 --- a/test/core/tests/current_singleton.py +++ b/test/core/tests/current_singleton.py @@ -100,6 +100,7 @@ def step_all(self): def check_results(self, flow, checker): run = checker.get_run() from metaflow import get_namespace + checker_namespace = get_namespace() if run is None: # very basic sanity check for CLI @@ -129,7 +130,7 @@ def check_results(self, flow, checker): assert_equals( task.data.run_obj[task.data.step_name].id, task.data.step_name ) - # Restore the original namespace back + # Restore the original namespace back for these tests namespace(checker_namespace) assert_equals(run.data.run_obj.pathspec, run.pathspec) assert_equals(run.data.project_names, {"current_singleton"})