Skip to content

Add watchdog for stalled Ray initialization - #2016

Open
kalectory wants to merge 1 commit into
NovaSky-AI:mainfrom
kalectory:neil/ray-init-timeout
Open

Add watchdog for stalled Ray initialization#2016
kalectory wants to merge 1 commit into
NovaSky-AI:mainfrom
kalectory:neil/ray-init-timeout

Conversation

@kalectory

Copy link
Copy Markdown

Summary

  • add an opt-in watchdog around Ray driver initialization
  • log the configured deadline and successful initialization duration
  • dump all Python thread stacks and terminate the driver when initialization exceeds the deadline

Motivation

Trajectory TCLI cold starts XIDs 1042972 and 1042974 both blocked indefinitely while the driver CoreWorker registered with the local raylet. The API process stayed healthy and create_model remained pending until its caller timed out after 3600 seconds. Because ray.init is synchronous, timing out only the caller leaves the unusable engine alive.

Set SKYRL_RAY_INIT_TIMEOUT_IN_S to enable the watchdog. The default remains disabled so existing deployments do not change behavior.

Test plan

  • uvx --from ruff==0.11.9 ruff check skyrl/train/utils/utils.py tests/train/utils/test_initialize_ray.py
  • uvx --from black==24.10.0 black --check skyrl/train/utils/utils.py tests/train/utils/test_initialize_ray.py
  • uv run --isolated --extra dev --extra fsdp pytest -q tests/train/utils/test_initialize_ray.py

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a watchdog mechanism to terminate the driver process if ray.init() hangs or fails to make progress within a configurable timeout specified by the SKYRL_RAY_INIT_TIMEOUT_IN_S environment variable. The review feedback suggests several robustness improvements: wrapping faulthandler.dump_traceback in a try-except block to guarantee process termination even if traceback dumping fails, safely parsing the timeout environment variable to prevent crashes on invalid values, and handling NaN timeout values to avoid raising exceptions in the threading library.

Comment on lines +43 to +45
logger.error(f"ray.init() did not complete within {timeout_s:g}s; terminating the driver")
faulthandler.dump_traceback(file=sys.stderr, all_threads=True)
os._exit(1)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

If sys.stderr is redirected or replaced with a custom stream that does not support a file descriptor (e.g., in certain application servers, container environments, or testing frameworks), faulthandler.dump_traceback will raise a RuntimeError or AttributeError.

If an exception is raised here, the watchdog thread will crash before executing os._exit(1), leaving the stalled process hanging indefinitely.

Wrap the traceback dump in a try...except block to ensure that the driver is guaranteed to terminate even if the traceback dump fails.

Suggested change
logger.error(f"ray.init() did not complete within {timeout_s:g}s; terminating the driver")
faulthandler.dump_traceback(file=sys.stderr, all_threads=True)
os._exit(1)
logger.error(f"ray.init() did not complete within {timeout_s:g}s; terminating the driver")
try:
faulthandler.dump_traceback(file=sys.stderr, all_threads=True)
except Exception as e:
logger.error(f"Failed to dump traceback: {e}")
os._exit(1)

# log_to_driver=True allows training progress from skyrl_entrypoint to reach stdout.
# Infrastructure logs (vLLM, workers) are redirected to log file via os.dup2 in their init.
ray.init(runtime_env={"env_vars": env_vars}, log_to_driver=True)
ray_init_timeout_s = float(os.environ.get("SKYRL_RAY_INIT_TIMEOUT_IN_S", "0"))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Parsing the environment variable SKYRL_RAY_INIT_TIMEOUT_IN_S directly with float() without handling potential parsing errors can cause the driver to crash with a ValueError if the variable is set to an empty string or an invalid non-numeric value (which is common in templated deployments).

Safely parse the environment variable with a try...except ValueError block and default to 0.0 (disabled) with a warning log.

Suggested change
ray_init_timeout_s = float(os.environ.get("SKYRL_RAY_INIT_TIMEOUT_IN_S", "0"))
ray_init_timeout_env = os.environ.get("SKYRL_RAY_INIT_TIMEOUT_IN_S", "0")
try:
ray_init_timeout_s = float(ray_init_timeout_env) if ray_init_timeout_env.strip() else 0.0
except ValueError:
logger.warning(
f"Invalid SKYRL_RAY_INIT_TIMEOUT_IN_S value: {ray_init_timeout_env!r}. "
"Disabling the watchdog (timeout=0)."
)
ray_init_timeout_s = 0.0

Comment on lines +37 to +38
if timeout_s <= 0:
return completed

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

If timeout_s is NaN (e.g., if parsed from a malformed float representation), the condition timeout_s <= 0 will evaluate to False. Subsequently, calling completed.wait(timeout_s) with NaN will raise a ValueError in Python's threading library, crashing the watchdog thread.

Add a check for math.isnan(timeout_s) to safely handle this edge case.

Suggested change
if timeout_s <= 0:
return completed
if timeout_s <= 0 or math.isnan(timeout_s):
return completed

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant