1# Copyright 2021 Google LLC 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14"""Provides the `serve_slices` operation. 15 16This wraps the generated op and ensures that necessary shared libraries 17are loaded. 18""" 19 20import tensorflow as tf 21 22from fcp.tensorflow import _serve_slices_op 23from fcp.tensorflow import gen_serve_slices_py 24 25_serve_slices_so = tf.load_op_library( 26 tf.compat.v1.resource_loader.get_path_to_datafile('./_serve_slices_op.so')) 27 28 29def _to_tensor_list(list_of_python_values, dtype=None): 30 return [ 31 tf.convert_to_tensor(subvalue, dtype=dtype) 32 for subvalue in list_of_python_values 33 ] 34 35 36def serve_slices(callback_token, server_val, max_key, select_fn_initialize_op, 37 select_fn_server_val_input_tensor_names, 38 select_fn_key_input_tensor_name, 39 select_fn_filename_input_tensor_name, 40 select_fn_target_tensor_name): 41 """Calls into a preregistered `callback_token` to serve slices of a value. 42 43 In addition to the arguments to this function, `serve_slices` requires that 44 a TensorFlow graph containing a selection function (`select_fn`) be provided 45 to the server running `serve_slices`. `serve_slices` is responsible for 46 providing the server with the names of the placeholder tensor inputs to the 47 selection function (`select_fn_X_input_tensor_names`, 48 `select_fn_key_input_tensor_name`, and `select_fn_filename_input_tensor_name`) 49 and the target tensor to evalate to ensure that the slice is written to the 50 provided filename (`select_fn_target_tensor_name`). 51 52 Args: 53 callback_token: An string ID corresponding to a callback registered with the 54 `register_serve_slices_callback` function. This function will be invoked 55 when `serve_slices` is called. 56 server_val: A list of arbitrary-typed tensors from which slices may be 57 generated using `select_fn`. These tensors must be passed into the 58 `select_fn` by writing them to the placeholder tensors named by 59 `select_fn_server_val_input_names`, which must contain exactly one tensor 60 name for each tensor in `server_val`. 61 max_key: An integer indicating the maxiumum slice index which may be 62 requested. Slice indices start at zero and may go up to `max_key` 63 (inclusive). 64 select_fn_initialize_op: An op to run before each call to `select_fn` in 65 order to reinitialize any state `select_fn` may contain. 66 select_fn_server_val_input_tensor_names: A list of names of the tensors that 67 make up the `server_val` portion of the inputs to `select_fn`. Must be the 68 same length as the number of tensors in `server_val`. 69 select_fn_key_input_tensor_name: The name of the tensor that is the `key` 70 input to `select_fn`. 71 select_fn_filename_input_tensor_name: The name of the placeholder tensor 72 that is the `filename` input to `select_fn`. The `filename` is used to 73 specify where the resulting slice should be written. 74 select_fn_target_tensor_name: The name of the `target` tensor to run which 75 will result in `select_fn`'s output being written to `filename`. 76 77 Returns: 78 A string identifier given by the underlying callback which can be used by 79 clients to access the generated slices. 80 """ 81 return gen_serve_slices_py.serve_slices( 82 callback_token=tf.convert_to_tensor(callback_token, dtype=tf.string), 83 server_val=_to_tensor_list(server_val), 84 max_key=tf.convert_to_tensor(max_key, dtype=tf.int32), 85 select_fn_initialize_op=tf.convert_to_tensor( 86 select_fn_initialize_op, dtype=tf.string), 87 select_fn_server_val_input_tensor_names=_to_tensor_list( 88 select_fn_server_val_input_tensor_names, dtype=tf.string), 89 select_fn_key_input_tensor_name=tf.convert_to_tensor( 90 select_fn_key_input_tensor_name, dtype=tf.string), 91 select_fn_filename_input_tensor_name=tf.convert_to_tensor( 92 select_fn_filename_input_tensor_name, dtype=tf.string), 93 select_fn_target_tensor_name=tf.convert_to_tensor( 94 select_fn_target_tensor_name, dtype=tf.string)) 95 96 97def register_serve_slices_callback(callback): 98 """Registers a callback to be invoked by the `ServeSlices` op.""" 99 def callback_adapter(callback_token, server_val, *args): 100 # Convert the serialized TensorProtos to ndarrays. 101 tensor_proto = tf.make_tensor_proto(0) 102 converted_server_val = [ 103 tf.make_ndarray(tensor_proto.FromString(val)) for val in server_val 104 ] 105 return callback(callback_token, converted_server_val, *args) 106 107 return _serve_slices_op.register_serve_slices_callback(callback_adapter) 108