# Copyright 2020 Uber Technologies, Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import functools
import queue

from horovod.common.exceptions import HorovodInternalError, HostsUpdatedInterrupt
from horovod.runner.elastic.worker import HostUpdateResult, WorkerNotificationManager


notification_manager = WorkerNotificationManager()


class State(object):
    """State representation used for tracking in memory state across workers.

    Args:
        bcast_object: Function used to broadcast a variable from rank 0 to the other workers.
        get_rank: Function that returns the current rank of this worker.
    """
    def __init__(self, bcast_object, get_rank):
        self._bcast_object = bcast_object
        self._rank = get_rank
        self._host_messages = queue.Queue()
        self._last_updated_timestamp = 0
        self._reset_callbacks = []

    def register_reset_callbacks(self, callbacks):
        """Register callbacks that will be invoked following a reset event (worker added or removed).

        For example, a common use of a reset callback would be to update the learning rate scale with the
        new number of workers.

        Args:
            callbacks: list of functions to execute.
        """
        self._reset_callbacks.extend(callbacks)

    def on_reset(self):
        self._host_messages = queue.Queue()
        self.reset()
        for callback in self._reset_callbacks:
            callback()

    def on_hosts_updated(self, timestamp, update_res):
        self._host_messages.put((timestamp, update_res))

    def commit(self):
        """Commits all modifications to state tracked by this object to host memory.

        This call will also check for any changes to known hosts, and raise a `HostsUpdatedInterrupt`
        if any were detected.

        Because commits are a heavy operation involving data copy (potentially from GPU to host), it is
        recommended to consider committing less frequently than once per batch. This allows users to tradeoff
        between per-batch execution time and lost training steps in the event of a worker failure.
        """
        self.save()
        self.check_host_updates()

    def check_host_updates(self):
        """Checks that a notification has been sent indicating that hosts can be added or will be removed.

        Raises a `HostsUpdatedInterrupt` if such a notification has been received.
        """
        # Iterate through the update messages sent from the server. If the update timestamp
        # is greater than the last update timestamp, then trigger a HostsUpdatedException.
        last_updated_timestamp = prev_timestamp = self._last_updated_timestamp
        all_update = HostUpdateResult.no_update
        while not self._host_messages.empty():
            timestamp, update = self._host_messages.get()
            if timestamp > last_updated_timestamp:
                last_updated_timestamp = timestamp
                all_update |= update

        # In order to ensure all workers raise the exception at the same time, we need to sync
        # the updated state across all the workers.
        # TODO(travis): this should be a max allreduce to account for changes in rank 0
        prev_timestamp, self._last_updated_timestamp, all_update = \
            self._bcast_object((prev_timestamp, last_updated_timestamp, all_update))

        # At this point, updated state is globally consistent across all ranks.
        if self._last_updated_timestamp > prev_timestamp:
            raise HostsUpdatedInterrupt(all_update == HostUpdateResult.removed)


    def save(self):
        """Saves state to host memory."""
        raise NotImplementedError()

    def restore(self):
        """Restores the last committed state, undoing any uncommitted modifications."""
        raise NotImplementedError()

    def sync(self):
        """Synchronize state across workers."""
        raise NotImplementedError()

    def reset(self):
        """Reset objects and variables following a reset event (before synchronization)."""
        pass


class ObjectState(State):
    """State for simple Python objects.

    Every object is specified as a keyword argument, and will be assigned as an attribute.

    Args:
        bcast_object: Horovod broadcast object function used to sync state dictionary.
        get_rank: Horovod rank function used to identify is this process is the coordinator.
        kwargs: Properties to sync, will be exposed as attributes of the object.
    """
    def __init__(self, bcast_object, get_rank, **kwargs):
        self._bcast_object = bcast_object
        self._saved_state = kwargs
        self._set_attrs()
        super(ObjectState, self).__init__(bcast_object=bcast_object, get_rank=get_rank)

    def save(self):
        new_state = {}
        for attr in self._saved_state.keys():
            new_state[attr] = getattr(self, attr)
        self._saved_state = new_state

    def restore(self):
        self._set_attrs()

    def sync(self):
        if self._saved_state:
            self._saved_state = self._bcast_object(self._saved_state)
            self._set_attrs()

    def _set_attrs(self):
        for attr, value in self._saved_state.items():
            setattr(self, attr, value)


def run_fn(func, reset):
    @functools.wraps(func)
    def wrapper(state, *args, **kwargs):
        notification_manager.init()
        notification_manager.register_listener(state)
        skip_sync = False

        try:
            while True:
                if not skip_sync:
                    state.sync()

                try:
                    return func(state, *args, **kwargs)
                except HorovodInternalError:
                    state.restore()
                    skip_sync = False
                except HostsUpdatedInterrupt as e:
                    skip_sync = e.skip_sync

                reset()
                state.on_reset()
        finally:
            notification_manager.remove_listener(state)
    return wrapper

