diff --git a/tests/v1/tracing/test_tracing.py b/tests/v1/tracing/test_tracing.py index 2b450a6299..1b7b243c9d 100644 --- a/tests/v1/tracing/test_tracing.py +++ b/tests/v1/tracing/test_tracing.py @@ -7,6 +7,8 @@ import time from opentelemetry.sdk.environment_variables import OTEL_EXPORTER_OTLP_TRACES_INSECURE from vllm import LLM, SamplingParams +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.platforms import current_platform from vllm.tracing import SpanAttributes # Import shared fixtures from the tracing conftest @@ -23,6 +25,11 @@ def test_traces( ): with monkeypatch.context() as m: m.setenv(OTEL_EXPORTER_OTLP_TRACES_INSECURE, "true") + if current_platform.is_rocm(): + # The fake OTLP server starts gRPC worker threads before the engine + # core is launched. On ROCm CI, forking while those threads are + # active can segfault in gRPC during engine startup or teardown. + m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") sampling_params = SamplingParams( temperature=0.01, @@ -30,58 +37,70 @@ def test_traces( max_tokens=256, ) model = "facebook/opt-125m" - llm = LLM( - model=model, - otlp_traces_endpoint=FAKE_TRACE_SERVER_ADDRESS, - gpu_memory_utilization=0.3, - disable_log_stats=False, - ) - prompts = ["This is a short prompt"] - outputs = llm.generate(prompts, sampling_params=sampling_params) - print(f"test_traces outputs is : {outputs}") + llm = None + try: + llm = LLM( + model=model, + otlp_traces_endpoint=FAKE_TRACE_SERVER_ADDRESS, + gpu_memory_utilization=0.3, + disable_log_stats=False, + ) + prompts = ["This is a short prompt"] + outputs = llm.generate(prompts, sampling_params=sampling_params) + print(f"test_traces outputs is : {outputs}") - # Wait for the "llm_request" span to be exported. - # The BatchSpanProcessor batches spans and exports them periodically, - # so we need to wait specifically for the llm_request span to appear. - timeout = 15 - deadline = time.time() + timeout - llm_request_spans = [] - while time.time() < deadline: - all_spans = trace_service.get_all_spans() - llm_request_spans = [s for s in all_spans if s["name"] == "llm_request"] - if llm_request_spans: - break - time.sleep(0.5) + # Wait for the "llm_request" span to be exported. + # The BatchSpanProcessor batches spans and exports them periodically, + # so we need to wait specifically for the llm_request span to appear. + timeout = 15 + deadline = time.time() + timeout + llm_request_spans = [] + while time.time() < deadline: + all_spans = trace_service.get_all_spans() + llm_request_spans = [s for s in all_spans if s["name"] == "llm_request"] + if llm_request_spans: + break + time.sleep(0.5) - assert len(llm_request_spans) == 1, ( - f"Expected exactly 1 'llm_request' span, but got {len(llm_request_spans)}. " - f"All span names: {[s['name'] for s in all_spans]}" - ) + assert len(llm_request_spans) == 1, ( + f"Expected exactly 1 'llm_request' span, but got " + f"{len(llm_request_spans)}. " + f"All span names: {[s['name'] for s in all_spans]}" + ) - attributes = llm_request_spans[0]["attributes"] - # assert attributes.get(SpanAttributes.GEN_AI_RESPONSE_MODEL) == model - assert attributes.get(SpanAttributes.GEN_AI_REQUEST_ID) == outputs[0].request_id - assert ( - attributes.get(SpanAttributes.GEN_AI_REQUEST_TEMPERATURE) - == sampling_params.temperature - ) - assert ( - attributes.get(SpanAttributes.GEN_AI_REQUEST_TOP_P) == sampling_params.top_p - ) - assert ( - attributes.get(SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS) - == sampling_params.max_tokens - ) - assert attributes.get(SpanAttributes.GEN_AI_REQUEST_N) == sampling_params.n - assert attributes.get(SpanAttributes.GEN_AI_USAGE_PROMPT_TOKENS) == len( - outputs[0].prompt_token_ids - ) - completion_tokens = sum(len(o.token_ids) for o in outputs[0].outputs) - assert ( - attributes.get(SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS) - == completion_tokens - ) + attributes = llm_request_spans[0]["attributes"] + # assert attributes.get(SpanAttributes.GEN_AI_RESPONSE_MODEL) == model + assert ( + attributes.get(SpanAttributes.GEN_AI_REQUEST_ID) + == outputs[0].request_id + ) + assert ( + attributes.get(SpanAttributes.GEN_AI_REQUEST_TEMPERATURE) + == sampling_params.temperature + ) + assert ( + attributes.get(SpanAttributes.GEN_AI_REQUEST_TOP_P) + == sampling_params.top_p + ) + assert ( + attributes.get(SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS) + == sampling_params.max_tokens + ) + assert attributes.get(SpanAttributes.GEN_AI_REQUEST_N) == sampling_params.n + assert attributes.get(SpanAttributes.GEN_AI_USAGE_PROMPT_TOKENS) == len( + outputs[0].prompt_token_ids + ) + completion_tokens = sum(len(o.token_ids) for o in outputs[0].outputs) + assert ( + attributes.get(SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS) + == completion_tokens + ) - assert attributes.get(SpanAttributes.GEN_AI_LATENCY_TIME_IN_QUEUE) > 0 - assert attributes.get(SpanAttributes.GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN) > 0 - assert attributes.get(SpanAttributes.GEN_AI_LATENCY_E2E) > 0 + assert attributes.get(SpanAttributes.GEN_AI_LATENCY_TIME_IN_QUEUE) > 0 + assert attributes.get(SpanAttributes.GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN) > 0 + assert attributes.get(SpanAttributes.GEN_AI_LATENCY_E2E) > 0 + finally: + if llm is not None: + shutdown_timeout = 60.0 if current_platform.is_rocm() else 5.0 + llm.llm_engine.engine_core.shutdown(timeout=shutdown_timeout) + cleanup_dist_env_and_memory()