From 64860b487578c6cfac0855cfb45aea89ee2457af Mon Sep 17 00:00:00 2001 From: savin Date: Thu, 28 Sep 2023 08:03:20 -0700 Subject: [PATCH] fix issue --- metaflow/plugins/pypi/conda_decorator.py | 24 ++++++++++++++++------ metaflow/plugins/pypi/conda_environment.py | 12 +++++++---- metaflow/plugins/pypi/pypi_decorator.py | 22 ++++++++------------ 3 files changed, 34 insertions(+), 24 deletions(-) diff --git a/metaflow/plugins/pypi/conda_decorator.py b/metaflow/plugins/pypi/conda_decorator.py index 04c8cfbadac..5bd58d69916 100644 --- a/metaflow/plugins/pypi/conda_decorator.py +++ b/metaflow/plugins/pypi/conda_decorator.py @@ -79,6 +79,8 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge if self.attributes["disabled"] is not None else super_attributes["disabled"] ) + + # Set default for `disabled` argument. if not self.attributes["disabled"]: self.attributes["disabled"] = False # Set Python interpreter to user's Python if necessary. @@ -86,12 +88,12 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge self.attributes["python"] = platform.python_version() # CPython! # Take care of `disabled` argument. - if self.attributes["disabled"]: - _step = next(step for step in self.flow if step.name == self.step) - _step.decorators[:] = [ - deco for deco in _step.decorators if deco.name not in ["conda", "pypi"] - ] - del self.attributes["disabled"] + # if self.attributes["disabled"]: + # _step = next(step for step in self.flow if step.name == self.step) + # # _step.decorators[:] = [ + # # deco for deco in _step.decorators if deco.name not in ["conda", "pypi"] + # # ] + # del self.attributes["disabled"] # @conda uses a conda environment to create a virtual environment. # The conda environment can be created through micromamba. @@ -126,6 +128,8 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge environment.set_local_root(LocalStorage.get_datastore_root_from_config(logger)) def runtime_init(self, flow, graph, package, run_id): + if self.attributes["disabled"]: + return # Create a symlink to metaflow installed outside the virtual environment. self.metaflow_dir = tempfile.TemporaryDirectory(dir="/tmp") os.symlink( @@ -169,6 +173,8 @@ def runtime_init(self, flow, graph, package, run_id): def runtime_task_created( self, task_datastore, task_id, split_index, input_paths, is_cloned, ubf_context ): + if self.attributes["disabled"]: + return self.interpreter = ( self.environment.interpreter(self.step) if not any( @@ -194,6 +200,8 @@ def task_pre_step( ubf_context, inputs, ): + if self.attributes["disabled"]: + return # Add Python interpreter's parent to the path to ensure that any non-pythonic # dependencies in the virtual environment are visible to the user code. # sys.executable points to the Python interpreter in the virtual environment @@ -212,6 +220,8 @@ def task_pre_step( def runtime_step_cli( self, cli_args, retry_count, max_user_code_retries, ubf_context ): + if self.attributes["disabled"]: + return # TODO: Support unbounded for-each # TODO: Check what happens when PYTHONPATH is defined via @environment # Ensure local installation of Metaflow is visible to user code @@ -229,6 +239,8 @@ def runtime_step_cli( cli_args.entrypoint[0] = self.interpreter def runtime_finished(self, exception): + if self.attributes["disabled"]: + return self.metaflow_dir.cleanup() diff --git a/metaflow/plugins/pypi/conda_environment.py b/metaflow/plugins/pypi/conda_environment.py index 4fc64707490..0de1df3c38a 100644 --- a/metaflow/plugins/pypi/conda_environment.py +++ b/metaflow/plugins/pypi/conda_environment.py @@ -183,10 +183,14 @@ def get_environment(self, step): for decorator in step.decorators: # @conda decorator is guaranteed to exist thanks to self.decospecs if decorator.name in ["conda", "pypi"]: - environment[decorator.name] = dict(decorator.attributes) - # environment can be empty for @conda/@pypi(disabled=True) - if not environment: - return {} + # handle @conda/@pypi(disabled=True) + if decorator.attributes["disabled"]: + return {} + environment[decorator.name] = { + k: decorator.attributes[k] + for k in decorator.attributes + if k != "disabled" + } # TODO: Support dependencies for `--metadata`. # TODO: Introduce support for `--telemetry` as a follow up. # Certain packages are required for metaflow runtime to function correctly. diff --git a/metaflow/plugins/pypi/pypi_decorator.py b/metaflow/plugins/pypi/pypi_decorator.py index 9c0b8efd9aa..65d43790627 100644 --- a/metaflow/plugins/pypi/pypi_decorator.py +++ b/metaflow/plugins/pypi/pypi_decorator.py @@ -48,13 +48,14 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge else super_attributes["disabled"] ) - # Take care of `disabled` argument. - if self.attributes["disabled"]: - _step = next(step for step in self.flow if step.name == self.step) - _step.decorators[:] = [ - deco for deco in _step.decorators if deco.name not in ["conda", "pypi"] - ] - del self.attributes["disabled"] + # # Take care of `disabled` argument. + # if self.attributes["disabled"]: + # _step = next(step for step in self.flow if step.name == self.step) + # # _step.decorators[:] = [ + # # deco for deco in _step.decorators if deco.name not in ["conda", "pypi"] + # # ] + + # del self.attributes["disabled"] # At the moment, @pypi uses a conda environment as a virtual environment. This # is to ensure that we can have a dedicated Python interpreter within the @@ -88,13 +89,6 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge ) ) - # TODO: This code snippet can be done away with by altering the constructor of - # MetaflowEnvironment. A good first-task exercise. - # Avoid circular import - from metaflow.plugins.datastores.local_storage import LocalStorage - - environment.set_local_root(LocalStorage.get_datastore_root_from_config(logger)) - class PyPIFlowDecorator(FlowDecorator): """