# Path: tensorflowonspark/gpu_info.py
# MAX_RETRIES = 3           #: Maximum retries to allocate GPUs
# AS_STRING = 'string'
# AS_LIST = 'list'
# def is_gpu_available():
# def get_gpus(num_gpu=1, worker_index=-1, format=AS_STRING):
#   def parse_gpu(gpu_str):
# 
# Path: tensorflowonspark/reservation.py
# TFOS_SERVER_PORT = "TFOS_SERVER_PORT"
# TFOS_SERVER_HOST = "TFOS_SERVER_HOST"
# BUFSIZE = 1024
# MAX_RETRIES = 3
#       CONNECTIONS = []
# class Reservations:
# class MessageSocket(object):
# class Server(MessageSocket):
# class Client(MessageSocket):
#   def __init__(self, required):
#   def add(self, meta):
#   def done(self):
#   def get(self):
#   def remaining(self):
#   def receive(self, sock):
#   def send(self, sock, msg):
#   def __init__(self, count):
#   def await_reservations(self, sc, status={}, timeout=600):
#   def _handle_message(self, sock, msg):
#   def start(self):
#     def _listen(self, sock):
#   def get_server_ip(self):
#   def get_server_ports(self):
#   def start_listening_socket(self):
#   def stop(self):
#   def __init__(self, server_addr):
#   def _request(self, msg_type, msg_data=None):
#   def close(self):
#   def register(self, reservation):
#   def get_reservations(self):
#   def await_reservations(self):
#   def request_stop(self):
# 
# Path: tensorflowonspark/TFSparkNode.py
# class TFSparkNode(object):
#   """Low-level functions used by the high-level TFCluster APIs to manage cluster state.
# 
#   **This class is not intended for end-users (see TFNode for end-user APIs)**.
# 
#   For cluster management, this wraps the per-node cluster logic as Spark RDD mapPartitions functions, where the RDD is expected to be
#   a "nodeRDD" of the form: ``nodeRDD = sc.parallelize(range(num_executors), num_executors)``.
# 
#   For data feeding, this wraps the feeding logic as Spark RDD mapPartitions functions on a standard "dataRDD".
# 
#   This also manages a reference to the TFManager "singleton" per executor.  Since Spark can spawn more than one python-worker
#   per executor, this will reconnect to the "singleton" instance as needed.
#   """
#   mgr = None                #: TFManager instance
#   cluster_id = None         #: Unique ID for a given TensorFlowOnSpark cluster, used for invalidating state for new clusters.


import argparse
import os
import random
import unittest
from tensorflowonspark import gpu_info, reservation, TFSparkNode
from unittest.mock import patch



class TFSparkNodeTest(unittest.TestCase):
  def setUp(self):
    self.server = reservation.Server(1)
    self.server_addr = self.server.start()

    self.parser = argparse.ArgumentParser()

    self.default_fn = lambda args, ctx: print("{}:{} args: {}".format(ctx.job_name, ctx.task_index, args))
    self.job_name = 'chief'
    self.task_index = 0
    self.cluster_meta = {
      'id': random.getrandbits(64),
      'cluster_template': {self.job_name: [0]},
      'num_executors': 1,
      'default_fs': 'file://',
      'working_dir': '.',
      'server_addr': self.server_addr
    }
    self.tensorboard = False
    self.log_dir = None
    self.queues = ['input']
    self.background = False

  def tearDown(self):

# Path: tensorflowonspark/gpu_info.py
# MAX_RETRIES = 3           #: Maximum retries to allocate GPUs
# AS_STRING = 'string'
# AS_LIST = 'list'
# def is_gpu_available():
# def get_gpus(num_gpu=1, worker_index=-1, format=AS_STRING):
#   def parse_gpu(gpu_str):
# 
# Path: tensorflowonspark/reservation.py
# TFOS_SERVER_PORT = "TFOS_SERVER_PORT"
# TFOS_SERVER_HOST = "TFOS_SERVER_HOST"
# BUFSIZE = 1024
# MAX_RETRIES = 3
#       CONNECTIONS = []
# class Reservations:
# class MessageSocket(object):
# class Server(MessageSocket):
# class Client(MessageSocket):
#   def __init__(self, required):
#   def add(self, meta):
#   def done(self):
#   def get(self):
#   def remaining(self):
#   def receive(self, sock):
#   def send(self, sock, msg):
#   def __init__(self, count):
#   def await_reservations(self, sc, status={}, timeout=600):
#   def _handle_message(self, sock, msg):
#   def start(self):
#     def _listen(self, sock):
#   def get_server_ip(self):
#   def get_server_ports(self):
#   def start_listening_socket(self):
#   def stop(self):
#   def __init__(self, server_addr):
#   def _request(self, msg_type, msg_data=None):
#   def close(self):
#   def register(self, reservation):
#   def get_reservations(self):
#   def await_reservations(self):
#   def request_stop(self):
# 
# Path: tensorflowonspark/TFSparkNode.py
# class TFSparkNode(object):
#   """Low-level functions used by the high-level TFCluster APIs to manage cluster state.
# 
#   **This class is not intended for end-users (see TFNode for end-user APIs)**.
# 
#   For cluster management, this wraps the per-node cluster logic as Spark RDD mapPartitions functions, where the RDD is expected to be
#   a "nodeRDD" of the form: ``nodeRDD = sc.parallelize(range(num_executors), num_executors)``.
# 
#   For data feeding, this wraps the feeding logic as Spark RDD mapPartitions functions on a standard "dataRDD".
# 
#   This also manages a reference to the TFManager "singleton" per executor.  Since Spark can spawn more than one python-worker
#   per executor, this will reconnect to the "singleton" instance as needed.
#   """
#   mgr = None                #: TFManager instance
#   cluster_id = None         #: Unique ID for a given TensorFlowOnSpark cluster, used for invalidating state for new clusters.


# Path: tests/test_TFSparkNode.py
import argparse
import os
import random
import unittest
from tensorflowonspark import gpu_info, reservation, TFSparkNode
from unittest.mock import patch




class TFSparkNodeTest(unittest.TestCase):
  def setUp(self):
    self.server = reservation.Server(1)
    self.server_addr = self.server.start()

    self.parser = argparse.ArgumentParser()

    self.default_fn = lambda args, ctx: print("{}:{} args: {}".format(ctx.job_name, ctx.task_index, args))
    self.job_name = 'chief'
    self.task_index = 0
    self.cluster_meta = {
      'id': random.getrandbits(64),
      'cluster_template': {self.job_name: [0]},
      'num_executors': 1,
      'default_fs': 'file://',
      'working_dir': '.',
      'server_addr': self.server_addr
    }
    self.tensorboard = False
    self.log_dir = None
    self.queues = ['input']
    self.background = False

  def tearDown(self):

    client = reservation.Client(self.server_addr)