chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,429 @@
|
||||
import abc
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple
|
||||
|
||||
from ray._private.thirdparty.tabulate.tabulate import tabulate
|
||||
from ray.util.annotations import Deprecated, DeveloperAPI
|
||||
from ray.widgets import Template
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Syncing period for syncing checkpoints between nodes or to cloud.
|
||||
DEFAULT_SYNC_PERIOD = 300
|
||||
|
||||
# Default sync timeout after which syncing processes are aborted
|
||||
DEFAULT_SYNC_TIMEOUT = 1800
|
||||
|
||||
|
||||
@Deprecated
|
||||
@dataclass
|
||||
class SyncConfig:
|
||||
sync_period: int = DEFAULT_SYNC_PERIOD
|
||||
sync_timeout: int = DEFAULT_SYNC_TIMEOUT
|
||||
sync_artifacts: bool = False
|
||||
sync_artifacts_on_checkpoint: bool = True
|
||||
|
||||
def _repr_html_(self) -> str:
|
||||
"""Generate an HTML representation of the SyncConfig."""
|
||||
return Template("scrollableTable.html.j2").render(
|
||||
table=tabulate(
|
||||
{
|
||||
"Setting": ["Sync period", "Sync timeout"],
|
||||
"Value": [self.sync_period, self.sync_timeout],
|
||||
},
|
||||
tablefmt="html",
|
||||
showindex=False,
|
||||
headers="keys",
|
||||
),
|
||||
max_height="none",
|
||||
)
|
||||
|
||||
|
||||
class _BackgroundProcess:
|
||||
def __init__(self, fn: Callable):
|
||||
self._fn = fn
|
||||
self._process = None
|
||||
self._result = {}
|
||||
self._start_time = float("-inf")
|
||||
|
||||
@property
|
||||
def is_running(self):
|
||||
return self._process and self._process.is_alive()
|
||||
|
||||
@property
|
||||
def start_time(self):
|
||||
return self._start_time
|
||||
|
||||
def start(self, *args, **kwargs):
|
||||
if self.is_running:
|
||||
return False
|
||||
|
||||
self._result = {}
|
||||
|
||||
def entrypoint():
|
||||
try:
|
||||
result = self._fn(*args, **kwargs)
|
||||
except Exception as e:
|
||||
self._result["exception"] = e
|
||||
return
|
||||
|
||||
self._result["result"] = result
|
||||
|
||||
self._process = threading.Thread(target=entrypoint)
|
||||
self._process.daemon = True
|
||||
self._process.start()
|
||||
self._start_time = time.time()
|
||||
|
||||
def wait(self, timeout: Optional[float] = None) -> Any:
|
||||
"""Waits for the background process to finish running. Waits until the
|
||||
background process has run for at least `timeout` seconds, counting from
|
||||
the time when the process was started."""
|
||||
if not self._process:
|
||||
return None
|
||||
|
||||
time_remaining = None
|
||||
if timeout:
|
||||
elapsed = time.time() - self.start_time
|
||||
time_remaining = max(timeout - elapsed, 0)
|
||||
|
||||
self._process.join(timeout=time_remaining)
|
||||
|
||||
if self._process.is_alive():
|
||||
self._process = None
|
||||
raise TimeoutError(
|
||||
f"{getattr(self._fn, '__name__', str(self._fn))} did not finish "
|
||||
f"running within the timeout of {timeout} seconds."
|
||||
)
|
||||
|
||||
self._process = None
|
||||
|
||||
exception = self._result.get("exception")
|
||||
if exception:
|
||||
raise exception
|
||||
|
||||
result = self._result.get("result")
|
||||
|
||||
self._result = {}
|
||||
return result
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class Syncer(abc.ABC):
|
||||
"""Syncer class for synchronizing data between Ray nodes and remote (cloud) storage.
|
||||
|
||||
This class handles data transfer for two cases:
|
||||
|
||||
1. Synchronizing data such as experiment state snapshots from the driver to
|
||||
cloud storage.
|
||||
2. Synchronizing data such as trial checkpoints from remote trainables to
|
||||
cloud storage.
|
||||
|
||||
Synchronizing tasks are usually asynchronous and can be awaited using ``wait()``.
|
||||
The base class implements a ``wait_or_retry()`` API that will retry a failed
|
||||
sync command.
|
||||
|
||||
The base class also exposes an API to only kick off syncs every ``sync_period``
|
||||
seconds.
|
||||
|
||||
Args:
|
||||
sync_period: The minimum time in seconds between sync operations, as
|
||||
used by ``sync_up/down_if_needed``.
|
||||
sync_timeout: The maximum time to wait for a sync process to finish before
|
||||
issuing a new sync operation. Ex: should be used by ``wait`` if launching
|
||||
asynchronous sync tasks.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sync_period: float = DEFAULT_SYNC_PERIOD,
|
||||
sync_timeout: float = DEFAULT_SYNC_TIMEOUT,
|
||||
):
|
||||
self.sync_period = sync_period
|
||||
self.sync_timeout = sync_timeout
|
||||
self.last_sync_up_time = float("-inf")
|
||||
self.last_sync_down_time = float("-inf")
|
||||
|
||||
@abc.abstractmethod
|
||||
def sync_up(
|
||||
self, local_dir: str, remote_dir: str, exclude: Optional[List] = None
|
||||
) -> bool:
|
||||
"""Synchronize local directory to remote directory.
|
||||
|
||||
This function can spawn an asynchronous process that can be awaited in
|
||||
``wait()``.
|
||||
|
||||
Args:
|
||||
local_dir: Local directory to sync from.
|
||||
remote_dir: Remote directory to sync up to. This is an URI
|
||||
(``protocol://remote/path``).
|
||||
exclude: Pattern of files to exclude, e.g.
|
||||
``["*/checkpoint_*]`` to exclude trial checkpoints.
|
||||
|
||||
Returns:
|
||||
True if sync process has been spawned, False otherwise.
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
def sync_down(
|
||||
self, remote_dir: str, local_dir: str, exclude: Optional[List] = None
|
||||
) -> bool:
|
||||
"""Synchronize remote directory to local directory.
|
||||
|
||||
This function can spawn an asynchronous process that can be awaited in
|
||||
``wait()``.
|
||||
|
||||
Args:
|
||||
remote_dir: Remote directory to sync down from. This is an URI
|
||||
(``protocol://remote/path``).
|
||||
local_dir: Local directory to sync to.
|
||||
exclude: Pattern of files to exclude, e.g.
|
||||
``["*/checkpoint_*]`` to exclude trial checkpoints.
|
||||
|
||||
Returns:
|
||||
True if sync process has been spawned, False otherwise.
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
def delete(self, remote_dir: str) -> bool:
|
||||
"""Delete directory on remote storage.
|
||||
|
||||
This function can spawn an asynchronous process that can be awaited in
|
||||
``wait()``.
|
||||
|
||||
Args:
|
||||
remote_dir: Remote directory to delete. This is an URI
|
||||
(``protocol://remote/path``).
|
||||
|
||||
Returns:
|
||||
True if sync process has been spawned, False otherwise.
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def retry(self):
|
||||
"""Retry the last sync up, sync down, or delete command.
|
||||
|
||||
You should implement this method if you spawn asynchronous syncing
|
||||
processes.
|
||||
"""
|
||||
pass
|
||||
|
||||
def wait(self, timeout: Optional[float] = None):
|
||||
"""Wait for asynchronous sync command to finish.
|
||||
|
||||
You should implement this method if you spawn asynchronous syncing
|
||||
processes. This method should timeout after the asynchronous command
|
||||
has run for `sync_timeout` seconds and raise a `TimeoutError`.
|
||||
"""
|
||||
pass
|
||||
|
||||
def sync_up_if_needed(
|
||||
self, local_dir: str, remote_dir: str, exclude: Optional[List] = None
|
||||
) -> bool:
|
||||
"""Syncs up if time since last sync up is greater than sync_period.
|
||||
|
||||
Args:
|
||||
local_dir: Local directory to sync from.
|
||||
remote_dir: Remote directory to sync up to. This is an URI
|
||||
(``protocol://remote/path``).
|
||||
exclude: Pattern of files to exclude, e.g.
|
||||
``["*/checkpoint_*]`` to exclude trial checkpoints.
|
||||
|
||||
Returns:
|
||||
The result of ``sync_up`` if a sync was triggered, otherwise ``None``.
|
||||
"""
|
||||
now = time.time()
|
||||
if now - self.last_sync_up_time >= self.sync_period:
|
||||
result = self.sync_up(
|
||||
local_dir=local_dir, remote_dir=remote_dir, exclude=exclude
|
||||
)
|
||||
self.last_sync_up_time = now
|
||||
return result
|
||||
|
||||
def sync_down_if_needed(
|
||||
self, remote_dir: str, local_dir: str, exclude: Optional[List] = None
|
||||
):
|
||||
"""Syncs down if time since last sync down is greater than sync_period.
|
||||
|
||||
Args:
|
||||
remote_dir: Remote directory to sync down from. This is an URI
|
||||
(``protocol://remote/path``).
|
||||
local_dir: Local directory to sync to.
|
||||
exclude: Pattern of files to exclude, e.g.
|
||||
``["*/checkpoint_*]`` to exclude trial checkpoints.
|
||||
|
||||
Returns:
|
||||
The result of ``sync_down`` if a sync was triggered, otherwise ``None``.
|
||||
"""
|
||||
now = time.time()
|
||||
if now - self.last_sync_down_time >= self.sync_period:
|
||||
result = self.sync_down(
|
||||
remote_dir=remote_dir, local_dir=local_dir, exclude=exclude
|
||||
)
|
||||
self.last_sync_down_time = now
|
||||
return result
|
||||
|
||||
def wait_or_retry(self, max_retries: int = 2, backoff_s: int = 5):
|
||||
assert max_retries > 0
|
||||
last_error_traceback = None
|
||||
for i in range(max_retries + 1):
|
||||
try:
|
||||
self.wait()
|
||||
except Exception as e:
|
||||
attempts_remaining = max_retries - i
|
||||
|
||||
# If we're out of retries, then save the full traceback of the last
|
||||
# error and show it when raising an exception.
|
||||
if attempts_remaining == 0:
|
||||
last_error_traceback = traceback.format_exc()
|
||||
break
|
||||
|
||||
logger.error(
|
||||
f"The latest sync operation failed with the following error: "
|
||||
f"{repr(e)}\n"
|
||||
f"Retrying {attempts_remaining} more time(s) after sleeping "
|
||||
f"for {backoff_s} seconds..."
|
||||
)
|
||||
time.sleep(backoff_s)
|
||||
self.retry()
|
||||
continue
|
||||
# Succeeded!
|
||||
return
|
||||
raise RuntimeError(
|
||||
f"Failed sync even after {max_retries} retries. "
|
||||
f"The latest sync failed with the following error:\n{last_error_traceback}"
|
||||
)
|
||||
|
||||
def reset(self):
|
||||
self.last_sync_up_time = float("-inf")
|
||||
self.last_sync_down_time = float("-inf")
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
def _repr_html_(self) -> str:
|
||||
return
|
||||
|
||||
|
||||
class _BackgroundSyncer(Syncer):
|
||||
"""Syncer using a background process for asynchronous file transfer."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sync_period: float = DEFAULT_SYNC_PERIOD,
|
||||
sync_timeout: float = DEFAULT_SYNC_TIMEOUT,
|
||||
):
|
||||
super(_BackgroundSyncer, self).__init__(
|
||||
sync_period=sync_period, sync_timeout=sync_timeout
|
||||
)
|
||||
self._sync_process = None
|
||||
self._current_cmd = None
|
||||
|
||||
def _should_continue_existing_sync(self):
|
||||
"""Returns whether a previous sync is still running within the timeout."""
|
||||
return (
|
||||
self._sync_process
|
||||
and self._sync_process.is_running
|
||||
and time.time() - self._sync_process.start_time < self.sync_timeout
|
||||
)
|
||||
|
||||
def _launch_sync_process(self, sync_command: Tuple[Callable, Dict]):
|
||||
"""Waits for the previous sync process to finish,
|
||||
then launches a new process that runs the given command."""
|
||||
if self._sync_process:
|
||||
try:
|
||||
self.wait()
|
||||
except Exception:
|
||||
logger.warning(
|
||||
f"Last sync command failed with the following error:\n"
|
||||
f"{traceback.format_exc()}"
|
||||
)
|
||||
|
||||
self._current_cmd = sync_command
|
||||
self.retry()
|
||||
|
||||
def sync_up(
|
||||
self, local_dir: str, remote_dir: str, exclude: Optional[List] = None
|
||||
) -> bool:
|
||||
if self._should_continue_existing_sync():
|
||||
logger.debug(
|
||||
f"Last sync still in progress, "
|
||||
f"skipping sync up of {local_dir} to {remote_dir}"
|
||||
)
|
||||
return False
|
||||
|
||||
sync_up_cmd = self._sync_up_command(
|
||||
local_path=local_dir, uri=remote_dir, exclude=exclude
|
||||
)
|
||||
self._launch_sync_process(sync_up_cmd)
|
||||
|
||||
return True
|
||||
|
||||
def _sync_up_command(
|
||||
self, local_path: str, uri: str, exclude: Optional[List] = None
|
||||
) -> Tuple[Callable, Dict]:
|
||||
raise NotImplementedError
|
||||
|
||||
def sync_down(
|
||||
self, remote_dir: str, local_dir: str, exclude: Optional[List] = None
|
||||
) -> bool:
|
||||
if self._should_continue_existing_sync():
|
||||
logger.warning(
|
||||
f"Last sync still in progress, "
|
||||
f"skipping sync down of {remote_dir} to {local_dir}"
|
||||
)
|
||||
return False
|
||||
|
||||
sync_down_cmd = self._sync_down_command(uri=remote_dir, local_path=local_dir)
|
||||
self._launch_sync_process(sync_down_cmd)
|
||||
|
||||
return True
|
||||
|
||||
def _sync_down_command(self, uri: str, local_path: str) -> Tuple[Callable, Dict]:
|
||||
raise NotImplementedError
|
||||
|
||||
def delete(self, remote_dir: str) -> bool:
|
||||
if self._should_continue_existing_sync():
|
||||
logger.warning(
|
||||
f"Last sync still in progress, skipping deletion of {remote_dir}"
|
||||
)
|
||||
return False
|
||||
|
||||
delete_cmd = self._delete_command(uri=remote_dir)
|
||||
self._launch_sync_process(delete_cmd)
|
||||
|
||||
return True
|
||||
|
||||
def _delete_command(self, uri: str) -> Tuple[Callable, Dict]:
|
||||
raise NotImplementedError
|
||||
|
||||
def wait(self, timeout: Optional[float] = None):
|
||||
if self._sync_process:
|
||||
try:
|
||||
self._sync_process.wait(timeout=timeout or self.sync_timeout)
|
||||
except Exception as e:
|
||||
raise e
|
||||
finally:
|
||||
# Regardless of whether the sync process succeeded within the timeout,
|
||||
# clear the sync process so a new one can be created.
|
||||
self._sync_process = None
|
||||
|
||||
def retry(self):
|
||||
if not self._current_cmd:
|
||||
raise RuntimeError("No sync command set, cannot retry.")
|
||||
cmd, kwargs = self._current_cmd
|
||||
self._sync_process = _BackgroundProcess(cmd)
|
||||
self._sync_process.start(**kwargs)
|
||||
|
||||
def __getstate__(self):
|
||||
state = self.__dict__.copy()
|
||||
state["_sync_process"] = None
|
||||
return state
|
||||
Reference in New Issue
Block a user