mirror of
https://github.com/wassname/ray.git
synced 2026-07-12 09:43:55 +08:00
* [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:
committed by
GitHub
parent
bb03e2499b
commit
3f22448834
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user