def main(is_ps, run_id, rank, world_size, cluster_spec, batch_size, replicas_to_aggregate, light_target=False):
    logging.info('Initial.')
    job_name = ('ps' if is_ps else 'worker')
    cluster = tf.train.ClusterSpec(cluster_spec)
    gpu_options = tf.GPUOptions(allow_growth=True, per_process_gpu_memory_fraction=0.2)
    session_conf = tf.ConfigProto(gpu_options=gpu_options, allow_soft_placement=True, log_device_placement=False)
    server = tf.train.Server(cluster, job_name=job_name, task_index=rank, config=session_conf)
    if is_ps:
        server.join()
    else:
        device_fn = tf.train.replica_device_setter(ps_tasks=None, ps_device='/job:ps', worker_device='/job:{}/task:{}/device:GPU:{}'.format(job_name, rank, rank), merge_devices=True, cluster=cluster, ps_ops=None, ps_strategy=None)
        with tf.Graph().as_default():
            with tf.device(device_fn):
                data_loader = DatasetCifar(dataset='cifar-10', dataset_root='/datasets', batch_size=batch_size, world_size=world_size, rank=rank, seed=42, tf_dtype=tf.float32)
                (train_op, loss, metrics, hooks) = define_graph(data_loader.inputs, data_loader.labels, data_loader.training, batch_size, replicas_to_aggregate)
                local_init_op = tf.group(tf.local_variables_initializer(), data_loader.train_init_op, data_loader.validation_init_op)
                scaffold = tf.train.Scaffold(init_op=None, init_feed_dict=None, init_fn=None, ready_op=None, ready_for_local_init_op=None, local_init_op=local_init_op)
                lr_tensor_name = tf.get_default_graph().get_tensor_by_name('learning_rate:0')
            with tf.train.MonitoredTrainingSession(config=session_conf, master=server.target, scaffold=scaffold, is_chief=(rank == 0), checkpoint_dir=None, save_checkpoint_secs=None, save_summaries_steps=None, stop_grace_period_secs=5, hooks=hooks) as sess:
                logging.info('Begin training.')
                final_epoch = 164
                if light_target:
                    goal = task1_time_to_accuracy_light_goal()
                else:
                    goal = task1_time_to_accuracy_goal()
                tracker = Tracker(metrics, run_id, rank, goal=goal)
                tracker.start()
                for i_epoch in range(final_epoch):
                    logging.debug('=> Epoch {}'.format(i_epoch))
                    train_round(sess, data_loader.train_init_op, train_op, loss, metrics, batch_size, data_loader.num_batches_per_epoch_for_train, tracker, lr_tensor=lr_tensor_name, lr_scheduler_level='epoch')
                    validation_round(sess, data_loader.validation_init_op, loss, metrics, batch_size, data_loader.num_batches_per_epoch_for_eval, tracker)
                    tracker.epoch_end()
                    if tracker.goal_reached:
                        print('Goal Reached!')
                        return
            logging.info('Finish.')
def main(is_ps, run_id, rank, world_size, cluster_spec, batch_size, replicas_to_aggregate, light_target=False):
    logging.info('Initial.')
    job_name = ('ps' if is_ps else 'worker')
    cluster = tf.train.ClusterSpec(cluster_spec)
    gpu_options = tf.GPUOptions(allow_growth=True, per_process_gpu_memory_fraction=0.2)
    session_conf = tf.ConfigProto(gpu_options=gpu_options, allow_soft_placement=True, log_device_placement=False)
    server = tf.train.Server(cluster, job_name=job_name, task_index=rank, config=session_conf)
    if is_ps:
        server.join()
    else:
        device_fn = tf.train.replica_device_setter(ps_tasks=None, ps_device='/job:ps', worker_device='/job:{}/task:{}/device:GPU:{}'.format(job_name, rank, rank), merge_devices=True, cluster=cluster, ps_ops=None, ps_strategy=None)
        with tf.Graph().as_default():
            with tf.device(device_fn):
                data_loader = DatasetCifar(dataset='cifar-10', dataset_root='/datasets', batch_size=batch_size, world_size=world_size, rank=rank, seed=42, tf_dtype=tf.float32)
                (train_op, loss, metrics, hooks) = define_graph(data_loader.inputs, data_loader.labels, data_loader.training, batch_size, replicas_to_aggregate)
                local_init_op = tf.group(tf.local_variables_initializer(), data_loader.train_init_op, data_loader.validation_init_op)
                scaffold = tf.train.Scaffold(init_op=None, init_feed_dict=None, init_fn=None, ready_op=None, ready_for_local_init_op=None, local_init_op=local_init_op)
                lr_tensor_name = tf.get_default_graph().get_tensor_by_name('learning_rate:0')
            with tf.train.MonitoredTrainingSession(config=session_conf, master=server.target, scaffold=scaffold, is_chief=(rank == 0), checkpoint_dir=None, save_checkpoint_secs=None, save_summaries_steps=None, stop_grace_period_secs=5, hooks=hooks) as sess:
                logging.info('Begin training.')
                final_epoch = 164
                if light_target:
                    goal = task1_time_to_accuracy_light_goal()
                else:
                    goal = task1_time_to_accuracy_goal()
                tracker = Tracker(metrics, run_id, rank, goal=goal)
                tracker.start()
                for i_epoch in range(final_epoch):
                    logging.debug('=> Epoch {}'.format(i_epoch))
                    train_round(sess, data_loader.train_init_op, train_op, loss, metrics, batch_size, data_loader.num_batches_per_epoch_for_train, tracker, lr_tensor=lr_tensor_name, lr_scheduler_level='epoch')
                    validation_round(sess, data_loader.validation_init_op, loss, metrics, batch_size, data_loader.num_batches_per_epoch_for_eval, tracker)
                    tracker.epoch_end()
                    if tracker.goal_reached:
                        print('Goal Reached!')
                        return
            logging.info('Finish.')