# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# This notebook shows how to use the simple model serving functionality found in `tfe.serving`.

# # Setup

# +
from collections import OrderedDict

import tensorflow as tf
import tf_encrypted as tfe
# -

# Turn on TFE profling flags since we want to inspect with TensorBoard.

# +
TENSORBOARD_DIR = "/tmp/tensorboard"

tfe.set_tfe_trace_flag(True)
tfe.set_tfe_events_flag(True)
tfe.set_log_directory(TENSORBOARD_DIR)

# !rm -rf {TENSORBOARD_DIR}
# -

# # Protocol
#
# We first configure the protocol we will be using, as well as the servers on which we want to run it.
#
# Note that the configuration is saved to file as we will be needing it in the client as well.

# +
players = OrderedDict([
    ('server0', 'localhost:4000'),
    ('server1', 'localhost:4001'),
    ('server2', 'localhost:4002'),
])

config = tfe.RemoteConfig(players)
config.save('/tmp/tfe.config')
# -

tfe.set_config(config)
tfe.set_protocol(tfe.protocol.Pond())

# ## Launching servers
#
# Before actually serving the computation below we need to launch TensorFlow servers in new processes. Run the following in three different terminals. You may have to allow Python to accept incoming connections.

for player_name in players.keys():
    print("python -m tf_encrypted.player --config /tmp/tfe.config {}".format(player_name))

# # Computation
#
# We then define the computation we want to run. These will happen on private tensors on the servers defined above.
#
# Note that the only single-tensor computations are currently supported.

# +
input_shape = (5, 5)
output_shape = (5, 5)

def computation(x):
    return x * x


# -

# We can then set up a new `tfe.serving.QueueServer` to serve this computation.

server = tfe.serving.QueueServer(
    input_shape=input_shape,
    output_shape=output_shape,
    computation_fn=computation)

# # Serving
#
# With all of the above in place we can finally connect to our servers, push our graph to them, and start serving computations.

# +
sess = tfe.Session(disable_optimizations=True)

tf.Session.reset(sess.target)


# +
def step_fn():
    print("Next")

server.run(
    sess,
    num_steps=5,
    step_fn=step_fn)
# -

# At this point we switch to the client notebook to run computations.



StarcoderdataJupyterScriptsDedupFiltered