From 9c31d3b76ee4ee48924b56a9b057d6984c00f99f Mon Sep 17 00:00:00 2001 From: Sakari Ikonen Date: Tue, 22 Aug 2023 12:04:12 +0300 Subject: [PATCH] prevent tracing module loading during conda bootstrapping process. --- metaflow/metaflow_config.py | 4 ++++ metaflow/plugins/pypi/conda_environment.py | 6 +++++- metaflow/tracing/__init__.py | 5 ++++- 3 files changed, 13 insertions(+), 2 deletions(-) diff --git a/metaflow/metaflow_config.py b/metaflow/metaflow_config.py index 09ea9529ea..a0f04e0850 100644 --- a/metaflow/metaflow_config.py +++ b/metaflow/metaflow_config.py @@ -402,6 +402,10 @@ OTEL_ENDPOINT = from_conf("OTEL_ENDPOINT") ZIPKIN_ENDPOINT = from_conf("ZIPKIN_ENDPOINT") CONSOLE_TRACE_ENABLED = from_conf("CONSOLE_TRACE_ENABLED", False) +# internal env used for preventing the tracing module from loading during Conda bootstrapping. +BOOTSTRAPPING_CONDA_ENVIRONMENT = bool( + os.environ.get("BOOTSTRAPPING_CONDA_ENVIRONMENT", False) +) # MAX_ATTEMPTS is the maximum number of attempts, including the first # task, retries, and the final fallback task and its retries. diff --git a/metaflow/plugins/pypi/conda_environment.py b/metaflow/plugins/pypi/conda_environment.py index c497d79fce..6d21d5c9c2 100644 --- a/metaflow/plugins/pypi/conda_environment.py +++ b/metaflow/plugins/pypi/conda_environment.py @@ -320,7 +320,11 @@ def bootstrap_commands(self, step_name, datastore_type): if id_: return [ "echo 'Bootstrapping virtual environment...'", - 'python -m metaflow.plugins.pypi.bootstrap "%s" %s "%s" linux-64' + # We have to prevent the tracing module from loading, + # as the bootstrapping process uses the internal S3 client which would fail to import tracing + # due to the required dependencies being bundled into the conda environment, + # which is yet to be initialized at this point. + 'BOOTSTRAPPING_CONDA_ENVIRONMENT=True python -m metaflow.plugins.pypi.bootstrap "%s" %s "%s" linux-64' % (self.flow.name, id_, self.datastore_type), "echo 'Environment bootstrapped.'", "export PATH=$PATH:$(pwd)/micromamba", diff --git a/metaflow/tracing/__init__.py b/metaflow/tracing/__init__.py index 72da6682d0..aadb7793c6 100644 --- a/metaflow/tracing/__init__.py +++ b/metaflow/tracing/__init__.py @@ -2,6 +2,7 @@ OTEL_ENDPOINT, ZIPKIN_ENDPOINT, CONSOLE_TRACE_ENABLED, + BOOTSTRAPPING_CONDA_ENVIRONMENT, ) from functools import wraps import contextlib @@ -49,7 +50,9 @@ def wrapper_func(*args, **kwargs): return wrapper_func -if CONSOLE_TRACE_ENABLED or OTEL_ENDPOINT or ZIPKIN_ENDPOINT: +if not BOOTSTRAPPING_CONDA_ENVIRONMENT and ( + CONSOLE_TRACE_ENABLED or OTEL_ENDPOINT or ZIPKIN_ENDPOINT +): try: from .tracing_modules import *