/test/registered/unit/observability/test_trace.py

"""Unit tests for trace.py — no server, no model loading."""

import os

from sglang.test.ci.ci_register import register_cpu_ci

register_cpu_ci(est_time=6, suite="base-a-test-cpu")

import threading
import unittest
from unittest.mock import patch

import sglang.srt.observability.trace as mod
from sglang.srt.observability.trace import (
    SpanAttributes,
    TraceEvent,
    TraceNullContext,
    TraceReqContext,
    TraceSliceContext,
    TraceThreadInfo,
    extract_trace_headers,
    get_global_trace_level,
    get_global_tracing_enabled,
    process_tracing_init,
    set_global_trace_level,
    trace_set_thread_info,
)

try:
    from opentelemetry import trace as otel_trace
    from opentelemetry.sdk.trace import TracerProvider

    from sglang.srt.observability.trace import get_otlp_span_exporter

    _has_otel = True
except ImportError:
    _has_otel = False

# Access the private module-level function (avoid name mangling inside classes).
_get_host_id = getattr(mod, "_get_host_id")


class TestTraceFunctions(unittest.TestCase):
    def test_extract_trace_headers(self):
        headers = {"traceparent": "abc", "tracestate": "xyz", "other": "skip"}
        result = extract_trace_headers(headers)
        self.assertEqual(result, {"traceparent": "abc", "tracestate": "xyz"})

    def test_extract_trace_headers_missing(self):
        self.assertEqual(extract_trace_headers({}), {})

    def test_set_global_trace_level(self):
        from sglang.srt.runtime_context import get_resources

        orig = get_resources().trace_level
        try:
            set_global_trace_level(5)
            self.assertEqual(get_global_trace_level(), 5)
        finally:
            get_resources().trace_level = orig

    def test_global_trace_level_env_var(self):
        # The level lives on ctx.resources and is seeded lazily from the env
        # on first read after a reset (no module reload involved).
        from sglang.srt.runtime_context import get_resources

        orig = get_resources().trace_level
        try:
            with patch.dict(os.environ, {"SGLANG_TRACE_LEVEL": "2"}):
                get_resources().trace_level = None
                self.assertEqual(get_global_trace_level(), 2)
            get_resources().trace_level = None  # SGLANG_TRACE_LEVEL unset → 3
            self.assertEqual(get_global_trace_level(), 3)
        finally:
            get_resources().trace_level = orig

    def test_get_global_tracing_enabled(self):
        self.assertEqual(get_global_tracing_enabled(), mod.opentelemetry_initialized)

    def test_get_cur_time_ns(self):
        ts = mod.get_cur_time_ns()
        self.assertIsInstance(ts, int)
        self.assertGreater(ts, 0)


class TestTraceNullContext(unittest.TestCase):
    def test_null_object_pattern(self):
        ctx = TraceNullContext()
        self.assertFalse(ctx.tracing_enable)
        # Any attribute access returns self
        self.assertIs(ctx.some_method, ctx)
        # Callable returns self
        self.assertIs(ctx("arg1", key="val"), ctx)
        # Chaining works
        self.assertIs(ctx.foo.bar.baz(1, 2, 3), ctx)


class TestSpanAttributes(unittest.TestCase):
    def test_constants_exist(self):
        self.assertEqual(SpanAttributes.GEN_AI_LATENCY_E2E, "gen_ai.latency.e2e")
        self.assertIsInstance(SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS, str)


