Args: server (tf.train.Server): Returns: tf.train.SessionCreator
(server)
| 6 | |
| 7 | |
| 8 | def get_distributed_session_creator(server): |
| 9 | """ |
| 10 | Args: |
| 11 | server (tf.train.Server): |
| 12 | |
| 13 | Returns: |
| 14 | tf.train.SessionCreator |
| 15 | """ |
| 16 | |
| 17 | server_def = server.server_def |
| 18 | is_chief = (server_def.job_name == 'worker') and (server_def.task_index == 0) |
| 19 | |
| 20 | init_op = tf.global_variables_initializer() |
| 21 | local_init_op = tf.local_variables_initializer() |
| 22 | ready_op = tf.report_uninitialized_variables() |
| 23 | ready_for_local_init_op = tf.report_uninitialized_variables(tf.global_variables()) |
| 24 | sm = tf.train.SessionManager( |
| 25 | local_init_op=local_init_op, |
| 26 | ready_op=ready_op, |
| 27 | ready_for_local_init_op=ready_for_local_init_op, |
| 28 | graph=tf.get_default_graph()) |
| 29 | |
| 30 | # to debug wrong variable collection |
| 31 | # from pprint import pprint |
| 32 | # print("GLOBAL:") |
| 33 | # pprint([(k.name, k.device) for k in tf.global_variables()]) |
| 34 | # print("LOCAL:") |
| 35 | # pprint([(k.name, k.device) for k in tf.local_variables()]) |
| 36 | |
| 37 | class _Creator(tf.train.SessionCreator): |
| 38 | def create_session(self): |
| 39 | if is_chief: |
| 40 | return sm.prepare_session(master=server.target, init_op=init_op) |
| 41 | else: |
| 42 | tf.logging.set_verbosity(tf.logging.INFO) # print message about uninitialized vars |
| 43 | ret = sm.wait_for_session(master=server.target) |
| 44 | tf.logging.set_verbosity(tf.logging.WARN) |
| 45 | return ret |
| 46 | |
| 47 | return _Creator() |