import json import sys import pytest from pydantic import ValidationError from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( AlignerPerStepPlan, AlignerPlan, ) from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import ( PositionalSeqId, TokenAlignerPlan, TokenAlignerSeqInfo, TokenAlignerStepAux, TokenLocator, ) from sglang.srt.debug_utils.comparator.aligner.unsharder.types import ( AxisInfo, ConcatParams, UnsharderPlan, ) from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis, TokenLayout from sglang.srt.debug_utils.comparator.output_types import ( ErrorLog, NonTensorComparisonRecord, SkipComparisonRecord, SummaryRecord, TensorComparisonRecord, parse_record_json, ) from sglang.srt.debug_utils.comparator.tensor_comparator.types import ( DiffInfo, TensorInfo, TensorStats, ) from sglang.srt.debug_utils.comparator.utils import Pair, _check_equal_lengths from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="default", nightly=True) class TestCheckEqualLengths: def test_all_equal(self): _check_equal_lengths(a=[1, 2], b=[3, 4]) def test_empty_lists(self): _check_equal_lengths(a=[], b=[]) def test_mismatch_raises(self): with pytest.raises(ValueError, match="Length mismatch"): _check_equal_lengths(a=[1, 2], b=[3]) class TestTokenAlignerStepAux: def test_valid(self): aux = TokenAlignerStepAux( input_ids=[10, 20, 30], positions=[0, 1, 2], seq_lens=[2, 1], seq_ids=[ PositionalSeqId(step=0, seq_index=0), PositionalSeqId(step=0, seq_index=1), ], ) assert len(aux.input_ids) == 3 def test_token_length_mismatch(self): with pytest.raises(ValueError, match="Length mismatch"): TokenAlignerStepAux( input_ids=[10, 20, 30], positions=[0, 1], seq_lens=[2, 1], seq_ids=[ PositionalSeqId(step=0, seq_index=0), PositionalSeqId(step=0, seq_index=1), ], ) def test_seq_length_mismatch(self): with pytest.raises(ValueError, match="Length mismatch"): TokenAlignerStepAux( input_ids=[10, 20, 30], positions=[0, 1, 2], seq_lens=[2, 1], seq_ids=[PositionalSeqId(step=0, seq_index=0)], ) def test_sum_seq_lens_mismatch(self): with pytest.raises(ValueError, match="sum\\(seq_lens\\)"): TokenAlignerStepAux( input_ids=[10, 20, 30], positions=[0, 1, 2], seq_lens=[1, 1], seq_ids=[ PositionalSeqId(step=0, seq_index=0), PositionalSeqId(step=0, seq_index=1), ], ) class TestTokenAlignerSeqInfo: def test_valid(self): info = TokenAlignerSeqInfo( input_ids=[10, 20, 30], positions=[0, 1, 2], locator=TokenLocator(steps=[0, 0, 1], token_index_in_step=[0, 1, 0]), ) assert len(info.input_ids) == 3 def test_length_mismatch(self): with pytest.raises(ValidationError): TokenAlignerSeqInfo( input_ids=[10, 20, 30], positions=[0, 1, 2], locator=TokenLocator(steps=[0, 0], token_index_in_step=[0, 1, 0]), ) def test_positions_not_sequential(self): with pytest.raises(ValidationError, match="positions must be"): TokenAlignerSeqInfo( input_ids=[10, 20, 30], positions=[0, 2, 1], locator=TokenLocator(steps=[0, 0, 1], token_index_in_step=[0, 1, 0]), ) class TestTokenAlignerPlan: def test_valid(self): plan = TokenAlignerPlan( locators=Pair( x=TokenLocator(steps=[0, 0, 1], token_index_in_step=[0, 1, 0]), y=TokenLocator(steps=[0, 1, 1], token_index_in_step=[0, 0, 1]), ), layouts=Pair(x=TokenLayout.T, y=TokenLayout.T), ) assert len(plan.locators.x.steps) == 3 def test_length_mismatch(self): with pytest.raises(ValidationError, match="Length mismatch"): TokenAlignerPlan( locators=Pair( x=TokenLocator(steps=[0, 0], token_index_in_step=[0, 1]), y=TokenLocator(steps=[0, 1, 1], token_index_in_step=[0, 0, 1]), ), layouts=Pair(x=TokenLayout.T, y=TokenLayout.T), ) class TestSummaryRecord: def test_valid(self): record = SummaryRecord(total=10, passed=7, failed=2, skipped=1) assert record.total == 10 def test_total_mismatch(self): with pytest.raises(ValidationError, match="total=10"): SummaryRecord(total=10, passed=5, failed=2, skipped=1) class TestAxisInfo: def test_valid(self): info = AxisInfo(axis_rank=0, axis_size=4) assert info.axis_rank == 0 def test_axis_size_zero(self): with pytest.raises(ValidationError, match="axis_size must be > 0"): AxisInfo(axis_rank=0, axis_size=0) def test_axis_size_negative(self): with pytest.raises(ValidationError, match="axis_size must be > 0"): AxisInfo(axis_rank=0, axis_size=-1) def test_axis_rank_negative(self): with pytest.raises(ValidationError, match="axis_rank must be in"): AxisInfo(axis_rank=-1, axis_size=4) def test_axis_rank_too_large(self): with pytest.raises(ValidationError, match="axis_rank must be in"): AxisInfo(axis_rank=4, axis_size=4) def test_axis_rank_equals_size_minus_one(self): info = AxisInfo(axis_rank=3, axis_size=4) assert info.axis_rank == 3 def _make_tensor_info() -> TensorInfo: return TensorInfo( shape=[4, 4], dtype="float32", stats=TensorStats(mean=0.0, abs_mean=0.8, std=1.0, min=-2.0, max=2.0), ) def _make_diff_info(*, passed: bool) -> DiffInfo: return DiffInfo( rel_diff=0.001, max_abs_diff=0.01, mean_abs_diff=0.005, max_diff_coord=[0, 0], baseline_at_max=1.0, target_at_max=1.01, diff_threshold=1e-3, passed=passed, ) def _make_comparison_record( *, diff: DiffInfo | None, errors: list | None = None, ) -> TensorComparisonRecord: ti: TensorInfo = _make_tensor_info() return TensorComparisonRecord( name="t", baseline=ti, target=ti, unified_shape=[4, 4], shape_mismatch=False, diff=diff, errors=errors or [], ) class TestOutputRecordCategories: def test_skip_record_with_errors_is_failed(self) -> None: record = SkipComparisonRecord( name="t", reason="test", errors=[ErrorLog(category="c", message="m")], ) assert record.category == "failed" def test_skip_record_no_warnings_is_skipped(self) -> None: record = SkipComparisonRecord(name="t", reason="test") assert record.category == "skipped" def test_comparison_record_diff_none_is_failed(self) -> None: record: TensorComparisonRecord = _make_comparison_record(diff=None) assert record.category == "failed" def test_comparison_record_passed_with_errors_is_failed(self) -> None: record: TensorComparisonRecord = _make_comparison_record( diff=_make_diff_info(passed=True), errors=[ErrorLog(category="c", message="m")], ) assert record.category == "failed" def test_comparison_record_passed_no_warnings_is_passed(self) -> None: record: TensorComparisonRecord = _make_comparison_record( diff=_make_diff_info(passed=True), ) assert record.category == "passed" def test_non_tensor_record_equal_is_passed(self) -> None: record = NonTensorComparisonRecord( name="sm_scale", baseline_value="0.125", target_value="0.125", baseline_type="float", target_type="float", values_equal=True, ) assert record.category == "passed" def test_non_tensor_record_different_is_failed(self) -> None: record = NonTensorComparisonRecord( name="sm_scale", baseline_value="0.125", target_value="0.25", baseline_type="float", target_type="float", values_equal=False, ) assert record.category == "failed" def test_non_tensor_record_with_errors_is_failed(self) -> None: record = NonTensorComparisonRecord( name="sm_scale", baseline_value="0.125", target_value="0.125", baseline_type="float", target_type="float", values_equal=True, errors=[ErrorLog(category="c", message="m")], ) assert record.category == "failed" def test_non_tensor_record_json_roundtrip(self) -> None: record = NonTensorComparisonRecord( name="sm_scale", baseline_value="0.125", target_value="0.25", baseline_type="float", target_type="float", values_equal=False, ) json_str: str = record.model_dump_json() roundtripped = parse_record_json(json_str) assert isinstance(roundtripped, NonTensorComparisonRecord) assert roundtripped.name == "sm_scale" assert roundtripped.values_equal is False assert roundtripped.baseline_value == "0.125" assert roundtripped.target_value == "0.25" def test_non_tensor_record_text_format_equal(self) -> None: record = NonTensorComparisonRecord( name="sm_scale", baseline_value="0.125", target_value="0.125", baseline_type="float", target_type="float", values_equal=True, ) text: str = record.to_text() assert "sm_scale" in text assert "[equal]" in text def test_non_tensor_record_text_format_different(self) -> None: record = NonTensorComparisonRecord( name="sm_scale", baseline_value="0.125", target_value="0.25", baseline_type="float", target_type="float", values_equal=False, ) text: str = record.to_text() assert "baseline" in text assert "target" in text def _make_aligner_plan() -> AlignerPlan: unsharder = UnsharderPlan( axis=ParallelAxis.TP, params=ConcatParams(dim_name="h"), groups=[[0, 1]], ) return AlignerPlan( per_step_plans=Pair( x=[ AlignerPerStepPlan( step=0, input_object_indices=[0, 1], sub_plans=[unsharder] ) ], y=[ AlignerPerStepPlan( step=0, input_object_indices=[0, 1], sub_plans=[unsharder] ) ], ), ) class TestAlignerPlanInTensorComparisonRecord: def test_comparison_record_with_aligner_plan(self) -> None: plan: AlignerPlan = _make_aligner_plan() record: TensorComparisonRecord = _make_comparison_record( diff=_make_diff_info(passed=True), ) record_with_plan = record.model_copy(update={"aligner_plan": plan}) assert record_with_plan.aligner_plan is not None assert record_with_plan.aligner_plan.per_step_plans.x[0].step == 0 def test_aligner_plan_json_roundtrip(self) -> None: plan: AlignerPlan = _make_aligner_plan() record: TensorComparisonRecord = _make_comparison_record( diff=_make_diff_info(passed=True), ) record_with_plan = record.model_copy(update={"aligner_plan": plan}) json_str: str = record_with_plan.model_dump_json() parsed = json.loads(json_str) assert "aligner_plan" in parsed assert ( parsed["aligner_plan"]["per_step_plans"]["x"][0]["sub_plans"][0]["type"] == "unsharder" ) roundtripped: TensorComparisonRecord = parse_record_json(json_str) assert roundtripped.aligner_plan is not None assert ( roundtripped.aligner_plan.per_step_plans.x[0].sub_plans[0].type == "unsharder" ) def test_comparison_record_without_aligner_plan(self) -> None: record: TensorComparisonRecord = _make_comparison_record( diff=_make_diff_info(passed=True), ) json_str: str = record.model_dump_json() roundtripped: TensorComparisonRecord = parse_record_json(json_str) assert roundtripped.aligner_plan is None def test_aligner_plan_text_format(self) -> None: plan: AlignerPlan = _make_aligner_plan() record: TensorComparisonRecord = _make_comparison_record( diff=_make_diff_info(passed=True), ) record_with_plan = record.model_copy(update={"aligner_plan": plan}) text: str = record_with_plan.to_text() assert "Aligner Plan:" in text assert "unsharder" in text if __name__ == "__main__": sys.exit(pytest.main([__file__]))