nit: empty line after this function.

Use type annotations? def my_tensor_function(torch.Tensor a, torch.Tensor b)?

Use type annotations?

type annotation failed in CI python 2 build, although we skipped python 2 tests, but the helper function still need to build in python 2

Can we make my_class hold a really complex object with tensors inside? This would be a nice test case. One example would be something like: Tuple[Tensor, List[List[Tensor]], Dict[Tuple[Tensor, List[Tensor]], List[Tensor]]

We should add some python docs to all public methods in this file (ones without _).

nit: Can we use six.PY3 instead? Also, IIRC, we only support py3 right? Do we still need this check?

yes, because when internal_rpc_util module is imported by other files, the global variables will be checked at build time. it will fail to build in python2 if not adding the check

why six.PY3 is better? trying to be consistent with other rpc files that are using sys.version_info

I just find six.PY3 to be nicer since its much more concise that doing comparisons on sys.version_info.

Can tensor_reducer and serialize be private?

Shouldn't _dispatch_table be global here?

It's not immediately clear why this is needed. I see it's related to the reducer in the dispatch table. Would it be possible to turn the dispatch table into a first class object and associate some methods with it? E.g. _dispatch_table.clear_send_tables() or something? Then we can keep all the logic related to pickling/unpickling in one place vs. mucking with globals.

@pietern in order to achieve that we'd need a per-serialization dispatch table (or at least a thread-local dispatch table). I don't know what's the relative cost of copying the dispatch table, but I'd rather have a global dispatch table and a thread-local pickling context object instead.

===

On a separate note, I believe that there is a corner case not covered here: it is possible for a custom reducer to call an RPC. In that case you'll have nested calls to serialize on the same thread. Example:

class MyCustomClass:
    def __getstate__(self):
        dist.rpc(my_random_remote_function_called_inside_of_serialize, ..., to='w1')
        return (...)

m = MyCustomClass()
rpc(my_function_taking_a_custom_class_instance, args=(m, ), 'w1')

In order to cover this corner case you'll need to save a possible current _thread_local_tensor_tables.send_tables and restore it after the call. Even if we go with alternative solutions such as suggested by Pieter we'll have to deal with this case.

One thing that you didn't seem to have tested is UDFs that return a complex value containing tensors inside. We should test that as well as this will exercise a different serialization / deserialization codepath.

any reason to make this range(10) instead of say range(100)? just making it more likely to exercise race conditions.

@pietern in order to achieve that we'd need a per-serialization dispatch table (or at least a thread-local dispatch table). I don't know what's the relative cost of copying the dispatch table, but I'd rather have a global dispatch table and a thread-local pickling context object instead.

===

On a separate note, I believe that there is a corner case not covered here: it is possible for a custom reducer to call an RPC. In that case you'll have nested calls to serialize on the same thread. Example:

class MyCustomClass:
    def __getstate__(self):
        dist.rpc(my_random_remote_function_called_inside_of_serialize, ..., to='w1')
        return (...)

m = MyCustomClass()
rpc(my_function_taking_a_custom_class_instance, args=(m, ), 'w1')

In order to cover this corner case you'll need to save a possible current _thread_local_tensor_tables.send_tables and restore it after the call. Even if we go with alternative solutions such as suggested by Pieter we'll have to deal with this case.

I believe nested deserialization could happen in the same thread too -- if for example setstate in a class actually calls into rpc -- see comment above.

I believe you can simplify this with :

  old_send_tables = _thread_local_tensor_tables.send_tables if hasattr('send_tables', _thread_local_tensor_tqbles) else None

no need for should_restore.

Use python docstrings for the class and member functions instead of multiple lines starting with '#'.

python docstring

nit: make this _InternalRPCPickler to indicate its private.

I just find six.PY3 to be nicer since its much more concise that doing comparisons on sys.version_info.

Question: IIUC, OSS does not build Python files. Do you see errors in internal tests?

I see errors in OSS tests

nit: sort imports in ascii order