def CreateValue(self, request, context): """Creates a value embedded in the executor. Args: request: An instance of `executor_pb2.CreateValueRequest`. context: An instance of `grpc.ServicerContext`. Returns: An instance of `executor_pb2.CreateValueResponse`. """ py_typecheck.check_type(request, executor_pb2.CreateValueRequest) try: value, value_type = (executor_service_utils.deserialize_value( request.value)) value_id = str(uuid.uuid4()) future_val = asyncio.run_coroutine_threadsafe( self._executor.create_value(value, value_type), self._event_loop) with self._lock: self._values[value_id] = future_val return executor_pb2.CreateValueResponse( value_ref=executor_pb2.ValueRef(id=value_id)) except (ValueError, TypeError) as err: logging.error(traceback.format_exc()) context.set_code(grpc.StatusCode.INVALID_ARGUMENT) context.set_details(str(err)) return executor_pb2.CreateValueResponse()
def get_value(self, value_id): response = self._stub.Compute( executor_pb2.ComputeRequest(value_ref=executor_pb2.ValueRef( id=value_id))) py_typecheck.check_type(response, executor_pb2.ComputeResponse) value, _ = executor_service_utils.deserialize_value(response.value) return value
async def _compute(self, value_ref): py_typecheck.check_type(value_ref, executor_pb2.ValueRef) request = executor_pb2.ComputeRequest(value_ref=value_ref) response = self._stub.Compute(request) py_typecheck.check_type(response, executor_pb2.ComputeResponse) value, _ = executor_service_utils.deserialize_value(response.value) return value
def test_serialize_deserialize_sequence_of_namedtuples(self): test_tuple_type = collections.namedtuple('TestTuple', ['a', 'b', 'c']) def make_test_tuple(x): return test_tuple_type(a=x * 2, b=tf.cast(x, tf.int32), c=tf.cast(x - 1, tf.float32)) ds = tf.data.Dataset.range(5).map(make_test_tuple) element_type = computation_types.NamedTupleType([ ('a', tf.int64), ('b', tf.int32), ('c', tf.float32), ]) sequence_type = computation_types.SequenceType(element=element_type) value_proto, value_type = executor_service_utils.serialize_value( ds, sequence_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(value_type, sequence_type) y, type_spec = executor_service_utils.deserialize_value(value_proto) self.assertEqual(type_spec, sequence_type) actual_values = self.evaluate([y_val for y_val in y]) expected_values = [ test_tuple_type(a=x * 2, b=x, c=x - 1.) for x in range(5) ] for actual, expected in zip(actual_values, expected_values): self.assertAllClose(actual, expected)
def get_value(self, value_id: str): """Retrieves a value using the `Compute` endpoint.""" response = self._stub.Compute( executor_pb2.ComputeRequest(value_ref=executor_pb2.ValueRef( id=value_id))) py_typecheck.check_type(response, executor_pb2.ComputeResponse) value, _ = executor_service_utils.deserialize_value(response.value) return value
def test_serialize_deserialize_nested_tuple_value_without_names(self): x = tuple([10, 20]) x_type = computation_types.to_type(tuple([tf.int32, tf.int32])) value_proto, value_type = executor_service_utils.serialize_value(x, x_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), '<int32,int32>') y, type_spec = executor_service_utils.deserialize_value(value_proto) self.assertEqual(str(type_spec), str(x_type)) self.assertCountEqual(y, (10, 20))
def test_serialize_deserialize_tensor_value_with_nontrivial_shape(self): x = tf.constant([10, 20, 30]).numpy() value_proto, value_type = executor_service_utils.serialize_value( x, computation_types.TensorType(tf.int32, [3])) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), 'int32[3]') y, type_spec = executor_service_utils.deserialize_value(value_proto) self.assertEqual(str(type_spec), 'int32[3]') self.assertTrue(np.array_equal(x, y))
def test_serialize_deserialize_tensor_value_with_different_dtype(self): x = tf.constant(10.0).numpy() value_proto, value_type = (executor_service_utils.serialize_value( x, computation_types.TensorType(tf.int32))) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), 'int32') y, type_spec = executor_service_utils.deserialize_value(value_proto) self.assertEqual(str(type_spec), 'int32') self.assertEqual(y, 10)
def test_serialize_deserialize_sequence_of_scalars(self): ds = tf.data.Dataset.range(5).map(lambda x: x * 2) value_proto, value_type = executor_service_utils.serialize_value( ds, computation_types.SequenceType(tf.int64)) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), 'int64*') y, type_spec = executor_service_utils.deserialize_value(value_proto) self.assertEqual(str(type_spec), 'int64*') self.assertAllEqual([y_val for y_val in y], [x * 2 for x in range(5)])
def test_serialize_deserialize_federated_at_server(self): x = 10 x_type = type_factory.at_server(tf.int32) value_proto, value_type = executor_service_utils.serialize_value( x, x_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), 'int32@SERVER') y, type_spec = executor_service_utils.deserialize_value(value_proto) self.assertEqual(str(type_spec), str(x_type)) self.assertEqual(y, 10)
def test_serialize_deserialize_federated_at_clients(self): x = [10, 20] x_type = type_factory.at_clients(tf.int32) value_proto, value_type = executor_service_utils.serialize_value( x, x_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), '{int32}@CLIENTS') y, type_spec = executor_service_utils.deserialize_value(value_proto) self.assertEqual(str(type_spec), str(x_type)) self.assertEqual(y, [10, 20])
def test_serialize_deserialize_tensor_value(self): x = tf.constant(10.0).numpy() type_spec = computation_types.TensorType(tf.as_dtype(x.dtype), x.shape) value_proto, value_type = executor_service_utils.serialize_value( x, type_spec) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), 'float32') y, type_spec = executor_service_utils.deserialize_value(value_proto) self.assertEqual(str(type_spec), 'float32') self.assertTrue(np.array_equal(x, y))
async def _compute(self, value_ref): py_typecheck.check_type(value_ref, executor_pb2.ValueRef) request = executor_pb2.ComputeRequest(value_ref=value_ref) if not self._bidi_stream: response = self._stub.Compute(request) else: response = (await self._bidi_stream.send_request( executor_pb2.ExecuteRequest(compute=request))).compute py_typecheck.check_type(response, executor_pb2.ComputeResponse) value, _ = executor_service_utils.deserialize_value(response.value) return value
def test_serialize_deserialize_computation_value(self): @computations.tf_computation def comp(): return tf.constant(10) value_proto, value_type = executor_service_utils.serialize_value(comp) self.assertEqual(value_proto.WhichOneof('value'), 'computation') self.assertEqual(str(value_type), '( -> int32)') comp, type_spec = executor_service_utils.deserialize_value(value_proto) self.assertIsInstance(comp, computation_pb2.Computation) self.assertEqual(str(type_spec), '( -> int32)')
def test_serialize_deserialize_sequence_of_nested_structures(self): test_tuple_type = collections.namedtuple('TestTuple', ['u', 'v']) def _make_nested_tf_structure(x): return collections.OrderedDict([ ('b', tf.cast(x, tf.int32)), ('a', tuple([ x, test_tuple_type(x * 2, x * 3), collections.OrderedDict([('x', x**2), ('y', x**3)]) ])), ]) ds = tf.data.Dataset.range(5).map(_make_nested_tf_structure) element_type = computation_types.NamedTupleType([ ('b', tf.int32), ('a', computation_types.NamedTupleType([ (None, tf.int64), (None, test_tuple_type(tf.int64, tf.int64)), (None, computation_types.NamedTupleType([('x', tf.int64), ('y', tf.int64)])), ])), ]) sequence_type = computation_types.SequenceType(element=element_type) value_proto, value_type = executor_service_utils.serialize_value( ds, sequence_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(value_type, sequence_type) y, type_spec = executor_service_utils.deserialize_value(value_proto) # These aren't the same because ser/de destroys the PyContainer type_spec.check_equivalent_to(sequence_type) def _build_expected_structure(x): return collections.OrderedDict([ ('b', x), ('a', tuple([ x, test_tuple_type(x * 2, x * 3), collections.OrderedDict([('x', x**2), ('y', x**3)]) ])), ]) actual_values = self.evaluate([y_val for y_val in y]) expected_values = [_build_expected_structure(x) for x in range(5)] for actual, expected in zip(actual_values, expected_values): self.assertEqual(type(actual), type(expected)) self.assertAllClose(actual, expected)
def test_serialize_deserialize_nested_tuple_value_with_names(self): x = collections.OrderedDict([('a', 10), ('b', [20, 30]), ('c', collections.OrderedDict([('d', 40)]))]) x_type = computation_types.to_type( collections.OrderedDict([('a', tf.int32), ('b', [tf.int32, tf.int32]), ('c', collections.OrderedDict([('d', tf.int32) ]))])) value_proto, value_type = executor_service_utils.serialize_value(x, x_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), '<a=int32,b=<int32,int32>,c=<d=int32>>') y, type_spec = executor_service_utils.deserialize_value(value_proto) self.assertEqual(str(type_spec), str(x_type)) self.assertTrue(str(y), '<a=10,b=<20,30>,c=<d=40>>')
def test_serialize_deserialize_sequence_of_tuples(self): ds = tf.data.Dataset.range(5).map(lambda x: (x * 2, tf.cast( x, tf.int32), tf.cast(x - 1, tf.float32))) value_proto, value_type = executor_service_utils.serialize_value( ds, computation_types.SequenceType(element=(tf.int64, tf.int32, tf.float32))) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), '<int64,int32,float32>*') y, type_spec = executor_service_utils.deserialize_value(value_proto) self.assertEqual(str(type_spec), '<int64,int32,float32>*') self.assertAllEqual(self.evaluate([y_val for y_val in y]), [(x * 2, x, x - 1.) for x in range(5)])
def CreateValue( self, request: executor_pb2.CreateValueRequest, context: grpc.ServicerContext, ) -> executor_pb2.CreateValueResponse: """Creates a value embedded in the executor.""" py_typecheck.check_type(request, executor_pb2.CreateValueRequest) try: value, value_type = ( executor_service_utils.deserialize_value(request.value)) value_id = str(uuid.uuid4()) future_val = asyncio.run_coroutine_threadsafe( self._executor.create_value(value, value_type), self._event_loop) with self._lock: self._values[value_id] = future_val return executor_pb2.CreateValueResponse( value_ref=executor_pb2.ValueRef(id=value_id)) except (ValueError, TypeError) as err: _set_invalid_arg_err(context, err) return executor_pb2.CreateValueResponse()
def CreateValue( self, request: executor_pb2.CreateValueRequest, context: grpc.ServicerContext, ) -> executor_pb2.CreateValueResponse: """Creates a value embedded in the executor.""" py_typecheck.check_type(request, executor_pb2.CreateValueRequest) try: with tracing.span('ExecutorService.CreateValue', 'deserialize_value'): value, value_type = ( executor_service_utils.deserialize_value(request.value)) value_id = str(uuid.uuid4()) coro = self._executor.create_value(value, value_type) future_val = self._run_coro_threadsafe_with_tracing(coro) with self._lock: self._values[value_id] = future_val return executor_pb2.CreateValueResponse( value_ref=executor_pb2.ValueRef(id=value_id)) except (ValueError, TypeError) as err: _set_invalid_arg_err(context, err) return executor_pb2.CreateValueResponse()
def CreateValue(self, request, context): """Creates a value embedded in the executor. Args: request: An instance of `executor_pb2.CreateValueRequest`. context: An instance of `grpc.ServicerContext`. Returns: An instance of `executor_pb2.CreateValueResponse`. """ py_typecheck.check_type(request, executor_pb2.CreateValueRequest) try: value, value_type = (executor_service_utils.deserialize_value( request.value)) value_id = str(uuid.uuid4()) future_val = self._executor.create_value(value, value_type) with self._lock: self._values[value_id] = future_val return executor_pb2.CreateValueResponse( value_ref=executor_pb2.ValueRef(id=value_id)) except (ValueError, TypeError): context.set_code(grpc.StatusCode.INVALID_ARGUMENT) return executor_pb2.CreateValueResponse()