def test_saving_string(self):
filename = self.temp()
with PackageExporter(filename, verbose=False) as he:
src = dedent("""\
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_scriptmodule_only_necessary_code(self):
"""
Test to verify when saving multiple packages with same CU
that packages don't include unnecessary torchscript code files.
The TorchVision code should only be saved in the package that
relies on it.
"""
from package_a.test_module import ModWithTensor
class ModWithTorchVision(torch.nn.Module):
def __init__(self, name: str):
super().__init__()
self.tvmod = resnet18()
def forward(self, input):
return input * 4
scripted_mod_0 = torch.jit.script(ModWithTorchVision("foo"))
scripted_mod_1 = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
buffer_0 = BytesIO()
with PackageExporter(buffer_0, verbose=False) as e:
e.save_pickle("res", "mod1.pkl", scripted_mod_0)
buffer_0.seek(0)
importer_0 = importer = PackageImporter(buffer_0)
buffer_1 = BytesIO()
with PackageExporter(buffer_1, verbose=False) as e:
e.save_pickle("res", "mod1.pkl", scripted_mod_1)
buffer_1.seek(0)
importer_1 = PackageImporter(buffer_1)
self.assertTrue("torchvision" in str(importer_0.file_structure()))
self.assertFalse("torchvision" in str(importer_1.file_structure()))
def test_package_then_fx(self):
from package_a.test_module import SimpleTest
model = SimpleTest()
f = BytesIO()
with PackageExporter(f) as pe:
pe.intern("**")
pe.save_pickle("model", "model.pkl", model)
f.seek(0)
pi = PackageImporter(f)
loaded = pi.load_pickle("model", "model.pkl")
traced = symbolic_trace(loaded)
input = torch.rand(2, 3)
self.assertEqual(loaded(input), traced(input))
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_package_script_class(self):
import package_a.fake_script_class as fake
buffer = BytesIO()
with PackageExporter(buffer, verbose=False) as pe:
pe.save_module(fake.__name__)
buffer.seek(0)
package_importer = PackageImporter(buffer)
loaded = package_importer.import_module(fake.__name__)
input = torch.tensor(1)
self.assertTrue(
torch.allclose(fake.uses_script_class(input),
loaded.uses_script_class(input)))
def test_load_shared_tensors(self):
"""
Test tensors shared across eager and ScriptModules on load
are the same.
"""
from package_a.test_module import ModWithTensor, ModWithTwoSubmodsAndTensor
shared_tensor = torch.ones(3, 3)
scripted_mod_0 = torch.jit.script(ModWithTensor(shared_tensor))
scripted_mod_1 = torch.jit.script(ModWithTensor(shared_tensor))
mod1 = ModWithTwoSubmodsAndTensor(shared_tensor, scripted_mod_0,
scripted_mod_1)
self.assertEqual(
shared_tensor.storage()._cdata,
scripted_mod_0.tensor.storage()._cdata,
)
self.assertEqual(
shared_tensor.storage()._cdata,
scripted_mod_1.tensor.storage()._cdata,
)
buffer = BytesIO()
with PackageExporter(buffer) as e:
e.intern("**")
e.save_pickle("res", "mod1.pkl", mod1)
buffer.seek(0)
importer = PackageImporter(buffer)
loaded_mod_1 = importer.load_pickle("res", "mod1.pkl")
self.assertEqual(
loaded_mod_1.tensor.storage()._cdata,
loaded_mod_1.sub_mod_0.tensor.storage()._cdata,
)
self.assertEqual(
loaded_mod_1.tensor.storage()._cdata,
loaded_mod_1.sub_mod_1.tensor.storage()._cdata,
)
loaded_mod_1.tensor.add_(torch.ones(3, 3))
self.assertTrue(
torch.allclose(loaded_mod_1.tensor, loaded_mod_1.sub_mod_0.tensor))
self.assertTrue(
torch.allclose(loaded_mod_1.tensor, loaded_mod_1.sub_mod_1.tensor))
def test_pickle_mocked(self):
import package_a.subpackage
obj = package_a.subpackage.PackageASubpackageObject()
obj2 = package_a.PackageAObject(obj)
buffer = BytesIO()
with PackageExporter(buffer, verbose=False) as he:
he.mock(include="package_a.subpackage")
he.save_pickle("obj", "obj.pkl", obj2)
buffer.seek(0)
hi = PackageImporter(buffer)
with self.assertRaises(NotImplementedError):
hi.load_pickle("obj", "obj.pkl")
def test_file_structure_has_file(self):
"""
Test Folder's has_file() method.
"""
buffer = BytesIO()
with PackageExporter(buffer, verbose=False) as he:
import package_a.subpackage
obj = package_a.subpackage.PackageASubpackageObject()
he.save_pickle("obj", "obj.pkl", obj)
export_file_structure = he.file_structure()
self.assertTrue(
export_file_structure.has_file("package_a/subpackage.py"))
self.assertFalse(
export_file_structure.has_file("package_a/subpackage"))
def test_save_module_binary(self):
f = BytesIO()
with PackageExporter(f) as he:
import module_a
import package_a
he.save_module(module_a.__name__)
he.save_module(package_a.__name__)
f.seek(0)
hi = PackageImporter(f)
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_save_scriptmodule_file(self):
"""
Test basic saving of ScriptModule in file.
"""
from package_a.test_module import ModWithTensor
scripted_mod = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
filename = self.temp()
with PackageExporter(filename) as e:
e.save_pickle("res", "mod.pkl", scripted_mod)
importer = PackageImporter(filename)
loaded_mod = importer.load_pickle("res", "mod.pkl")
input = torch.rand(1, 2, 3)
self.assertEqual(loaded_mod(input), scripted_mod(input))
def test_dunder_package_present(self):
"""
The attribute '__torch_package__' should be populated on imported modules.
"""
import package_a.subpackage
buffer = BytesIO()
obj = package_a.subpackage.PackageASubpackageObject()
with PackageExporter(buffer, verbose=False) as pe:
pe.save_pickle("obj", "obj.pkl", obj)
buffer.seek(0)
pi = PackageImporter(buffer)
mod = pi.import_module("package_a.subpackage")
self.assertTrue(hasattr(mod, "__torch_package__"))
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)
def test_invalid_import(self):
"""An incorrectly-formed import should raise a PackagingError."""
buffer = BytesIO()
with self.assertRaises(PackagingError) as e:
with PackageExporter(buffer) as exporter:
# This import will fail to load.
exporter.save_source_string("foo", "from ........ import lol")
self.assertEqual(
str(e.exception),
dedent("""
* Dependency resolution failed.
foo
Context: attempted relative import beyond top-level package
"""),
)
def test_deny_glob(self):
"""
Test marking packages as "deny" using globs instead of package names.
"""
filename = self.temp()
with self.assertRaisesRegex(
DeniedModuleError,
'required during packaging but has been explicitly blocklisted'
):
with PackageExporter(filename, verbose=False) as exporter:
exporter.deny(['package_a.*', 'module_*'])
exporter.save_source_string(
'test_module', """\
import package_a.subpackage
import module_a
""")
def test_different_package_interface(self):
"""Test a case where the interface defined in the package is
different than the one defined in the loading environment, to make
sure TorchScript can distinguish between the two.
"""
# Import one version of the interface
import package_a.fake_interface as fake
# Simulate a package that contains a different version of the
# interface, with the exact same name.
buffer = BytesIO()
with PackageExporter(buffer) as pe:
pe.save_source_string(
fake.__name__,
dedent(
"""\
import torch
from torch import Tensor
@torch.jit.interface
class ModuleInterface(torch.nn.Module):
def one(self, inp1: Tensor) -> Tensor:
pass
class ImplementsInterface(torch.nn.Module):
def one(self, inp1: Tensor) -> Tensor:
return inp1 + 1
class UsesInterface(torch.nn.Module):
proxy_mod: ModuleInterface
def __init__(self):
super().__init__()
self.proxy_mod = ImplementsInterface()
def forward(self, input: Tensor) -> Tensor:
return self.proxy_mod.one(input)
"""
),
)
buffer.seek(0)
package_importer = PackageImporter(buffer)
diff_fake = package_importer.import_module(fake.__name__)
# We should be able to script successfully.
torch.jit.script(diff_fake.UsesInterface())
def test_save_scriptmodule(self):
"""
Test basic saving of ScriptModule.
"""
from package_a.test_module import ModWithTensor
scripted_mod = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
buffer = BytesIO()
with PackageExporter(buffer) as e:
e.save_pickle("res", "mod.pkl", scripted_mod)
buffer.seek(0)
importer = PackageImporter(buffer)
loaded_mod = importer.load_pickle("res", "mod.pkl", map_location="cpu")
input = torch.rand(1, 2, 3)
self.assertEqual(loaded_mod(input), scripted_mod(input))
def test_loading_has_record(self):
"""
Test DirectoryReader's has_record().
"""
import package_a # noqa: F401
filename = self.temp()
with PackageExporter(filename) as e:
e.save_module("package_a")
zip_file = zipfile.ZipFile(filename, "r")
with TemporaryDirectory() as temp_dir:
zip_file.extractall(path=temp_dir)
dir_importer = PackageImporter(Path(temp_dir) / Path(filename).name)
self.assertTrue(dir_importer.zip_reader.has_record("package_a/__init__.py"))
self.assertFalse(dir_importer.zip_reader.has_record("package_a"))
def test_inspect_class(self):
"""Should be able to retrieve source for a packaged class."""
import package_a.subpackage
buffer = BytesIO()
obj = package_a.subpackage.PackageASubpackageObject()
with PackageExporter(buffer, verbose=False) as pe:
pe.save_pickle('obj', 'obj.pkl', obj)
buffer.seek(0)
pi = PackageImporter(buffer)
packaged_class = pi.import_module('package_a.subpackage').PackageASubpackageObject
regular_class = package_a.subpackage.PackageASubpackageObject
packaged_src = inspect.getsourcelines(packaged_class)
regular_src = inspect.getsourcelines(regular_class)
self.assertEqual(packaged_src, regular_src)
def test_deny_glob(self):
"""
Test marking packages as "deny" using globs instead of package names.
"""
buffer = BytesIO()
with self.assertRaises(PackagingError):
with PackageExporter(buffer, verbose=False) as exporter:
exporter.deny(["package_a.*", "module_*"])
exporter.save_source_string(
"test_module",
dedent(
"""\
import package_a.subpackage
import module_a
"""
),
)
def test_extern(self):
buffer = BytesIO()
with PackageExporter(buffer, verbose=False) as he:
he.extern(["package_a.subpackage", "module_a"])
he.save_source_string("foo", "import package_a.subpackage; import module_a")
buffer.seek(0)
hi = PackageImporter(buffer)
import module_a
import package_a.subpackage
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)
def test_importer_access(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 importlib
import torch_package_importer as 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_os_path_edge_case(self):
"""
Both 'os' and 'os.path' should be able to be imported into a package.
"""
import package_a.os_imports
buffer = BytesIO()
mod = package_a.os_imports.Module()
with PackageExporter(buffer, verbose=False) as pe:
pe.intern("**")
pe.save_pickle("obj", "obj.pkl", mod)
buffer.seek(0)
pi = PackageImporter(buffer)
mod = pi.load_pickle("obj", "obj.pkl")
mod()
def test_unique_module_names(self):
import package_a.subpackage
obj = package_a.subpackage.PackageASubpackageObject()
obj2 = package_a.PackageAObject(obj)
f1 = self.temp()
with PackageExporter(f1, verbose=False) as pe:
pe.save_pickle("obj", "obj.pkl", obj2)
importer1 = PackageImporter(f1)
loaded1 = importer1.load_pickle("obj", "obj.pkl")
importer2 = PackageImporter(f1)
loaded2 = importer2.load_pickle("obj", "obj.pkl")
# Modules from loaded packages should not shadow the names of modules.
# See mangling.md for more info.
self.assertNotEqual(type(obj2).__module__, type(loaded1).__module__)
self.assertNotEqual(type(loaded1).__module__, type(loaded2).__module__)
def test_package_fx_simple(self):
class SimpleTest(torch.nn.Module):
def forward(self, x):
return torch.relu(x + 3.0)
st = SimpleTest()
traced = symbolic_trace(st)
f = BytesIO()
with PackageExporter(f) as pe:
pe.save_pickle("model", "model.pkl", traced)
f.seek(0)
pi = PackageImporter(f)
loaded_traced = pi.load_pickle("model", "model.pkl")
input = torch.rand(2, 3)
self.assertTrue(torch.allclose(loaded_traced(input), traced(input)))
def test_mock(self):
filename = self.temp()
with PackageExporter(filename, verbose=False) as he:
he.mock(['package_a.subpackage', 'module_a'])
he.save_module('package_a')
he.require_module('package_a.subpackage')
he.require_module('module_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()
def test_save_module_with_module_object(self):
"""
Test that save_module works with a module object
instead of a module name.
"""
buffer = BytesIO()
with PackageExporter(buffer, verbose=False) as he:
import module_a
he.save_module(module_a)
buffer.seek(0)
hi = PackageImporter(buffer)
module_a_i = hi.import_module("module_a")
self.assertEqual(module_a_i.result, "module_a")
self.assertIsNot(module_a, module_a_i)
def test_loading_module(self):
"""
Test basic saving and loading of a packages from a DirectoryReader.
"""
import package_a
filename = self.temp()
with PackageExporter(filename) as e:
e.save_module("package_a")
zip_file = zipfile.ZipFile(filename, "r")
with TemporaryDirectory() as temp_dir:
zip_file.extractall(path=temp_dir)
dir_importer = PackageImporter(Path(temp_dir) / Path(filename).name)
dir_mod = dir_importer.import_module("package_a")
self.assertEqual(dir_mod.result, package_a.result)
def test_exporting_mismatched_code(self):
"""
If an object with the same qualified name is loaded from different
packages, the user should get an error if they try to re-save the
object with the wrong package's source code.
"""
import package_a.subpackage
obj = package_a.subpackage.PackageASubpackageObject()
obj2 = package_a.PackageAObject(obj)
f1 = self.temp()
with PackageExporter(f1, verbose=False) as pe:
pe.intern("**")
pe.save_pickle("obj", "obj.pkl", obj2)
importer1 = PackageImporter(f1)
loaded1 = importer1.load_pickle("obj", "obj.pkl")
importer2 = PackageImporter(f1)
loaded2 = importer2.load_pickle("obj", "obj.pkl")
f2 = self.temp()
def make_exporter():
pe = PackageExporter(f2,
verbose=False,
importer=[importer1, sys_importer])
# Ensure that the importer finds the 'PackageAObject' defined in 'importer1' first.
return pe
# This should fail. The 'PackageAObject' type defined from 'importer1'
# is not necessarily the same 'obj2's version of 'PackageAObject'.
pe = make_exporter()
with self.assertRaises(pickle.PicklingError):
pe.save_pickle("obj", "obj.pkl", obj2)
# This should also fail. The 'PackageAObject' type defined from 'importer1'
# is not necessarily the same as the one defined from 'importer2'
pe = make_exporter()
with self.assertRaises(pickle.PicklingError):
pe.save_pickle("obj", "obj.pkl", loaded2)
# This should succeed. The 'PackageAObject' type defined from
# 'importer1' is a match for the one used by loaded1.
pe = make_exporter()
pe.save_pickle("obj", "obj.pkl", loaded1)
def test_dunder_package_works_from_package(self):
"""
The attribute '__torch_package__' should be accessible from within
the module itself, so that packaged code can detect whether it's
being used in a packaged context or not.
"""
import package_a.use_dunder_package as mod
buffer = BytesIO()
with PackageExporter(buffer, verbose=False) as pe:
pe.save_module(mod.__name__)
buffer.seek(0)
pi = PackageImporter(buffer)
imported_mod = pi.import_module(mod.__name__)
self.assertTrue(imported_mod.is_from_package())
self.assertFalse(mod.is_from_package())
def test_importer_access(self):
buffer = BytesIO()
with PackageExporter(buffer, verbose=False) as he:
he.save_text("main", "main", "my string")
he.save_binary("main", "main_binary", "my string".encode("utf-8"))
src = dedent("""\
import importlib
import torch_package_importer as resources
t = resources.load_text('main', 'main')
b = resources.load_binary('main', 'main_binary')
""")
he.save_source_string("main", src, is_package=True)
buffer.seek(0)
hi = PackageImporter(buffer)
m = hi.import_module("main")
self.assertEqual(m.t, "my string")
self.assertEqual(m.b, "my string".encode("utf-8"))
