Repository: tensorflow
Issue: Share layer weights between processes in cluster
Labels: stat:awaiting tensorflower,type:feature,comp:apis,TF 2.10

### Issue type

Feature Request

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

source

### TensorFlow version

2.10.0

### Custom code

No

### OS platform and distribution

Linux Ubuntu 18.04

### Mobile device

_No response_

### Python version

3.10.11

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

11.2

### GPU model and memory

_No response_

### Current behavior?

Hey guys, I have a question regarding sharing layer weights between processes of a cluster.

I am in the flied of Deep Reinforcement Learning, where often there are multiple actors, generating experience (samples) and a central learner, who updates the network weights based on the gathered experience. For this to work efficient actors and learner are running in different processes on a single machine or possible on multiple machines. `tf.distribute` contains a lot functionalities useful for distributed Supervised Learning, but due to the different nature of DRL, `tf.distribute` can hardly be applied for this field.

Is there currently a way, or possibly a planned feature, to share trainable variables between TensorFlow instances running on different processes of a cluster? This is already done under the hood in `tf.distribute.experimental.MultiWorkerMirroredStrategy` and `tf.distribute.experimental.ParameterServerStrategy`, but as I mentioned above, are these distribution strategies not really suitable for DRL.

In the code below I create a cluster of processes using `tf.config.experimental_connect_to_cluster`, since running eager code and/or `tf.function` directly on a `tf.distribute.Server` is not possible. Experience is send to the learner via a shared `tf.queue.FIFOQueue` and now I would like to share the layer weights between the actor and learner.

### Standalone code to reproduce the issue

```shell
from absl import app
import multiprocessing as mp
import tensorflow as tf
from tensorflow.keras.layers import Dense


def run_actor_process(cluster_spec : tf.train.ClusterSpec) -> None:
    """
    Runs an actor which gathers experience and enqueues it to the learner.

    Args:
        cluster_spec (tf.distribute.Server):        The cluster specification the actor belongs to.
    """

    # only use first GPU
    physical_devices = tf.config.list_physical_devices('GPU')
    tf.config.set_visible_devices([physical_devices[0]], 'GPU')
    tf.config.experimental.set_memory_growth(physical_devices[0], True)

    # connect process to cluster
    tf.config.experimental_connect_to_cluster(cluster_spec, job_name='actor', task_index=0, protocol='grpc')

    # use shared queue on actor
    # NOTE: This needs to be done before any real TensorFlow operation is executed, otherwise the shared queue will not work.
    #       Not sure why this happens.
    with tf.device('/job:learner/replica:0/task:0/CPU:0'):
        queue = tf.queue.FIFOQueue(1, [tf.float32, tf.float32], shapes=[(5, ), (2, )], shared_name='trajectory_buffer')

    # NOTE: Ideally this layer's weights and biases would be shared between the processes.
    dense_layer = Dense(units=2,
                        activation=None,
                        use_bias=True,
                        kernel_initializer='glorot_uniform',
                        bias_initializer='zeros',
                        name='model_dense')

    @tf.function
    def run_actor(input_tensor : tf.Tensor):

        output_tensor = dense_layer(input_tensor)
        _ = queue.enqueue([tf.squeeze(input_tensor), tf.squeeze(output_tensor)])
        return output_tensor

    index = 1
    while True:

        observation = tf.range(index, index + 5, delta=1, dtype=tf.float32, name='range')
        observation = tf.expand_dims(observation, 0)

        output = run_actor(observation)

        print(f'Actor enqueued [{observation[0].numpy()}, {output[0].numpy()}].')
        index += 1


def run_server_process(cluster_spec : tf.distribute.Server) -> None:
    """
    Runs a TensorFlow distributed server.

    Args:
        cluster_spec (tf.distribute.Server):        The cluster specification the server belongs to.
    """

    # disable GPUs for TensorFlow session and server
    tf.config.set_visible_devices([], 'GPU')
    server_config = tf.compat.v1.ConfigProto(device_count={'GPU': 0})

    unused_server = tf.distribute.Server(cluster_spec,
                                         job_name='ps',
                                         task_index=0,
                                         protocol='grpc',
                                         config=server_config)
    unused_server.join()

def main(unused_argv):

    mp.set_start_method('spawn')

    # only use first GPU
    physical_devices = tf.config.list_physical_devices('GPU')
    tf.config.set_visible_devices([physical_devices[0]], 'GPU')
    tf.config.experimental.set_memory_growth(physical_devices[0], True)

    # define training cluster
    cluster_spec = tf.train.ClusterSpec({'ps' : ['localhost:7000'],
                                         'learner': ['localhost:7001'],
                                         'actor' : ['localhost:7002']})

    # run the actor in a separate process
    actor_process = mp.Process(target=run_actor_process, args=(cluster_spec, ), daemon=True)
    actor_process.start()

    # start parameter sever to initialize cluster
    # NOTE: This running server is solely needed to mimic a running cluster which the actor and learner can
    #       connect to using 'tf.config.experimental_connect_to_cluster'.
    server_process = mp.Process(target=run_server_process, args=(cluster_spec, ), daemon=True)
    server_process.start()

    # NOTE: The learner should not see the actor, otherwise the learner will wait until the actor is connected
    #       and the actor will wait until the leaner is connected, resulting in a deadlock.
    learner_cluster_spec = tf.train.ClusterSpec({'ps' : ['localhost:7000'],
                                                 'learner': ['localhost:7001']})

    # connect process to cluster
    tf.config.experimental_connect_to_cluster(learner_cluster_spec,
                                              job_name='learner',
                                              task_index=0,
                                              protocol='grpc')

    # use shared queue on learner
    # NOTE: This needs to be done before any real TensorFlow operation is executed, otherwise the shared queue will not work.
    #       Not sure why this happens.
    with tf.device('/job:learner/replica:0/task:0/CPU:0'):
        queue = tf.queue.FIFOQueue(1, [tf.float32, tf.float32], shapes=[(5, ), (2, )], shared_name='trajectory_buffer')

    # NOTE: Ideally this layer's weights and biases would be shared between the processes.
    dense_layer = Dense(units=2,
                        activation=None,
                        use_bias=True,
                        kernel_initializer='glorot_uniform',
                        bias_initializer='zeros',
                        name='model_dense')

    @tf.function
    def run_learner():

        dequeued = queue.dequeue_many(1)
        return dequeued, dense_layer(dequeued[0])

    while True:

        dequeued, learner_result = run_learner()

        # NOTE: If layer weights would be shared between learner and actor 'dequeued[1][0]' and 'learner_result[0]' would be equal.
        print(f'Learner dequeued [{dequeued[0][0].numpy()}, {dequeued[1][0].numpy()}] with result {learner_result[0].numpy()}.')


if __name__ == "__main__":
    try:
        app.run(main)
    except SystemExit:
        pass
```


### Relevant log output

_No response_