torch.package - a way to package models and code (pytorch#45015)

Summary:
Pull Request resolved: pytorch#45015

torch.package allows you to write packages of code, pickled python data, and
arbitrary binary and text resources into a self-contained package.

torch.package.PackageExporter writes the packages and
torch.package.PackageImporter reads them.

The importers can load this code in a hermetic way, such that code is loaded
from the package rather than the normal python import system. This allows
for the packaging of PyTorch model code and data so that it can be run
on a server or used in the future for transfer learning.

The code contained in packages is copied file-by-file from the original
source when it is created, and the file format is a specially organized
zip file. Future users of the package can unzip the package, and edit the code
in order to perform custom modifications to it.

The importer for packages ensures that code in the module can only be loaded from
within the package, except for modules explicitly listed as external using :method:`extern_module`.
The file `extern_modules` in the zip archive lists all the modules that a package externally depends on.
This prevents "implicit" dependencies where the package runs locally because it is importing
a locally-installed package, but then fails when the package is copied to another machine.

Test Plan: Imported from OSS

Reviewed By: SplitInfinity

Differential Revision: D23824337

Pulled By: zdevito

fbshipit-source-id: 1247c34ba9b656f9db68a83e31f2a0fbe3bea6bd

Author: zdevito

test/module_a.py

@@ -0,0 +1 @@
+result = 'module_a'

test/namespace_b/subpackage.py

@@ -0,0 +1,7 @@
+result = 'package_a'
+
+class PackageAObject:
+    __slots__ = ['obj']
+
+    def __init__(self, obj):
+        self.obj = obj

test/package_a/__init__.py

@@ -0,0 +1,3 @@
+result = 'package_a.subpackage'
+class PackageASubpackageObject:
+    pass

test/package_a/subpackage.py

@@ -89,7 +89,8 @@
     'test_determination',
     'test_futures',
     'test_fx',
-    'test_functional_autograd_benchmark'
+    'test_functional_autograd_benchmark',
+    'test_package',
 ]
 
 WINDOWS_BLOCKLIST = [

test/run_test.py

@@ -0,0 +1,309 @@
+from unittest import main, skipIf
+from torch.testing._internal.common_utils import TestCase, IS_WINDOWS
+from tempfile import NamedTemporaryFile
+from torch.package import PackageExporter, PackageImporter
+from pathlib import Path
+from tempfile import TemporaryDirectory
+import torch
+from sys import version_info
+
+try:
+    from torchvision.models import resnet18
+    HAS_TORCHVISION = True
+except ImportError:
+    HAS_TORCHVISION = False
+skipIfNoTorchVision = skipIf(not HAS_TORCHVISION, "no torchvision")
+
+
+
+packaging_directory = Path(__file__).parent
+
+class PackagingTest(TestCase):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._temporary_files = []
+
+    def temp(self):
+        t = NamedTemporaryFile()
+        name = t.name
+        if IS_WINDOWS:
+            t.close()  # can't read an open file in windows
+        else:
+            self._temporary_files.append(t)
+        return name
+
+    def tearDown(self):
+        for t in self._temporary_files:
+            t.close()
+        self._temporary_files = []
+
+    def test_saving_source(self):
+        filename = self.temp()
+        with PackageExporter(filename, verbose=False) as he:
+            he.save_source_file('foo', str(packaging_directory / 'module_a.py'))
+            he.save_source_file('foodir', str(packaging_directory / 'package_a'))
+        hi = PackageImporter(filename)
+        foo = hi.import_module('foo')
+        s = hi.import_module('foodir.subpackage')
+        self.assertEqual(foo.result, 'module_a')
+        self.assertEqual(s.result, 'package_a.subpackage')
+
+    def test_saving_string(self):
+        filename = self.temp()
+        with PackageExporter(filename, verbose=False) as he:
+            src = """\
+import math
+the_math = math
+"""
+            he.save_source_string('my_mod', src)
+        hi = PackageImporter(filename)
+        m = hi.import_module('math')
+        import math
+        self.assertIs(m, math)
+        my_mod = hi.import_module('my_mod')
+        self.assertIs(my_mod.math, math)
+
+    def test_save_module(self):
+        filename = self.temp()
+        with PackageExporter(filename, verbose=False) as he:
+            import module_a
+            import package_a
+            he.save_module(module_a.__name__)
+            he.save_module(package_a.__name__)
+        hi = PackageImporter(filename)
+        module_a_i = hi.import_module('module_a')
+        self.assertEqual(module_a_i.result, 'module_a')
+        self.assertIsNot(module_a, module_a_i)
+        package_a_i = hi.import_module('package_a')
+        self.assertEqual(package_a_i.result, 'package_a')
+        self.assertIsNot(package_a_i, package_a)
+
+    def test_pickle(self):
+        import package_a.subpackage
+        obj = package_a.subpackage.PackageASubpackageObject()
+        obj2 = package_a.PackageAObject(obj)
+
+        filename = self.temp()
+        with PackageExporter(filename, verbose=False) as he:
+            he.save_pickle('obj', 'obj.pkl', obj2)
+        hi = PackageImporter(filename)
+
+        # check we got dependencies
+        sp = hi.import_module('package_a.subpackage')
+        # check we didn't get other stuff
+        with self.assertRaises(ImportError):
+            hi.import_module('module_a')
+
+        obj_loaded = hi.load_pickle('obj', 'obj.pkl')
+        self.assertIsNot(obj2, obj_loaded)
+        self.assertIsInstance(obj_loaded.obj, sp.PackageASubpackageObject)
+        self.assertIsNot(package_a.subpackage.PackageASubpackageObject, sp.PackageASubpackageObject)
+
+    def test_resources(self):
+        filename = self.temp()
+        with PackageExporter(filename, verbose=False) as he:
+            he.save_text('main', 'main', "my string")
+            he.save_binary('main', 'main_binary', "my string".encode('utf-8'))
+            src = """\
+import resources
+t = resources.load_text('main', 'main')
+b = resources.load_binary('main', 'main_binary')
+"""
+            he.save_source_string('main', src, is_package=True)
+        hi = PackageImporter(filename)
+        m = hi.import_module('main')
+        self.assertEqual(m.t, "my string")
+        self.assertEqual(m.b, "my string".encode('utf-8'))
+
+    def test_extern(self):
+        filename = self.temp()
+        with PackageExporter(filename, verbose=False) as he:
+            he.extern_modules(['package_a.subpackage', 'module_a'])
+            he.save_module('package_a')
+        hi = PackageImporter(filename)
+        import package_a.subpackage
+        import module_a
+
+        module_a_im = hi.import_module('module_a')
+        hi.import_module('package_a.subpackage')
+        package_a_im = hi.import_module('package_a')
+
+        self.assertIs(module_a, module_a_im)
+        self.assertIsNot(package_a, package_a_im)
+        self.assertIs(package_a.subpackage, package_a_im.subpackage)
+
+    @skipIf(version_info.major < 3 or version_info.minor < 7, 'mock uses __getattr__ a 3.7 feature')
+    def test_mock(self):
+        filename = self.temp()
+        with PackageExporter(filename, verbose=False) as he:
+            he.mock_modules(['package_a.subpackage', 'module_a'])
+            he.save_module('package_a')
+        hi = PackageImporter(filename)
+        import package_a.subpackage
+        _ = package_a.subpackage
+        import module_a
+        _ = module_a
+
+        m = hi.import_module('package_a.subpackage')
+        r = m.result
+        with self.assertRaisesRegex(NotImplementedError, 'was mocked out'):
+            r()
+
+    @skipIf(version_info.major < 3 or version_info.minor < 7, 'mock uses __getattr__ a 3.7 feature')
+    def test_custom_requires(self):
+        filename = self.temp()
+
+        class Custom(PackageExporter):
+            def require_module(self, name, dependencies):
+                if name == 'module_a':
+                    self.mock_module('module_a')
+                elif name == 'package_a':
+                    self.save_source_string('package_a', 'import module_a\nresult = 5\n')
+                else:
+                    raise NotImplementedError('wat')
+
+        with Custom(filename, verbose=False) as he:
+            he.save_source_string('main', 'import package_a\n')
+
+        hi = PackageImporter(filename)
+        hi.import_module('module_a').should_be_mocked
+        bar = hi.import_module('package_a')
+        self.assertEqual(bar.result, 5)
+
+    @skipIfNoTorchVision
+    def test_resnet(self):
+        resnet = resnet18()
+
+        f1 = self.temp()
+
+        # create a package that will save it along with its code
+        with PackageExporter(f1, verbose=False) as e:
+            # put the pickled resnet in the package, by default
+            # this will also save all the code files references by
+            # the objects in the pickle
+            e.save_pickle('model', 'model.pkl', resnet)
+
+        # we can now load the saved model
+        i = PackageImporter(f1)
+        r2 = i.load_pickle('model', 'model.pkl')
+
+        # test that it works
+        input = torch.rand(1, 3, 224, 224)
+        ref = resnet(input)
+        self.assertTrue(torch.allclose(r2(input), ref))
+
+        # functions exist also to get at the private modules in each package
+        torchvision = i.import_module('torchvision')
+
+        f2 = self.temp()
+        # if we are doing transfer learning we might want to re-save
+        # things that were loaded from a package
+        with PackageExporter(f2, verbose=False) as e:
+            # We need to tell the exporter about any modules that
+            # came from imported packages so that it can resolve
+            # class names like torchvision.models.resnet.ResNet
+            # to their source code.
+
+            e.importers.insert(0, i.import_module)
+
+            # e.importers is a list of module importing functions
+            # that by default contains importlib.import_module.
+            # it is searched in order until the first success and
+            # that module is taken to be what torchvision.models.resnet
+            # should be in this code package. In the case of name collisions,
+            # such as trying to save a ResNet from two different packages,
+            # we take the first thing found in the path, so only ResNet objects from
+            # one importer will work. This avoids a bunch of name mangling in
+            # the source code. If you need to actually mix ResNet objects,
+            # we suggest reconstructing the model objects using code from a single package
+            # using functions like save_state_dict and load_state_dict to transfer state
+            # to the correct code objects.
+            e.save_pickle('model', 'model.pkl', r2)
+
+        i2 = PackageImporter(f2)
+        r3 = i2.load_pickle('model', 'model.pkl')
+        self.assertTrue(torch.allclose(r3(input), ref))
+
+        # test we can load from a directory
+        import zipfile
+        zf = zipfile.ZipFile(f1, 'r')
+
+        with TemporaryDirectory() as td:
+            zf.extractall(path=td)
+            iz = PackageImporter(str(Path(td) / Path(f1).name))
+            r4 = iz.load_pickle('model', 'model.pkl')
+            self.assertTrue(torch.allclose(r4(input), ref))
+
+    @skipIfNoTorchVision
+    def test_model_save(self):
+
+        # This example shows how you might package a model
+        # so that the creator of the model has flexibility about
+        # how they want to save it but the 'server' can always
+        # use the same API to load the package.
+
+        # The convension is for each model to provide a
+        # 'model' package with a 'load' function that actual
+        # reads the model out of the archive.
+
+        # How the load function is implemented is up to the
+        # the packager.
+
+        # get our normal torchvision resnet
+        resnet = resnet18()
+
+
+        f1 = self.temp()
+        # Option 1: save by pickling the whole model
+        # + single-line, similar to torch.jit.save
+        # - more difficult to edit the code after the model is created
+        with PackageExporter(f1, verbose=False) as e:
+            e.save_pickle('model', 'pickled', resnet)
+            # note that this source is the same for all models in this approach
+            # so it can be made part of an API that just takes the model and
+            # packages it with this source.
+            src = """\
+import resources # gives you access to the importer from within the package
+
+# server knows to call model.load() to get the model,
+# maybe in the future it passes options as arguments by convension
+def load():
+    return resources.load_pickle('model', 'pickled')
+        """
+            e.save_source_string('model', src, is_package=True)
+
+        f2 = self.temp()
+        # Option 2: save with state dict
+        # - more code to write to save/load the model
+        # + but this code can be edited later to adjust adapt the model later
+        with PackageExporter(f2, verbose=False) as e:
+            e.save_pickle('model', 'state_dict', resnet.state_dict())
+            src = """\
+import resources # gives you access to the importer from within the package
+from torchvision.models.resnet import resnet18
+def load():
+    # if you want, you can later edit how resnet is constructed here
+    # to edit the model in the package, while still loading the original
+    # state dict weights
+    r = resnet18()
+    state_dict = resources.load_pickle('model', 'state_dict')
+    r.load_state_dict(state_dict)
+    return r
+        """
+            e.save_source_string('model', src, is_package=True)
+
+
+
+        # regardless of how we chose to package, we can now use the model in a server in the same way
+        input = torch.rand(1, 3, 224, 224)
+        results = []
+        for m in [f1, f2]:
+            importer = PackageImporter(m)
+            the_model = importer.import_module('model').load()
+            r = the_model(input)
+            results.append(r)
+
+        self.assertTrue(torch.allclose(*results))
+
+if __name__ == '__main__':
+    main()

test/test_package.py

@@ -0,0 +1,2 @@
+from .importer import PackageImporter
+from .exporter import PackageExporter

torch/package/__init__.py

@@ -0,0 +1,78 @@
+from pickle import _Pickler, _getattribute, whichmodule, _extension_registry, _compat_pickle  # type: ignore
+from pickle import GLOBAL, STACK_GLOBAL, EXT1, EXT2, EXT4, PicklingError
+from struct import pack
+
+class CustomImportPickler(_Pickler):
+    def __init__(self, import_module, *args, **kwargs):
+        self.import_module = import_module
+        super().__init__(*args, **kwargs)
+
+    def save_global(self, obj, name=None):
+        # unfortunately the pickler code is factored in a way that
+        # forces us to copy/paste this function. The only change is marked
+        # CHANGED below.
+        write = self.write
+        memo = self.memo
+
+        if name is None:
+            name = getattr(obj, '__qualname__', None)
+        if name is None:
+            name = obj.__name__
+
+        module_name = whichmodule(obj, name)
+        try:
+            # CHANGED: self.import_module rather than
+            # __import__
+            module = self.import_module(module_name)
+            obj2, parent = _getattribute(module, name)
+        except (ImportError, KeyError, AttributeError):
+            raise PicklingError(
+                "Can't pickle %r: it's not found as %s.%s" %
+                (obj, module_name, name)) from None
+        else:
+            if obj2 is not obj:
+                raise PicklingError(
+                    "Can't pickle %r: it's not the same object as %s.%s" %
+                    (obj, module_name, name))
+
+        if self.proto >= 2:
+            code = _extension_registry.get((module_name, name))
+            if code:
+                assert code > 0
+                if code <= 0xff:
+                    write(EXT1 + pack("<B", code))
+                elif code <= 0xffff:
+                    write(EXT2 + pack("<H", code))
+                else:
+                    write(EXT4 + pack("<i", code))
+                return
+        lastname = name.rpartition('.')[2]
+        if parent is module:
+            name = lastname
+        # Non-ASCII identifiers are supported only with protocols >= 3.
+        if self.proto >= 4:
+            self.save(module_name)
+            self.save(name)
+            write(STACK_GLOBAL)
+        elif parent is not module:
+            self.save_reduce(getattr, (parent, lastname))
+        elif self.proto >= 3:
+            write(GLOBAL + bytes(module_name, "utf-8") + b'\n' +
+                  bytes(name, "utf-8") + b'\n')
+        else:
+            if self.fix_imports:
+                r_name_mapping = _compat_pickle.REVERSE_NAME_MAPPING
+                r_import_mapping = _compat_pickle.REVERSE_IMPORT_MAPPING
+                if (module_name, name) in r_name_mapping:
+                    module_name, name = r_name_mapping[(module_name, name)]
+                elif module_name in r_import_mapping:
+                    module_name = r_import_mapping[module_name]
+            try:
+                write(GLOBAL + bytes(module_name, "ascii") + b'\n' +
+                      bytes(name, "ascii") + b'\n')
+            except UnicodeEncodeError:
+                raise PicklingError(
+                    "can't pickle global identifier '%s.%s' using "
+                    "pickle protocol %i" % (module, name, self.proto)) from None
+
+        self.memoize(obj)

torch/package/_custom_import_pickler.py

@@ -0,0 +1,83 @@
+import _warnings
+import os.path
+# note: implementations 
+# copied from cpython's import code
+
+
+# _zip_searchorder defines how we search for a module in the Zip
+# archive: we first search for a package __init__, then for
+# non-package .pyc, and .py entries. The .pyc entries
+# are swapped by initzipimport() if we run in optimized mode. Also,
+# '/' is replaced by path_sep there.
+
+_zip_searchorder = (
+    ('/__init__.py', True),
+    ('.py', False),
+)
+
+# Replace any occurrences of '\r\n?' in the input string with '\n'.
+# This converts DOS and Mac line endings to Unix line endings.
+def _normalize_line_endings(source):
+    source = source.replace(b'\r\n', b'\n')
+    source = source.replace(b'\r', b'\n')
+    return source
+
+def _resolve_name(name, package, level):
+    """Resolve a relative module name to an absolute one."""
+    bits = package.rsplit('.', level - 1)
+    if len(bits) < level:
+        raise ValueError('attempted relative import beyond top-level package')
+    base = bits[0]
+    return '{}.{}'.format(base, name) if name else base
+
+def _sanity_check(name, package, level):
+    """Verify arguments are "sane"."""
+    if not isinstance(name, str):
+        raise TypeError('module name must be str, not {}'.format(type(name)))
+    if level < 0:
+        raise ValueError('level must be >= 0')
+    if level > 0:
+        if not isinstance(package, str):
+            raise TypeError('__package__ not set to a string')
+        elif not package:
+            raise ImportError('attempted relative import with no known parent '
+                              'package')
+    if not name and level == 0:
+        raise ValueError('Empty module name')
+
+def _calc___package__(globals):
+    """Calculate what __package__ should be.
+
+    __package__ is not guaranteed to be defined or could be set to None
+    to represent that its proper value is unknown.
+
+    """
+    package = globals.get('__package__')
+    spec = globals.get('__spec__')
+    if package is not None:
+        if spec is not None and package != spec.parent:
+            _warnings.warn("__package__ != __spec__.parent "
+                           f"({package!r} != {spec.parent!r})",
+                           ImportWarning, stacklevel=3)
+        return package
+    elif spec is not None:
+        return spec.parent
+    else:
+        _warnings.warn("can't resolve package from __spec__ or __package__, "
+                       "falling back on __name__ and __path__",
+                       ImportWarning, stacklevel=3)
+        package = globals['__name__']
+        if '__path__' not in globals:
+            package = package.rpartition('.')[0]
+    return package
+
+def _normalize_path(path):
+    """Normalize a path by ensuring it is a string.
+
+    If the resulting string contains path separators, an exception is raised.
+    """
+    parent, file_name = os.path.split(path)
+    if parent:
+        raise ValueError('{!r} must be only a file name'.format(path))
+    else:
+        return file_name

torch/package/_importlib.py

@@ -0,0 +1,39 @@
+
+_magic_methods = ['__subclasscheck__', '__hex__', '__rmul__',
+                  '__float__', '__idiv__', '__setattr__', '__div__', '__invert__',
+                  '__nonzero__', '__rshift__',
+                  '__eq__', '__pos__', '__round__',
+                  '__rand__', '__or__', '__complex__', '__divmod__',
+                  '__len__', '__reversed__', '__copy__', '__reduce__',
+                  '__deepcopy__', '__rdivmod__', '__rrshift__', '__ifloordiv__',
+                  '__hash__', '__iand__', '__xor__', '__isub__', '__oct__',
+                  '__ceil__', '__imod__', '__add__', '__truediv__',
+                  '__unicode__', '__le__', '__delitem__', '__sizeof__', '__sub__',
+                  '__ne__', '__pow__', '__bytes__', '__mul__',
+                  '__itruediv__', '__bool__', '__iter__', '__abs__',
+                  '__gt__', '__iadd__', '__enter__',
+                  '__floordiv__', '__call__', '__neg__',
+                  '__and__', '__ixor__', '__getitem__', '__exit__', '__cmp__',
+                  '__getstate__', '__index__', '__contains__', '__floor__', '__lt__', '__getattr__',
+                  '__mod__', '__trunc__', '__delattr__', '__instancecheck__', '__setitem__', '__ipow__',
+                  '__ilshift__', '__long__', '__irshift__', '__imul__',
+                  '__lshift__', '__dir__', '__ge__', '__int__', '__ior__']
+
+
+class MockedObject:
+    _name: str
+
+    def __init__(self, name):
+        self.__dict__['_name'] = name
+
+    def __repr__(self):
+        return f"MockedObject({self._name})"
+
+
+def install_method(method_name):
+    def _not_implemented(self, *args, **kwargs):
+        raise NotImplementedError(f"Object '{self._name}' was mocked out during packaging but it is being used in {method_name}")
+    setattr(MockedObject, method_name, _not_implemented)
+
+for method_name in _magic_methods:
+    install_method(method_name)

torch/package/_mock.py

@@ -0,0 +1,48 @@
+import torch
+from glob import glob
+import os.path
+from typing import List, Any
+
+_storages : List[Any] = [
+    torch.DoubleStorage,
+    torch.FloatStorage,
+    torch.LongStorage,
+    torch.IntStorage,
+    torch.ShortStorage,
+    torch.CharStorage,
+    torch.ByteStorage,
+    torch.BoolStorage,
+]
+_dtype_to_storage = {
+    data_type(0).dtype: data_type for data_type in _storages
+}
+
+# because get_storage_from_record returns a tensor!?
+class _HasStorage(object):
+    def __init__(self, storage):
+        self._storage = storage
+
+    def storage(self):
+        return self._storage
+
+
+class MockZipReader(object):
+    def __init__(self, directory):
+        self.directory = directory
+
+    def get_record(self, name):
+        filename = f'{self.directory}/{name}'
+        with open(filename, 'rb') as f:
+            return f.read()
+
+    def get_storage_from_record(self, name, numel, dtype):
+        storage = _dtype_to_storage[dtype]
+        filename = f'{self.directory}/{name}'
+        return _HasStorage(storage.from_file(filename=filename, size=numel))
+
+    def get_all_records(self, ):
+        files = []
+        for filename in glob(f'{self.directory}/**', recursive=True):
+            if not os.path.isdir(filename):
+                files.append(filename[len(self.directory) + 1:])
+        return files

torch/package/_mock_zipreader.py

@@ -0,0 +1,42 @@
+from typing import List, Optional, Tuple
+import ast
+from ._importlib import _resolve_name
+
+class _ExtractModuleReferences(ast.NodeVisitor):
+    """
+    Extract the list of global variables a block of code will read and write
+    """
+
+    @classmethod
+    def run(cls, src: str, package: str) -> List[Tuple[str, Optional[str]]]:
+        visitor = cls(package)
+        tree = ast.parse(src)
+        visitor.visit(tree)
+        return list(visitor.references.keys())
+
+    def __init__(self, package):
+        super().__init__()
+        self.package = package
+        self.references = {}
+
+    def _absmodule(self, module_name: str, level: int) -> str:
+        if level > 0:
+            return _resolve_name(module_name, self.package, level)
+        return module_name
+
+    def visit_Import(self, node):
+        for alias in node.names:
+            self.references[(alias.name, None)] = True
+
+    def visit_ImportFrom(self, node):
+        name = self._absmodule(node.module, 0 if node.level is None else node.level)
+        for alias in node.names:
+            # from my_package import foo
+            # foo may be a module, so we have to add it to the list of
+            # potential references, if import of it fails, we will ignore it
+            if alias.name != '*':
+                self.references[(name, alias.name)] = True
+            else:
+                self.references[(name, None)] = True
+
+find_files_source_depends_on = _ExtractModuleReferences.run

torch/package/exporter.py

@@ -821,7 +821,7 @@ def restore_location(storage, location):
     return restore_location
 
 
-def _load(zip_file, map_location, pickle_module, **pickle_load_args):
+def _load(zip_file, map_location, pickle_module, pickle_file='data.pkl', **pickle_load_args):
     restore_location = _get_restore_location(map_location)
 
     loaded_storages = {}
@@ -847,7 +847,7 @@ def persistent_load(saved_id):
         return storage
 
     # Load the data (which may in turn use `persistent_load` to load tensors)
-    data_file = io.BytesIO(zip_file.get_record('data.pkl'))
+    data_file = io.BytesIO(zip_file.get_record(pickle_file))
     unpickler = pickle_module.Unpickler(data_file, **pickle_load_args)
     unpickler.persistent_load = persistent_load
     result = unpickler.load()