# ...
if FLAGS.job_name == 'ps':
    with tf.device('/job:ps/task:0'):
        counter = tf.Variable(0, name='counter', trainable=False)
    with tf.Session(server.target):
        sess.run(tf.global_variables_initializer())

elif FLAGS.job_name == 'worker':

    with tf.device('/job:ps/task:0'):
        counter = tf.Variable(0, name='counter', trainable=False)

    with tf.Session(server.target):
      sess.run(counter.assign_add(1))