Repository: linkedin/shiv

Issue:
large zipapp in a memory limited docker container fails to bootstrap due to using max workers for compiling site packages
I have a microservice zipapp with a lot of dependencies, which fails to bootstrap in a memory limited docker container:
``` traceback
Traceback (most recent call last):
  File "/usr/local/lib/python3.6/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/usr/local/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/microservice/app.pyz/__main__.py", line 3, in <module>
  File "/microservice/app.pyz/_bootstrap/__init__.py", line 106, in bootstrap
  File "/microservice/app.pyz/_bootstrap/__init__.py", line 80, in extract_site_packages
  File "/usr/local/lib/python3.6/compileall.py", line 86, in compile_dir
    success = min(results, default=True)
  File "/usr/local/lib/python3.6/concurrent/futures/process.py", line 366, in _chain_from_iterable_of_lists
    for element in iterable:
  File "/usr/local/lib/python3.6/concurrent/futures/_base.py", line 586, in result_iterator
    yield fs.pop().result()
  File "/usr/local/lib/python3.6/concurrent/futures/_base.py", line 432, in result
    return self.__get_result()
  File "/usr/local/lib/python3.6/concurrent/futures/_base.py", line 384, in __get_result
    raise self._exception
concurrent.futures.process.BrokenProcessPool: A process in the process pool was terminated abruptly while the future was running or pending.
```
I am wondering/asking if it would be possible to limit the number of workers used to compile the site packages in https://github.com/linkedin/shiv/blob/master/src/shiv/bootstrap/__init__.py#L80 by specifying it at zipapp creation time or even completely disable pre-compilation (I guess, at least in my case a lot of modules from some dependencies are not used anyway)?

compiling with max workers(*):
![mprof-zipapp-memory](https://user-images.githubusercontent.com/513316/44577990-bf384580-a792-11e8-9b3c-46fb4a10f266.png)

no pre-compiling:
![mprof-zipapp-memory-no-compile](https://user-images.githubusercontent.com/513316/44578008-c8c1ad80-a792-11e8-8760-e13695d3d984.png)

In my case (zipapp in docker to save image size) compilation does not work well, it takes a lot of memory and extends the bootstrap phase unnecessary.

\* - profiling done with https://pypi.org/project/memory_profiler/

Discussion:

Patch:
diff --git a/setup.cfg b/setup.cfg
index c504615..0e35d76 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -4,7 +4,6 @@ commit = True
 tag = True
 
 [flake8]
-ignore = C901
 max-line-length = 160
 max-complexity = 12
 exclude = docs test/package/setup.py
diff --git a/setup.py b/setup.py
index 43da6db..7fe9d53 100644
--- a/setup.py
+++ b/setup.py
@@ -12,12 +12,20 @@ import venv
 from pathlib import Path
 from setuptools.command import easy_install
 
-requirements = [
+install_requires = [
     'click==6.7',
-    'pip>=9.0.1',
-    'importlib_resources>=0.4',
+    'pip>=9.0.3',
 ]
 
+extras_require = {
+    ":python_version<'3.7'": ["importlib_resources>=0.4"]
+}
+
+if int(setuptools.__version__.split('.')[0]) < 18:
+    extras_require = {}
+    if sys.version_info < (3, 7):
+        install_requires.append('importlib_resources>=0.4')
+
 # The following template and classmethod are copied from
 # fast entry points, Copyright (c) 2016, Aaron Christianson
 # https://github.com/ninjaaron/fast-entry_point
@@ -110,7 +118,8 @@ setuptools.setup(
     url="https://github.com/linkedin/shiv",
     packages=setuptools.find_packages('src'),
     package_dir={'': 'src'},
-    install_requires=requirements,
+    install_requires=install_requires,
+    extras_require=extras_require,
     entry_points={
         'console_scripts': [
             'shiv = shiv.cli:main',
@@ -124,5 +133,6 @@ setuptools.setup(
         'Programming Language :: Python :: 3 :: Only',
         'Programming Language :: Python :: 3',
         'Programming Language :: Python :: 3.6',
+        'Programming Language :: Python :: 3.7',
     ],
 )
diff --git a/src/shiv/bootstrap/__init__.py b/src/shiv/bootstrap/__init__.py
index 02488ba..9eab5ad 100644
--- a/src/shiv/bootstrap/__init__.py
+++ b/src/shiv/bootstrap/__init__.py
@@ -64,7 +64,7 @@ def cache_path(archive, root_dir, build_id):
     return root / f"{name}_{build_id}"
 
 
-def extract_site_packages(archive, target_path):
+def extract_site_packages(archive, target_path, compile_pyc, compile_workers=0):
     """Extract everything in site-packages to a specified path.
 
     :param ZipFile archive: The zipfile object we are bootstrapping from.
@@ -76,8 +76,8 @@ def extract_site_packages(archive, target_path):
         if filename.startswith("site-packages"):
             archive.extract(filename, target_path_tmp)
 
-    # compile pyc with stderr silenced
-    compileall.compile_dir(target_path_tmp, quiet=2, workers=0)
+    if compile_pyc:
+        compileall.compile_dir(target_path_tmp, quiet=2, workers=compile_workers)
 
     # atomic move
     shutil.move(str(target_path_tmp), str(target_path))
@@ -85,7 +85,7 @@ def extract_site_packages(archive, target_path):
 
 def _first_sitedir_index():
     for index, part in enumerate(sys.path):
-        if Path(part).stem == 'site-packages':
+        if Path(part).stem == "site-packages":
             return index
 
 
@@ -103,7 +103,7 @@ def bootstrap():
 
     # determine if first run or forcing extract
     if not site_packages.exists() or env.force_extract:
-        extract_site_packages(archive, site_packages.parent)
+        extract_site_packages(archive, site_packages.parent, env.compile_pyc, env.compile_workers)
 
     # get sys.path's length
     length = len(sys.path)
diff --git a/src/shiv/bootstrap/environment.py b/src/shiv/bootstrap/environment.py
index 5825437..1afc3da 100644
--- a/src/shiv/bootstrap/environment.py
+++ b/src/shiv/bootstrap/environment.py
@@ -2,47 +2,55 @@
 This module contains the ``Environment`` object, which combines settings decided at build time with
 overrides defined at runtime (via environment variables).
 """
-import copy
 import json
 import os
 
 from pathlib import Path
 
 
+def str_bool(v):
+    if not isinstance(v, bool):
+        return str(v).lower() in ("yes", "true", "t", "1")
+    return v
+
+
 class Environment:
     INTERPRETER = "SHIV_INTERPRETER"
     ENTRY_POINT = "SHIV_ENTRY_POINT"
     MODULE = "SHIV_MODULE"
     ROOT = "SHIV_ROOT"
     FORCE_EXTRACT = "SHIV_FORCE_EXTRACT"
+    COMPILE_PYC = "SHIV_COMPILE_PYC"
+    COMPILE_WORKERS = "SHIV_COMPILE_WORKERS"
 
     def __init__(
         self,
         build_id=None,
         entry_point=None,
         always_write_cache=False,
+        compile_pyc=True,
     ):
         self.build_id = build_id
         self.always_write_cache = always_write_cache
 
         # properties
         self._entry_point = entry_point
+        self._compile_pyc = compile_pyc
 
     @classmethod
     def from_json(cls, json_data):
         return Environment(**json.loads(json_data))
 
     def to_json(self):
-        d = copy.copy(self.__dict__)
-        del d["_entry_point"]
-        d["entry_point"] = self.entry_point
-        return json.dumps(d)
+        return json.dumps(
+            # we strip the leading underscores to retain properties (such as _entry_point)
+            {key.lstrip("_"): value for key, value in self.__dict__.items()}
+        )
 
     @property
     def entry_point(self):
         return os.environ.get(
-            self.ENTRY_POINT,
-            os.environ.get(self.MODULE, self._entry_point),
+            self.ENTRY_POINT, os.environ.get(self.MODULE, self._entry_point)
         )
 
     @property
@@ -56,4 +64,15 @@ class Environment:
 
     @property
     def force_extract(self):
-        return bool(os.environ.get(self.FORCE_EXTRACT, self.always_write_cache))
+        return str_bool(os.environ.get(self.FORCE_EXTRACT, self.always_write_cache))
+
+    @property
+    def compile_pyc(self):
+        return str_bool(os.environ.get(self.COMPILE_PYC, self._compile_pyc))
+
+    @property
+    def compile_workers(self):
+        try:
+            return int(os.environ.get(self.COMPILE_WORKERS, 0))
+        except ValueError:
+            return 0
diff --git a/src/shiv/cli.py b/src/shiv/cli.py
index d375559..8bd16a3 100644
--- a/src/shiv/cli.py
+++ b/src/shiv/cli.py
@@ -1,8 +1,12 @@
-import importlib_resources  # type: ignore
 import shutil
 import sys
 import uuid
 
+try:
+    import importlib.resources as importlib_resources  # type: ignore
+except ImportError:
+    import importlib_resources  # type: ignore
+
 from configparser import ConfigParser
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -71,6 +75,11 @@ def copy_bootstrap(bootstrap_target: Path) -> None:
     default=True,
     help="Whether or not to compress your zip.",
 )
+@click.option(
+    "--compile-pyc/--no-compile-pyc",
+    default=True,
+    help="Whether or not to compile pyc files during initial bootstrap.",
+)
 @click.argument("pip_args", nargs=-1, type=click.UNPROCESSED)
 def main(
     output_file: str,
@@ -78,6 +87,7 @@ def main(
     console_script: Optional[str],
     python: Optional[str],
     compressed: bool,
+    compile_pyc: bool,
     pip_args: List[str],
 ) -> None:
     """
@@ -110,9 +120,7 @@ def main(
         site_packages.mkdir(parents=True, exist_ok=True)
 
         # install deps into staged site-packages
-        pip.install(
-            ["--target", str(site_packages)] + list(pip_args),
-        )
+        pip.install(["--target", str(site_packages)] + list(pip_args))
 
         # if entry_point is a console script, get the callable
         if entry_point is None and console_script is not None:
@@ -123,8 +131,7 @@ def main(
 
         # create runtime environment metadata
         env = Environment(
-            build_id=str(uuid.uuid4()),
-            entry_point=entry_point,
+            build_id=str(uuid.uuid4()), entry_point=entry_point, compile_pyc=compile_pyc
         )
 
         Path(working_path, "environment.json").write_text(env.to_json())
diff --git a/tox.ini b/tox.ini
index 94f264d..421b619 100644
--- a/tox.ini
+++ b/tox.ini
@@ -1,5 +1,5 @@
 [tox]
-envlist = py36
+envlist = py36, py37
 
 [testenv]
 passenv = TRAVIS TRAVIS_*


Test patch:
diff --git a/test/test_bootstrap.py b/test/test_bootstrap.py
index 1f77034..7e33357 100644
--- a/test/test_bootstrap.py
+++ b/test/test_bootstrap.py
@@ -89,6 +89,18 @@ class TestEnvironment:
         with env_var('SHIV_FORCE_EXTRACT', '1'):
             assert env.force_extract is True
 
+        assert env.compile_pyc is True
+        with env_var("SHIV_COMPILE_PYC", "False"):
+            assert env.compile_pyc is False
+
+        assert env.compile_workers == 0
+        with env_var("SHIV_COMPILE_WORKERS", "1"):
+            assert env.compile_workers == 1
+
+        # ensure that non-digits are ignored
+        with env_var("SHIV_COMPILE_WORKERS", "one bazillion"):
+            assert env.compile_workers == 0
+
     def test_serialize(self):
         env = Environment()
         env_as_json = env.to_json()
diff --git a/test/test_builder.py b/test/test_builder.py
index 2eac558..dabe641 100644
--- a/test/test_builder.py
+++ b/test/test_builder.py
@@ -25,12 +25,14 @@ def tmp_write_prefix(interpreter):
 
 
 class TestBuilder:
+
     @pytest.mark.parametrize(
-        'interpreter,expected', [
-            ('/usr/bin/python', b'#!/usr/bin/python\n'),
-            ('/usr/bin/env python', b'#!/usr/bin/env python\n'),
-            ('/some/other/path/python -sE', b'#!/some/other/path/python -sE\n'),
-        ]
+        "interpreter,expected",
+        [
+            ("/usr/bin/python", b"#!/usr/bin/python\n"),
+            ("/usr/bin/env python", b"#!/usr/bin/env python\n"),
+            ("/some/other/path/python -sE", b"#!/some/other/path/python -sE\n"),
+        ],
     )
     def test_file_prefix(self, interpreter, expected):
         assert tmp_write_prefix(interpreter) == expected
@@ -41,24 +43,25 @@ class TestBuilder:
 
     def test_create_archive(self, sp):
         with tempfile.TemporaryDirectory() as tmpdir:
-            target = Path(tmpdir, 'test.zip')
+            target = Path(tmpdir, "test.zip")
 
             # create an archive
-            create_archive(sp, target, sys.executable, 'code:interact')
+            create_archive(sp, target, sys.executable, "code:interact")
 
             # create one again (to ensure we overwrite)
-            create_archive(sp, target, sys.executable, 'code:interact')
+            create_archive(sp, target, sys.executable, "code:interact")
 
             assert zipfile.is_zipfile(str(target))
 
             with pytest.raises(ZipAppError):
-                create_archive(sp, target, sys.executable, 'alsjdbas,,,')
+                create_archive(sp, target, sys.executable, "alsjdbas,,,")
 
-    @pytest.mark.skipif(os.name == 'nt',
-                        reason='windows has no concept of execute permissions')
+    @pytest.mark.skipif(
+        os.name == "nt", reason="windows has no concept of execute permissions"
+    )
     def test_archive_permissions(self, sp):
         with tempfile.TemporaryDirectory() as tmpdir:
-            target = Path(tmpdir, 'test.zip')
-            create_archive(sp, target, sys.executable, 'code:interact')
+            target = Path(tmpdir, "test.zip")
+            create_archive(sp, target, sys.executable, "code:interact")
 
             assert target.stat().st_mode & UGOX == UGOX
diff --git a/test/test_cli.py b/test/test_cli.py
index 2d24138..7cc46cd 100644
--- a/test/test_cli.py
+++ b/test/test_cli.py
@@ -45,12 +45,13 @@ class TestCLI:
         # assert we got the correct reason
         assert strip_header(result.output) == DISALLOWED_PIP_ARGS.format(arg=arg, reason=reason)
 
-    def test_hello_world(self, tmpdir, runner, package_location):
+    @pytest.mark.parametrize('compile_option', ["--compile-pyc", "--no-compile-pyc"])
+    def test_hello_world(self, tmpdir, runner, package_location, compile_option):
 
         with tempfile.TemporaryDirectory(dir=tmpdir) as tmpdir:
             output_file = Path(tmpdir, 'test.pyz')
 
-            result = runner(['-e', 'hello:main', '-o', str(output_file), str(package_location)])
+            result = runner(['-e', 'hello:main', '-o', str(output_file), str(package_location), compile_option])
 
             # check that the command successfully completed
             assert result.exit_code == 0