Skip to content

Commit

Permalink
call .freeze() instead of .build() in ipex launcher and add PEX_IPEX_…
Browse files Browse the repository at this point in the history
…SKIP_EXECUTION env var (#10530)

### Problem

The "ipex" execution method for pex files from #8793 wastes a lot of time zipping up the 3rdparty modules it just hydrated when it is first executed. This is unnecessary, as pex files can be executed as unzipped directories.

### Solution

- Call `.freeze()` instead of `.build()` in `ipex_launcher.py`.
- Add the capability to skip actually executing the pex after hydration if `PEX_IPEX_SKIP_EXECUTION` is set in the subprocess environment.

### Result

On `1.25.x-twtr`:
```bash
> ./pants binary.py --generate-ipex examples/src/python/example/tensorflow_custom_op:show-tf-version && \
  time PEX_PROFILE_FILENAME=ipex.prof ./dist/show-tf-version.ipex
...
41.51s user 12.65s system 86% cpu 1:02.38 total
```

On this PR:
```bash
> ./pants binary.py --generate-ipex examples/src/python/example/tensorflow_custom_op:show-tf-version && \
  time PEX_PROFILE_FILENAME=ipex.prof ./dist/show-tf-version.ipex
...
16.29s user 8.48s system 77% cpu 32.000 total
```
  • Loading branch information
cosmicexplorer authored Aug 6, 2020
1 parent bf9afc2 commit de31c57
Show file tree
Hide file tree
Showing 3 changed files with 36 additions and 18 deletions.
28 changes: 16 additions & 12 deletions src/python/pants/backend/python/subsystems/ipex/ipex_launcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ def modify_pex_info(pex_info, **kwargs):
return PexInfo.from_json(json.dumps(new_info))


def _hydrate_pex_file(self, hydrated_pex_file):
def _hydrate_pex_file(self, hydrated_pex_dir):
# We extract source files into a temporary directory before creating the pex.
td = tempfile.mkdtemp()

Expand Down Expand Up @@ -86,24 +86,28 @@ def _hydrate_pex_file(self, hydrated_pex_file):
for resolved_dist in resolved_distributions:
bootstrap_builder.add_distribution(resolved_dist.distribution)

bootstrap_builder.build(hydrated_pex_file, bytecode_compile=False)
# We call .freeze() instead of .build() here in order to avoid zipping up all the large 3rdparty
# modules we just added to the chroot.
# NB: Bytecode compilation can take an extremely long time for large 3rdparty modules.
bootstrap_builder.freeze(bytecode_compile=False)
os.rename(src=bootstrap_builder.path(), dst=hydrated_pex_dir)


def main(self):
filename_base, ext = os.path.splitext(self)

# If the ipex (this pex) is already named '.pex', ensure the output filename doesn't collide by
# inserting an intermediate '.ipex'!
if ext == ".pex":
hydrated_pex_file = "{filename_base}.ipex.pex".format(filename_base=filename_base)
else:
hydrated_pex_file = "{filename_base}.pex".format(filename_base=filename_base)
# Incorporate the code hash into the output unpacked pex directory in order to:
# (a) avoid execing an out of date hydrated ipex,
# (b) avoid collisions with other similarly-named (i)pex files in the same directory!
code_hash = PexInfo.from_pex(self).code_hash
hydrated_pex_dir = "{}-{}{}".format(filename_base, code_hash, ext)

if not os.path.exists(hydrated_pex_file):
_log("Hydrating {} to {}...".format(self, hydrated_pex_file))
_hydrate_pex_file(self, hydrated_pex_file)
if not os.path.exists(hydrated_pex_dir):
_log("Hydrating {} to {}...".format(self, hydrated_pex_dir))
_hydrate_pex_file(self, hydrated_pex_dir)

os.execv(sys.executable, [sys.executable, hydrated_pex_file] + sys.argv[1:])
if "PEX_IPEX_SKIP_EXECUTION" not in os.environ:
os.execv(sys.executable, [sys.executable, hydrated_pex_dir] + sys.argv[1:])


if __name__ == "__main__":
Expand Down
2 changes: 2 additions & 0 deletions tests/python/pants_test/backend/python/tasks/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -140,9 +140,11 @@ python_tests(
name = 'python_binary_create',
sources = ['test_python_binary_create.py'],
dependencies = [
'3rdparty/python:pex',
'src/python/pants/backend/python/tasks',
'src/python/pants/base:run_info',
'src/python/pants/build_graph',
'src/python/pants/util:contextutil',
':python_task_test_base',
],
tags = {"partially_type_checked"},
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,14 @@
from textwrap import dedent

from colors import blue
from pex.pex_info import PexInfo

from pants.backend.python.tasks.gather_sources import GatherSources
from pants.backend.python.tasks.python_binary_create import PythonBinaryCreate
from pants.backend.python.tasks.select_interpreter import SelectInterpreter
from pants.base.run_info import RunInfo
from pants.build_graph.register import build_file_aliases as register_core
from pants.util.contextutil import environment_as
from pants_test.backend.python.tasks.python_task_test_base import PythonTaskTestBase


Expand Down Expand Up @@ -152,7 +154,7 @@ def main():
binary, expected_output="Hello World!\n", expected_shebang=b"#!/usr/bin/env python2\n"
)

def test_generate_ipex_ansicolors(self):
def _assert_generate_ipex_ansicolors(self, expected_output: str):
self.create_python_requirement_library(
"3rdparty/ipex", "ansicolors", requirements=["ansicolors"]
)
Expand All @@ -176,11 +178,21 @@ def test_generate_ipex_ansicolors(self):
self.set_options(generate_ipex=True)
dist_dir = os.path.join(self.build_root, "dist")

self._assert_pex(
binary, expected_output=blue("i just lazy-loaded the ansicolors dependency!") + "\n"
)
self._assert_pex(binary, expected_output=expected_output)

dehydrated_ipex_file = os.path.join(dist_dir, "bin.ipex")
assert os.path.isfile(dehydrated_ipex_file)
hydrated_pex_output_file = os.path.join(dist_dir, "bin.pex")
assert os.path.isfile(hydrated_pex_output_file)

filename_base, ext = os.path.splitext(dehydrated_ipex_file)
code_hash = PexInfo.from_pex(dehydrated_ipex_file).code_hash
hydrated_pex_dir = f"{filename_base}-{code_hash}{ext}"
assert os.path.isdir(hydrated_pex_dir)

def test_generate_ipex_ansicolors(self):
self._assert_generate_ipex_ansicolors(
blue("i just lazy-loaded the ansicolors dependency!") + "\n"
)

def test_generate_ipex_skip_execution(self):
with environment_as(PEX_IPEX_SKIP_EXECUTION="true"):
self._assert_generate_ipex_ansicolors(expected_output="")

0 comments on commit de31c57

Please sign in to comment.