# __get_host_id
class TestGetHostId(unittest.TestCase):
    def test_from_machine_id_file(self):
        with (
            patch("os.path.exists", return_value=True),
            patch(
                "builtins.open",
                unittest.mock.mock_open(read_data="abc123\n"),
            ),
        ):
            self.assertEqual(_get_host_id(), "abc123")

    def test_from_machine_id_file_error(self):
        """Falls back to MAC address when file read fails."""
        with (
            patch("os.path.exists", return_value=True),
            patch("builtins.open", side_effect=IOError("read error")),
        ):
            result = _get_host_id()
            self.assertIsInstance(result, str)
            self.assertGreater(len(result), 0)

    def test_from_mac_address(self):
        with (
            patch("os.path.exists", return_value=False),
            patch("uuid.getnode", return_value=0x112233445566),
        ):
            result = _get_host_id()
            self.assertIsInstance(result, str)
            self.assertGreater(len(result), 0)

    def test_unknown_fallback(self):
        with (
            patch("os.path.exists", return_value=False),
            patch("uuid.getnode", return_value=0),
        ):
            self.assertEqual(_get_host_id(), "unknown")


@unittest.skipUnless(_has_otel, "opentelemetry not installed")
class TestGetOtlpSpanExporter(unittest.TestCase):
    def test_grpc_default(self):

        with patch.dict(os.environ, {}, clear=False):
            os.environ.pop("OTEL_EXPORTER_OTLP_TRACES_PROTOCOL", None)
            exporter = get_otlp_span_exporter("localhost:4317")
        self.assertIsNotNone(exporter)

    def test_http_protobuf(self):

        with patch.dict(
            os.environ, {"OTEL_EXPORTER_OTLP_TRACES_PROTOCOL": "http/protobuf"}
        ):
            exporter = get_otlp_span_exporter("http://localhost:4318/v1/traces")
        self.assertIsNotNone(exporter)

    def test_invalid_protocol(self):

        with patch.dict(os.environ, {"OTEL_EXPORTER_OTLP_TRACES_PROTOCOL": "invalid"}):
            with self.assertRaises(ValueError):
                get_otlp_span_exporter("localhost:4317")


class TestProcessTracingInit(unittest.TestCase):
    def test_raises_without_otel(self):

        orig = mod.opentelemetry_imported
        mod.opentelemetry_imported = False
        try:
            with self.assertRaises(RuntimeError):
                process_tracing_init("localhost:4317", "test")
        finally:
            mod.opentelemetry_imported = orig


