Commit df0fe1cf authored by Mitar's avatar Mitar
Browse files

Make sure lists are pickable, too.

parent 341c276e
Pipeline #15524235 passed with stage
in 2 minutes and 46 seconds
......@@ -94,14 +94,20 @@ class SumPrimitive(base.SingletonOutputMixin[Inputs, Outputs, None, Hyperparams]
if DOCKER_KEY not in self.docker_containers:
raise ValueError("Docker key '{docker_key}' missing among provided Docker containers.".format(docker_key=DOCKER_KEY))
def _convert_value(self, value: typing.Any) -> typing.Union[numpy.ndarray, typing.List, typing.Any]:
# Server does not know about container types, just standard numpy arrays and lists.
if isinstance(value, container.ndarray):
return value.view(numpy.ndarray)
elif isinstance(value, container.List):
return [self._convert_value(v) for v in value]
return value
def produce(self, *, inputs: Inputs, timeout: float = None, iterations: int = None) -> base.CallResult[Outputs]:
# In the future, we should store here data in Arrow format into
# Plasma store and just pass an ObjectId of data over HTTP.
if isinstance(inputs, container.ndarray):
value = inputs.view(numpy.ndarray)
data = pickle.dumps(value)
data = pickle.dumps(inputs)
value = self._convert_value(inputs)
data = pickle.dumps(value)
# TODO: Retry if connection fails.
# This connection can sometimes fail because the service inside a Docker container
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment