sgl-project--sglang/test/manual/test_ray_engine.py

"""Integration tests for RayEngine and Ray HTTP server (requires GPU + Ray).

Tests the Ray actor scheduler backend:
  - Offline inference via Engine(use_ray=True) inside a Ray actor on a placement group
  - Data parallel (DP) and DP attention support
  - Custom placement_group and SGLANG_RAY_BUNDLE_INDICES for fine-grained bundle control
  - Error paths in RayEngine._launch_scheduler_processes()
  - HTTP server launched via --use-ray flag

Usage:
    # 1-GPU tests
    python -m pytest test/manual/test_ray_engine.py::TestRayEngineOfflineTP1 -v -s
    python -m pytest test/manual/test_ray_engine.py::TestRayEngineErrors -v -s
    python -m pytest test/manual/test_ray_engine.py::TestRayHTTPServerTP1 -v -s
    python -m pytest test/manual/test_ray_engine.py::TestRayEnginePlacementGroupErrors -v -s

    # 2-GPU tests
    python -m pytest test/manual/test_ray_engine.py::TestRayEngineOfflineTP2 -v -s
    python -m pytest test/manual/test_ray_engine.py::TestRayEngineOfflinePP2 -v -s
    python -m pytest test/manual/test_ray_engine.py::TestRayEngineOfflineDP2 -v -s
    python -m pytest test/manual/test_ray_engine.py::TestRayEngineOfflineDPAttention -v -s
    python -m pytest test/manual/test_ray_engine.py::TestRayEnginePlacementGroup -v -s
"""

from __future__ import annotations

import os
import time
import unittest

import torch

from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST

# Allow overriding the model via env var for environments without gated access
_MODEL = os.environ.get("SGLANG_TEST_MODEL", DEFAULT_SMALL_MODEL_NAME_FOR_TEST)

# DP attention requires a model whose num_kv_heads divides evenly across the
# attention-TP dimension.  Qwen2.5-0.5B (kv_heads=2, attn_heads=14) hits a
# shape mismatch in the KV cache, so we use a larger model here.
_DP_ATTN_MODEL = os.environ.get("SGLANG_TEST_DP_ATTN_MODEL", "Qwen/Qwen3-8B")

try:
    import ray
    from ray.runtime_env import RuntimeEnv
    from ray.util.placement_group import placement_group
    from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

    # Prevent Ray from overriding CUDA_VISIBLE_DEVICES so that all GPUs
    # remain visible inside actors regardless of num_gpus allocation.
    _env_vars = {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"}
    if os.environ.get("HF_TOKEN"):
        _env_vars["HF_TOKEN"] = os.environ["HF_TOKEN"]
    _RAY_RUNTIME_ENV = RuntimeEnv(env_vars=_env_vars)
    _has_ray = True
except ImportError:
    _has_ray = False
    _RAY_RUNTIME_ENV = None


_NUM_GPUS = torch.cuda.device_count()

_SAMPLING_PARAMS = {"max_new_tokens": 32, "temperature": 0.0}

_PROMPTS = [
    "The capital of France is",
    "Explain quantum computing in simple terms:",
    "Write a haiku about programming:",
    "What is 2 + 2?",
]


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _create_engine_on_pg(
    tp_size, pp_size=1, dp_size=1, model=_MODEL, extra_kwargs=None
):
    """Create an EngineActor on a placement group and wait for it to be ready.

    Returns (engine_actor, placement_group).
    """

    @ray.remote
    class EngineActor:
        def __init__(self, **kwargs):
            from sglang.srt.ray.engine import RayEngine

            self.engine = RayEngine(**kwargs)

        def is_ready(self):
            return True

        def generate(self, prompt, sampling_params):
            return self.engine.generate(prompt=prompt, sampling_params=sampling_params)

        def shutdown(self):
            if self.engine:
                self.engine.shutdown()
                self.engine = None

    enable_dp_attention = (extra_kwargs or {}).get("enable_dp_attention", False)
    if enable_dp_attention:
        # DP attention folds DP into TP — total GPUs = tp_size * pp_size
        total_gpus = tp_size * pp_size
    else:
        total_gpus = dp_size * tp_size * pp_size
    pg = placement_group(
        [{"CPU": 1, "GPU": total_gpus}],
        strategy="STRICT_PACK",
    )
    ray.get(pg.ready())

    kwargs = dict(
        model_path=model,
        tp_size=tp_size,
        pp_size=pp_size,
        dp_size=dp_size,
    )
    if extra_kwargs:
        kwargs.update(extra_kwargs)

    actor = EngineActor.options(
        num_cpus=1,
        num_gpus=0,
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=pg,
            placement_group_bundle_index=0,
        ),
    ).remote(**kwargs)

    ray.get(actor.is_ready.remote(), timeout=600)
    return actor, pg


