Re-Revert "[Core] zero-copy serializer for pytorch (#12344)" (#12478)

* [Core] zero-copy serializer for pytorch (#12344)

* zero-copy serializer for pytorch

* address possible bottleneck

* add tests & device support

(cherry picked from commit 0a505ca83d)

* add environmental variables

* update doc
This commit is contained in:
Siyuan (Ryans) Zhuang
2020-11-30 11:43:03 -08:00
committed by GitHub
parent bb03e2499b
commit 3f22448834
5 changed files with 111 additions and 1 deletions
+2
View File
@@ -26,6 +26,7 @@ from ray._raylet import (
MessagePackSerializedObject,
RawSerializedObject,
)
from ray import serialization_addons
logger = logging.getLogger(__name__)
@@ -155,6 +156,7 @@ class SerializationContext:
# Because objects have default __reduce__ method, we only need to
# treat ObjectRef specifically.
self._register_cloudpickle_reducer(ray.ObjectRef, object_ref_reducer)
serialization_addons.apply(self)
def _register_cloudpickle_reducer(self, cls, reducer):
pickle.CloudPickler.dispatch[cls] = reducer