wehub-resource-sync
126 lines
3.8 KiB
Python
126 lines
3.8 KiB
Python
import copy
|
|
import logging
|
|
import os
|
|
import queue
|
|
import threading
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
|
|
from ray.air.constants import _ERROR_REPORT_TIMEOUT
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def is_nan(value):
|
|
return np.isnan(value)
|
|
|
|
|
|
def is_nan_or_inf(value):
|
|
return is_nan(value) or np.isinf(value)
|
|
|
|
|
|
class StartTraceback(Exception):
|
|
"""These exceptions (and their tracebacks) can be skipped with `skip_exceptions`"""
|
|
|
|
pass
|
|
|
|
|
|
class StartTracebackWithWorkerRank(StartTraceback):
|
|
def __init__(self, worker_rank: int) -> None:
|
|
super().__init__()
|
|
self.worker_rank = worker_rank
|
|
|
|
def __reduce__(self):
|
|
return (self.__class__, (self.worker_rank,))
|
|
|
|
|
|
def skip_exceptions(exc: Optional[Exception]) -> Exception:
|
|
"""Skip all contained `StartTracebacks` to reduce traceback output.
|
|
|
|
Returns a shallow copy of the exception with all `StartTracebacks` removed.
|
|
|
|
If the RAY_AIR_FULL_TRACEBACKS environment variable is set,
|
|
the original exception (not a copy) is returned.
|
|
"""
|
|
should_not_shorten = bool(int(os.environ.get("RAY_AIR_FULL_TRACEBACKS", "0")))
|
|
|
|
if should_not_shorten:
|
|
return exc
|
|
|
|
if isinstance(exc, StartTraceback):
|
|
# If this is a StartTraceback, skip
|
|
return skip_exceptions(exc.__cause__)
|
|
|
|
# Perform a shallow copy to prevent recursive __cause__/__context__.
|
|
new_exc = copy.copy(exc).with_traceback(exc.__traceback__)
|
|
|
|
# Make sure nested exceptions are properly skipped.
|
|
cause = getattr(exc, "__cause__", None)
|
|
if cause:
|
|
new_exc.__cause__ = skip_exceptions(cause)
|
|
|
|
return new_exc
|
|
|
|
|
|
def exception_cause(exc: Optional[Exception]) -> Optional[Exception]:
|
|
if not exc:
|
|
return None
|
|
|
|
return getattr(exc, "__cause__", None)
|
|
|
|
|
|
class RunnerThread(threading.Thread):
|
|
"""Supervisor thread that runs your script."""
|
|
|
|
def __init__(self, *args, error_queue, **kwargs):
|
|
threading.Thread.__init__(self, *args, **kwargs)
|
|
self._error_queue = error_queue
|
|
self._ret = None
|
|
|
|
def _propagate_exception(self, e: BaseException):
|
|
try:
|
|
# report the error but avoid indefinite blocking which would
|
|
# prevent the exception from being propagated in the unlikely
|
|
# case that something went terribly wrong
|
|
self._error_queue.put(e, block=True, timeout=_ERROR_REPORT_TIMEOUT)
|
|
except queue.Full:
|
|
logger.critical(
|
|
(
|
|
"Runner Thread was unable to report error to main "
|
|
"function runner thread. This means a previous error "
|
|
"was not processed. This should never happen."
|
|
)
|
|
)
|
|
|
|
def run(self):
|
|
try:
|
|
self._ret = self._target(*self._args, **self._kwargs)
|
|
except StopIteration:
|
|
logger.debug(
|
|
(
|
|
"Thread runner raised StopIteration. Interpreting it as a "
|
|
"signal to terminate the thread without error."
|
|
)
|
|
)
|
|
except SystemExit as e:
|
|
# Do not propagate up for graceful termination.
|
|
if e.code == 0:
|
|
logger.debug(
|
|
(
|
|
"Thread runner raised SystemExit with error code 0. "
|
|
"Interpreting it as a signal to terminate the thread "
|
|
"without error."
|
|
)
|
|
)
|
|
else:
|
|
# If non-zero exit code, then raise exception to main thread.
|
|
self._propagate_exception(e)
|
|
except BaseException as e:
|
|
# Propagate all other exceptions to the main thread.
|
|
self._propagate_exception(e)
|
|
|
|
def join(self, timeout=None):
|
|
super(RunnerThread, self).join(timeout)
|
|
return self._ret
|