Repository: pytorch
Issue: torch.package: re-export docs suggestions don't work + other sharp edges
Labels: oncall: package/deploy,imported

## 📚 Documentation

This page: https://pytorch.org/docs/stable/package.html#re-export-an-imported-object, seems to have some errors. 

Take the following self-contained script:

```python
import torch
from torch import nn 
import tempfile

EXTERN = [
    "__main__"
]

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 10)
        self.relu = nn.ReLU()

    def forward(self, data):
        return self.linear(self.relu(data))
    
    def package(self, path, verbose=False, extern=[], **kwargs):
        package_name = type(self).__name__
        resource_name = f"{type(self).__name__}.pth"
  
        with torch.package.PackageExporter(path, verbose=verbose, **kwargs) as exp:
            exp.extern(EXTERN + extern)
            exp.save_pickle(package_name, resource_name, self)
        return path

    @classmethod
    def load_package(cls, path):
        package_name = cls.__name__
        resource_name = f"{cls.__name__}.pth"

        imp = torch.package.PackageImporter(path)
        model = imp.load_pickle(package_name, resource_name, "cpu")

        return model, imp

data = torch.randn(1, 10)

with tempfile.NamedTemporaryFile(suffix=".pth") as f:
    model = Model()
    print(1, model(data))

    model.package(f.name)
    model2, imp = Model.load_package(f.name)
    print(2, model2(data))

    model2.package(f.name, importer=(imp,))
    model3, imp = Model.load_package(f.name)
    print(3, model3(data))
```

which works, and has the following output:

```
1 tensor([[-0.1936,  0.3821, -0.6370, -0.1138, -1.1114,  0.6543, -0.3422, -0.5319,
         -0.1984,  0.7432]], grad_fn=<AddmmBackward>)
2 tensor([[-0.1936,  0.3821, -0.6370, -0.1138, -1.1114,  0.6543, -0.3422, -0.5319,
         -0.1984,  0.7432]], grad_fn=<AddmmBackward>)
3 tensor([[-0.1936,  0.3821, -0.6370, -0.1138, -1.1114,  0.6543, -0.3422, -0.5319,
         -0.1984,  0.7432]], grad_fn=<AddmmBackward>)
```

The first note is that the the re-export docs say to do the following:

```python
with PackageExporter(f2, importer=(importer, sys_importer)) as exporter:
    exporter.save_pickle("model", "model.pkl", obj)
```

I was assuming that `sys_importer` can be obtained via 
```python
from torch.package.importer import _SysImporter as sys_importer
```

but it's unclear. So if you do as suggested:

```python
model2.package(f.name, importer=(imp, sys_importer))
```

This happens:

```
1 tensor([[ 0.1494, -0.4079, -0.0159, -0.9270,  0.0571,  0.4879,  0.2713,  1.0109,
         -0.1130, -0.2018]], grad_fn=<AddmmBackward>)
2 tensor([[ 0.1494, -0.4079, -0.0159, -0.9270,  0.0571,  0.4879,  0.2713,  1.0109,
         -0.1130, -0.2018]], grad_fn=<AddmmBackward>)
Traceback (most recent call last):
  File "scripts/prem/repro.py", line 48, in <module>
    model2.package(f.name, importer=(imp, sys_importer))
  File "scripts/prem/repro.py", line 25, in package
    exp.save_pickle(package_name, resource_name, self)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/site-packages/torch/package/package_exporter.py", line 497, in save_pickle
    pickler.dump(obj)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 487, in dump
    self.save(obj)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 603, in save
    self.save_reduce(obj=obj, *rv)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 687, in save_reduce
    save(cls)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 1128, in save_type
    return self.save_global(obj)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/site-packages/torch/package/_package_pickler.py", line 32, in save_global
    module_name, name = self.importer.get_name(obj, name)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/site-packages/torch/package/importer.py", line 89, in get_name
    orig_module_name = self.whichmodule(obj, name)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/site-packages/torch/package/importer.py", line 206, in whichmodule
    module_name = importer.whichmodule(obj, name)
TypeError: whichmodule() missing 1 required positional argument: 'name'
```

so that should probably be corrected in the docs.

The second sharp edge not mentioned in the docs is this:

```python
EXTERN = [
  "__main__"
] 
```

If I change this to:

```python
EXTERN = []
```

then this happens:

```
1 tensor([[ 0.3872,  0.0951,  1.1489,  1.0664,  0.1613,  0.1006, -0.6773, -0.2391,
         -1.5084,  0.3271]], grad_fn=<AddmmBackward>)
Traceback (most recent call last):
  File "scripts/prem/repro.py", line 42, in <module>
    model.package(f.name)
  File "scripts/prem/repro.py", line 23, in package
    exp.save_pickle(package_name, resource_name, self)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/site-packages/torch/package/package_exporter.py", line 755, in __exit__
    self.close()
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/site-packages/torch/package/package_exporter.py", line 848, in close
    self._execute_dependency_graph()
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/site-packages/torch/package/package_exporter.py", line 784, in _execute_dependency_graph
    self._validate_dependency_graph()
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/site-packages/torch/package/package_exporter.py", line 771, in _validate_dependency_graph
    raise PackagingError(self.dependency_graph)
torch.package.package_exporter.PackagingError:
* Module did not match against any action pattern. Extern, mock, or intern it.
    __main__
```

so I'm finding I always have to manually add `__main__` to the extern.

Finally, this is a very tiny thing, but if I try to do something more clever, like saving the self.importer into the object so that the save/load works like so:

```python
    def package(self, path, verbose=False, extern=[], **kwargs):
        package_name = type(self).__name__
        resource_name = f"{type(self).__name__}.pth"

        if hasattr(self, "importer"):
            kwargs["importer"] = (self.importer,)
            # Don't re-save self.importer!
            del self.importer
  
        with torch.package.PackageExporter(path, verbose=verbose, **kwargs) as exp:
            exp.extern(EXTERN + extern)
            exp.save_pickle(package_name, resource_name, self)
        return path

    @classmethod
    def load_package(cls, path):
        package_name = cls.__name__
        resource_name = f"{cls.__name__}.pth"

        imp = torch.package.PackageImporter(path)
        model = imp.load_pickle(package_name, resource_name, "cpu")
        model.importer = imp

        return model
```

I *need* that `del self.importer` line, otherwise this happens:

```
1 tensor([[ 0.6387, -0.3141, -0.1326, -0.1238, -0.0272, -0.2990, -0.0568, -0.1781,
          0.0821, -0.1556]], grad_fn=<AddmmBackward>)
2 tensor([[ 0.6387, -0.3141, -0.1326, -0.1238, -0.0272, -0.2990, -0.0568, -0.1781,
          0.0821, -0.1556]], grad_fn=<AddmmBackward>)
Traceback (most recent call last):
  File "scripts/prem/repro.py", line 52, in <module>
    model2.package(f.name)
  File "scripts/prem/repro.py", line 28, in package
    exp.save_pickle(package_name, resource_name, self)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/site-packages/torch/package/package_exporter.py", line 497, in save_pickle
    pickler.dump(obj)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 487, in dump
    self.save(obj)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 603, in save
    self.save_reduce(obj=obj, *rv)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 717, in save_reduce
    save(state)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 971, in save_dict
    self._batch_setitems(obj.items())
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 997, in _batch_setitems
    save(v)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 603, in save
    self.save_reduce(obj=obj, *rv)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 717, in save_reduce
    save(state)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 971, in save_dict
    self._batch_setitems(obj.items())
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 997, in _batch_setitems
    save(v)
  File "/home/admin/miniconda3/envs/wav2wav-v2/lib/python3.8/pickle.py", line 578, in save
    rv = reduce(self.proto)
TypeError: cannot pickle 'torch._C.PyTorchFileReader' object
```