[train] Add worker_process_setup_hook to set mp start method to spawn

[SumanthRH]

What does this PR do?

Adds skyrl.utils.worker_setup.worker_setup_fn which sets the multiprocessing start method to 'spawn', and wires it into ray.init via the worker_process_setup_hookruntime_env key. Includes unit tests verifying the hook is applied in Ray workers.

We've previously seen many examples where ray <> fork interact in weird ways. There's code in the skyrl_entrypoint task (torch dataloader) as well as in other ray worker processes (ex: megatron workers) that rely on multiprocessing. Using worker setup hook provides us with a consistent way to handle this for all Ray worker processes.

Test plan:

• CPU tests
• One GPU test that uses ray : test_model_wrapper.py
• New CPU tests for worker setup function

---

[gemini-code-assist]

Code Review

This pull request introduces a worker_process_setup_hook to consistently set the multiprocessing start method to spawn for all Ray workers, which is a good practice to avoid issues with fork in a Ray environment. The changes are well-implemented, including a new utility function, its integration into ray.init, and corresponding unit tests. My feedback focuses on improving the robustness of the new tests to ensure they don't have side effects on other tests.

[gemini-code-assist]

These tests modify the global multiprocessing start method but don't restore its original state. This can lead to side effects and flaky tests, as the state change persists for subsequent tests in the same process. To improve test isolation, it's best practice to save the original state before the test and restore it in a finally block. This also makes the tests more portable across different operating systems.

def test_worker_setup_fn_sets_spawn():
    """Test that worker_setup_fn sets the mp start method to spawn."""
    if 'fork' not in multiprocessing.get_all_start_methods():
        pytest.skip("fork start method not supported")

    original_method = multiprocessing.get_start_method(allow_none=True)
    try:
        # Reset to a different state first
        multiprocessing.set_start_method("fork", force=True)
        worker_setup_fn()
        assert multiprocessing.get_start_method() == "spawn"
    finally:
        if original_method:
            multiprocessing.set_start_method(original_method, force=True)


def test_worker_setup_fn_idempotent():
    """Test that calling worker_setup_fn twice doesn't raise."""
    original_method = multiprocessing.get_start_method(allow_none=True)
    try:
        multiprocessing.set_start_method("spawn", force=True)
        worker_setup_fn()  # should not raise
        assert multiprocessing.get_start_method() == "spawn"
    finally:
        if original_method:
            multiprocessing.set_start_method(original_method, force=True)

[devin-ai-integration]

✅ Devin Review: No Issues Found

Devin Review analyzed this PR and found no potential bugs to report.

View in Devin Review to see 3 additional findings.