Support method decorator for tagging and add minimalistic comparator in dumper (#19559)
This commit is contained in:
@@ -332,7 +332,6 @@ class TestEntrypointGroupingRaw:
|
||||
enable=True,
|
||||
dir=str(side_dir),
|
||||
exp_name=_FIXED_EXP_NAME,
|
||||
enable_http_server=False,
|
||||
)
|
||||
)
|
||||
dumper.__dict__["_static_meta"] = {"world_rank": 0, "world_size": 1}
|
||||
@@ -1110,7 +1109,6 @@ class TestEntrypointAlignment:
|
||||
enable=True,
|
||||
dir=str(d),
|
||||
exp_name=_FIXED_EXP_NAME,
|
||||
enable_http_server=False,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -1167,7 +1165,6 @@ class TestEntrypointAlignment:
|
||||
enable=True,
|
||||
dir=str(sglang_dir),
|
||||
exp_name=_FIXED_EXP_NAME,
|
||||
enable_http_server=False,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -1226,7 +1223,6 @@ class TestEntrypointAlignment:
|
||||
enable=True,
|
||||
dir=str(megatron_dir),
|
||||
exp_name=_FIXED_EXP_NAME,
|
||||
enable_http_server=False,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -1468,9 +1464,7 @@ def _assert_single_comparison_passed(records: list[AnyRecord]) -> ComparisonReco
|
||||
|
||||
|
||||
def _make_dumper(directory: Path) -> _Dumper:
|
||||
return _Dumper(
|
||||
config=DumperConfig(enable=True, dir=str(directory), enable_http_server=False)
|
||||
)
|
||||
return _Dumper(config=DumperConfig(enable=True, dir=str(directory)))
|
||||
|
||||
|
||||
def _create_dumps(
|
||||
@@ -1528,7 +1522,6 @@ def _create_non_tensor_rank_dump(
|
||||
enable=True,
|
||||
dir=str(directory),
|
||||
exp_name=_FIXED_EXP_NAME,
|
||||
enable_http_server=False,
|
||||
)
|
||||
)
|
||||
dumper.__dict__["_static_meta"] = {"world_rank": rank, "world_size": 1}
|
||||
@@ -1611,7 +1604,6 @@ def _create_rank_dump(
|
||||
enable=True,
|
||||
dir=str(directory),
|
||||
exp_name=_FIXED_EXP_NAME,
|
||||
enable_http_server=False,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -1951,7 +1943,6 @@ class TestEntrypointThdCpZigzag:
|
||||
enable=True,
|
||||
dir=str(sglang_dir),
|
||||
exp_name=_FIXED_EXP_NAME,
|
||||
enable_http_server=False,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -403,79 +403,5 @@ class TestAlignerPlanInComparisonRecord:
|
||||
assert "unsharder" 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 TestAlignerPlanInComparisonRecord:
|
||||
def test_comparison_record_with_aligner_plan(self) -> None:
|
||||
plan: AlignerPlan = _make_aligner_plan()
|
||||
record: ComparisonRecord = _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: ComparisonRecord = _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: ComparisonRecord = 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: ComparisonRecord = _make_comparison_record(
|
||||
diff=_make_diff_info(passed=True),
|
||||
)
|
||||
json_str: str = record.model_dump_json()
|
||||
roundtripped: ComparisonRecord = 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: ComparisonRecord = _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__]))
|
||||
|
||||
Reference in New Issue
Block a user