import logging
import tensorflow as tf

from . import train
from .flags import FLAGS
from .file import read_files



def main(model):
  def run(_):
    cluster = tf.train.ClusterSpec({
      "ps": FLAGS.ps_hosts,
      "worker": FLAGS.worker_hosts,
    })

    server = tf.train.Server(cluster,
                             job_name=FLAGS.job_name,
                             task_index=FLAGS.task_index)

    if FLAGS.job_name == "ps":
      server.join()
    elif FLAGS.job_name == "worker":
      with tf.device(tf.train.replica_device_setter(
          worker_device="/job:worker/task:{}".format(FLAGS.task_index),
          cluster=cluster)):
        global_step = train.global_step()

        inputs = read_files(FLAGS.file_glob, FLAGS.file_format)
        train_op, _ = model(*inputs[1:])

        saver = tf.train.Saver(max_to_keep=2**16)
        summary_op = tf.merge_all_summaries()
        init_op = tf.initialize_all_variables()

      sv = tf.train.Supervisor(is_chief=(FLAGS.task_index == 0),
                               logdir=FLAGS.log_dir,
                               init_op=init_op,
                               summary_op=summary_op,
                               saver=saver,
                               global_step=global_step)

      with sv.managed_session(server.target) as sess, sess.as_default():
        step = global_step.eval()
        while not sv.should_stop() and step < FLAGS.num_epochs: # TODO: num_epochs != num_steps
          logging.info("#epochs = {}", step)
          _, step = sess.run([train_op, global_step])
        sv.saver.save(sess, sv.save_path, step)
      sv.stop()
    else:
      raise ValueError("Invalid job_name: {}".format(FLAGS.job_name))

  tf.app.run(run)

import logging
import tensorflow as tf

from . import train
from .flags import FLAGS
from .file import read_files



def main(model):
  def run(_):
    cluster = tf.train.ClusterSpec({
      "ps": FLAGS.ps_hosts,
      "worker": FLAGS.worker_hosts,
    })

    server = tf.train.Server(cluster,
                             job_name=FLAGS.job_name,
                             task_index=FLAGS.task_index)

    if FLAGS.job_name == "ps":
      server.join()
    elif FLAGS.job_name == "worker":
      with tf.device(tf.train.replica_device_setter(
          worker_device="/job:worker/task:{}".format(FLAGS.task_index),
          cluster=cluster)):
        global_step = train.global_step()

        inputs = read_files(FLAGS.file_glob, FLAGS.file_format)
        train_op, _ = model(*inputs[1:])

        saver = tf.train.Saver(max_to_keep=2**16)
        summary_op = tf.merge_all_summaries()
        init_op = tf.initialize_all_variables()

      sv = tf.train.Supervisor(is_chief=(FLAGS.task_index == 0),
                               logdir=FLAGS.log_dir,
                               init_op=init_op,
                               summary_op=summary_op,
                               saver=saver,
                               global_step=global_step)

      with sv.managed_session(server.target) as sess, sess.as_default():
        step = global_step.eval()
        while not sv.should_stop() and step < FLAGS.num_epochs: # TODO: num_epochs != num_steps
          logging.info("#epochs = %d", step)
          _, step = sess.run([train_op, global_step])
        sv.saver.save(sess, sv.save_path, step)
      sv.stop()
    else:
      raise ValueError("Invalid job_name: {}".format(FLAGS.job_name))

  tf.app.run(run)

Fix logging