Share variable across devices in Tensorflow 2 cluster

Question

I would like to use Tensorflow in order to share a variable across devices. In my use case, I want to share a queue stored on a parameter server that can be filled by all the workers. Here is my code to do this with a simple variable:

# run.py
import tensorflow as tf
import os
import json
import argparse
import time

parser = argparse.ArgumentParser()
parser.add_argument('--type', type=str, default='ps', required=False)
parser.add_argument('--index', type=int, default=0, required=False)
args = parser.parse_args()

cluster_dict = {
    "cluster": {
        "worker": ['localhost:50848'],
        "ps": ['localhost:50849'] 
    },
    "task": {"type": args.type, "index": args.index}
}


os.environ["GRPC_FAIL_FAST"] = "use_caller"
os.environ["TF_CONFIG"] = json.dumps(cluster_dict)
cluster_resolver = tf.distribute.cluster_resolver.TFConfigClusterResolver()

server = tf.distribute.Server(
    cluster_resolver.cluster_spec(), job_name=cluster_resolver.task_type,
    task_index=cluster_resolver.task_id, protocol="grpc"
)
# Initializes tf server internally.
strategy = tf.distribute.experimental.ParameterServerStrategy(cluster_resolver)


with tf.device('/job:ps/task:0'):
    v = tf.Variable(
        1, name='test')

if cluster_resolver.task_type == "worker":
    while True:
        with tf.device('/job:ps/task:0'):
            # : ok
            print(v.device)
            # Increments variable on the worker
            v.assign_add(1)
            # v is incremented: ok
            print('v:', v)

            time.sleep(1)
elif cluster_resolver.task_type == "ps":
    with tf.device('/job:ps/task:0'):
        while True:
            # : ok
            print(v.device)
            # v should be incremented by the worker
            print(v)
            # v=1: not working
            time.sleep(1)

else:
    pass

I then launch the worker and the parameter server using:

python run.py --type ps

and

python run.py --type worker

However the code in its current form does not work: the variable stored in the parameter server is updated by the worker but not visible by the parameter server. Note that I use the distribution strategy to initialize the server in the backend (the device is not found if I comment the line where the strategy is initialized).

My use case is similar to: https://github.com/tensorflow/tensorflow/issues/43983.

Do you know how to do this?

I am using Tensorflow 2.4.1 on Windows 10.

Share variable across devices in Tensorflow 2 cluster

0 Answers