Implementation:Tensorflow Serving Remote Predict Op Kernel Test
| Knowledge Sources | |
|---|---|
| Domains | Testing, Prediction, Experimental |
| Last Updated | 2026-02-13 00:00 GMT |
Overview
Test suite validating the experimental RemotePredict TensorFlow op kernel which makes RPC calls to a remote prediction service.
Description
This test file validates the RemotePredictOp kernel, an experimental TensorFlow custom op that enables making remote prediction RPC calls from within a TensorFlow graph. The tests use a MockPredictionService to simulate remote prediction behavior and a MockRpc class for RPC lifecycle management. The mock service routes behavior based on model name, supporting good model, bad model, tensor content checking, and proto field checking scenarios.
The RunRemotePredict helper function constructs a TF graph with the TfServingRemotePredict op, sets up input tensor aliases and tensors, configures output types, and executes using a ClientSession.
Key areas tested include:
- Successful remote prediction round-trip
- RPC error handling (fail on error vs return status)
- Tensor content encoding verification
- Proto field encoding verification
Usage
Run these tests to validate changes to the experimental remote predict op kernel, RPC integration, or tensor serialization for remote prediction.
Code Reference
Source Location
- Repository: Tensorflow_Serving
- File: tensorflow_serving/experimental/tensorflow/ops/remote_predict/kernels/remote_predict_op_kernel_test.cc
- Lines: 1-240
Test Fixture
class MockPredictionService {
public:
static absl::Status Create(const string& target_address,
std::unique_ptr<MockPredictionService>* service);
absl::StatusOr<MockRpc*> CreateRpc(absl::Duration max_rpc_deadline);
void Predict(MockRpc* rpc, PredictRequest* request,
PredictResponse* response,
std::function<void(absl::Status status)> callback);
static constexpr char kGoodModel[] = "good_model";
static constexpr char kBadModel[] = "bad_model";
static constexpr char kGoodModelCheckTensorContent[] =
"good_model_check_tensor_content";
static constexpr char kGoodModelCheckProtoField[] =
"good_model_check_proto_field";
};
REGISTER_KERNEL_BUILDER(
Name("TfServingRemotePredict").Device(DEVICE_CPU),
RemotePredictOp<MockPredictionService>);
Build Target
bazel test //tensorflow_serving/experimental/tensorflow/ops/remote_predict/kernels:remote_predict_op_kernel_test
Test Coverage
Key Test Cases
| Test Name | Category | Description |
|---|---|---|
| TestSimple | Integration | Tests successful remote prediction with input/output tensor mapping |
| TestRpcError | Error Handling | Tests RPC error propagation (fail_on_rpc_error=true) |
| TestRpcErrorReturnStatus | Error Handling | Tests RPC error returned as status tensor (fail_on_rpc_error=false) |
| TestTensorContent | Serialization | Tests tensor_content encoding in request protos |
| TestProtoField | Serialization | Tests proto field (int_val) encoding in request protos |
Usage Examples
Test Pattern
TEST(RemotePredictTest, TestSimple) {
std::vector<Tensor> outputs;
TF_ASSERT_OK(RunRemotePredict(
/*model_name=*/MockPredictionService::kGoodModel, &outputs));
ASSERT_EQ(4, outputs.size());
// Check status code is 0 (success) and no error message
EXPECT_EQ(0, outputs[0].scalar<int>()());
EXPECT_EQ("", outputs[1].scalar<tensorflow::tstring>()());
// Check output tensors match inputs
test::ExpectTensorEqual<int>(outputs[2], test::AsTensor<int>({1, 2}));
test::ExpectTensorEqual<int>(outputs[3], test::AsTensor<int>({3, 4}));
}