horovod/common/elastic.py
# Copyright 2020 Uber Technologies, Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import functools
import queue

from horovod.common.exceptions import HorovodInternalError, HostsUpdatedInterrupt
from horovod.runner.elastic.worker import HostUpdateResult, WorkerNotificationManager


notification_manager = WorkerNotificationManager()


class State(object):
    """State representation used for tracking in memory state across workers.

    Args:
        bcast_object: Function used to broadcast a variable from rank 0 to the other workers.
        get_rank: Function that returns the current rank of this worker.
    """
    def __init__(self, bcast_object, get_rank):
        self._bcast_object = bcast_object
        self._rank = get_rank
        self._host_messages = queue.Queue()
        self._last_updated_timestamp = 0
        self._reset_callbacks = []

    def register_reset_callbacks(self, callbacks):
        """Register callbacks that will be invoked following a reset event (worker added or removed).

        For example, a common use of a reset callback would be to update the learning rate scale with the
        new number of workers.

        Args:
            callbacks: list of functions to execute.
        """
        self._reset_callbacks.extend(callbacks)

    def on_reset(self):
        self._host_messages = queue.Queue()
        self.reset()
        for callback in self._reset_callbacks:
            callback()

    def on_hosts_updated(self, timestamp, update_res):
        self._host_messages.put((timestamp, update_res))

    def commit(self):
        """Commits all modifications to state tracked by this object to host memory.

        This call will also check for any changes to known hosts, and raise a `HostsUpdatedInterrupt`
        if any were detected.

        Because commits are a heavy operation involving data copy (potentially from GPU to host), it is
        recommended to consider committing less frequently than once per batch. This allows users to tradeoff
        between per-batch execution time and lost training steps in the event of a worker failure.
        """
        self.save()
        self.check_host_updates()

    def check_host_updates(self):
        """Checks that a notification has been sent indicating that hosts can be added or will be removed.

        Raises a `HostsUpdatedInterrupt` if such a notification has been received.
        """
        # Iterate through the update messages sent from the server. If the update timestamp
        # is greater than the last update timestamp, then trigger a HostsUpdatedException.
        last_updated_timestamp = prev_timestamp = self._last_updated_timestamp
        all_update = HostUpdateResult.no_update
        while not self._host_messages.empty():
            timestamp, update = self._host_messages.get()
            if timestamp > last_updated_timestamp:
                last_updated_timestamp = timestamp
                all_update |= update

        # In order to ensure all workers raise the exception at the same time, we need to sync
        # the updated state across all the workers.
        # TODO(travis): this should be a max allreduce to account for changes in rank 0
        prev_timestamp, self._last_updated_timestamp, all_update = \
            self._bcast_object((prev_timestamp, last_updated_timestamp, all_update))

        # At this point, updated state is globally consistent across all ranks.
        if self._last_updated_timestamp > prev_timestamp:
            raise HostsUpdatedInterrupt(all_update == HostUpdateResult.removed)


    def save(self):
        """Saves state to host memory."""
        raise NotImplementedError()

    def restore(self):
        """Restores the last committed state, undoing any uncommitted modifications."""
        raise NotImplementedError()

    def sync(self):
        """Synchronize state across workers."""
        raise NotImplementedError()

    def reset(self):
        """Reset objects and variables following a reset event (before synchronization)."""
        pass


class ObjectState(State):
    """State for simple Python objects.

    Every object is specified as a keyword argument, and will be assigned as an attribute.

    Args:
        bcast_object: Horovod broadcast object function used to sync state dictionary.
        get_rank: Horovod rank function used to identify is this process is the coordinator.
        kwargs: Properties to sync, will be exposed as attributes of the object.
    """
    def __init__(self, bcast_object, get_rank, **kwargs):
        self._bcast_object = bcast_object
        self._saved_state = kwargs
        self._set_attrs()
        super(ObjectState, self).__init__(bcast_object=bcast_object, get_rank=get_rank)

    def save(self):
        new_state = {}
        for attr in self._saved_state.keys():
            new_state[attr] = getattr(self, attr)
        self._saved_state = new_state

    def restore(self):
        self._set_attrs()

    def sync(self):
        if self._saved_state:
            self._saved_state = self._bcast_object(self._saved_state)
            self._set_attrs()

    def _set_attrs(self):
        for attr, value in self._saved_state.items():
            setattr(self, attr, value)


