# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import unittest from unittest.mock import MagicMock, patch import paddle from paddle.distributed.flex_checkpoint.dcp.key_validation import ( AOAMappingEntry, AOASliceMapping, KeyValidationResult, ShapeMismatchInfo, _append_src_lines, _build_aoa_mappings, _classify_mappings, _describe_ops, _emit, _format_key_list, _format_pattern_groups, _format_slice_range, _get_signature, _group_by_signature, _group_keys_adaptive, _print_aoa_report, _print_standard_report, _slice_covers_full, _try_fold_src_keys, validate_and_report_keys_aoa, validate_and_report_keys_standard, ) from paddle.distributed.flex_checkpoint.dcp.metadata import ( LocalTensorIndex, LocalTensorMetadata, Metadata, ) class TestSliceCoversFull(unittest.TestCase): def test_covers_full(self): sl = (slice(0, 4), slice(0, 8)) self.assertTrue(_slice_covers_full(sl, (4, 8))) def test_not_covers_partial(self): sl = (slice(0, 2), slice(0, 8)) self.assertFalse(_slice_covers_full(sl, (4, 8))) def test_not_covers_mismatched_dims(self): sl = (slice(0, 4),) self.assertFalse(_slice_covers_full(sl, (4, 8))) def test_non_zero_start(self): sl = (slice(1, 4), slice(0, 8)) self.assertFalse(_slice_covers_full(sl, (4, 8))) class TestFormatSliceRange(unittest.TestCase): def test_basic(self): src_sl = (slice(0, 4), slice(0, 8)) dst_sl = (slice(0, 4), slice(0, 8)) result = _format_slice_range(src_sl, dst_sl) self.assertIn("0:4", result) self.assertIn("0:8", result) self.assertIn("->", result) def test_partial_slices(self): src_sl = (slice(2, 6),) dst_sl = (slice(0, 4),) result = _format_slice_range(src_sl, dst_sl) self.assertIn("2:6", result) self.assertIn("0:4", result) class TestTryFoldSrcKeys(unittest.TestCase): def test_fold_consecutive(self): keys = [f"model.experts.{i}.weight" for i in range(8)] result = _try_fold_src_keys(keys) self.assertIsNotNone(result) self.assertIn("{0..7}", result) def test_no_fold_different_patterns(self): keys = ["model.a.weight", "model.b.weight"] result = _try_fold_src_keys(keys) self.assertIsNone(result) def test_no_fold_multiple_varying_positions(self): keys = ["layer.0.expert.0.w", "layer.1.expert.1.w"] result = _try_fold_src_keys(keys) self.assertIsNone(result) def test_single_key(self): result = _try_fold_src_keys(["a.0.b"]) self.assertIsNone(result) def test_empty(self): result = _try_fold_src_keys([]) self.assertIsNone(result) class TestDescribeOps(unittest.TestCase): def test_single_with_permute(self): entry = AOAMappingEntry( dst_key="a.weight", dst_global_shape=(4, 8), slice_mappings=[ AOASliceMapping( "b.weight", (slice(0, 4), slice(0, 8)), (slice(0, 4), slice(0, 8)), ["[1, 0]"], ) ], ) result = _describe_ops(entry) self.assertIn("permute([1, 0])", result) def test_concat_with_cast(self): entry = AOAMappingEntry( dst_key="a.weight", dst_global_shape=(8, 4), slice_mappings=[ AOASliceMapping( "b.weight", (slice(0, 4), slice(0, 4)), (slice(0, 4), slice(0, 4)), ["bfloat16"], ), AOASliceMapping( "c.weight", (slice(0, 4), slice(0, 4)), (slice(4, 8), slice(0, 4)), ["bfloat16"], ), ], ) result = _describe_ops(entry) self.assertIn("concat", result) self.assertIn("cast(bfloat16)", result) def test_no_ops(self): entry = AOAMappingEntry( dst_key="a.weight", dst_global_shape=(4, 8), slice_mappings=[ AOASliceMapping( "b.weight", (slice(0, 4), slice(0, 8)), (slice(0, 4), slice(0, 8)), None, ) ], ) result = _describe_ops(entry) self.assertEqual(result, "") def test_empty_slice_mappings(self): entry = AOAMappingEntry( dst_key="a.weight", dst_global_shape=(4,), slice_mappings=[] ) result = _describe_ops(entry) self.assertEqual(result, "") class TestClassifyMappings(unittest.TestCase): def _make_entry(self, dst_key, src_key, pp=None, multi_src=False): if multi_src: sms = [ AOASliceMapping(src_key, (slice(0, 4),), (slice(0, 4),), pp), AOASliceMapping( src_key + ".2", (slice(0, 4),), (slice(4, 8),), pp ), ] else: sms = [AOASliceMapping(src_key, (slice(0, 4),), (slice(0, 4),), pp)] return AOAMappingEntry( dst_key=dst_key, dst_global_shape=(8,), slice_mappings=sms ) def test_rename_only(self): entry = self._make_entry("model.layers.2.w", "model.layers.0.w") rename, transform, struct = _classify_mappings([entry]) self.assertEqual(len(rename), 1) self.assertEqual(len(transform), 0) self.assertEqual(len(struct), 0) def test_with_transform(self): entry = self._make_entry( "model.layers.2.w", "model.layers.0.w", ["[1, 0]"] ) rename, transform, struct = _classify_mappings([entry]) self.assertEqual(len(rename), 0) self.assertEqual(len(transform), 1) self.assertEqual(len(struct), 0) def test_structural_multi_src(self): entry = self._make_entry( "model.layers.2.qkv", "model.layers.0.q", multi_src=True ) rename, transform, struct = _classify_mappings([entry]) self.assertEqual(len(rename), 0) self.assertEqual(len(transform), 0) self.assertEqual(len(struct), 1) def test_structural_different_pattern(self): entry = self._make_entry("model.decoder.0.w", "model.encoder.0.w") rename, transform, struct = _classify_mappings([entry]) self.assertEqual(len(struct), 1) class TestGroupBySignature(unittest.TestCase): def test_same_signature_grouped(self): entries = [] for i in range(4): entries.append( AOAMappingEntry( dst_key=f"model.layers.{i}.w", dst_global_shape=(4,), slice_mappings=[ AOASliceMapping( f"src.layers.{i}.w", (slice(0, 4),), (slice(0, 4),), ["[1, 0]"], ) ], ) ) groups = _group_by_signature(entries) self.assertEqual(len(groups), 1) self.assertEqual(len(next(iter(groups.values()))), 4) def test_different_signatures(self): e1 = AOAMappingEntry( dst_key="model.layers.0.w", dst_global_shape=(4,), slice_mappings=[ AOASliceMapping( "src.layers.0.w", (slice(0, 4),), (slice(0, 4),), None ) ], ) e2 = AOAMappingEntry( dst_key="model.layers.0.qkv", dst_global_shape=(12,), slice_mappings=[ AOASliceMapping( "src.layers.0.q", (slice(0, 4),), (slice(0, 4),), None ), AOASliceMapping( "src.layers.0.k", (slice(0, 4),), (slice(4, 8),), None ), ], ) groups = _group_by_signature([e1, e2]) self.assertEqual(len(groups), 2) class TestGetSignature(unittest.TestCase): def test_digits_normalized(self): entry = AOAMappingEntry( dst_key="model.layers.5.weight", dst_global_shape=(4,), slice_mappings=[ AOASliceMapping( "src.layers.5.weight", (slice(0, 4),), (slice(0, 4),), None ) ], ) sig = _get_signature(entry) self.assertIn("{N}", sig) self.assertNotIn("5", sig) class TestFormatPatternGroups(unittest.TestCase): def test_basic_output(self): entries = [ AOAMappingEntry( dst_key="model.layers.0.w", dst_global_shape=(4, 8), slice_mappings=[ AOASliceMapping( "src.layers.0.w", (slice(0, 4), slice(0, 8)), (slice(0, 4), slice(0, 8)), ["[1, 0]"], ) ], ) ] groups = {"sig1": entries} lines, next_idx = _format_pattern_groups(groups, "test", 1) self.assertTrue(any("Pattern #1" in l for l in lines)) self.assertEqual(next_idx, 2) def test_numbering_continues(self): e1 = [ AOAMappingEntry( dst_key="a.0.w", dst_global_shape=(4,), slice_mappings=[ AOASliceMapping( "b.0.w", (slice(0, 4),), (slice(0, 4),), None ) ], ) ] e2 = [ AOAMappingEntry( dst_key="c.0.w", dst_global_shape=(4,), slice_mappings=[ AOASliceMapping( "d.0.w", (slice(0, 4),), (slice(0, 4),), None ) ], ) ] groups = {"sig1": e1, "sig2": e2} lines, next_idx = _format_pattern_groups(groups, "test", 5) self.assertEqual(next_idx, 7) def test_max_patterns_truncation(self): import paddle.distributed.flex_checkpoint.dcp.key_validation as kv old = kv._MAX_PATTERNS_SHOWN kv._MAX_PATTERNS_SHOWN = 2 try: groups = {} for i in range(5): groups[f"sig{i}"] = [ AOAMappingEntry( dst_key=f"x.{i}.w", dst_global_shape=(4,), slice_mappings=[ AOASliceMapping( f"y.{i}.w", (slice(0, 4),), (slice(0, 4),), None ) ], ) ] lines, _ = _format_pattern_groups(groups, "test", 1) self.assertTrue(any("more" in l for l in lines)) finally: kv._MAX_PATTERNS_SHOWN = old class TestAppendSrcLines(unittest.TestCase): def test_few_srcs(self): sms = [ AOASliceMapping("a.w", (slice(0, 4),), (slice(0, 4),), None), AOASliceMapping("b.w", (slice(0, 4),), (slice(4, 8),), None), ] lines = [] _append_src_lines(lines, sms) self.assertEqual(len(lines), 2) self.assertIn("SRC:", lines[0]) self.assertIn("+", lines[1]) def test_many_srcs_foldable(self): sms = [ AOASliceMapping( f"experts.{i}.w", (slice(0, 4),), (slice(i * 4, (i + 1) * 4),), None, ) for i in range(10) ] lines = [] _append_src_lines(lines, sms) # Should fold into single line with ×N self.assertTrue(any("\u00d7" in l for l in lines)) def test_many_srcs_not_foldable(self): sms = [ AOASliceMapping( "src_alpha.w", (slice(0, 4),), (slice(0, 4),), None ), AOASliceMapping("src_beta.w", (slice(0, 4),), (slice(4, 8),), None), AOASliceMapping( "src_gamma.w", (slice(0, 4),), (slice(8, 12),), None ), AOASliceMapping( "src_delta.w", (slice(0, 4),), (slice(12, 16),), None ), AOASliceMapping( "src_epsilon.w", (slice(0, 4),), (slice(16, 20),), None ), AOASliceMapping( "src_zeta.w", (slice(0, 4),), (slice(20, 24),), None ), ] lines = [] _append_src_lines(lines, sms) # Should show first 2, ..., last 1 self.assertTrue(any("more" in l for l in lines)) class TestGroupKeysAdaptive(unittest.TestCase): def test_basic_grouping(self): keys = [ "model.layers.0.weight", "model.layers.1.weight", "model.layers.2.weight", "model.embed.weight", ] groups = _group_keys_adaptive(keys) self.assertEqual(len(groups), 2) # layers.* grouped together layer_group = [g for g in groups.values() if len(g) == 3] self.assertEqual(len(layer_group), 1) def test_no_digits(self): keys = ["model.weight", "model.bias"] groups = _group_keys_adaptive(keys) self.assertEqual(len(groups), 2) class TestFormatKeyList(unittest.TestCase): def test_few_keys(self): keys = {"a.w", "b.w", "c.w"} lines = _format_key_list(keys) self.assertEqual(len(lines), 3) def test_many_keys_grouped(self): keys = {f"model.layers.{i}.weight" for i in range(100)} lines = _format_key_list(keys) # Should be grouped and folded self.assertTrue(len(lines) < 100) self.assertTrue(any("[" in l for l in lines)) def test_empty(self): lines = _format_key_list(set()) self.assertEqual(lines, []) class TestEmit(unittest.TestCase): @patch("paddle.distributed.flex_checkpoint.dcp.key_validation.logger") def test_normal_output(self, mock_logger): lines = ["line1", "line2", "line3"] _emit(lines) self.assertEqual(mock_logger.info.call_count, 3) @patch("paddle.distributed.flex_checkpoint.dcp.key_validation.logger") def test_truncation(self, mock_logger): import paddle.distributed.flex_checkpoint.dcp.key_validation as kv old_max = kv._MAX_TOTAL_LINES kv._MAX_TOTAL_LINES = 5 try: lines = ["x"] * 20 _emit(lines) # 5 lines + 1 truncation msg = 6 self.assertEqual(mock_logger.info.call_count, 6) finally: kv._MAX_TOTAL_LINES = old_max class TestPrintStandardReport(unittest.TestCase): @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit") def test_all_matched(self, mock_emit): result = KeyValidationResult() _print_standard_report(result, "/tmp/ckpt", 100) lines = mock_emit.call_args[0][0] self.assertTrue(any("[OK]" in l for l in lines)) @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit") def test_with_missing_and_unexpected(self, mock_emit): result = KeyValidationResult( missing_keys={"a.w", "b.w"}, unexpected_keys={"c.w"}, shape_mismatches=[ShapeMismatchInfo("d.w", (4, 8), (4, 16))], ) _print_standard_report(result, "/tmp/ckpt", 100) lines = mock_emit.call_args[0][0] self.assertTrue(any("Missing" in l for l in lines)) self.assertTrue(any("Unexpected" in l for l in lines)) self.assertTrue(any("Shape" in l for l in lines)) self.assertTrue(any("Matched: 98/100" in l for l in lines)) @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit") def test_shape_mismatch_truncation(self, mock_emit): import paddle.distributed.flex_checkpoint.dcp.key_validation as kv old = kv._MAX_SHAPE_MISMATCHES kv._MAX_SHAPE_MISMATCHES = 2 try: mismatches = [ ShapeMismatchInfo(f"k{i}", (4,), (8,)) for i in range(5) ] result = KeyValidationResult( missing_keys={"x"}, shape_mismatches=mismatches ) _print_standard_report(result, "/tmp/ckpt", 10) lines = mock_emit.call_args[0][0] self.assertTrue(any("and 3 more" in l for l in lines)) finally: kv._MAX_SHAPE_MISMATCHES = old class TestPrintAoaReport(unittest.TestCase): @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit") def test_all_resolved(self, mock_emit): mappings = [ AOAMappingEntry( dst_key="a.w", dst_global_shape=(4,), slice_mappings=[ AOASliceMapping("b.w", (slice(0, 4),), (slice(0, 4),), None) ], is_identity=False, ), ] result = KeyValidationResult() _print_aoa_report(result, mappings, set(), "/tmp/ckpt") lines = mock_emit.call_args[0][0] self.assertTrue(any("[OK]" in l for l in lines)) @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit") def test_with_missing(self, mock_emit): mappings = [ AOAMappingEntry( "a.w", (4,), [AOASliceMapping("b.w", (slice(0, 4),), (slice(0, 4),), None)], ) ] result = KeyValidationResult( missing_keys={"c.w"}, unexpected_keys={"d.w"} ) _print_aoa_report(result, mappings, {"removed.w"}, "/tmp/ckpt") lines = mock_emit.call_args[0][0] self.assertTrue(any("Missing" in l for l in lines)) self.assertTrue(any("Unexpected" in l for l in lines)) self.assertTrue(any("Removed" in l for l in lines)) @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit") def test_randomly_initialized_keys(self, mock_emit): mappings = [] result = KeyValidationResult( randomly_initialized_keys={"init.w", "init.b"} ) _print_aoa_report(result, mappings, set(), "/tmp/ckpt") lines = mock_emit.call_args[0][0] self.assertTrue(any("Initialized (2)" in l for l in lines)) @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit") def test_removed_keys_truncation(self, mock_emit): mappings = [] removed = {f"removed.key.{i}" for i in range(10)} result = KeyValidationResult() _print_aoa_report(result, mappings, removed, "/tmp/ckpt") lines = mock_emit.call_args[0][0] self.assertTrue(any("more" in l for l in lines)) class TestBuildAoaMappings(unittest.TestCase): def test_basic(self): engine = MagicMock() td1 = MagicMock() td1.shape = [4, 8] td1.slices = [ ( "src.w", (slice(0, 4), slice(0, 8)), (slice(0, 4), slice(0, 8)), None, ) ] td2 = MagicMock() td2.shape = [8, 8] td2.slices = [ ( "src.q", (slice(0, 4), slice(0, 8)), (slice(0, 4), slice(0, 8)), ["[1, 0]"], ), ( "src.k", (slice(0, 4), slice(0, 8)), (slice(4, 8), slice(0, 8)), ["[1, 0]"], ), ] ov = MagicMock() ov.items.return_value = sorted({"dst.qkv": td2, "dst.w": td1}.items()) engine.output_vars = ov results = _build_aoa_mappings(engine) self.assertEqual(len(results), 2) qkv = next(r for r in results if r.dst_key == "dst.qkv") self.assertFalse(qkv.is_identity) self.assertEqual(len(qkv.slice_mappings), 2) def test_identity_detection(self): engine = MagicMock() td = MagicMock() td.shape = [4, 8] td.slices = [ ( "same.key", (slice(0, 4), slice(0, 8)), (slice(0, 4), slice(0, 8)), None, ) ] ov = MagicMock() ov.items.return_value = [("same.key", td)] engine.output_vars = ov results = _build_aoa_mappings(engine) self.assertEqual(len(results), 1) self.assertTrue(results[0].is_identity) def test_none_tensor_desc_skipped(self): engine = MagicMock() ov = MagicMock() ov.items.return_value = [("a", None), ("b", None)] engine.output_vars = ov results = _build_aoa_mappings(engine) self.assertEqual(len(results), 0) class TestValidateAndReportKeysStandard(unittest.TestCase): def _make_metadata(self, keys_shapes): """keys_shapes: dict of {key: shape_tuple}""" storage_metadata = {} state_dict_metadata = {} for key, shape in keys_shapes.items(): idx = LocalTensorIndex( tensor_key=key, global_offset=tuple([0] * len(shape)), replica_id=0, ) storage_metadata[idx] = f"{key}.distcp" state_dict_metadata[key] = [ LocalTensorMetadata( global_offset=tuple([0] * len(shape)), local_shape=shape, dtype="float32", global_shape=shape, ) ] return Metadata( state_dict_metadata=state_dict_metadata, storage_metadata=storage_metadata, ) @patch("paddle.distributed.get_rank", return_value=0) @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit") def test_all_match(self, mock_emit, mock_rank): metadata = self._make_metadata({"w1": (4, 8), "w2": (4, 8)}) state_dict = { "w1": paddle.zeros([4, 8]), "w2": paddle.zeros([4, 8]), } result = validate_and_report_keys_standard( [metadata], {"w1", "w2"}, None, False, "/tmp/ckpt", state_dict ) self.assertEqual(len(result.missing_keys), 0) self.assertEqual(len(result.unexpected_keys), 0) @patch("paddle.distributed.get_rank", return_value=0) @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit") def test_missing_keys(self, mock_emit, mock_rank): metadata = self._make_metadata({"w1": (4,)}) state_dict = { "w1": paddle.zeros([4]), "w2": paddle.zeros([4]), } result = validate_and_report_keys_standard( [metadata], {"w1", "w2"}, None, False, "/tmp/ckpt", state_dict ) self.assertIn("w2", result.missing_keys) @patch("paddle.distributed.get_rank", return_value=0) @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit") def test_unexpected_keys(self, mock_emit, mock_rank): metadata = self._make_metadata({"w1": (4,), "w2": (4,), "w3": (4,)}) state_dict = {"w1": paddle.zeros([4])} result = validate_and_report_keys_standard( [metadata], {"w1"}, None, False, "/tmp/ckpt", state_dict ) self.assertIn("w2", result.unexpected_keys) self.assertIn("w3", result.unexpected_keys) @patch("paddle.distributed.get_rank", return_value=0) @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit") def test_shape_mismatch(self, mock_emit, mock_rank): metadata = self._make_metadata({"w1": (4, 8)}) state_dict = {"w1": paddle.zeros([4, 16])} result = validate_and_report_keys_standard( [metadata], {"w1"}, None, False, "/tmp/ckpt", state_dict ) self.assertEqual(len(result.shape_mismatches), 1) self.assertEqual(result.shape_mismatches[0].src_global_shape, (4, 8)) self.assertEqual(result.shape_mismatches[0].dst_global_shape, (4, 16)) @patch("paddle.distributed.get_rank", return_value=0) @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit") def test_replica_id_filtered(self, mock_emit, mock_rank): """Keys with replica_id != 0 should be filtered out.""" storage_metadata = { LocalTensorIndex( tensor_key="w1", global_offset=(0,), replica_id=0 ): "f1", LocalTensorIndex( tensor_key="w2", global_offset=(0,), replica_id=1 ): "f2", } metadata = Metadata( state_dict_metadata={ "w1": [LocalTensorMetadata((0,), (4,), "float32", (4,))] }, storage_metadata=storage_metadata, ) state_dict = {"w1": paddle.zeros([4])} result = validate_and_report_keys_standard( [metadata], {"w1"}, None, False, "/tmp/ckpt", state_dict ) self.assertEqual(len(result.unexpected_keys), 0) @patch( "paddle.distributed.flex_checkpoint.dcp.key_validation._get_rank", return_value=1, ) @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit") @patch("paddle.distributed.all_gather_object") def test_non_rank0_no_print(self, mock_gather, mock_emit, mock_rank): metadata = self._make_metadata({"w1": (4,)}) state_dict = {"w1": paddle.zeros([4])} def gather_side_effect(out_list, obj, group=None): out_list.clear() out_list.append(obj) mock_gather.side_effect = gather_side_effect validate_and_report_keys_standard( [metadata], {"w1"}, None, True, "/tmp/ckpt", state_dict ) mock_emit.assert_not_called() class TestValidateAndReportKeysAoa(unittest.TestCase): def _make_mock_engine(self): engine = MagicMock() td1 = MagicMock() td1.shape = [4, 8] td1.slices = [ ( "src.w1", (slice(0, 4), slice(0, 8)), (slice(0, 4), slice(0, 8)), None, ) ] td2 = MagicMock() td2.shape = [8, 8] td2.slices = [ ( "src.q", (slice(0, 4), slice(0, 8)), (slice(0, 4), slice(0, 8)), ["[1, 0]"], ), ( "src.k", (slice(0, 4), slice(0, 8)), (slice(4, 8), slice(0, 8)), ["[1, 0]"], ), ] # output_vars: need .items() for _build_aoa_mappings and iteration for values() ov = MagicMock() ov.items.return_value = sorted({"dst.w1": td1, "dst.qkv": td2}.items()) ov.values.return_value = [td1, td2] ov.__iter__ = lambda self: iter({"dst.w1": td1, "dst.qkv": td2}) ov.__getitem__ = lambda self, k: {"dst.w1": td1, "dst.qkv": td2}[k] engine.output_vars = ov engine.need_add_output_vars = ["dst.init"] engine.need_remove_input_vars = ["src.removed"] engine.input_vars = MagicMock() engine.input_vars.keys.return_value = [ "src.w1", "src.q", "src.k", "src.removed", "src.leftover", ] engine.context = MagicMock() engine.context.get_all_dst_state_keys.return_value = { "dst.w1", "dst.qkv", "dst.init", } return engine @patch("paddle.distributed.get_rank", return_value=0) @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit") def test_all_resolved(self, mock_emit, mock_rank): engine = self._make_mock_engine() metadata = MagicMock() result = validate_and_report_keys_aoa(engine, metadata, "/tmp/ckpt") # dst.w1 and dst.qkv are covered; dst.init is randomly initialized self.assertEqual(len(result.missing_keys), 0) # src.leftover not consumed and not removed self.assertIn("src.leftover", result.unexpected_keys) self.assertIn("dst.init", result.randomly_initialized_keys) @patch("paddle.distributed.get_rank", return_value=0) @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit") def test_truly_missing(self, mock_emit, mock_rank): engine = self._make_mock_engine() # Add a dst key that is NOT covered engine.context.get_all_dst_state_keys = lambda: { "dst.w1", "dst.qkv", "dst.init", "dst.missing", } metadata = MagicMock() result = validate_and_report_keys_aoa(engine, metadata, "/tmp/ckpt") self.assertIn("dst.missing", result.missing_keys) @patch("paddle.distributed.get_rank", return_value=1) @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit") def test_non_rank0_no_print(self, mock_emit, mock_rank): engine = self._make_mock_engine() metadata = MagicMock() validate_and_report_keys_aoa(engine, metadata, "/tmp/ckpt") mock_emit.assert_not_called() class TestColorHelpers(unittest.TestCase): def test_no_color(self): from paddle.distributed.flex_checkpoint.dcp.key_validation import _C self.assertEqual(_C.green("test"), "test") self.assertEqual(_C.yellow("test"), "test") self.assertEqual(_C.red("test"), "test") self.assertEqual(_C.cyan("test"), "test") if __name__ == "__main__": unittest.main()