def _cleanup(actor, pg):
    """Shutdown engine actor and remove placement group."""
    try:
        ray.get(actor.shutdown.remote(), timeout=60)
    except Exception:
        pass
    try:
        ray.util.remove_placement_group(pg)
    except Exception:
        pass


# ---------------------------------------------------------------------------
# Tests: Offline TP=1
# ---------------------------------------------------------------------------


@unittest.skipUnless(_has_ray, "ray is not installed")
@unittest.skipUnless(_NUM_GPUS >= 1, "requires at least 1 GPU")
class TestRayEngineOfflineTP1(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        if not ray.is_initialized():
            ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
        cls.actor, cls.pg = _create_engine_on_pg(tp_size=1)

    @classmethod
    def tearDownClass(cls):
        _cleanup(cls.actor, cls.pg)
        ray.shutdown()

    def test_offline_generate(self):
        result = ray.get(
            self.actor.generate.remote("The capital of France is", _SAMPLING_PARAMS)
        )
        self.assertIn("text", result)
        self.assertGreater(len(result["text"]), 0)
        print(f"Generated: {result['text'][:200]}")

    def test_batch_generate(self):
        for prompt in _PROMPTS:
            result = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
            self.assertIn("text", result)
            self.assertGreater(len(result["text"]), 0, f"Empty output for: {prompt}")

    def test_deterministic(self):
        prompt = "The meaning of life is"
        r1 = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
        r2 = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
        self.assertEqual(r1["text"], r2["text"])


# ---------------------------------------------------------------------------
# Tests: Offline TP=2
# ---------------------------------------------------------------------------


@unittest.skipUnless(_has_ray, "ray is not installed")
@unittest.skipUnless(_NUM_GPUS >= 2, "requires at least 2 GPUs")
class TestRayEngineOfflineTP2(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        if not ray.is_initialized():
            ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
        cls.actor, cls.pg = _create_engine_on_pg(tp_size=2)

    @classmethod
    def tearDownClass(cls):
        _cleanup(cls.actor, cls.pg)
        ray.shutdown()

    def test_offline_generate_tp2(self):
        result = ray.get(
            self.actor.generate.remote("The capital of France is", _SAMPLING_PARAMS)
        )
        self.assertIn("text", result)
        self.assertGreater(len(result["text"]), 0)
        print(f"Generated (TP=2): {result['text'][:200]}")

    def test_batch_generate_tp2(self):
        for prompt in _PROMPTS:
            result = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
            self.assertIn("text", result)
            self.assertGreater(len(result["text"]), 0, f"Empty output for: {prompt}")


# ---------------------------------------------------------------------------
# Tests: Offline PP=2
# ---------------------------------------------------------------------------


@unittest.skipUnless(_has_ray, "ray is not installed")
@unittest.skipUnless(_NUM_GPUS >= 2, "requires at least 2 GPUs")
class TestRayEngineOfflinePP2(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        if not ray.is_initialized():
            ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
        cls.actor, cls.pg = _create_engine_on_pg(tp_size=1, pp_size=2)

    @classmethod
    def tearDownClass(cls):
        _cleanup(cls.actor, cls.pg)
        ray.shutdown()

    def test_offline_generate_pp2(self):
        result = ray.get(
            self.actor.generate.remote("The capital of France is", _SAMPLING_PARAMS)
        )
        self.assertIn("text", result)
        self.assertGreater(len(result["text"]), 0)
        print(f"Generated (PP=2): {result['text'][:200]}")

    def test_batch_generate_pp2(self):
        for prompt in _PROMPTS:
            result = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
            self.assertIn("text", result)
            self.assertGreater(len(result["text"]), 0, f"Empty output for: {prompt}")


# ---------------------------------------------------------------------------
# Tests: Error paths
# ---------------------------------------------------------------------------


@unittest.skipUnless(_has_ray, "ray is not installed")
@unittest.skipUnless(_NUM_GPUS >= 2, "requires at least 2 GPUs")
class TestRayEngineOfflineDP2(unittest.TestCase):
    """Test Ray engine with dp_size=2, tp_size=1."""

    @classmethod
    def setUpClass(cls):
        if not ray.is_initialized():
            ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
        cls.actor, cls.pg = _create_engine_on_pg(tp_size=1, dp_size=2)

    @classmethod
    def tearDownClass(cls):
        _cleanup(cls.actor, cls.pg)
        ray.shutdown()

    def test_offline_generate_dp2(self):
        result = ray.get(
            self.actor.generate.remote("The capital of France is", _SAMPLING_PARAMS)
        )
        self.assertIn("text", result)
        self.assertGreater(len(result["text"]), 0)
        print(f"Generated (DP=2): {result['text'][:200]}")

    def test_batch_generate_dp2(self):
        for prompt in _PROMPTS:
            result = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
            self.assertIn("text", result)
            self.assertGreater(len(result["text"]), 0, f"Empty output for: {prompt}")


# ---------------------------------------------------------------------------
# Tests: Offline DP Attention (dp=2, tp=2)
# ---------------------------------------------------------------------------


@unittest.skipUnless(_has_ray, "ray is not installed")
@unittest.skipUnless(_NUM_GPUS >= 2, "requires at least 2 GPUs")
class TestRayEngineOfflineDPAttention(unittest.TestCase):
    """Test Ray engine with dp_size=2, tp_size=2, enable_dp_attention=True."""

    @classmethod
    def setUpClass(cls):
        if not ray.is_initialized():
            ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
        cls.actor, cls.pg = _create_engine_on_pg(
            tp_size=2,
            dp_size=2,
            model=_DP_ATTN_MODEL,
            extra_kwargs={
                "enable_dp_attention": True,
                "disable_cuda_graph": True,
                "port": 31500,
            },
        )

    @classmethod
    def tearDownClass(cls):
        _cleanup(cls.actor, cls.pg)
        ray.shutdown()

    def test_offline_generate_dp_attention(self):
        result = ray.get(
            self.actor.generate.remote("The capital of France is", _SAMPLING_PARAMS)
        )
        self.assertIn("text", result)
        self.assertGreater(len(result["text"]), 0)
        print(f"Generated (DP-Attention): {result['text'][:200]}")

    def test_batch_generate_dp_attention(self):
        for prompt in _PROMPTS:
            result = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
            self.assertIn("text", result)
            self.assertGreater(len(result["text"]), 0, f"Empty output for: {prompt}")


# ---------------------------------------------------------------------------
# Tests: Error paths
# ---------------------------------------------------------------------------


@unittest.skipUnless(_has_ray, "ray is not installed")
@unittest.skipUnless(_NUM_GPUS >= 1, "requires at least 1 GPU")
class TestRayEngineErrors(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        if not ray.is_initialized():
            ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)

    @classmethod
    def tearDownClass(cls):
        ray.shutdown()

    def test_missing_placement_group_raises(self):
        """RayEngine without a placement group should raise RuntimeError."""

        @ray.remote(num_gpus=1)
        def _try_create_without_pg():
            from sglang.srt.ray.engine import RayEngine

            try:
                RayEngine(
                    model_path=_MODEL,
                    tp_size=1,
                    use_ray=True,
                )
                return None
            except RuntimeError as e:
                return str(e)

        error_msg = ray.get(_try_create_without_pg.remote(), timeout=120)
        self.assertIsNotNone(
            error_msg, "Expected RuntimeError but RayEngine created OK"
        )
        self.assertIn("placement group", error_msg.lower())


# ---------------------------------------------------------------------------
# Tests: HTTP server
# ---------------------------------------------------------------------------


@unittest.skipUnless(_has_ray, "ray is not installed")
@unittest.skipUnless(_NUM_GPUS >= 1, "requires at least 1 GPU")
class TestRayHTTPServerTP1(unittest.TestCase):
    """Test the Ray HTTP server path (launch_server.py --use-ray).

    Launches the server inside a Ray task on a placement group (mirrors
    examples/anyscale/driver_online.py) and sends HTTP requests to it.
    """

    @classmethod
    def setUpClass(cls):
        import requests as req_lib

        if not ray.is_initialized():
            ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)

        cls.port = 30100
        cls.pg = placement_group(
            [{"CPU": 1, "GPU": 1}],
            strategy="STRICT_PACK",
        )
        ray.get(cls.pg.ready())

        pg_strategy = PlacementGroupSchedulingStrategy(
            placement_group=cls.pg,
            placement_group_bundle_index=0,
        )

        # Resolve the node IP where the server will run
        @ray.remote(num_cpus=0, num_gpus=0)
        def _get_ip():
            return ray.util.get_node_ip_address()

        cls.node_ip = ray.get(_get_ip.options(scheduling_strategy=pg_strategy).remote())
        cls.base_url = f"http://{cls.node_ip}:{cls.port}"

        # Launch server as a Ray task (blocks until server exits)
        @ray.remote
        def _launch(**kwargs):
            from sglang.srt.ray.http_server import launch_server
            from sglang.srt.server_args import ServerArgs

            launch_server(ServerArgs(**kwargs))

        cls.server_ref = _launch.options(
            num_cpus=1,
            num_gpus=0,
            scheduling_strategy=pg_strategy,
        ).remote(
            model_path=_MODEL,
            tp_size=1,
            port=cls.port,
            host="0.0.0.0",
            use_ray=True,
        )

        # Wait for health check
        t0 = time.time()
        timeout = 600
        healthy = False
        while time.time() - t0 < timeout:
            ready, _ = ray.wait([cls.server_ref], timeout=0)
            if ready:
                try:
                    ray.get(cls.server_ref)
                except Exception as e:
                    raise RuntimeError(f"Server task crashed: {e}") from e
                raise RuntimeError("Server task exited before becoming healthy")
            try:
                if req_lib.get(f"{cls.base_url}/health", timeout=5).status_code == 200:
                    healthy = True
                    break
            except req_lib.exceptions.RequestException:
                pass
            time.sleep(3)

        if not healthy:
            ray.cancel(cls.server_ref, force=True)
            raise RuntimeError(f"Server did not become healthy within {timeout}s")

    @classmethod
    def tearDownClass(cls):
        try:
            ray.cancel(cls.server_ref, force=True)
        except Exception:
            pass
        try:
            ray.util.remove_placement_group(cls.pg)
        except Exception:
            pass
        ray.shutdown()

    def test_health_endpoint(self):
        import requests

        resp = requests.get(f"{self.base_url}/health", timeout=10)
        self.assertEqual(resp.status_code, 200)

    def test_generate_endpoint(self):
        import requests

        resp = requests.post(
            f"{self.base_url}/generate",
            json={
                "text": "The capital of France is",
                "sampling_params": _SAMPLING_PARAMS,
            },
            timeout=60,
        )
        resp.raise_for_status()
        data = resp.json()
        self.assertIn("text", data)
        self.assertGreater(len(data["text"]), 0)
        print(f"HTTP response: {data['text'][:200]}")

    def test_generate_multiple(self):
        import requests

        for prompt in _PROMPTS:
            resp = requests.post(
                f"{self.base_url}/generate",
                json={
                    "text": prompt,
                    "sampling_params": _SAMPLING_PARAMS,
                },
                timeout=60,
            )
            resp.raise_for_status()
            data = resp.json()
            self.assertIn("text", data)
            self.assertGreater(len(data["text"]), 0, f"Empty output for: {prompt}")


# ---------------------------------------------------------------------------
# Tests: Custom placement_group and SGLANG_RAY_BUNDLE_INDICES
# ---------------------------------------------------------------------------


@unittest.skipUnless(_has_ray, "ray is not installed")
@unittest.skipUnless(_NUM_GPUS >= 2, "requires at least 2 GPUs")
class TestRayEnginePlacementGroup(unittest.TestCase):
    """Test RayEngine with custom placement_group and SGLANG_RAY_BUNDLE_INDICES."""

    @classmethod
    def setUpClass(cls):
        if not ray.is_initialized():
            ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)

    @classmethod
    def tearDownClass(cls):
        ray.shutdown()

    def test_custom_pg_dp1_tp2(self):
        """Test custom placement_group with dp_size=1, tp_size=2."""
        from sglang.srt.ray.engine import RayEngine

        pg = placement_group([{"GPU": 1}] * 2, strategy="STRICT_PACK")
        ray.get(pg.ready())

        engine = RayEngine(
            model_path=_MODEL,
            tp_size=2,
            placement_group=pg,
            use_ray=True,
        )

        result = engine.generate("The capital of France is", _SAMPLING_PARAMS)
        self.assertIn("text", result)
        self.assertGreater(len(result["text"]), 0)
        print(f"Generated (dp=1, tp=2, custom PG): {result['text'][:200]}")

        engine.shutdown()
        ray.util.remove_placement_group(pg)

    def test_bundle_indices_dp1_tp2(self):
        """Test SGLANG_RAY_BUNDLE_INDICES with dp_size=1, tp_size=2."""
        from sglang.srt.ray.engine import RayEngine

        os.environ["SGLANG_RAY_BUNDLE_INDICES"] = "0,1"

        pg = placement_group([{"GPU": 1}] * 2, strategy="STRICT_PACK")
        ray.get(pg.ready())

        engine = RayEngine(
            model_path=_MODEL,
            tp_size=2,
            placement_group=pg,
            use_ray=True,
        )

        result = engine.generate("The capital of France is", _SAMPLING_PARAMS)
        self.assertIn("text", result)
        self.assertGreater(len(result["text"]), 0)
        print(f"Generated (dp=1, tp=2, indices=0,1): {result['text'][:200]}")

        engine.shutdown()
        ray.util.remove_placement_group(pg)
        del os.environ["SGLANG_RAY_BUNDLE_INDICES"]

    def test_custom_pg_dp2_tp1(self):
        """Test custom placement_group with dp_size=2, tp_size=1."""
        from sglang.srt.ray.engine import RayEngine

        pg = placement_group([{"GPU": 1}] * 2, strategy="STRICT_PACK")
        ray.get(pg.ready())

        engine = RayEngine(
            model_path=_MODEL,
            tp_size=1,
            dp_size=2,
            placement_group=pg,
            use_ray=True,
        )

        result = engine.generate("The capital of France is", _SAMPLING_PARAMS)
        self.assertIn("text", result)
        self.assertGreater(len(result["text"]), 0)
        print(f"Generated (dp=2, tp=1, custom PG): {result['text'][:200]}")

        engine.shutdown()
        ray.util.remove_placement_group(pg)

    def test_bundle_indices_skip_bundle(self):
        """Test skipping unhealthy GPU by using bundle_indices."""
        from sglang.srt.ray.engine import RayEngine

        os.environ["SGLANG_RAY_BUNDLE_INDICES"] = "1"  # Skip bundle 0

        pg = placement_group([{"GPU": 1}] * 2, strategy="STRICT_PACK")
        ray.get(pg.ready())

        engine = RayEngine(
            model_path=_MODEL,
            tp_size=1,
            placement_group=pg,
            use_ray=True,
        )

        result = engine.generate("The capital of France is", _SAMPLING_PARAMS)
        self.assertIn("text", result)
        self.assertGreater(len(result["text"]), 0)
        print(f"Generated (tp=1, skip bundle 0): {result['text'][:200]}")

        engine.shutdown()
        ray.util.remove_placement_group(pg)
        del os.environ["SGLANG_RAY_BUNDLE_INDICES"]

    def test_custom_pg_dp_attention(self):
        """Test custom placement_group with enable_dp_attention=True."""
        from sglang.srt.ray.engine import RayEngine

        pg = placement_group([{"GPU": 1}] * 2, strategy="STRICT_PACK")
        ray.get(pg.ready())

        engine = RayEngine(
            model_path=_DP_ATTN_MODEL,
            tp_size=2,
            placement_group=pg,
            use_ray=True,
            enable_dp_attention=True,
        )

        result = engine.generate("The capital of France is", _SAMPLING_PARAMS)
        self.assertIn("text", result)
        self.assertGreater(len(result["text"]), 0)
        print(f"Generated (dp attention, tp=2, custom PG): {result['text'][:200]}")

        engine.shutdown()
        ray.util.remove_placement_group(pg)


@unittest.skipUnless(_has_ray, "ray is not installed")
@unittest.skipUnless(_NUM_GPUS >= 1, "requires at least 1 GPU")
class TestRayEnginePlacementGroupErrors(unittest.TestCase):
    """Test error handling for placement_group and bundle indices."""

    @classmethod
    def setUpClass(cls):
        if not ray.is_initialized():
            ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)

    @classmethod
    def tearDownClass(cls):
        ray.shutdown()

    def test_multi_gpu_bundle_raises_error(self):
        """Custom PG with multi-GPU bundles should raise an error."""

        @ray.remote(num_gpus=0)
        def _try_multi_gpu_bundle():
            pg = placement_group([{"GPU": 2}], strategy="STRICT_PACK")
            ray.get(pg.ready())

            from sglang.srt.ray.engine import RayEngine

            try:
                RayEngine(
                    model_path=_MODEL,
                    tp_size=2,
                    placement_group=pg,
                    use_ray=True,
                )
                return None
            except Exception as e:
                return str(e)
            finally:
                ray.util.remove_placement_group(pg)

        error_msg = ray.get(_try_multi_gpu_bundle.remote(), timeout=120)
        self.assertIsNotNone(error_msg)
        self.assertIn("exactly 1 GPU per bundle", error_msg)

    def test_invalid_bundle_index_raises_error(self):
        """SGLANG_RAY_BUNDLE_INDICES with invalid index should raise an error."""

        @ray.remote(num_gpus=0)
        def _try_invalid_bundle_index():
            import os

            os.environ["SGLANG_RAY_BUNDLE_INDICES"] = "0,10"

            pg = placement_group([{"GPU": 1}] * 2, strategy="STRICT_PACK")
            ray.get(pg.ready())

            from sglang.srt.ray.engine import RayEngine

            try:
                RayEngine(
                    model_path=_MODEL,
                    tp_size=2,
                    placement_group=pg,
                    use_ray=True,
                )
                return None
            except Exception as e:
                return str(e)
            finally:
                os.environ.pop("SGLANG_RAY_BUNDLE_INDICES", None)
                ray.util.remove_placement_group(pg)

        error_msg = ray.get(_try_invalid_bundle_index.remote(), timeout=120)
        self.assertIsNotNone(error_msg)
        self.assertIn("out of range", error_msg)


if __name__ == "__main__":
    unittest.main()