diff --git a/metaflow/client/core.py b/metaflow/client/core.py index 04c48596b1..812e861bae 100644 --- a/metaflow/client/core.py +++ b/metaflow/client/core.py @@ -269,13 +269,19 @@ def __init__( _parent: Optional["MetaflowObject"] = None, _namespace_check: bool = True, _current_namespace: Optional[str] = None, + _current_metadata: Optional[str] = None, ): self._metaflow = Metaflow() self._parent = _parent self._path_components = None self._attempt = attempt self._current_namespace = _current_namespace or get_namespace() + self._current_metadata = _current_metadata or get_metadata() self._namespace_check = _namespace_check + + # Set the metadata being passed.. + metadata(self._current_metadata) + # If the current namespace is False, we disable checking for namespace for this # and all children objects. Not setting namespace_check to False has the consequence # of preventing access to children objects after the namespace changes @@ -394,6 +400,7 @@ def __iter__(self) -> Iterator["MetaflowObject"]: _current_namespace=( self._current_namespace if self._namespace_check else None ), + _current_metadata=self._current_metadata, ) for obj in unfiltered_children ), @@ -509,6 +516,7 @@ def __getitem__(self, id: str) -> "MetaflowObject": _current_namespace=( self._current_namespace if self._namespace_check else None ), + _current_metadata=self._current_metadata, ) else: raise KeyError(id) @@ -540,16 +548,17 @@ def _unpickle_284(self, data): ) def _unpickle_2124(self, data): - if len(data) != 4: + if len(data) != 5: raise MetaflowInternalError( "Unexpected size of array: {}".format(len(data)) ) - pathspec, attempt, ns, namespace_check = data + pathspec, attempt, ns, namespace_check, metadata = data self.__init__( pathspec=pathspec, attempt=attempt, _namespace_check=namespace_check, _current_namespace=ns, + _current_metadata=metadata, ) _UNPICKLE_FUNC = {"2.8.4": _unpickle_284, "2.12.4": _unpickle_2124} @@ -579,6 +588,7 @@ def __setstate__(self, state): attempt=state.get("_attempt", None), _namespace_check=state.get("_namespace_check", False), _current_namespace=None, + _current_metadata=None, ) def __getstate__(self): @@ -601,6 +611,7 @@ def __getstate__(self): self._attempt, self._current_namespace, self._namespace_check, + self._current_metadata, ], } diff --git a/metaflow/runner/metaflow_runner.py b/metaflow/runner/metaflow_runner.py index 161d1706d6..02434d6431 100644 --- a/metaflow/runner/metaflow_runner.py +++ b/metaflow/runner/metaflow_runner.py @@ -6,7 +6,7 @@ from typing import Dict, Iterator, Optional, Tuple -from metaflow import Run, metadata +from metaflow import Run from .utils import handle_timeout, clear_and_set_os_environ from .subprocess_manager import CommandManager, SubprocessManager @@ -284,11 +284,9 @@ def __get_executing_run(self, tfp_runner_attribute, command_obj): # Set the environment variables to what they were before the run executed. clear_and_set_os_environ(self.old_env) - # Set the correct metadata from the runner_attribute file corresponding to this run. - metadata_for_flow = content.get("metadata") - metadata(metadata_for_flow) - - run_object = Run(pathspec, _namespace_check=False) + run_object = Run( + pathspec, _namespace_check=False, _current_metadata=content.get("metadata") + ) return ExecutingRun(self, command_obj, run_object) def run(self, **kwargs) -> ExecutingRun: