Skip to content

/ pytorch

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Jitted function cannot be pickled #61210

Open
twoertwein opened this issue Jul 3, 2021 · 3 comments · May be fixed by #61381
Open

Jitted function cannot be pickled #61210

twoertwein opened this issue Jul 3, 2021 · 3 comments · May be fixed by #61381
Assignees
@SiddeshSambasivam
Labels
good first issue oncall: jit OSS contribution wanted
Projects
JIT Triage

Comments

@twoertwein
Copy link
Contributor

@twoertwein twoertwein commented Jul 3, 2021

🚀 Feature

Support to pickle a jitted function (or at least throw a TypeError when using protocol 0 and 1).

Motivation

Trying to pickle a jitted function either raises TypeError: cannot pickle 'torch._C.ScriptFunction' object when protocol>1 or far worse when using protocol=0 or protocol=1 python 3.9.5 dies with:

terminate called after throwing an instance of 'std::runtime_error'
  what():  instance allocation failed: new instance has no pybind11-registered base types
Aborted (core dumped)

Example:

import pickle

import torch

@torch.jit.script
def fun(x: torch.Tensor) -> torch.Tensor:
    return x

pickle.dumps(fun, protocol=0)

Pitch

Implement pickle support for all pickle protocols (or raise a TypeError when used with protocl<2).

Additional context

I encountered this when using a jitted function with dask dask/dask#7861
@github-actions github-actions bot added this to Need triage in JIT Triage Jul 3, 2021
@gmagogsfm
Copy link
Contributor

@gmagogsfm gmagogsfm commented Jul 6, 2021

Thanks for reporting this issue.

I think your suggestion makes sense. Given that complexity of this task is not high, I will leave it as a task for OSS contribution.

Here are some pointers:

  1. A scripted function is a ScriptFunction object defined in:

    pytorch/torch/jit/_script.py

    Line 52 in 05c1e5b

    ScriptFunction = torch._C.ScriptFunction
  2. We can add a __getstate__ method to ScriptFunction. The method should do nothing more than throwing an exception similar to the one here:

    pytorch/torch/jit/_script.py

    Line 436 in 05c1e5b

    raise pickle.PickleError("ScriptClasses cannot be pickled")
  3. A test should also be added to cover new exception
@gmagogsfm gmagogsfm moved this from Need triage to Pending in JIT Triage Jul 6, 2021
@jrbourbeau jrbourbeau mentioned this issue Jul 6, 2021
Closed
@SiddeshSambasivam
Copy link

@SiddeshSambasivam SiddeshSambasivam commented Jul 6, 2021

Hi @gmagogsfm ! I'm more than happy to help with this task.
@gmagogsfm
Copy link
Contributor

@gmagogsfm gmagogsfm commented Jul 6, 2021

Hi @gmagogsfm ! I'm more than happy to help with this task.

Awesome! Please see my previous comment for pointers on implementing this improvement. Feel free to ask here if you see any unexpected issue.
@SiddeshSambasivam SiddeshSambasivam linked a pull request that will close this issue Jul 7, 2021
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@SiddeshSambasivam SiddeshSambasivam

Labels
good first issue oncall: jit OSS contribution wanted
Projects
JIT Triage
  
Pending
Milestone
No milestone
Linked pull requests

Successfully merging a pull request may close this issue.

Added ScriptFunction pkl exception
4 participants
@gmagogsfm @facebook-github-bot @twoertwein @SiddeshSambasivam