# 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
  def test_gpu_default(self, mock_available, mock_get_gpus):
    """Default to one GPU if not explicitly requested"""

    mock_available.return_value = True
    mock_get_gpus.return_value = ['0']

    tf_args = self.parser.parse_args([])

    map_fn = TFSparkNode.run(self.default_fn, tf_args, self.cluster_meta, self.tensorboard, self.log_dir, self.queues, self.background)
    map_fn([0])
    self.assertEqual(os.environ['CUDA_VISIBLE_DEVICES'], '0')
    mock_get_gpus.assert_called_with(1, 0, format=gpu_info.AS_LIST)

  @patch('tensorflowonspark.TFSparkNode._get_cluster_spec')
  @patch('tensorflowonspark.gpu_info.get_gpus')
  @patch('tensorflowonspark.gpu_info.is_gpu_available')
  def test_gpu_cluster_spec(self, mock_available, mock_get_gpus, mock_get_spec):
    """Request GPU when multiple TF nodes land on same executor"""

    mock_available.return_value = True
    mock_get_gpus.return_value = ['0']
    mock_get_spec.return_value = {'chief': ['1.1.1.1:2222'], 'worker': ['1.1.1.1:2223', '1.1.1.1:2224', '2.2.2.2:2222']}

    self.cluster_meta['cluster_template'] = {'chief': [0], 'worker': [1, 2, 3]}
    self.parser.add_argument("--num_gpus", help="number of gpus to use", type=int)
    tf_args = self.parser.parse_args(["--num_gpus", "1"])
    print("tf_args:", tf_args)

    map_fn = TFSparkNode.run(self.default_fn, tf_args, self.cluster_meta, self.tensorboard, self.log_dir, self.queues, self.background)
    map_fn([2])     # worker:1

# 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

  def test_gpu_default(self, mock_available, mock_get_gpus):
    """Default to one GPU if not explicitly requested"""

    mock_available.return_value = True
    mock_get_gpus.return_value = ['0']

    tf_args = self.parser.parse_args([])

    map_fn = TFSparkNode.run(self.default_fn, tf_args, self.cluster_meta, self.tensorboard, self.log_dir, self.queues, self.background)
    map_fn([0])
    self.assertEqual(os.environ['CUDA_VISIBLE_DEVICES'], '0')
    mock_get_gpus.assert_called_with(1, 0, format=gpu_info.AS_LIST)

  @patch('tensorflowonspark.TFSparkNode._get_cluster_spec')
  @patch('tensorflowonspark.gpu_info.get_gpus')
  @patch('tensorflowonspark.gpu_info.is_gpu_available')
  def test_gpu_cluster_spec(self, mock_available, mock_get_gpus, mock_get_spec):
    """Request GPU when multiple TF nodes land on same executor"""

    mock_available.return_value = True
    mock_get_gpus.return_value = ['0']
    mock_get_spec.return_value = {'chief': ['1.1.1.1:2222'], 'worker': ['1.1.1.1:2223', '1.1.1.1:2224', '2.2.2.2:2222']}

    self.cluster_meta['cluster_template'] = {'chief': [0], 'worker': [1, 2, 3]}
    self.parser.add_argument("--num_gpus", help="number of gpus to use", type=int)
    tf_args = self.parser.parse_args(["--num_gpus", "1"])
    print("tf_args:", tf_args)

    map_fn = TFSparkNode.run(self.default_fn, tf_args, self.cluster_meta, self.tensorboard, self.log_dir, self.queues, self.background)
    map_fn([2])     # worker:1

    mock_get_gpus.assert_called_with(1, 2, format=gpu_info.AS_LIST)