def run_fn(func, reset):
    @functools.wraps(func)
    def wrapper(state, *args, **kwargs):
        notification_manager.init()
        notification_manager.register_listener(state)
        skip_sync = False

        try:
            while True:
                try:
                    if not skip_sync:
                        state.sync()

                    return func(state, *args, **kwargs)
                except HorovodInternalError:
                    state.restore()
                    skip_sync = False
                except HostsUpdatedInterrupt as e:
                    skip_sync = e.skip_sync

                reset()
                state.on_reset()
        finally:
            notification_manager.remove_listener(state)
    return wrapper

horovod/common/elastic.py
diff --git a/horovod/common/elastic.py b/horovod/common/elastic.py
index 84512a299f..f7096eb7ee 100644
--- a/horovod/common/elastic.py
+++ b/horovod/common/elastic.py
@@ -157,10 +157,10 @@ def wrapper(state, *args, **kwargs):
 
         try:
             while True:
-                if not skip_sync:
-                    state.sync()
-
                 try:
+                    if not skip_sync:
+                        state.sync()
+
                     return func(state, *args, **kwargs)
                 except HorovodInternalError:
                     state.restore()

【Elastic Horovod】Should we catch exceptions for state.sync()？
**Environment:**
1. Framework: (TensorFlow, Keras, PyTorch, MXNet): Pytorch
2. Framework version: 1.6.0
3. Horovod version: 0.21.3
4. MPI version: 4.0.3
5. CUDA version: 10.2
6. NCCL version: 2.7.6
7. Python version: 3.6

