Repository: ray-project/ray

Issue:
Time to initialize a policy grows linearly with the number of agents
<!--
General questions should be asked on the mailing list ray-dev@googlegroups.com.
Questions about how to use Ray should be asked on
[StackOverflow](https://stackoverflow.com/questions/tagged/ray).

Before submitting an issue, please fill out the following form.
-->

### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Linux Ubuntu 18.04
- **Ray installed from (source or binary)**: Binary
- **Ray version**: 0.7.4
- **Python version**: 3.7.4
- **Exact command to reproduce**: N/A

<!--
You can obtain the Ray version with

python -c "import ray; print(ray.__version__)"
-->

### Describe the problem
<!-- Describe the problem clearly here. -->
I noticed that in multi agent settings, the time to initialize a policy per agent increases as more agents are initialized. In the sample output I provided below, you can see that the time to initialize a single DynamicTFPolicy grows from 4.6 seconds to 15.3 seconds from the first agent to the tenth agent created. Line 291 of `rllib/policy/dynamic_tf_policy.py` is
```python
self._sess.run(tf.global_variables_initializer())
```
which I believe will run one time for each agent initialized. If I'm not mistaken, this means that every variable in the computation graph is being initialized each time that we initialize a DynamicTFPolicy. If initializing a DynamicTFPolicy adds new variables to the computation graph (as I believe it does), this would explain why the time to initialize a DynamicTFPolicy grows over time: We are initializing every variable in the computation graph, and the computation graph is growing. My question is, why does line 291 run a global variables initializer? Is there a reason for this that I can't see inside this method? How hard would it be to modify this to only initialize variables in the individual policy that we care to initialize?

I'm asking this because as detailed in #5753, I'm trying to modify rllib to allow initialization and removal of policies during training. The overhead incurred by this initialization quickly slows the training script down enough to be useless. Also, if anyone knows what the resource bottleneck is for policy initialization, that would be very helpful to know for when we're picking new hardware. Does it need a ton of cores to run in parallel, or more memory, or a bigger GPU or more GPUs or something? Thanks.


### Source code / logs
<!-- Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem. -->

```
(pytorch) root@e3a955e42cae:~/bees/bees# python trainer.py settings/settings.json
2019-10-23 12:38:04,168 WARNING worker.py:1426 -- WARNING: Not updating worker name since `setproctitle` is not installed. Install this with `pip install setproctitle` (or ray[debug]) to enable monitoring of worker processes.
2019-10-23 12:38:04,169 INFO resource_spec.py:205 -- Starting Ray with 3.52 GiB memory available for workers and up to 1.78 GiB for objects. You can adjust these settings with ray.init(memory=<bytes>, object_store_memory=<bytes>).
2019-10-23 12:38:05,354 INFO trainer.py:344 -- Tip: set 'eager': true or the --eager flag to enable TensorFlow eager execution
2019-10-23 12:38:05,754 WARNING ppo.py:149 -- Using the simple minibatch optimizer. This will significantly reduce performance, consider simple_optimizer=False.
DTFP: 4.604122s
DTFP: 4.856234s
DTFP: 5.630484s
DTFP: 6.850456s
DTFP: 7.856700s
DTFP: 9.624164s
DTFP: 10.894944s
DTFP: 12.129192s
DTFP: 14.210247s
DTFP: 15.342738s
```

Line 130 in `tf_policy_template.py`  (modified to print debug output above)
```python
            t = time.time()

            DynamicTFPolicy.__init__(
                self,
                obs_space,
                action_space,
                config,
                loss_fn,
                stats_fn=stats_fn,
                grad_stats_fn=grad_stats_fn,
                before_loss_init=before_loss_init_wrapper,
                make_model=make_model,
                action_sampler_fn=action_sampler_fn,
                existing_model=existing_model,
                existing_inputs=existing_inputs,
                get_batch_divisibility_req=get_batch_divisibility_req,
                obs_include_prev_action_reward=obs_include_prev_action_reward)

            print("DTFP: %fs" % (time.time() - t))
```

Snippet of trainer script used.
```python
# pylint: disable=invalid-name
if __name__ == "__main__":
    ray.init()

    # Get ``settings`` file for now.
    settings_file = sys.argv[1]
    with open(settings_file, "r") as f:
        settings = json.load(f)

    env_config = settings["env"]
    time_steps = env_config["time_steps"]

    space_env = create_env(settings)
    env = create_env(settings)

    # Register environment
    register_env("world", lambda _: env)

    # Build environment instance to get ``obs_space``.
    obs_space = space_env.observation_space
    act_space = space_env.action_space

    # You can also have multiple policies per trainer, but here we just
    # show one each for PPO and DQN.
    policies: Dict[str, Tuple[Any, gym.Space, gym.Space, Dict[Any, Any]]] = {
        "0": (PPOTFPolicy, obs_space, act_space, {}),
        "1": (PPOTFPolicy, obs_space, act_space, {}),
        "2": (PPOTFPolicy, obs_space, act_space, {}),
        "3": (PPOTFPolicy, obs_space, act_space, {}),
        "4": (PPOTFPolicy, obs_space, act_space, {}),
        "5": (PPOTFPolicy, obs_space, act_space, {}),
        "6": (PPOTFPolicy, obs_space, act_space, {}),
        "7": (PPOTFPolicy, obs_space, act_space, {}),
        "8": (PPOTFPolicy, obs_space, act_space, {}),
        "9": (PPOTFPolicy, obs_space, act_space, {}),
    }

    def policy_mapping_fn(agent_id: int) -> str:
        """ Returns the given agent's policy identifier. """
        return str(agent_id)

    ppo_trainer = PPOTrainer(
        env="bee_world",
        config={
            "multiagent": {
                "policies": policies,
                "policy_mapping_fn": policy_mapping_fn,
                "policies_to_train": ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"],
            },
            "simple_optimizer": True,
            # Disable filters, otherwise we would need to synchronize those
            # as well to the DQN agent.
            "observation_filter": "NoFilter",
            "num_workers": 2,
            "num_gpus": 1,
            "train_batch_size": 2,
            "sample_batch_size": 1,
            "sgd_minibatch_size": 2,
        },
    )
```





Discussion:
I am a collaborator on the project with @brendanxwhitaker, and I wanted to add another peculiar thing about the initialization code in DynamicTFPolicy. As pointed out above, line 291 of `rllib/policy/dynamic_tf_policy.py` is
```python
self._sess.run(tf.global_variables_initializer())
```
However, line 350 of the same initialization method in the same file is the exact same line, so that initialization runs twice. This is briefly touched on in a comment on line 290
```python
# postprocessing might depend on variable init, so run it first here
```
This is obviously a bit redundant but seems necessary unless something else in the method changes. My question about this is, under what conditions is the postprocessing dependent on the variable initialization, and how hard would it be to change this method so that the initialization is not redundant? 
Thanks for looking into this... agree we should try to fix the redundant initialization. One possible solution may be to initialize uninitialized variables. It looks like TF doesn't have a way to do this out of the box, but there are some workarounds: 

https://stackoverflow.com/questions/35164529/in-tensorflow-is-there-any-way-to-just-initialize-uninitialised-variables

Maybe one of these would work?

> This is obviously a bit redundant but seems necessary unless something else in the method changes. My question about this is, under what conditions is the postprocessing dependent on the variable initialization, and how hard would it be to change this method so that the initialization is not redundant?

This happens in for example PPO, where the postprocessor is accessing the value function. If you comment it out, then PPO will crash during policy creation trying to read uninitialized weights.
Okay, so just to be clear, when you say redundant, are you referring to calling `tf.global_variables_initializer` instead of only initializing the uninitialized variables, or are you referring to there being two of those calls for every call to the init function?

>Maybe one of these would work?

That seems like it would. Do you think it's possible to just create a variable scope that only contains the new, uninitialized variables and then just initialize all of those? Or maybe just add them to a list as they're created and pass a reference back up to the init function?

Also, is it expected behavior for the initialization time to be growing with each new policy added? Is this something unique to PPO, or is it going to be observed no matter what algorithm we use?

Are there other algorithms better suited for large numbers of individual policies running on a single machine?
The former (which includes the latter, though double init is less of an
issue).

On Fri, Oct 25, 2019, 5:11 AM Brendan Whitaker <notifications@github.com>
wrote:

> Okay, so just to be clear, when you say redundant, are you referring to
> calling tf.global_variables_initializer and initializing everything on
> each of those calls, or are you referring to there being two of those calls
> for every call to the init function?
>
> —
> You are receiving this because you were assigned.
> Reply to this email directly, view it on GitHub
> <https://github.com/ray-project/ray/issues/5982?email_source=notifications&email_token=AAADUSWGLFRZ2T4UX4BPOODQQLO6JA5CNFSM4JEA3DT2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOECIFA6Y#issuecomment-546328699>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AAADUSS6IPZ3UYTEBZEYLVLQQLO6JANCNFSM4JEA3DTQ>
> .
>

Looking through the stack overflow page, none of those answers are super clean but one of them is not as slow, with only one sess.run to find which variables are uninitialized. We are going to try a workaround and see how the performance is. Also, I am curious about what Brendan asked about regarding the initialization time growing. Do you know if this is something we should expect from initialization, that the initialization time will be larger for the last agents initialized than for the first, even if we implement one of the workarounds above? Thanks.
@mtcrawshaw I don't think it should be slowing down over time once the redundant inits are removed, if so that might be a bug in TensorFlow.
So we have implemented a couple of the workarounds from the stack overflow page that you mentioned, and we noticed some weird behavior from RLlib that we wanted to ask you about. In policy/dynamic_tf_policy.py, we replaced

```python
self._sess.run(tf.global_variables_initializer())
```

on lines 291 and 350 with

```python
initialize_uninitialized(self._sess)
```

where ``initialize_uninitialized`` defined by

```python
def initialize_uninitialized(sess):
    bytestr_var_names = sess.run(tf.report_uninitialized_variables())
    global_vars = tf.global_variables()
    strnames = [name.decode("utf-8") for name in bytestr_var_names]
    uninit_vars = [
        v
        for v in tf.global_variables()
        if v.name.split(":")[0]
        in set(strnames)
    ]
    initializer = tf.variables_initializer(uninit_vars)
    sess.run(initializer)
```

However, we still were getting an increase in the initialization time as the number of initialized agents grew. We threw a couple of print statements in ``initialize_uninitialized`` to get the names and number of uninitialized variables, modified version below.

```python
def initialize_uninitialized(sess):
    bytestr_var_names = sess.run(tf.report_uninitialized_variables())
    global_vars = tf.global_variables()
    strnames = [name.decode("utf-8") for name in bytestr_var_names]
    uninit_vars = [
        v
        for v in tf.global_variables()
        if v.name.split(":")[0]
        in set(strnames)
    ]
    print("Names of uninitialized variables: %s" % str([v.name for v in uninit_vars]))
    print("Number of uninitvars: %d" % len(uninit_vars))
    print("Number of names from report: %d" % len(bytestr_var_names))
    print("Number of variables: %d" % len(global_vars))
    initializer = tf.variables_initializer(uninit_vars)
    sess.run(initializer)
```

I won't provide the output here because it is quite long, but it shows that the variables which get initialized during the initialization of the first agent show up as uninitialized during the initialization of the second agent! The variables corresponding to the first and second agent are then initialized, but during the initialization of the third agent, they again reappear as uninitialized. My question is, do you have any idea why these variables might be getting "de-initialized" at any point during the initialization process?
That's really weird, perhaps try printing out the current TF graph id?
Maybe there are new graphs being created each time. I can't think of any
code that would be uninitializing variables otherwise.

On Tue, Oct 29, 2019, 5:57 AM Michael Crawshaw <notifications@github.com>
wrote:

> So we have implemented a couple of the workarounds from the stack overflow
> page that you mentioned, and we noticed some weird behavior from RLlib that
> we wanted to ask you about. In policy/dynamic_tf_policy.py, we replaced
>
> self._sess.run(tf.global_variables_initializer())
>
> with a call to initialize_uninitialized(self._sess), where
> initialize_uninitialized defined by
>
> def initialize_uninitialized(sess):
>     bytestr_var_names = sess.run(tf.report_uninitialized_variables())
>     global_vars = tf.global_variables()
>     strnames = [name.decode("utf-8") for name in bytestr_var_names]
>     uninit_vars = [
>         v
>         for v in tf.global_variables()
>         if v.name.split(":")[0]
>         in set(strnames)
>     ]
>     initializer = tf.variables_initializer(uninit_vars)
>     sess.run(initializer)
>
> However, we still were getting an increase in the initialization time as
> the number of initialized agents grew. We threw a couple of print
> statements in initialize_uninitialized to get the names and number of
> uninitialized variables, modified version below.
>
> def initialize_uninitialized(sess):
>     bytestr_var_names = sess.run(tf.report_uninitialized_variables())
>     global_vars = tf.global_variables()
>     strnames = [name.decode("utf-8") for name in bytestr_var_names]
>     uninit_vars = [
>         v
>         for v in tf.global_variables()
>         if v.name.split(":")[0]
>         in set(strnames)
>     ]
>     print("Names of uninitialized variables: %s" % str([v.name for v in uninit_vars]))
>     print("Number of uninitvars: %d" % len(uninit_vars))
>     print("Number of names from report: %d" % len(bytestr_var_names))
>     print("Number of variables: %d" % len(global_vars))
>     initializer = tf.variables_initializer(uninit_vars)
>     sess.run(initializer)
>
> I won't provide the output here because it is quite long, but it shows
> that the variables which get initialized during the initialization of the
> first agent show up as uninitialized during the initialization of the
> second agent! The variables corresponding to the first and second agent are
> then initialized, but during the initialization of the third agent, they
> again reappear as uninitialized. My question is, do you have any idea why
> these variables might be getting "de-initialized" at any point during the
> initialization process?
>
> —
> You are receiving this because you were assigned.
> Reply to this email directly, view it on GitHub
> <https://github.com/ray-project/ray/issues/5982?email_source=notifications&email_token=AAADUSUTNO3MSFXURACF44TQRAXM3A5CNFSM4JEA3DT2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOECQMCAA#issuecomment-547406080>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AAADUSQ56CLKJAL5S577XUTQRAXM3ANCNFSM4JEA3DTQ>
> .
>

Alright, so last week we were actually running that in a fork of RLLib. We ran the same tests in a fresh install, and found that the initialization time still grows, but that the number of uninitialized variables is constant (normal). 

We think this is entirely because of the (3?) calls to the global variable initializer. We'll investigate the graph id thing. 
I'm running into the same situation. As the agent initialization time grows linearly per-agent, this means that there is superlinear growth in total initialization time as the number of agents increases.

My agents all use the same type of custom policy, but each have their own instance of that policy. Although a single part of the graph does grow exponentially with the amount of agents, it is insignificant compared to the rest of of the parameters: a change of <20k parameters versus hundreds of thousands of other fixed parameters.

I ran cprofile/kcachegrind with `local_mode=True` to take a peek.

- 2 agents: Initialization takes 44 seconds
- 4 agents: Initialization takes 172 seconds
- 5 agents: Initialization takes 275 seconds

Respectively, their call graphs look like this:

![2_agent_profile](https://user-images.githubusercontent.com/5096835/81180011-be693680-8faa-11ea-9c6b-ff19aad3e788.png)
![4_agent_profile](https://user-images.githubusercontent.com/5096835/81180013-bf01cd00-8faa-11ea-9307-ff97c8dd53ef.png)
![5_agent_profile](https://user-images.githubusercontent.com/5096835/81180014-bf9a6380-8faa-11ea-860c-6d8960787fee.png)

As the number of agents grows, the 3 functions that dominate are:

1. _pywrap_tensorflow_internal.TF_SessionRun_wrapper
2. _pywrap_tensorflow_internal.ExtendSession
3. _pywrap_tensorflow_internal.TF_SessionMakeCallable

If anyone wants to dig into this, check out these [3 profiler files](https://github.com/ray-project/ray/files/4586880/profiler_files.zip). You can start one in KCachegrind with the following command:

`pyprof2calltree -i prof-2.out -k`
I experimented a bit, and I've found [this line](https://github.com/ray-project/ray/blob/master/rllib/optimizers/multi_gpu_optimizer.py#L139) `self.workers.local_worker().get_weights()` to be the slowest (makes up for 80-90% of initialization time with 5 agents), but only when it is called for the first time. Subsequent calls are far, far faster.

Somewhere down the call stack Ray then calls [this function](https://github.com/ray-project/ray/blob/master/python/ray/experimental/tf_utils.py#L157) for each policy, for each element of that policy. The tf graph is the same size for each agent, and grows linearly (in my case, 2 agents/80k elements, 5 agents/200k elements). Thus, you get a linearly increasing amount of eval calls to a linearly increasingly large graph = exponential growth.

I'm not sure whether the graph size increase is intentional: each policy is after all evaluated independently, as shown in these two lines in sampler.py {[1](https://github.com/ray-project/ray/blob/releases/0.8.5/rllib/evaluation/sampler.py#L590), [2](https://github.com/ray-project/ray/blob/releases/0.8.5/rllib/evaluation/sampler.py#L613)}. Is this done so that optimization is faster?
Here's a graph for memory use while evaluating/training on a small amount of steps (30k-ish) with respectively `[2,3,4,5]` agents. The longest flat lines in each segment are where get_weights() is being called.

![num_agents](https://user-images.githubusercontent.com/5096835/82066727-181ae080-96d0-11ea-8c28-17179c37423c.png)


Resolving patch:
diff --git a/python/ray/experimental/tf_utils.py b/python/ray/experimental/tf_utils.py
--- a/python/ray/experimental/tf_utils.py
+++ b/python/ray/experimental/tf_utils.py
@@ -161,10 +161,7 @@ def get_weights(self):
             Dictionary mapping variable names to their weights.
         """
         self._check_sess()
-        return {
-            k: v.eval(session=self.sess)
-            for k, v in self.variables.items()
-        }
+        return self.sess.run(self.variables)
 
     def set_weights(self, new_weights):
         """Sets the weights to new_weights.


Test patch: