tensorflowonspark.reservation module

This module contains client/server methods to manage node reservations during TFCluster startup.

class Client(server_addr)

Bases: tensorflowonspark.reservation.MessageSocket

Client to register and await node reservations.

Args:

server_addr

a tuple of (host, port) pointing to the Server.

await_reservations()

Poll until all reservations completed, then return cluster_info.

close()

Close the client socket.

get_reservations()

Get current list of reservations.

register(reservation)

Register reservation with server.

request_stop()

Request server stop.

server_addr = None

address of server

sock = None

socket to server TCP connection

class MessageSocket

Bases: object

Abstract class w/ length-prefixed socket send/receive functions.

receive(sock)

Receive a message on sock.

send(sock, msg)

Send msg to destination sock.

class Reservations(required)

Bases: object

Thread-safe store for node reservations.

Args:

required

expected number of nodes in the cluster.

add(meta)

Add a reservation.

Args:

meta

a dictonary of metadata about a node

done()

Returns True if the required number of reservations have been fulfilled.

get()

Get the list of current reservations.

remaining()

Get a count of remaining/unfulfilled reservations.

class Server(count)

Bases: tensorflowonspark.reservation.MessageSocket

Simple socket server with length-prefixed pickle messages.

Args:

count

expected number of nodes in the cluster.

await_reservations(sc, status={}, timeout=600)

Block until all reservations are received.

done = False

boolean indicating if server should be shutdown.

get_server_ip()

Returns the value of TFOS_SERVER_HOST environment variable (if set), otherwise defaults to current host/IP.

get_server_ports()

Returns a list of target ports as defined in the TFOS_SERVER_PORT environment (if set), otherwise defaults to 0 (any port).

TFOS_SERVER_PORT should be either a single port number or a range, e.g. ‘8888’ or ‘9997-9999’

reservations = None

List of reservations managed by this server.

start()

Start listener in a background thread

Returns:

address of the Server as a tuple of (host, port)

start_listening_socket()

Starts the registration server socket listener.

stop()

Stop the Server’s socket listener.