wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/python/distribute/experimental/rpc/BUILD
+64
View File
@@ -0,0 +1,64 @@
# Python bindings for RPC client and server ops.
load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
pytype_strict_library(
name = "rpc_ops",
srcs = [
"rpc_ops.py",
],
deps = [
"//tensorflow/distribute/experimental/rpc/kernels:gen_rpc_ops",
"//tensorflow/distribute/experimental/rpc/proto:tf_rpc_service_proto_py",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:none_tensor",
"//tensorflow/python/framework:type_spec",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/saved_model:nested_structure_coder",
"//tensorflow/python/types:core",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:tf_export",
],
)
tf_py_strict_test(
name = "rpc_ops_test",
size = "medium",
srcs = ["rpc_ops_test.py"],
shard_count = 7,
tags = [
"no_mac", # flaky, b/205156709
],
deps = [
":rpc_ops",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:nest",
"@pypi//portpicker",
],
)
tensorflow/python/distribute/experimental/rpc/rpc_ops.py
+505
View File
@@ -0,0 +1,505 @@
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Module to expose RPC APIs in tensorflow."""
from typing import Optional, Sequence, Union
import tensorflow.distribute.experimental.rpc.kernels.gen_rpc_ops as gen_rpc_ops
from tensorflow.distribute.experimental.rpc.proto import tf_rpc_service_pb2 as rpc_pb2
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.eager import function as tf_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import none_tensor
from tensorflow.python.framework import type_spec
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.saved_model import nested_structure_coder
from tensorflow.python.types import core as core_tf_types
from tensorflow.python.util import nest
from tensorflow.python.util.tf_export import tf_export
def get_output_specs_from_function(func: tf_function.ConcreteFunction):
output_specs = nest.map_structure(type_spec.type_spec_from_value,
func.structured_outputs)
output_specs_proto = nested_structure_coder.encode_structure(output_specs)
return output_specs_proto.SerializeToString()
def get_input_specs_from_function(func: tf_function.ConcreteFunction):
arg_specs, _ = func.structured_input_signature
arg_specs_proto = nested_structure_coder.encode_structure(arg_specs)
return arg_specs_proto.SerializeToString()
@tf_export("distribute.experimental.rpc.Server", v1=[])
class Server(object):
"""A Server base class for accepting RPCs for registered tf.functions.
Functions can be registered on the server and are exposed via RPCs.
"""
@staticmethod
def create(rpc_layer, address):
"""Create TF RPC server at given address.
Args:
rpc_layer: Communication layer between client and server. Only "grpc" rpc
layer is supported at the moment.
address: Address where RPC server is hosted.
Returns:
An instance of `tf.distribute.experimental.rpc.Server` class.
Raises:
A ValueError if rpc_layer other than "grpc" is used. Only GRPC
is supported at the moment.
Example usage:
>>> import portpicker
>>> @tf.function(input_signature=[
... tf.TensorSpec([], tf.int32),
... tf.TensorSpec([], tf.int32)])
... def remote_fn(a, b):
... return tf.add(a, b)
>>> port = portpicker.pick_unused_port()
>>> address = "localhost:{}".format(port)
>>> server = tf.distribute.experimental.rpc.Server.create("grpc", address)
>>> server.register("addition", remote_fn)
>>> server.start()
"""
if rpc_layer != "grpc":
raise ValueError("Only GRPC backend is supported at the moment.")
return GrpcServer(address=address)
def register(self, method_name: str,
func: Union[def_function.Function,
tf_function.ConcreteFunction]):
"""Method for registering tf.function on server.
Registered methods can be invoked remotely from clients.
Args:
method_name: Name of the tf.function. Clients use this method_name to make
RPCs.
func: A `tf.function` or ConcreteFunction to register.
"""
raise NotImplementedError("Please use create_server method to create a"
"concrete subclass of Server.")
def start(self):
"""Starts the RPC server on provided address.
Server listens for new requests from client, once it is started.
"""
raise NotImplementedError("Please use create_server method to create a"
"concrete subclass of Server.")
@tf_export("distribute.experimental.rpc.Client", v1=[])
class Client(object):
"""Client class for invoking RPCs to the server."""
@staticmethod
def create(rpc_layer, address, name="", timeout_in_ms=0):
"""Create TF RPC client to connect to the given address.
Args:
rpc_layer: Communication layer between client and server. Only "grpc" rpc
layer is supported at the moment.
address: Address of the server to connect the RPC client to.
name: Name of the RPC Client. You can create multiple clients connecting
to same server and distinguish them using different names.
timeout_in_ms: The default timeout to use for outgoing RPCs from client. 0
indicates no timeout. Exceeding timeout during RPC will raise
DeadlineExceeded error.
Returns:
An instance of `tf.distribute.experimental.rpc.Client` with the following
dynamically added methods for eagerly created clients:
* `Registered methods` e.g. multiply(**args):
If Client is created when executing eagerly, client will request the
list of registered methods from server during client creation.
The convenience methods for RPCs will be dynamically added to the
created Client instance.
For example, when a server has method "multiply" registered, the
client object created in eager mode will have 'multiply' method
available. Users can use client.multiply(..) to make RPC, instead of
client.call("multiply", ...)
Both "call" and "multiply" methods are non-blocking i.e. they return
a StatusOrResult object which should be used to wait for getting
value or error.
Along with the above, blocking versions of the registered
methods are also dynamically added to client instance.
e.g. multiply_blocking(**args). These methods block till the RPC is
finished and return response for successful RPC. Otherwise raise
exception.
These methods are not available when Client is created inside a
tf.function.
Raises:
A ValueError if rpc_layer other than "grpc" is used. Only GRPC
is supported at the moment.
A DeadlineExceeded exception in eager mode if timeout exceeds while
creating and listing client methods.
Example usage:
>>> # Have server already started.
>>> import portpicker
>>> @tf.function(input_signature=[
... tf.TensorSpec([], tf.int32),
... tf.TensorSpec([], tf.int32)])
... def remote_fn(a, b):
... return tf.add(a, b)
>>> port = portpicker.pick_unused_port()
>>> address = "localhost:{}".format(port)
>>> server = tf.distribute.experimental.rpc.Server.create("grpc", address)
>>> server.register("addition", remote_fn)
>>> server.start()
>>> # Start client
>>> client = tf.distribute.experimental.rpc.Client.create("grpc",
... address=address, name="test_client")
>>> a = tf.constant(2, dtype=tf.int32)
>>> b = tf.constant(3, dtype=tf.int32)
>>> result = client.call(
... args=[a, b],
... method_name="addition",
... output_specs=tf.TensorSpec((), tf.int32))
>>> if result.is_ok():
... result.get_value()
>>> result = client.addition(a, b)
>>> if result.is_ok():
... result.get_value()
>>> value = client.addition_blocking(a, b)
"""
if rpc_layer != "grpc":
raise ValueError("Only GRPC backend is supported at the moment.")
if context.executing_eagerly():
list_registered_methods = True
else:
list_registered_methods = False
return GrpcClient(
address=address,
name=name,
list_registered_methods=list_registered_methods,
timeout_in_ms=timeout_in_ms)
def call(self,
method_name: str,
args: Optional[Sequence[core_tf_types.Tensor]] = None,
output_specs=None,
timeout_in_ms=0):
"""Method for making RPC calls to remote server.
This invokes RPC to the server, executing the registered method_name
remotely.
Args:
method_name: Remote registered method to invoke
args: List of arguments for the registered method.
output_specs: Output specs for the output from method.
For example, if tf.function is: @tf.function(input_signature=[
tf.TensorSpec([], tf.int32), tf.TensorSpec([], tf.int32) ])
def multiply_fn(a, b): return tf.math.multiply(a, b)
output_spec is: tf.TensorSpec((), tf.int32) If you have access to TF
Function, the output specs can be generated
from tf.function by calling: output_specs =
tf.nest.map_structure(tf.type_spec_from_value,
tf_function.get_concrete_function().structured_outputs If output_specs
are not provided, flattened list of tensors will be returned in
response.
timeout_in_ms: Timeout for this call. If 0, default client timeout will be
used.
Returns:
An instance of `StatusOrResult` class with the following available
methods.
* `is_ok()`:
Returns True of RPC was successful.
* `get_error()`:
Returns TF error_code and error message for the RPC.
* `get_value()`:
Returns the returned value from remote TF function execution
when RPC is successful.
Calling any of the above methods will block till RPC is completed and
result is available.
"""
raise NotImplementedError("Must be implemented in inherited classes.")
class GrpcServer(Server):
"""GrpcServer object encapsulates a resource with GRPC server.
Functions can be registered locally and are exposed via RPCs.
Example:
```
server = rpc_ops.GrpcServer("host:port")
@tf.function
def add(a, b):
return a + b
server.register("add", add)
server.start()
```
"""
def __init__(self, address: str):
self._server_handle = gen_rpc_ops.rpc_server(address)
if context.executing_eagerly():
self._handle_deleter = resource_variable_ops.EagerResourceDeleter(
handle=self._server_handle, handle_device=self._server_handle.device)
else:
raise NotImplementedError("Please create the server outside tf.function.")
def register(self, method_name: str,
func: Union[def_function.Function,
tf_function.ConcreteFunction]):
"""Method for registering functions."""
if isinstance(func, def_function.Function):
if func.function_spec.arg_names:
if func.input_signature is None:
raise ValueError("Input signature not specified for the function.")
concrete_fn = func.get_concrete_function()
gen_rpc_ops.rpc_server_register(
self._server_handle,
method_name=method_name,
captured_inputs=concrete_fn.captured_inputs,
input_specs=get_input_specs_from_function(concrete_fn),
output_specs=get_output_specs_from_function(concrete_fn),
f=concrete_fn)
elif isinstance(func, tf_function.ConcreteFunction):
gen_rpc_ops.rpc_server_register(
self._server_handle,
method_name=method_name,
captured_inputs=func.captured_inputs,
input_specs=get_input_specs_from_function(func),
output_specs=get_output_specs_from_function(func),
f=func)
else:
# Python functions
# TODO(b/186762191): Add an implementation to support python functions.
raise ValueError("Only TF functions are supported with Register method")
def start(self):
"""Starts GRPC server."""
gen_rpc_ops.rpc_server_start(self._server_handle)
class GrpcClient(Client):
"""Client wrapper to connect to remote RPC server using GRPC.
If Client is created with (list_registered_methods=True):
1. Input and output specs for the methods till this point will be fetched from
Server.
2. convenience methods are added to invoke registered methods directly from
client.
For example:
For call a server method `add`
client.add(a, b) or client.add_async(a, b) can be used instead of
client.call(args=[a,b], output_specs=[..])
Prerequisite for using list_registered_methods=True:
1. Server should be already started with the registered methods.
2. Client must be created in Eager mode.
"""
def __init__(self,
address: str,
name: str = "",
list_registered_methods=False,
timeout_in_ms=0):
self._client_handle, methods = gen_rpc_ops.rpc_client(
shared_name=name,
server_address=address,
list_registered_methods=list_registered_methods,
timeout_in_ms=timeout_in_ms)
if context.executing_eagerly():
self._handle_deleter = resource_variable_ops.EagerResourceDeleter(
handle=self._client_handle, handle_device=self._client_handle.device)
else:
raise NotImplementedError(
"Client creation is supported only in eager mode.")
self._server_address = address
self._method_registry = {}
for method in methods.numpy():
m = rpc_pb2.RegisteredMethod()
m.ParseFromString(method)
output_specs = nested_structure_coder.decode_proto(m.output_specs)
input_specs = nested_structure_coder.decode_proto(m.input_specs)
self._method_registry[m.method] = output_specs
# TODO(ishark): Perhaps doc string can also be taken as input during
# function registration.
doc_string = "RPC Call for " + m.method + " method to server " + address
self._add_method(m.method, output_specs, input_specs, self._client_handle,
doc_string)
def _add_method(self, method_name, output_specs, input_specs, client_handle,
doc_string):
"""Method to add RPC methods to the client object."""
def validate_and_get_flat_inputs(*args):
if args is None:
args = []
if input_specs:
nest.assert_same_structure(args, input_specs)
flat_inputs = nest.flatten(args)
return flat_inputs
def call_wrapper(*args, timeout_in_ms=0):
status_or, deleter = gen_rpc_ops.rpc_call(
client_handle,
args=validate_and_get_flat_inputs(*args),
method_name=method_name,
timeout_in_ms=timeout_in_ms)
return StatusOrResult(status_or, deleter, output_specs)
def call_blocking_wrapper(*args, timeout_in_ms=0):
status_or, deleter = gen_rpc_ops.rpc_call(
client_handle,
args=validate_and_get_flat_inputs(*args),
method_name=method_name,
timeout_in_ms=timeout_in_ms)
status_or = StatusOrResult(status_or, deleter, output_specs)
if status_or.is_ok():
return status_or.get_value()
else:
error_code, error_msg = status_or.get_error()
raise errors.exception_type_from_error_code(error_code.numpy())(
None, None, error_msg.numpy())
setattr(self, method_name, call_wrapper)
call_wrapper.__doc__ = doc_string
blocking_method_name = method_name + "_blocking"
setattr(self, blocking_method_name, call_blocking_wrapper)
call_blocking_wrapper.__doc__ = doc_string
def call(self,
method_name: str,
args: Optional[Sequence[core_tf_types.Tensor]] = None,
output_specs=None,
timeout_in_ms=0):
"""Method to invoke remote registered functions on the connected server.
Server should be started before making an RPC Call.
Args:
method_name: Registered method to invoke on Server.
args: Input arguments for the method.
output_specs: Output specs for the output from method.
timeout_in_ms: Timeout for this call. If 0, default client timeout will be
used.
Returns:
StatusOrResult object. This function issues the RPC call to server, it
does not block for the duration of RPC. Please call is_ok, get_error or
get_value methods on the returned object to blocked till RPC finishes.
"""
if args is None:
args = []
status_or, deleter = gen_rpc_ops.rpc_call(
self._client_handle,
args=nest.flatten(args),
method_name=method_name,
timeout_in_ms=timeout_in_ms)
return StatusOrResult(status_or, deleter, output_specs)
class StatusOrResult(object):
"""Class representing result and status from RPC Call."""
def __init__(self, status_or, deleter, output_specs=None):
self._status_or = status_or
self._output_specs = output_specs
self._deleter = deleter
self._error_code: dtypes.int64 = None
self._error_message: dtypes.string = None
def _check_status(self):
if self._error_code is None:
self._error_code, self._error_message = gen_rpc_ops.rpc_check_status(
self._status_or)
def __del__(self):
# Make sure the resource is deleted in the same mode as it was created in.
if context.executing_eagerly():
with context.eager_mode():
gen_rpc_ops.delete_rpc_future_resource(
handle=self._status_or, deleter=self._deleter)
else:
with context.graph_mode():
gen_rpc_ops.delete_rpc_future_resource(
handle=self._status_or, deleter=self._deleter)
def is_ok(self):
"""Returns True if RPC is successful, otherwise returns False.
This call will block for RPC result.
"""
self._check_status()
return math_ops.equal(self._error_code,
constant_op.constant(0, dtype=dtypes.int64))
def get_error(self):
"""Returns (TF Error Code, Error Message) from RPC Response.
This call will block for RPC result.
"""
self._check_status()
return self._error_code, self._error_message
def get_value(self):
"""Returns the returned response value from RPC Call when RPC is successful.
The returned value is tensors in the output_specs format as returned from
the RPC call
This call will block for RPC result.
"""
self._check_status()
if self._output_specs is None or isinstance(self._output_specs,
none_tensor.NoneTensorSpec):
flat_output_dtypes = []
return_none = True
else:
return_none = False
flat_output_dtypes = [s.dtype for s in nest.flatten(self._output_specs)]
result = gen_rpc_ops.rpc_get_value(self._status_or, Tout=flat_output_dtypes)
if return_none:
return None
else:
return nest.pack_sequence_as(self._output_specs, result)
tensorflow/python/distribute/experimental/rpc/rpc_ops_test.py
+850
View File
@@ -0,0 +1,850 @@
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for rpc_ops.py."""
import threading
import time
import numpy as np
import portpicker
from tensorflow.python.distribute.experimental.rpc import rpc_ops
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function as eager_def_function
from tensorflow.python.framework import config
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_spec
from tensorflow.python.framework import test_util
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.util import nest
@test_util.with_eager_op_as_function
class RpcOpsTest(test.TestCase):
def setUp(self):
super(RpcOpsTest, self).setUp()
cpus = config.list_physical_devices("CPU")
# Set 2 virtual CPUs
config.set_logical_device_configuration(cpus[0], [
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration()
])
def test_generated_rpc_ops(self):
@eager_def_function.function(input_signature=[
tensor_spec.TensorSpec([], dtypes.int32),
tensor_spec.TensorSpec([], dtypes.int32)
])
def remote_fn(a, b):
return math_ops.multiply(a, b)
concrete_remote_fn = remote_fn.get_concrete_function()
a = variables.Variable(2, dtype=dtypes.int32)
b = variables.Variable(3, dtype=dtypes.int32)
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server_resource = rpc_ops.gen_rpc_ops.rpc_server(server_address=address)
rpc_ops.gen_rpc_ops.rpc_server_register(
server_resource,
f=concrete_remote_fn,
captured_inputs=concrete_remote_fn.captured_inputs,
output_specs=rpc_ops.get_output_specs_from_function(concrete_remote_fn),
method_name="multiply")
rpc_ops.gen_rpc_ops.rpc_server_start(server_resource)
client_handle, _ = rpc_ops.gen_rpc_ops.rpc_client(
server_address=address, timeout_in_ms=5000)
future_resource, deleter = rpc_ops.gen_rpc_ops.rpc_call(
client_handle, args=[a, b], method_name="multiply", timeout_in_ms=0)
error_code, _ = rpc_ops.gen_rpc_ops.rpc_check_status(future_resource)
self.assertAllEqual(error_code, 0)
self.assertAllEqual(
rpc_ops.gen_rpc_ops.rpc_get_value(future_resource, Tout=[dtypes.int32]),
[6])
resource_variable_ops.EagerResourceDeleter(
handle=server_resource, handle_device=server_resource.device)
resource_variable_ops.EagerResourceDeleter(
handle=client_handle, handle_device=client_handle.device)
rpc_ops.gen_rpc_ops.delete_rpc_future_resource(future_resource, deleter)
def test_exported_rpc_api_static_factory(self):
@eager_def_function.function(input_signature=[
tensor_spec.TensorSpec([], dtypes.int32),
tensor_spec.TensorSpec([], dtypes.int32)
])
def _remote_fn(a, b):
return math_ops.multiply(a, b)
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server_resource = rpc_ops.Server.create("grpc", address)
server_resource.register("multiply", _remote_fn)
server_resource.start()
client = rpc_ops.Client.create("grpc", address=address, name="test_client")
a = variables.Variable(2, dtype=dtypes.int32)
b = variables.Variable(3, dtype=dtypes.int32)
mul_or = client.call(
args=[a, b],
method_name="multiply",
output_specs=tensor_spec.TensorSpec((), dtypes.int32))
self.assertAllEqual(mul_or.is_ok(), True)
self.assertAllEqual(mul_or.get_value(), 6)
# Test empty client name
client1 = rpc_ops.Client.create("grpc", address)
mul_or = client1.call(
args=[a, b],
method_name="multiply",
output_specs=tensor_spec.TensorSpec((), dtypes.int32))
self.assertAllEqual(mul_or.is_ok(), True)
self.assertAllEqual(mul_or.get_value(), 6)
# Test without output_spec
mul_or = client1.multiply(a, b)
self.assertAllEqual(mul_or.is_ok(), True)
self.assertAllEqual(mul_or.get_value(), 6)
self.assertEqual(client1.multiply.__doc__,
"RPC Call for multiply method to server " + address)
def test_rpc_ops_call_method(self):
@eager_def_function.function(input_signature=[
tensor_spec.TensorSpec([], dtypes.int32),
tensor_spec.TensorSpec([], dtypes.int32)
])
def _remote_fn(a, b):
return math_ops.multiply(a, b)
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server_resource = rpc_ops.GrpcServer(address)
@eager_def_function.function(input_signature=[
tensor_spec.TensorSpec([], dtypes.int32),
tensor_spec.TensorSpec([], dtypes.int32)
])
def add_fn(a, b):
return math_ops.add(a, b)
# Register TF function
server_resource.register("multiply", _remote_fn)
# Register concrete Function
server_resource.register("add", add_fn.get_concrete_function())
server_resource.start()
client = rpc_ops.GrpcClient(address=address, name="test_client")
a = variables.Variable(2, dtype=dtypes.int32)
b = variables.Variable(3, dtype=dtypes.int32)
mul_or = client.call(
args=[a, b],
method_name="multiply",
output_specs=tensor_spec.TensorSpec((), dtypes.int32))
self.assertAllEqual(mul_or.is_ok(), True)
self.assertAllEqual(mul_or.get_value(), 6)
add_or = client.call(
args=[a, b],
method_name="add",
output_specs=tensor_spec.TensorSpec((), dtypes.int32))
self.assertAllEqual(add_or.is_ok(), True)
self.assertAllEqual(add_or.get_value(), 5)
# Test empty client name
client1 = rpc_ops.GrpcClient(address, list_registered_methods=True)
mul_or = client1.call(
args=[a, b],
method_name="multiply",
output_specs=tensor_spec.TensorSpec((), dtypes.int32))
self.assertAllEqual(mul_or.is_ok(), True)
self.assertAllEqual(mul_or.get_value(), 6)
def test_rpc_ops_non_blocking_convenience_methods(self):
@eager_def_function.function(input_signature=[
tensor_spec.TensorSpec([], dtypes.int32),
tensor_spec.TensorSpec([], dtypes.int32)
])
def _remote_fn(a, b):
return math_ops.multiply(a, b)
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server_resource = rpc_ops.GrpcServer(address)
# Register TF function
server_resource.register("multiply", _remote_fn)
server_resource.start()
a = variables.Variable(2, dtype=dtypes.int32)
b = variables.Variable(3, dtype=dtypes.int32)
client = rpc_ops.GrpcClient(address, list_registered_methods=True)
mul_or = client.multiply(a, b)
self.assertAllEqual(mul_or.is_ok(), True)
self.assertAllEqual(mul_or.get_value(), 6)
self.assertEqual(client.multiply.__doc__,
"RPC Call for multiply method to server " + address)
def test_rpc_ops_blocking_convenience_methods(self):
@eager_def_function.function(input_signature=[
tensor_spec.TensorSpec([], dtypes.int32),
tensor_spec.TensorSpec([], dtypes.int32)
])
def _remote_fn(a, b):
return math_ops.multiply(a, b)
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server_resource = rpc_ops.GrpcServer(address)
# Register TF function
server_resource.register("multiply", _remote_fn)
server_resource.start()
client = rpc_ops.GrpcClient(address, list_registered_methods=True)
a = variables.Variable(2, dtype=dtypes.int32)
b = variables.Variable(3, dtype=dtypes.int32)
self.assertAllEqual(client.multiply_blocking(a, b), 6)
self.assertEqual(
client.multiply_blocking.__doc__,
"RPC Call for multiply method to server " + address)
def test_output_specs(self):
@eager_def_function.function(
input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
def test_dict(val):
return {"key": val}
@eager_def_function.function(
input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
def is_positive(a):
if a > 0:
return True
return False
@eager_def_function.function(input_signature=[])
def do_nothing():
return []
@eager_def_function.function(
input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
def test_nested_structure(v):
return {"test": (v, [v, v]), "test1": (v,)}
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server_resource = rpc_ops.GrpcServer(address)
server_resource.register("test_dict", test_dict)
server_resource.register("is_positive", is_positive)
server_resource.register("test_nested_structure", test_nested_structure)
server_resource.register("do_nothing", do_nothing)
server_resource.start()
client = rpc_ops.GrpcClient(
address=address, name="test_client", list_registered_methods=True)
a = variables.Variable(2, dtype=dtypes.int32)
result_or = client.test_dict(a)
self.assertAllEqual(result_or.is_ok(), True)
nest.map_structure(self.assertAllEqual, result_or.get_value(), {"key": 2})
result_or = client.is_positive(a)
self.assertTrue(result_or.is_ok())
self.assertTrue(result_or.get_value())
result_or = client.test_nested_structure(a)
self.assertAllEqual(result_or.is_ok(), True)
nest.map_structure(self.assertAllEqual, result_or.get_value(), {
"test": (2, [2, 2]),
"test1": (2,)
})
result_or = client.do_nothing()
self.assertAllEqual(result_or.is_ok(), True)
self.assertAllEqual(result_or.get_value(), [])
def test_input_specs(self):
@eager_def_function.function(input_signature=[{
"a": tensor_spec.TensorSpec([], dtypes.int32),
"b": tensor_spec.TensorSpec([], dtypes.int32)
}])
def test_input_dict(value):
return math_ops.add(value["a"], value["b"])
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server_resource = rpc_ops.GrpcServer(address)
server_resource.register("test_input_dict", test_input_dict)
server_resource.start()
client = rpc_ops.GrpcClient(
address=address, name="test_client", list_registered_methods=True)
a = variables.Variable(2, dtype=dtypes.int32)
b = variables.Variable(3, dtype=dtypes.int32)
result_or = client.test_input_dict({"a": a, "b": b})
self.assertAllEqual(result_or.is_ok(), True)
self.assertAllEqual(result_or.get_value(), 5)
with self.assertRaises(TypeError):
client.test_input_dict([a, b])
def test_call_register_ordering(self):
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
# Create client succeeds before server start and registration
client = rpc_ops.GrpcClient(address)
# Create client with list_registered_methods fails before server is started.
with self.assertRaises(errors.DeadlineExceededError):
rpc_ops.GrpcClient(
address,
name="client1",
list_registered_methods=True,
timeout_in_ms=1)
v = variables.Variable(initial_value=0, dtype=dtypes.int64)
@eager_def_function.function(
input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
def assign_add(a):
v.assign_add(a)
@eager_def_function.function(input_signature=[])
def read_var():
return v.value()
server = rpc_ops.GrpcServer(address)
def start_server():
# Delay server start to test whether client creation also waits
# till server is up.
time.sleep(1)
server.register("assign_add", assign_add)
server.start()
t = threading.Thread(target=start_server)
t.start()
# Create same "client1" again should succeed.
client1_with_listed_methods = rpc_ops.GrpcClient(
address, name="client1", list_registered_methods=True)
result_or = client1_with_listed_methods.assign_add(
variables.Variable(2, dtype=dtypes.int64))
self.assertAllEqual(result_or.is_ok(), True)
result_or = client.call("assign_add",
[variables.Variable(2, dtype=dtypes.int64)])
self.assertAllEqual(result_or.is_ok(), True)
# Create client with registered methods
client2_with_listed_methods = rpc_ops.GrpcClient(
address=address, name="client2", list_registered_methods=True)
result_or = client2_with_listed_methods.assign_add(
variables.Variable(2, dtype=dtypes.int64))
self.assertAllEqual(result_or.is_ok(), True)
self.assertAllEqual(v, 6)
# Register new method after server started.
with self.assertRaisesRegex(
errors.FailedPreconditionError,
"All methods must be registered before starting the server"):
server.register("read_var", read_var)
def test_client_timeout(self):
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
@eager_def_function.function(input_signature=[
tensor_spec.TensorSpec([], dtypes.int32),
tensor_spec.TensorSpec([], dtypes.int32)
])
def add(a, b):
return math_ops.add(a, b)
server = rpc_ops.GrpcServer(address)
def start_server():
# Delay server start to simulate deadline exceeded for 1st RPC call
# response. Client waits till server is started, thus it can trigger
# deadline exceeded.
time.sleep(1)
server.register("add", add)
server.start()
t = threading.Thread(target=start_server)
t.start()
def ensure_server_is_ready(client):
server_ready = False
while not server_ready:
result_or = client.call(
"add", [constant_op.constant(20),
constant_op.constant(30)])
if result_or.is_ok():
server_ready = True
else:
error_code, _ = result_or.get_error()
if error_code == errors.UNAVAILABLE:
server_ready = False
else:
server_ready = True
return
# Create client with list_registered_methods fails before server is started.
with self.assertRaises(errors.DeadlineExceededError):
rpc_ops.GrpcClient(
address,
name="client1",
list_registered_methods=True,
timeout_in_ms=1)
# Create same client again should succeed with
# list_registered_methods=False. Default timeout for client is 1 ms.
client = rpc_ops.GrpcClient(
address, name="client1", list_registered_methods=False, timeout_in_ms=1)
ensure_server_is_ready(client)
# Make explicit RPC call, the timeout of 1 ms should lead to
# deadline exceeded error.
result_or = client.call(
"add", [constant_op.constant(20),
constant_op.constant(30)],
timeout_in_ms=1)
self.assertAllEqual(result_or.is_ok(), False)
error_code, error_message = result_or.get_error()
self.assertAllEqual(error_code, errors.DEADLINE_EXCEEDED, error_message)
# Specifying reasonable timeout for call should succeed.
result_or = client.call(
"add", [constant_op.constant(20),
constant_op.constant(30)],
timeout_in_ms=5000)
self.assertAllEqual(result_or.is_ok(), True)
error_code, _ = result_or.get_error()
# Test timeouts for convenience methods
# Restart server again with delay to simulate deadline exceeded.
del server
server = rpc_ops.GrpcServer(address)
t = threading.Thread(target=start_server)
t.start()
# Client with no default timeout.
client = rpc_ops.GrpcClient(
address, name="client2", list_registered_methods=True)
# Succeeds with reasonable timeout.
result_or = client.add(
constant_op.constant(20), constant_op.constant(30), timeout_in_ms=5000)
self.assertAllEqual(result_or.is_ok(), True)
def test_async_call_op_wrapper(self):
v = variables.Variable(initial_value=0, dtype=dtypes.int64)
@eager_def_function.function(
input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
def assign_add(a):
v.assign_add(a)
@eager_def_function.function(input_signature=[])
def read_var():
return v.value()
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server = rpc_ops.GrpcServer(address)
server.register("assign_add", assign_add)
server.register("read_var", read_var)
server.start()
client = rpc_ops.GrpcClient(address)
futures = []
for _ in range(10):
futures.append(
client.call("assign_add",
[variables.Variable(2, dtype=dtypes.int64)]))
for f in futures:
f.is_ok()
result_or = client.call(
"read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
self.assertAllEqual(result_or.is_ok(), True)
self.assertAllEqual(result_or.get_value(), [20])
def test_rpc_call_op_in_tf_function(self):
@eager_def_function.function(input_signature=[
tensor_spec.TensorSpec([], dtypes.int32),
tensor_spec.TensorSpec([], dtypes.int32)
])
def _remote_fn(a, b):
return math_ops.multiply(a, b)
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server_resource = rpc_ops.GrpcServer(address)
server_resource.register("remote_fn", _remote_fn)
server_resource.start()
client = rpc_ops.GrpcClient(address=address, name="test_client")
a = variables.Variable(2, dtype=dtypes.int32)
b = variables.Variable(3, dtype=dtypes.int32)
@eager_def_function.function
def call_fn():
result_or = client.call(
args=[a, b],
method_name="remote_fn",
output_specs=[tensor_spec.TensorSpec([], dtypes.int32)])
self.assertAllEqual(True, result_or.is_ok())
result = result_or.get_value()
self.assertEqual(len(result), 1) # Call returns a list(tensors)
# TODO(ishark): Shape for output tensor is unknown currently.
# Add attribute for capturing TensorSpec for output and enable
# check below:
# self.assertIsNotNone(result[0].shape.rank)
return result
self.assertAllEqual(call_fn(), [6])
def test_resource_deletion(self):
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server = rpc_ops.GrpcServer(address)
server_handle = server._server_handle
# Test Future resource deletion
v = variables.Variable(initial_value=0, dtype=dtypes.int64)
@eager_def_function.function(input_signature=[])
def read_var():
return v.value()
server.register("read_var", read_var)
server.start()
client = rpc_ops.GrpcClient(address)
client_handle = client._client_handle
# Check future resource deletion without calling get_value.
def _create_and_delete_rpc_future():
handle = client.call(
"read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
return handle._status_or
@eager_def_function.function
def _create_and_delete_rpc_future_fn():
handle = client.call(
"read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
return handle._status_or
for _ in range(2):
handle = _create_and_delete_rpc_future()
with self.assertRaises(errors.NotFoundError):
resource_variable_ops.destroy_resource_op(
handle, ignore_lookup_error=False)
for _ in range(2):
handle = _create_and_delete_rpc_future_fn()
with self.assertRaises(errors.NotFoundError):
resource_variable_ops.destroy_resource_op(
handle, ignore_lookup_error=False)
# Check future resource deletion with calling get_value.
def _create_and_delete_with_future():
handle = client.call(
"read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
status_or_handle = handle._status_or
handle.get_value()
return status_or_handle
# Check future resource deletion with calling get_value with tf.function.
@eager_def_function.function
def _create_and_delete_with_future_fn():
handle = client.call(
"read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
status_or_handle = handle._status_or
handle.get_value()
return status_or_handle
for _ in range(2):
resource_handle = _create_and_delete_with_future()
with self.assertRaises(errors.NotFoundError):
resource_variable_ops.destroy_resource_op(
resource_handle, ignore_lookup_error=False)
for _ in range(2):
resource_handle = _create_and_delete_with_future_fn()
with self.assertRaises(errors.NotFoundError):
resource_variable_ops.destroy_resource_op(
resource_handle, ignore_lookup_error=False)
# Test server client resource gets deleted.
del client
with self.assertRaises(errors.NotFoundError):
resource_variable_ops.destroy_resource_op(
client_handle, ignore_lookup_error=False)
# Test server server resource gets deleted.
del server
with self.assertRaises(errors.NotFoundError):
resource_variable_ops.destroy_resource_op(
server_handle, ignore_lookup_error=False)
def test_rpc_error(self):
v = variables.Variable(initial_value=0, dtype=dtypes.int64)
@eager_def_function.function(
input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
def assign_add(a):
v.assign_add(a)
@eager_def_function.function(input_signature=[])
def read_var():
return v.value()
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server = rpc_ops.GrpcServer(address)
server.register("assign_add", assign_add)
server.register("read_var", read_var)
server.start()
client = rpc_ops.GrpcClient(address, list_registered_methods=True)
# confirm it works as expected when arguments are passed.
result_or = client.call("assign_add",
[variables.Variable(2, dtype=dtypes.int64)])
self.assertAllEqual(result_or.is_ok(), True)
result_or = client.call(
"read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
self.assertAllEqual(result_or.is_ok(), True)
self.assertAllEqual(result_or.get_value(), [2])
result_or = client.assign_add(variables.Variable(2, dtype=dtypes.int64))
self.assertAllEqual(True, result_or.is_ok())
result_or = client.read_var()
self.assertAllEqual(True, result_or.is_ok())
self.assertAllEqual(result_or.get_value(), 4)
# Fails with invalid argument error when no arguments are passed.
result_or = client.call("assign_add")
self.assertAllEqual(result_or.is_ok(), False)
error_code, _ = result_or.get_error()
self.assertAllEqual(error_code, errors.INVALID_ARGUMENT)
del server
with self.assertRaises(errors.DeadlineExceededError):
_ = client.assign_add_blocking(
variables.Variable(2, dtype=dtypes.int64), timeout_in_ms=1)
def test_captured_inputs(self):
v = variables.Variable(initial_value=0, dtype=dtypes.int64)
@eager_def_function.function(
input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
def assign_add(a):
v.assign_add(a)
@eager_def_function.function(input_signature=[])
def read_var():
return v.value()
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server = rpc_ops.GrpcServer(address)
server.register("assign_add", assign_add)
server.register("read_var", read_var)
server.start()
client = rpc_ops.GrpcClient(address)
result_or = client.call("assign_add",
[variables.Variable(2, dtype=dtypes.int64)])
self.assertAllEqual(result_or.is_ok(), True)
result_or = client.call("assign_add",
[variables.Variable(2, dtype=dtypes.int64)])
self.assertAllEqual(result_or.is_ok(), True)
result_or = client.call(
"read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
self.assertAllEqual(result_or.is_ok(), True)
self.assertAllEqual(result_or.get_value(), [4])
def test_register_method_twice(self):
v = variables.Variable(initial_value=0, dtype=dtypes.int64)
@eager_def_function.function(
input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
def assign_add(a):
v.assign_add(a)
@eager_def_function.function(
input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
def assign(a):
v.assign(a)
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server = rpc_ops.GrpcServer(address)
server.register("assign", assign_add)
with self.assertRaisesRegex(errors.InvalidArgumentError,
"assign is already registered."):
# Reusing the same error name.
server.register("assign", assign)
def test_tf_function_register_without_input_signature(self):
v = variables.Variable(initial_value=0, dtype=dtypes.int64)
@eager_def_function.function
def assign(a):
v.assign(a)
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server = rpc_ops.GrpcServer(address)
with self.assertRaisesRegex(
ValueError, "Input signature not specified for the function."):
server.register("assign", assign)
# Register without input signature should work for functions without input
# args.
@eager_def_function.function
def read_var():
return v.value()
server.register("read_var", read_var)
def test_multi_device_resource(self):
elements = np.random.randint(100, size=[200])
with ops.device("/device:CPU:1"):
queue = data_flow_ops.FIFOQueue(200, dtypes.int64, shapes=[])
@eager_def_function.function()
def populate_queue():
queue.enqueue_many(elements)
queue.close()
with ops.device("/device:CPU:0"):
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server = rpc_ops.GrpcServer(address)
server.register("populate_queue", populate_queue)
server.start()
client = rpc_ops.GrpcClient(address, list_registered_methods=True)
client.populate_queue()
for e in elements:
self.assertAllEqual(e, queue.dequeue())
def test_queue_resource(self):
elements = np.random.randint(100, size=[200])
queue = data_flow_ops.FIFOQueue(200, dtypes.int64, shapes=[])
@eager_def_function.function()
def populate_queue():
queue.enqueue_many(elements)
queue.close()
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server = rpc_ops.GrpcServer(address)
server.register("populate_queue", populate_queue)
server.start()
client = rpc_ops.GrpcClient(address, list_registered_methods=True)
client.populate_queue()
for e in elements:
self.assertAllEqual(e, queue.dequeue())
def test_multi_device_resource_cpu(self):
with ops.device("/device:cpu:1"):
v = variables.Variable(initial_value=0, dtype=dtypes.int64)
@eager_def_function.function(
input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
def assign_add(a):
v.assign_add(a)
with ops.device("/device:CPU:0"):
port = portpicker.pick_unused_port()
address = "localhost:{}".format(port)
server = rpc_ops.GrpcServer(address)
server.register("assign_add", assign_add)
server.start()
client = rpc_ops.GrpcClient(address, list_registered_methods=True)
result_or = client.assign_add(variables.Variable(2, dtype=dtypes.int64))
self.assertAllEqual(result_or.is_ok(), True)
self.assertAllEqual(v, 2)
if __name__ == "__main__":
ops.enable_eager_execution()
test.main()