Reuse code of checking large pickles (#2291)

This commit is contained in:
Hao Chen
2018-06-29 10:51:23 +08:00
committed by Robert Nishihara
parent cd63804768
commit 20c0ecb522
3 changed files with 43 additions and 41 deletions
+10 -27
View File
@@ -30,7 +30,12 @@ import ray.signature
import ray.local_scheduler
import ray.plasma
import ray.ray_constants as ray_constants
from ray.utils import random_string, binary_to_hex, is_cython
from ray.utils import (
binary_to_hex,
check_oversized_pickle,
is_cython,
random_string,
)
SCRIPT_MODE = 0
WORKER_MODE = 1
@@ -653,18 +658,8 @@ class Worker(object):
else:
del function.__globals__[function.__name__]
if len(pickled_function) > ray_constants.PICKLE_OBJECT_WARNING_SIZE:
warning_message = ("Warning: The remote function {} has size {} "
"when pickled. It will be stored in Redis, "
"which could cause memory issues. This may "
"mean that the function definition uses a "
"large array or other object.".format(
function_name, len(pickled_function)))
ray.utils.push_error_to_driver(
self,
ray_constants.PICKLING_LARGE_OBJECT_PUSH_ERROR,
warning_message,
driver_id=self.task_driver_id.id())
check_oversized_pickle(pickled_function, function_name,
"remote function", self)
self.redis_client.hmset(
key, {
@@ -714,20 +709,8 @@ class Worker(object):
# we don't need to export it again.
return
if (len(pickled_function) >
ray_constants.PICKLE_OBJECT_WARNING_SIZE):
warning_message = ("Warning: The function {} has size {} when "
"pickled. It will be stored in Redis, "
"which could cause memory issues. This may "
"mean that the remote function definition "
"uses a large array or other object."
.format(function.__name__,
len(pickled_function)))
ray.utils.push_error_to_driver(
self,
ray_constants.PICKLING_LARGE_OBJECT_PUSH_ERROR,
warning_message,
driver_id=self.task_driver_id.id())
check_oversized_pickle(pickled_function, function.__name__,
"function", self)
# Run the function on all workers.
self.redis_client.hmset(