**Checklist:**
1. Did you search issues to find if somebody asked this question before? Yes.
2. If your question is about hang, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/running.rst)?
3. If your question is about docker, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/docker.rst)?
4. Did you check if you question is answered in the [troubleshooting guide]  (https://github.com/horovod/horovod/blob/master/docs/troubleshooting.rst)? Yes

**Bug report:**
When a new worker was added in host discovery script, old workers will sync their state to new one. But if any worker failed during state synchronization,  unfortunately, the elastic horovod task will fail and  it seems not play a role for elastic:
```
[0]<stderr>:[2021-06-21 21:35:05.743047: E /tmp/pip-req-build-4rhufbvy/horovod/common/operations.cc:640] Horovod background loop uncaught exception: [/pytorch/third_party/gloo/gloo/transport/tcp/pair.cc:575] Connection closed by peer [11.198.63.123]:50349
[0]<stdout>:[2021-06-21 21:35:05.773132: D /tmp/pip-req-build-4rhufbvy/horovod/common/operations.cc:652] [0]: Shutting down background thread
[0]<stderr>:Traceback (most recent call last):
[0]<stderr>:  File "/usr/local/lib64/python3.6/site-packages/horovod/torch/mpi_ops.py", line 960, in synchronize
[0]<stderr>:    mpi_lib.horovod_torch_wait_and_clear(handle)[0]<stderr>:RuntimeError: Horovod has been shut down. This was caused by an exception on one of the ranks or an attempt to allreduce, allgather or broadcast a tensor after one of the ranks finished execution. If the shutdown was caused by an exception, you should see the exception in the log before the first shutdown message.
[0]<stderr>:
[0]<stderr>:During handling of the above exception, another exception occurred:[0]<stderr>: 
[0]<stderr>:Traceback (most recent call last):[0]<stderr>:  File "pytorch_synthetic_benchmark_elastic.py", line 140, in <module>
[0]<stderr>:    run_benchmark(state)[0]<stderr>:  File "/usr/local/lib64/python3.6/site-packages/horovod/common/elastic.py", line 162, in wrapper
[0]<stderr>:    state.sync()
[0]<stderr>:  File "/usr/local/lib64/python3.6/site-packages/horovod/torch/elastic/state.py", line 62, in sync
[0]<stderr>:    handler.sync()[0]<stderr>:  File "/usr/local/lib64/python3.6/site-packages/horovod/torch/elastic/state.py", line 101, in sync
[0]<stderr>:    broadcast_parameters(self.value.state_dict(), root_rank=0)[0]<stderr>:  File "/usr/local/lib64/python3.6/site-packages/horovod/torch/functions.py", line 58, in broadcast_parameters
[0]<stderr>:    synchronize(handle)[0]<stderr>:  File "/usr/local/lib64/python3.6/site-packages/horovod/torch/mpi_ops.py", line 964, in synchronize
[0]<stderr>:    raise HorovodInternalError(e)
[0]<stderr>:horovod.common.exceptions.HorovodInternalError: Horovod has been shut down. This was caused by an exception on one of the ranks or an attempt to allreduce, allgather or broadcast a tensor after one of the ranks finished execution. If the shutdown was caused by an exception, you should see the exception in the log before the first shutdown message.
``` 

I think that was caused by this [code segment:](https://github.com/horovod/horovod/blob/139416965ab9aa5850baf96ec54ce35c58b05119/horovod/common/elastic.py#L161)
It works well for me when I fix code as follows
```
def run_fn(func, reset):
       ....
        try:
            while True:
                try:
                    # Here we also catch exceptions for state.sync().
                    if not skip_sync:
                         state.sync()
                    return func(state, *args, **kwargs)
                except HorovodInternalError:
                    state.restore()
                    skip_sync = False
                except HostsUpdatedInterrupt as e:
                    skip_sync = e.skip_sync

                reset()
                state.on_reset()
        finally:
            notification_manager.remove_listener(state)
    return wrapper
```


**Steps to reproduce.**
1. In order to easily reproduce the problem, we add one line  in `horovod/examples/elastic/pytorch/pytorch_synthetic_benchmark_elastic.py` as follows:
```
...
state.register_reset_callbacks([on_state_reset])
# Here we sleep 30s to keep old workers stay in state.sync() when a new worker
# was add in host-discovery-script.
time.sleep(30)
run_benchmark(state)
...
```
2. Run elastic horovod:
```
horovodrun -np 1 --host-discovery-script ./discovery_hosts.sh --network-interface eth1 --min-np 1 --log-level DEBUG --verbose  python3 pytorch_synthetic_benchmark_elastic.py --num-iters=1000
```
3. After some iteration passed, we add a new worker in host-discovery-script to raise `HostsUpdatedInterrupt`. The old workers will call `state.sync()` and hang in `state.sync()` for 30s as new worker will sleep 30s before `hvd.elastic.run`
4. When old worker was hang in `state.sync`,  we kill one old worker to raise `HorovodInternalError` . At this time the elastic horovod will fail. The content of stderr as follows:
```
[0]<stderr>:[2021-06-21 21:35:05.743047: E /tmp/pip-req-build-4rhufbvy/horovod/common/operations.cc:640] Horovod background loop uncaught exception: [/pytorch/third_party/gloo/gloo/transport/tcp/pair.cc:575] Connection closed by peer [11.198.63.123]:50349
[0]<stdout>:[2021-06-21 21:35:05.773132: D /tmp/pip-req-build-4rhufbvy/horovod/common/operations.cc:652] [0]: Shutting down background thread
[0]<stderr>:Traceback (most recent call last):
[0]<stderr>:  File "/usr/local/lib64/python3.6/site-packages/horovod/torch/mpi_ops.py", line 960, in synchronize
[0]<stderr>:    mpi_lib.horovod_torch_wait_and_clear(handle)[0]<stderr>:RuntimeError: Horovod has been shut down. This was caused by an exception on one of the ranks or an attempt to allreduce, allgather or broadcast a tensor after one of the ranks finished execution. If the shutdown was caused by an exception, you should see the exception in the log before the first shutdown message.
[0]<stderr>:
[0]<stderr>:During handling of the above exception, another exception occurred:[0]<stderr>: 
[0]<stderr>:Traceback (most recent call last):[0]<stderr>:  File "pytorch_synthetic_benchmark_elastic.py", line 140, in <module>
[0]<stderr>:    run_benchmark(state)[0]<stderr>:  File "/usr/local/lib64/python3.6/site-packages/horovod/common/elastic.py", line 162, in wrapper
[0]<stderr>:    state.sync()
[0]<stderr>:  File "/usr/local/lib64/python3.6/site-packages/horovod/torch/elastic/state.py", line 62, in sync
[0]<stderr>:    handler.sync()[0]<stderr>:  File "/usr/local/lib64/python3.6/site-packages/horovod/torch/elastic/state.py", line 101, in sync
[0]<stderr>:    broadcast_parameters(self.value.state_dict(), root_rank=0)[0]<stderr>:  File "/usr/local/lib64/python3.6/site-packages/horovod/torch/functions.py", line 58, in broadcast_parameters
[0]<stderr>:    synchronize(handle)[0]<stderr>:  File "/usr/local/lib64/python3.6/site-packages/horovod/torch/mpi_ops.py", line 964, in synchronize
[0]<stderr>:    raise HorovodInternalError(e)
[0]<stderr>:horovod.common.exceptions.HorovodInternalError: Horovod has been shut down. This was caused by an exception on one of the ranks or an attempt to allreduce, allgather or broadcast a tensor after one of the ranks finished execution. If the shutdown was caused by an exception, you should see the exception in the log before the first shutdown message.
``` 


