Repository: pex-tool/pex

Issue:
Pex fails when used in conjunction with Pytorch
The following code:

```
import torch
import torch.utils.data

class Test(torch.utils.data.Dataset):
    def __len__(self):
        return 10

    def __getitem__(self, idx):
        return torch.randn(5, 5)


if __name__ == '__main__':
  n_workers = 4
  data = Test()
  train_loader = torch.utils.data.DataLoader(data, batch_size=2, num_workers=n_workers)

  for item in train_loader:
    print(item)
    break
```

Runs fine with a python interpreter with pytorch and numpy installed:

```
$ python test.py

(0 ,.,.) =
 -0.4520 -1.2118  0.3244  0.7276 -0.8220
  1.9171  0.9921  0.5206  0.2886 -0.3736
  1.3648  1.0917  0.1983  0.0402 -1.2338
  0.3155 -0.6425  0.3849 -0.4376  1.2371
  0.5352 -0.3482  1.4171  0.4569 -0.3687

(1 ,.,.) =
  0.1180 -1.0528  0.6809  0.2318 -2.2028
  0.1399  0.5150 -0.5507 -0.7639 -0.8540
  1.9070  0.8260  0.8660 -0.1029 -0.1823
 -2.3825 -2.0044  1.1457 -0.9469 -0.4882
 -0.3794  0.6706 -0.8183 -0.4942 -1.7503
[torch.FloatTensor of size 2x5x5]
```

But fails with pex:

```
$ pex torch numpy -- test.py
Traceback (most recent call last):
  File "/private/var/folders/b5/hgpy98310bg_g3hfqngl759m0000gn/T/tmpxYzeas/.bootstrap/_pex/pex.py", line 367, in execute
    self._wrap_coverage(self._wrap_profiling, self._execute)
  File "/private/var/folders/b5/hgpy98310bg_g3hfqngl759m0000gn/T/tmpxYzeas/.bootstrap/_pex/pex.py", line 293, in _wrap_coverage
    runner(*args)
  File "/private/var/folders/b5/hgpy98310bg_g3hfqngl759m0000gn/T/tmpxYzeas/.bootstrap/_pex/pex.py", line 325, in _wrap_profiling
    runner(*args)
  File "/private/var/folders/b5/hgpy98310bg_g3hfqngl759m0000gn/T/tmpxYzeas/.bootstrap/_pex/pex.py", line 413, in _execute
    return self.execute_interpreter()
  File "/private/var/folders/b5/hgpy98310bg_g3hfqngl759m0000gn/T/tmpxYzeas/.bootstrap/_pex/pex.py", line 423, in execute_interpreter
    self.execute_content(name, content)
  File "/private/var/folders/b5/hgpy98310bg_g3hfqngl759m0000gn/T/tmpxYzeas/.bootstrap/_pex/pex.py", line 453, in execute_content
    exec_function(ast, globals())
  File "<exec_function>", line 4, in exec_function
  File "test_dl.py", line 18, in <module>
    for item in train_loader:
  File "/private/var/folders/b5/hgpy98310bg_g3hfqngl759m0000gn/T/tmpxYzeas/.deps/torch-0.4.0-cp27-none-macosx_10_6_x86_64.whl/torch/utils/data/dataloader.py", line 286, in __next__
    return self._process_next_batch(batch)
  File "/private/var/folders/b5/hgpy98310bg_g3hfqngl759m0000gn/T/tmpxYzeas/.deps/torch-0.4.0-cp27-none-macosx_10_6_x86_64.whl/torch/utils/data/dataloader.py", line 307, in _process_next_batch
    raise batch.exc_type(batch.exc_msg)
RuntimeError: Traceback (most recent call last):
  File "/private/var/folders/b5/hgpy98310bg_g3hfqngl759m0000gn/T/tmpxYzeas/.deps/torch-0.4.0-cp27-none-macosx_10_6_x86_64.whl/torch/utils/data/dataloader.py", line 57, in _worker_loop
    samples = collate_fn([dataset[i] for i in batch_indices])
  File "/private/var/folders/b5/hgpy98310bg_g3hfqngl759m0000gn/T/tmpxYzeas/.deps/torch-0.4.0-cp27-none-macosx_10_6_x86_64.whl/torch/utils/data/dataloader.py", line 113, in default_collate
    storage = batch[0].storage()._new_shared(numel)
  File "/private/var/folders/b5/hgpy98310bg_g3hfqngl759m0000gn/T/tmpxYzeas/.deps/torch-0.4.0-cp27-none-macosx_10_6_x86_64.whl/torch/storage.py", line 114, in _new_shared
    return cls._new_using_filename(size)
RuntimeError: error executing torch_shm_manager at "/private/var/folders/b5/hgpy98310bg_g3hfqngl759m0000gn/T/tmpxYzeas/.deps/torch-0.4.0-cp27-none-macosx_10_6_x86_64.whl/torch/lib/torch_shm_manager" at /Users/soumith/code/builder/wheel/pytorch-src/torch/lib/libshm/core.cpp:125
```

Relevant info from the PyTorch guys on what they're doing is here: https://pytorch.org/docs/stable/multiprocessing.html#file-system-file-system



Discussion:
Additional info: the above code was run on Mac and breaks. On Linux PyTorch uses a `file_descriptor` based sharing system which does work well.

However forced setting of the sharing strategy:
```
torch.multiprocessing.set_sharing_strategy('file_system')
```
makes it also fail on linux.

So the conclusion is that something in the `file_system` strategy of shared memory management is incompatible with pex.

