From 5683593c2676212d81ffa9b239dce9fe66215bed Mon Sep 17 00:00:00 2001 From: Sakari Ikonen Date: Mon, 3 Mar 2025 15:36:29 +0200 Subject: [PATCH 1/5] add pip patcher and rename pip.py to avoid name collision with pip module --- metaflow/plugins/pypi/conda_environment.py | 2 +- metaflow/plugins/pypi/pip_patcher.py | 15 +++++++++++++++ metaflow/plugins/pypi/{pip.py => pip_resolver.py} | 8 ++++++-- 3 files changed, 22 insertions(+), 3 deletions(-) create mode 100644 metaflow/plugins/pypi/pip_patcher.py rename metaflow/plugins/pypi/{pip.py => pip_resolver.py} (97%) diff --git a/metaflow/plugins/pypi/conda_environment.py b/metaflow/plugins/pypi/conda_environment.py index 292d1fb7535..6c5973b1270 100644 --- a/metaflow/plugins/pypi/conda_environment.py +++ b/metaflow/plugins/pypi/conda_environment.py @@ -55,7 +55,7 @@ def validate_environment(self, logger, datastore_type): # Initialize necessary virtual environments for all Metaflow tasks. # Use Micromamba for solving conda packages and Pip for solving pypi packages. from .micromamba import Micromamba - from .pip import Pip + from .pip_resolver import Pip print_lock = threading.Lock() diff --git a/metaflow/plugins/pypi/pip_patcher.py b/metaflow/plugins/pypi/pip_patcher.py new file mode 100644 index 00000000000..dda28cc9754 --- /dev/null +++ b/metaflow/plugins/pypi/pip_patcher.py @@ -0,0 +1,15 @@ +from unittest.mock import patch +import os +import sys + + +@patch("platform.machine", lambda: os.environ.get("PIP_PATCH_MACHINE", "x86_64")) +@patch("platform.system", lambda: os.environ.get("PIP_PATCH_SYSTEM", "Linux")) +def _main(args): + from pip import main + + main(args) + + +if __name__ == "__main__": + _main(sys.argv[1:]) diff --git a/metaflow/plugins/pypi/pip.py b/metaflow/plugins/pypi/pip_resolver.py similarity index 97% rename from metaflow/plugins/pypi/pip.py rename to metaflow/plugins/pypi/pip_resolver.py index 5aeb03ec5f4..3c24ee76e93 100644 --- a/metaflow/plugins/pypi/pip.py +++ b/metaflow/plugins/pypi/pip_resolver.py @@ -102,7 +102,10 @@ def solve(self, id_, packages, python, platform): else: cmd.append(f"{package}=={version}") try: - self._call(prefix, cmd) + env = {} + if platform == "linux-64": + env = {"PIP_PATCH_SYSTEM": "Linux", "PIP_PATCH_MACHINE": "x86_64"} + self._call(prefix, cmd, env) except PipPackageNotFound as ex: # pretty print package errors raise PipException( @@ -306,7 +309,8 @@ def _call(self, prefix, args, env=None, isolated=True): "run", "--prefix", prefix, - "pip3", + "python", + os.path.join(os.path.dirname(__file__), "pip_patcher.py"), "--disable-pip-version-check", "--no-color", ] From 345ce3e34b879b4512b34894203f27d4595c8b47 Mon Sep 17 00:00:00 2001 From: Sakari Ikonen Date: Mon, 3 Mar 2025 15:42:04 +0200 Subject: [PATCH 2/5] adding comment to pip patcher script explaining use case --- metaflow/plugins/pypi/pip_patcher.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/metaflow/plugins/pypi/pip_patcher.py b/metaflow/plugins/pypi/pip_patcher.py index dda28cc9754..e43dc380974 100644 --- a/metaflow/plugins/pypi/pip_patcher.py +++ b/metaflow/plugins/pypi/pip_patcher.py @@ -3,6 +3,9 @@ import sys +# Because Pip does not offer a direct way to set target platform_system and platform_machine values for resolving packages transitive dependencies, we need to instead +# manually patch the correct target architecture values for pip to be able to resolve the whole dependency tree successfully. +# This is necessary for packages that have conditional dependencies dependent on machine/system, e.g. Torch @patch("platform.machine", lambda: os.environ.get("PIP_PATCH_MACHINE", "x86_64")) @patch("platform.system", lambda: os.environ.get("PIP_PATCH_SYSTEM", "Linux")) def _main(args): From 1f125a023aeb8c150be2cd0968f064c9af628842 Mon Sep 17 00:00:00 2001 From: Sakari Ikonen Date: Mon, 3 Mar 2025 16:26:59 +0200 Subject: [PATCH 3/5] fix defaults for patching --- metaflow/plugins/pypi/pip_patcher.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/metaflow/plugins/pypi/pip_patcher.py b/metaflow/plugins/pypi/pip_patcher.py index e43dc380974..b9e1a6c13b1 100644 --- a/metaflow/plugins/pypi/pip_patcher.py +++ b/metaflow/plugins/pypi/pip_patcher.py @@ -6,8 +6,12 @@ # Because Pip does not offer a direct way to set target platform_system and platform_machine values for resolving packages transitive dependencies, we need to instead # manually patch the correct target architecture values for pip to be able to resolve the whole dependency tree successfully. # This is necessary for packages that have conditional dependencies dependent on machine/system, e.g. Torch -@patch("platform.machine", lambda: os.environ.get("PIP_PATCH_MACHINE", "x86_64")) -@patch("platform.system", lambda: os.environ.get("PIP_PATCH_SYSTEM", "Linux")) +@patch( + "platform.machine", lambda: os.environ.get("PIP_PATCH_MACHINE", os.uname().machine) +) +@patch( + "platform.system", lambda: os.environ.get("PIP_PATCH_SYSTEM", os.uname().sysname) +) def _main(args): from pip import main From 7dee484edd8b6b2e9b0652690851e398eaef7cc7 Mon Sep 17 00:00:00 2001 From: Sakari Ikonen Date: Mon, 3 Mar 2025 16:46:07 +0200 Subject: [PATCH 4/5] exit code fix for patcher script --- metaflow/plugins/pypi/pip_patcher.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/metaflow/plugins/pypi/pip_patcher.py b/metaflow/plugins/pypi/pip_patcher.py index b9e1a6c13b1..c4ac7ca2852 100644 --- a/metaflow/plugins/pypi/pip_patcher.py +++ b/metaflow/plugins/pypi/pip_patcher.py @@ -15,7 +15,8 @@ def _main(args): from pip import main - main(args) + exitcode = main(args) + sys.exit(exitcode) if __name__ == "__main__": From d787ea8a88da4eaeaeb41f10eec199515dddd4db Mon Sep 17 00:00:00 2001 From: Sakari Ikonen Date: Mon, 3 Mar 2025 16:49:09 +0200 Subject: [PATCH 5/5] note about pip deprecation --- metaflow/plugins/pypi/pip_patcher.py | 1 + 1 file changed, 1 insertion(+) diff --git a/metaflow/plugins/pypi/pip_patcher.py b/metaflow/plugins/pypi/pip_patcher.py index c4ac7ca2852..11cf2c314e9 100644 --- a/metaflow/plugins/pypi/pip_patcher.py +++ b/metaflow/plugins/pypi/pip_patcher.py @@ -13,6 +13,7 @@ "platform.system", lambda: os.environ.get("PIP_PATCH_SYSTEM", os.uname().sysname) ) def _main(args): + # TODO: Pip has deprecated using script wrappers for the cli. this will break in the future, and make patching the sys internals much harder from pip import main exitcode = main(args)