Repository: ray-project/ray
File: rllib/utils/checkpoints.py (Python)

Code before:
                    iter(
                        comp.foreach_actor(
                            _get_ip,
                            remote_actor_ids=[actor_to_use],
                        )
                    )
                )
                if not _result.ok:
                    raise _result.get()
                worker_ip_addr = _result.get()
                self_ip_addr = _get_ip()

                # Save the state to a temporary location on the `actor_to_use`'s
                # node.
                comp_state_ref = None
                if _state_provided:
                    comp_state_ref = ray.put(state.pop(comp_name))

                if worker_ip_addr == self_ip_addr:
                    comp.foreach_actor(
                        lambda w, _path=comp_path, _state=comp_state_ref, _use_msgpack=use_msgpack: (  # noqa
                            w.save_to_path(
                                _path,
                                state=(
                                    ray.get(_state)
                                    if _state is not None
                                    else w.get_state()
                                ),
                                use_msgpack=_use_msgpack,
                            )
                        ),
                        remote_actor_ids=[actor_to_use],
                    )
                else:
                    # Save the checkpoint to the temporary directory on the worker.
                    def _save(w, _state=comp_state_ref, _use_msgpack=use_msgpack):
                        import tempfile

                        # Create a temporary directory on the worker.
                        tmpdir = tempfile.mkdtemp()
                        w.save_to_path(
                            tmpdir,
                            state=(
                                ray.get(_state) if _state is not None else w.get_state()
                            ),
                            use_msgpack=_use_msgpack,
                        )
                        return tmpdir

                    _result = next(
                        iter(comp.foreach_actor(_save, remote_actor_ids=[actor_to_use]))

Review comment (suggestion):
From what I can see, passing on the filesystem to the children's save_to_path call is the only change in behaviour here? Lgtm

Code after:
                    iter(
                        comp.foreach_actor(
                            _get_ip,
                            remote_actor_ids=[actor_to_use],
                        )
                    )
                )
                if not _result.ok:
                    raise _result.get()
                worker_ip_addr = _result.get()
                self_ip_addr = _get_ip()

                # Save the state to a temporary location on the `actor_to_use`'s
                # node.
                comp_state_ref = None
                if _state_provided:
                    comp_state_ref = ray.put(state.pop(comp_name))

                # If worker_addr == self_addr, save directly to the path
                # provided by the user, make sure to use filesystem.
                if worker_ip_addr == self_ip_addr:
                    comp.foreach_actor(
                        lambda w, _path=comp_path, _filesystem=filesystem, _state=comp_state_ref, _use_msgpack=use_msgpack: (  # noqa
                            w.save_to_path(
                                path=_path,
                                filesystem=_filesystem,
                                state=(
                                    ray.get(_state)
                                    if _state is not None
                                    else w.get_state()
                                ),
                                use_msgpack=_use_msgpack,
                            )
                        ),
                        remote_actor_ids=[actor_to_use],
                    )
                # Transfer state files from the worker node to the head node
                else:
                    # Save the checkpoint to the temporary directory on the worker.
                    def _save(w, _state=comp_state_ref, _use_msgpack=use_msgpack):
                        import tempfile

                        # Create a temporary directory on the worker.
                        tmpdir = tempfile.mkdtemp()
                        w.save_to_path(
                            path=tmpdir,
                            state=(
                                ray.get(_state) if _state is not None else w.get_state()
                            ),
                            use_msgpack=_use_msgpack,
                        )