Adding `--not-zip-safe` doesn't fix the problem.
OK, root cause is `torch_shm_manager` is not executable in the pex environment. When I build with --not-zip-safe and `chmod +x /home/jsirois/.pex/install/torch-0.4.0-cp36-cp36m-manylinux1_x86_64.whl.40da15562268d7f03a1d135a85e2cc9ac9a46e8d/torch-0.4.0-cp36-cp36m-manylinux1_x86_64.whl/torch/lib/torch_shm_manager` after running once (so that the torch.pex and its internal wheels are exploded to the filesystem), things work. I ran a quick experiment using python's `ZipFile.extractall` and reproduced permissions not being preserved (using the `unzip` binary `torch_shm_manager` is executable after extraction). Here is the upstream bug: https://bugs.python.org/issue15795

I think we can work around this by subclassing ZipFile :/

And WOW - torch is a ~400MB wheel. That's getting to be the size of a docker image.

Resolving patch:
diff --git a/pex/archiver.py b/pex/archiver.py
--- a/pex/archiver.py
+++ b/pex/archiver.py
@@ -6,7 +6,7 @@
 import tarfile
 import zipfile
 
-from .common import safe_mkdtemp
+from .common import PermPreservingZipFile, safe_mkdtemp
 
 
 class Archiver(object):
@@ -19,7 +19,7 @@ class InvalidArchive(Error): pass
     '.tar.gz': (tarfile.TarFile.open, tarfile.ReadError),
     '.tar.bz2': (tarfile.TarFile.open, tarfile.ReadError),
     '.tgz': (tarfile.TarFile.open, tarfile.ReadError),
-    '.zip': (zipfile.ZipFile, zipfile.BadZipfile)
+    '.zip': (PermPreservingZipFile, zipfile.BadZipfile)
   }
 
   @classmethod
diff --git a/pex/common.py b/pex/common.py
--- a/pex/common.py
+++ b/pex/common.py
@@ -78,10 +78,27 @@ def teardown(self):
 _MKDTEMP_SINGLETON = MktempTeardownRegistry()
 
 
+class PermPreservingZipFile(zipfile.ZipFile):
+  """A ZipFile that works around https://bugs.python.org/issue15795"""
+
+  def _extract_member(self, member, targetpath, pwd):
+    result = super(PermPreservingZipFile, self)._extract_member(member, targetpath, pwd)
+    info = member if isinstance(member, zipfile.ZipInfo) else self.getinfo(member)
+    self._chmod(info, result)
+    return result
+
+  def _chmod(self, info, path):
+    # This magic works to extract perm bits from the 32 bit external file attributes field for
+    # unix-created zip files, for the layout, see:
+    #   https://www.forensicswiki.org/wiki/ZIP#External_file_attributes
+    attr = info.external_attr >> 16
+    os.chmod(path, attr)
+
+
 @contextlib.contextmanager
 def open_zip(path, *args, **kwargs):
   """A contextmanager for zip files.  Passes through positional and kwargs to zipfile.ZipFile."""
-  with contextlib.closing(zipfile.ZipFile(path, *args, **kwargs)) as zip:
+  with contextlib.closing(PermPreservingZipFile(path, *args, **kwargs)) as zip:
     yield zip
 
 


Test patch:
diff --git a/tests/test_common.py b/tests/test_common.py
--- a/tests/test_common.py
+++ b/tests/test_common.py
@@ -1,13 +1,14 @@
 # Copyright 2016 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
-
+import contextlib
 import errno
 import os
 from contextlib import contextmanager
 
 import pytest
 
-from pex.common import rename_if_empty
+from pex.common import PermPreservingZipFile, chmod_plus_x, rename_if_empty, touch
+from pex.testing import temporary_dir
 
 try:
   from unittest import mock
@@ -42,3 +43,46 @@ def test_rename_if_empty_enotempty():
 
 def test_rename_if_empty_eperm():
   rename_if_empty_test(errno.EPERM, expect_raises=OSError)
+
+
+def extract_perms(path):
+  return oct(os.stat(path).st_mode)
+
+
+@contextlib.contextmanager
+def zip_fixture():
+  with temporary_dir() as target_dir:
+    one = os.path.join(target_dir, 'one')
+    touch(one)
+
+    two = os.path.join(target_dir, 'two')
+    touch(two)
+    chmod_plus_x(two)
+
+    assert extract_perms(one) != extract_perms(two)
+
+    zip_file = os.path.join(target_dir, 'test.zip')
+    with contextlib.closing(PermPreservingZipFile(zip_file, 'w')) as zf:
+      zf.write(one, 'one')
+      zf.write(two, 'two')
+
+    yield zip_file, os.path.join(target_dir, 'extract'), one, two
+
+
+def test_perm_preserving_zipfile_extractall():
+  with zip_fixture() as (zip_file, extract_dir, one, two):
+    with contextlib.closing(PermPreservingZipFile(zip_file)) as zf:
+      zf.extractall(extract_dir)
+
+      assert extract_perms(one) == extract_perms(os.path.join(extract_dir, 'one'))
+      assert extract_perms(two) == extract_perms(os.path.join(extract_dir, 'two'))
+
+
+def test_perm_preserving_zipfile_extract():
+  with zip_fixture() as (zip_file, extract_dir, one, two):
+    with contextlib.closing(PermPreservingZipFile(zip_file)) as zf:
+      zf.extract('one', path=extract_dir)
+      zf.extract('two', path=extract_dir)
+
+      assert extract_perms(one) == extract_perms(os.path.join(extract_dir, 'one'))
+      assert extract_perms(two) == extract_perms(os.path.join(extract_dir, 'two'))