class TestTraceReqContextDisabled(unittest.TestCase):
    def setUp(self):
        self.orig = mod.opentelemetry_initialized
        mod.opentelemetry_initialized = False

    def tearDown(self):
        mod.opentelemetry_initialized = self.orig

    def test_init_disabled(self):
        ctx = TraceReqContext(rid="req-1")
        self.assertFalse(ctx.tracing_enable)
        self.assertFalse(ctx.is_tracing_enabled())

    def test_getstate_disabled(self):
        ctx = TraceReqContext(rid="req-1")
        state = ctx.__getstate__()
        self.assertEqual(state, {"tracing_enable": False})

    def test_setstate_disabled(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.__setstate__({"tracing_enable": True, "is_copy": False})
        # opentelemetry_initialized is False → tracing forced off
        self.assertFalse(ctx.tracing_enable)

        # Should not register anything


@unittest.skipUnless(_has_otel, "opentelemetry not installed")
class TestTraceReqContextEnabled(unittest.TestCase):
    def setUp(self):

        self.orig_initialized = mod.opentelemetry_initialized
        self.orig_tracer = mod.tracer
        self.orig_threads = mod.threads_info.copy()
        from sglang.srt.runtime_context import get_resources

        self.orig_level = get_resources().trace_level

        # Reset OTel global TracerProvider so set_tracer_provider works each test
        otel_trace._TRACER_PROVIDER_SET_ONCE._done = False
        otel_trace._TRACER_PROVIDER = None

        self.provider = TracerProvider()
        otel_trace.set_tracer_provider(self.provider)
        mod.opentelemetry_initialized = True
        mod.tracer = otel_trace.get_tracer("test")
        set_global_trace_level(3)

    def tearDown(self):
        mod.opentelemetry_initialized = self.orig_initialized
        mod.tracer = self.orig_tracer
        mod.threads_info.clear()
        mod.threads_info.update(self.orig_threads)
        from sglang.srt.runtime_context import get_resources

        get_resources().trace_level = self.orig_level

    def test_trace_set_thread_info(self):
        trace_set_thread_info("scheduler", tp_rank=0, dp_rank=0)

        pid = threading.get_native_id()
        self.assertIn(pid, mod.threads_info)
        self.assertEqual(mod.threads_info[pid].thread_label, "scheduler")

        # Second call for same thread is a no-op
        trace_set_thread_info("different_label")
        self.assertEqual(mod.threads_info[pid].thread_label, "scheduler")

    def test_module_filtering(self):
        """global_trace_modules gates only explicitly named modules."""
        orig_modules = mod.global_trace_modules
        mod.global_trace_modules = ["request"]
        try:
            # Default empty module_name is never filtered
            ctx = TraceReqContext(rid="req-1")
            self.assertTrue(ctx.tracing_enable)
            # Listed module is traced
            ctx = TraceReqContext(rid="req-1", module_name="request")
            self.assertTrue(ctx.tracing_enable)
            # Unlisted module is filtered out
            ctx = TraceReqContext(rid="req-1", module_name="mooncake")
            self.assertFalse(ctx.tracing_enable)
        finally:
            mod.global_trace_modules = orig_modules

    def test_full_lifecycle(self):
        """Start → slice_start → slice_end → finish."""
        ctx = TraceReqContext(rid="req-1", role="unified")
        self.assertTrue(ctx.tracing_enable)

        ctx.trace_req_start(ts=1000)
        self.assertEqual(ctx.start_time_ns, 1000)
        self.assertIsNotNone(ctx.root_span)
        self.assertIsNotNone(ctx.thread_context)

        ctx.trace_slice_start("prefill", level=1, ts=2000)
        self.assertEqual(len(ctx.thread_context.cur_slice_stack), 1)

        ctx.trace_slice_end("prefill", level=1, ts=3000)
        self.assertEqual(len(ctx.thread_context.cur_slice_stack), 0)
        self.assertIsNotNone(ctx.last_span_context)

        ctx.trace_req_finish(ts=4000, attrs={"tokens": 42})
        self.assertIsNone(ctx.root_span)

    def test_trace_req_start_with_bootstrap_room(self):
        ctx = TraceReqContext(rid="req-1", bootstrap_room=0xFF, role="prefill")
        ctx.trace_req_start(ts=1000)
        self.assertIsNotNone(ctx.root_span)
        ctx.trace_req_finish(ts=2000)

    def test_trace_slice_combined(self):
        """trace_slice() creates and ends a span in one call."""
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)

        s = TraceSliceContext(
            "decode",
            2000,
            end_time_ns=3000,
            level=1,
            attrs={"key": "val"},
            events=[TraceEvent("evt", 2500, {"e": 1})],
        )
        ctx.trace_slice(s)
        self.assertIsNotNone(ctx.last_span_context)
        ctx.trace_req_finish(ts=4000)

    def test_trace_slice_with_events_cache(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)

        # Add events to cache
        ctx.trace_event("schedule", level=1, ts=1500, attrs={"bid": "x"})
        self.assertEqual(len(ctx.events_cache), 1)

        # trace_slice_start + trace_slice_end flushes matching events
        ctx.trace_slice_start("prefill", level=1, ts=1200)
        ctx.trace_slice_end("prefill", level=1, ts=2000)
        self.assertEqual(len(ctx.events_cache), 0)

        ctx.trace_req_finish(ts=3000)

    def test_trace_slice_combined_with_events_cache(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)

        ctx.trace_event("evt", level=1, ts=1500)
        s = TraceSliceContext("decode", 1200, end_time_ns=2000, level=1)
        ctx.trace_slice(s)
        self.assertEqual(len(ctx.events_cache), 0)
        ctx.trace_req_finish(ts=3000)

    def test_trace_event_no_attrs(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        ctx.trace_event("evt", level=1, ts=1500, attrs=None)
        self.assertEqual(ctx.events_cache[0].attrs, {})
        ctx.trace_req_finish(ts=2000)

    def test_trace_slice_end_empty_stack(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        # End without start → warning, no crash
        ctx.trace_slice_end("missing", level=1, ts=2000)
        ctx.trace_req_finish(ts=3000)

    def test_trace_slice_end_name_mismatch(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        ctx.trace_slice_start("prefill", level=1, ts=1500)
        # Mismatched name → warning, slice popped
        ctx.trace_slice_end("wrong_name", level=1, ts=2000)
        self.assertEqual(len(ctx.thread_context.cur_slice_stack), 0)
        ctx.trace_req_finish(ts=3000)

    def test_trace_slice_end_with_attrs_and_thread_finish(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        ctx.trace_slice_start("dispatch", level=2, ts=1500)
        ctx.trace_slice_end(
            "dispatch",
            level=2,
            ts=2000,
            attrs={"key": "val"},
            thread_finish_flag=True,
        )
        # thread_finish_flag triggers abort → thread_context is None
        self.assertIsNone(ctx.thread_context)

    def test_trace_slice_combined_with_thread_finish(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        s = TraceSliceContext("dispatch", 1500, end_time_ns=2000, level=2)
        ctx.trace_slice(s, thread_finish_flag=True)
        self.assertIsNone(ctx.thread_context)

    def test_nested_slices(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        ctx.trace_slice_start("outer", level=1, ts=1500)
        ctx.trace_slice_start("inner", level=2, ts=1600)
        self.assertEqual(len(ctx.thread_context.cur_slice_stack), 2)
        ctx.trace_slice_end("inner", level=2, ts=1800)
        self.assertEqual(len(ctx.thread_context.cur_slice_stack), 1)
        ctx.trace_slice_end("outer", level=1, ts=2000)
        ctx.trace_req_finish(ts=3000)

    def test_nested_slice_with_last_span_context(self):
        """trace_slice uses last_span_context when slice stack is empty."""
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)

        # First slice sets last_span_context
        ctx.trace_slice_start("s1", level=1, ts=1500)
        ctx.trace_slice_end("s1", level=1, ts=2000)
        self.assertIsNotNone(ctx.last_span_context)

        # Second slice uses last_span_context as link
        ctx.trace_slice_start("s2", level=1, ts=2500)
        ctx.trace_slice_end("s2", level=1, ts=3000)

        # trace_slice also uses last_span_context
        s = TraceSliceContext("s3", 3500, end_time_ns=4000, level=1)
        ctx.trace_slice(s)

        ctx.trace_req_finish(ts=5000)

    def test_abort_with_unclosed_slices(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        ctx.trace_slice_start("s1", level=1, ts=1500)
        ctx.trace_slice_start("s2", level=2, ts=1600)
        ctx.abort(ts=2000)
        self.assertIsNone(ctx.thread_context)

    def test_abort_with_events_cache(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        ctx.trace_event("evt", level=1, ts=1500)
        ctx.abort(ts=2000)
        self.assertEqual(len(ctx.events_cache), 0)

    def test_abort_with_abort_info_dict(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        ctx.abort(ts=2000, abort_info={"reason": "cancelled"})
        self.assertIsNone(ctx.thread_context)

    def test_abort_with_base_finish_reason(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        from sglang.srt.managers.schedule_batch import FINISH_LENGTH

        abort_obj = FINISH_LENGTH(length=10)
        ctx.abort(ts=2000, abort_info=abort_obj)
        self.assertIsNone(ctx.thread_context)

    def test_check_fast_return_by_level(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        ctx.trace_level = 1  # instance-level, set at init from global
        # Level 2 > trace_level 1 → fast return
        ctx.trace_slice_start("s", level=2, ts=1500)
        self.assertEqual(len(ctx.thread_context.cur_slice_stack), 0)
        ctx.trace_level = 3
        ctx.trace_req_finish(ts=2000)

    def test_rebuild_thread_context(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        old_tc = ctx.thread_context
        ctx.rebuild_thread_context(ts=1500)
        self.assertIsNot(ctx.thread_context, old_tc)
        ctx.trace_req_finish(ts=2000)

    def test_getstate_enabled(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        state = ctx.__getstate__()
        self.assertTrue(state["tracing_enable"])
        self.assertEqual(state["rid"], "req-1")
        self.assertIn("root_span_context", state)
        ctx.trace_req_finish(ts=2000)

    def test_getstate_no_root_context(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        ctx.root_span_context = None
        state = ctx.__getstate__()
        self.assertFalse(state["tracing_enable"])
        ctx.root_span_context = True  # prevent __del__ issues
        ctx.trace_req_finish(ts=2000)

    def test_getstate_with_slice_stack(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        ctx.trace_slice_start("s1", level=1, ts=1500)
        state = ctx.__getstate__()
        self.assertIn("last_span_context", state)
        ctx.trace_req_finish(ts=2000)

    def test_setstate_enabled(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        state = ctx.__getstate__()
        ctx.trace_req_finish(ts=2000)

        ctx2 = TraceReqContext(rid="req-2")
        ctx2.__setstate__(state)
        self.assertTrue(ctx2.tracing_enable)
        self.assertTrue(ctx2.is_copy)
        self.assertIsNotNone(ctx2.root_span_context)

    def test_thread_context_with_tp_rank(self):
        """Covers tp_rank branch in __create_thread_context."""

        pid = threading.get_native_id()
        mod.threads_info[pid] = TraceThreadInfo(
            "host", pid, "sched", tp_rank=0, dp_rank=0, pp_rank=0
        )
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        self.assertIsNotNone(ctx.thread_context)
        ctx.trace_req_finish(ts=2000)

    def test_setstate_with_last_span_context(self):
        """Covers __setstate__ path where last_span_context is truthy."""
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        ctx.trace_slice_start("s1", level=1, ts=1500)
        ctx.trace_slice_end("s1", level=1, ts=2000)
        state = ctx.__getstate__()
        ctx.trace_req_finish(ts=3000)

        self.assertIsNotNone(state.get("last_span_context"))
        ctx2 = TraceReqContext(rid="req-2")
        ctx2.__setstate__(state)
        self.assertIsNotNone(ctx2.last_span_context)

    def test_events_cache_partial_match(self):
        """Events outside the slice time range stay in cache."""
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)

        ctx.trace_event("early", level=1, ts=500)
        ctx.trace_event("inside", level=1, ts=1500)
        ctx.trace_event("late", level=1, ts=5000)

        ctx.trace_slice_start("s", level=1, ts=1200)
        ctx.trace_slice_end("s", level=1, ts=2000)
        # "early" (500 < 1200) and "late" (5000 >= 2000) stay in cache
        self.assertEqual(len(ctx.events_cache), 2)
        ctx.trace_req_finish(ts=6000)

    def test_trace_slice_combined_events_partial_match(self):
        """Events outside slice range stay in cache for trace_slice method."""
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)

        ctx.trace_event("early", level=1, ts=500)
        ctx.trace_event("inside", level=1, ts=1500)

        s = TraceSliceContext("s", 1200, end_time_ns=2000, level=1)
        ctx.trace_slice(s)
        self.assertEqual(len(ctx.events_cache), 1)  # "early" stays
        ctx.trace_req_finish(ts=3000)

    def test_trace_slice_nested_parent(self):
        """trace_slice with parent from slice stack (not thread_span)."""
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)

        ctx.trace_slice_start("outer", level=1, ts=1500)
        s = TraceSliceContext("inner", 1600, end_time_ns=1800, level=2)
        ctx.trace_slice(s)
        ctx.trace_slice_end("outer", level=1, ts=2000)
        ctx.trace_req_finish(ts=3000)

    def test_del_triggers_abort(self):
        ctx = TraceReqContext(rid="req-1")
        ctx.trace_req_start(ts=1000)
        # __del__ calls abort
        ctx.__del__()
        self.assertIsNone(ctx.thread_context)


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