Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / standard_planner.py: 80%
284 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-25 04:27 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Standard planner implementations for checkpoint save and load."""
16from dataclasses import dataclass
17import dataclasses
18import pickle
19from typing import Any, Optional, Union
21from hyper_parallel.core.distributed_checkpoint.metadata import (
22 CHUNK_INFO,
23 Metadata,
24 MetadataIndex,
25 ChunkStorageMetadata,
26 ChunkInfo,
27 TensorStorageMetadata,
28 TensorProperties,
29 BytesStorageMetadata
30)
31from hyper_parallel.core.distributed_checkpoint.planner import (
32 SavePlan,
33 SavePlanner,
34 LoadPlan,
35 LoadPlanner,
36 WriteItem,
37 WriteItemType,
38 ReadItem,
39 LoadItemType
40)
41from hyper_parallel.core.distributed_checkpoint.reshard import infer_intersection
42from hyper_parallel.core.distributed_checkpoint.ragged_utils import (
43 create_ragged_write_items,
44 get_ragged_box_tensor,
45)
46from hyper_parallel.core.distributed_checkpoint.util import (
47 narrow_tensor_by_index,
48 chunk_to_area,
49 create_chunk_list_for_tensor,
50 remove_redundant_plans,
51 flatten_state_dict,
52 set_element,
53)
54from hyper_parallel.core.dtensor.dtensor import DTensor
55from hyper_parallel.core.dtensor.layout import Layout, infer_slice_area_by_layout
56from hyper_parallel.platform import get_platform
58platform = get_platform()
59Tensor = platform.Tensor
62@dataclass(frozen=True)
63class CachedSaveResult:
64 """Cached finalized save result keyed by planner cache namespace."""
66 final_plan: SavePlan
67 metadata: Metadata
70class StandardSavePlanner(SavePlanner):
71 """Standard implementation of SavePlanner for distributed checkpoint saving."""
73 _cached_save_result: dict[str, CachedSaveResult] = {}
75 def __init__(
76 self,
77 enable_plan_caching: bool = True,
78 remove_redundancy: bool = True,
79 save_to_minimum_rank: bool = False,
80 ):
81 self.state_dict: Optional[dict[str, Any]] = None
82 self.is_coordinator: bool = False
83 self.rank: int = 0
84 self.remove_redundancy: bool = remove_redundancy
85 self.save_to_minimum_rank: bool = save_to_minimum_rank
86 self.flatten_state_dict: bool = True
87 self._enable_plan_caching: bool = enable_plan_caching
88 self._default_enable_plan_caching: bool = enable_plan_caching
89 self._cached_plans_key: str = self.__class__.__name__
91 def configure_planner(self, state_dict: dict[str, Any], **kwargs) -> None:
92 """
93 Configure planner.
95 Args:
96 state_dict (dict[str, Any]): The state_dict to save.
97 **kwargs: Additional keyword arguments (e.g., is_coordinator, rank, remove_redundancy,
98 save_to_minimum_rank).
99 """
100 self.is_coordinator = kwargs.get("is_coordinator", False)
101 self.rank = kwargs.get("rank", 0)
102 self.remove_redundancy = kwargs.get("remove_redundancy", self.remove_redundancy)
103 self.save_to_minimum_rank = kwargs.get("save_to_minimum_rank", self.save_to_minimum_rank)
104 self.flatten_state_dict = kwargs.get("flatten_state_dict", True)
106 use_collectives = bool(kwargs.get("use_collectives", True))
107 self._enable_plan_caching = bool(
108 kwargs.get("enable_plan_caching", self._default_enable_plan_caching)
109 )
110 if not use_collectives:
111 self.remove_redundancy = False
112 self._enable_plan_caching = False
114 if self.flatten_state_dict:
115 state_dict, self.name_mapping = flatten_state_dict(state_dict)
116 self.state_dict = state_dict
117 if any(
118 isinstance(obj, DTensor)
119 and obj.layout is not None
120 and obj.layout.ragged_shard is not None
121 for obj in state_dict.values()
122 ):
123 self._enable_plan_caching = False
124 self._cached_plans_key = self._build_cache_key(state_dict)
126 def _build_cache_key(self, state_dict: dict[str, Any]) -> str:
127 """Build a stable cache namespace from sorted state_dict keys."""
128 return f"{self.__class__.__name__}:{'||'.join(state_dict.keys())}"
130 def build_local_plan(self) -> SavePlan:
131 """
132 Create local save plan.
134 Returns:
135 SavePlan: Local save plan containing WriteItems for this rank.
136 """
137 if self.state_dict is None:
138 raise RuntimeError("Planner not set up")
140 def compute_global_offsets(global_shape: tuple[int, ...], dtensor_layout: Layout) -> tuple[int, ...]:
141 """
142 Compute the offsets of local tensor in global tensor based on layout.
144 Args:
145 global_shape (tuple[int, ...]): Global shape of the tensor.
146 dtensor_layout (Layout): Layout of the DTensor.
148 Returns:
149 tuple[int, ...]: Tuple of offsets for each dimension.
150 """
151 if dtensor_layout is None:
152 # If layout is None, return all zeros (no sharding)
153 return tuple(0 for _ in global_shape)
155 # Validate layout attributes
156 if not hasattr(dtensor_layout, 'mesh_shape') or dtensor_layout.mesh_shape is None:
157 raise ValueError("Layout must have mesh_shape attribute")
158 if not hasattr(dtensor_layout, 'tensor_map') or dtensor_layout.tensor_map is None:
159 raise ValueError("Layout must have tensor_map attribute")
160 if not hasattr(dtensor_layout, 'rank_list') or dtensor_layout.rank_list is None:
161 raise ValueError("Layout must have rank_list attribute")
163 current_rank = self.rank
164 if current_rank not in dtensor_layout.rank_list:
165 raise ValueError(
166 f"Current rank {current_rank} not found in layout's rank_list {dtensor_layout.rank_list}")
168 inner_rank_id = dtensor_layout.rank_list.index(current_rank)
169 # Calculate slice area using infer_slice_area_by_rank
170 slice_area = infer_slice_area_by_layout(
171 dtensor_layout,
172 inner_rank_id,
173 global_shape,
174 )
175 # Extract offsets (start values) from slice_area
176 return tuple(start for start, _ in slice_area)
178 items = []
179 for fqn, obj in self.state_dict.items():
180 # Check if it's a DTensor
181 if isinstance(obj, DTensor):
182 if obj.layout is not None and obj.layout.ragged_shard is not None:
183 items.extend(create_ragged_write_items(fqn, obj))
184 continue
185 # Create write item for DTensor
186 local_tensor = obj.to_local()
187 layout = obj.layout
189 # Get chunk metadata with offsets
190 if layout:
191 offsets = compute_global_offsets(obj.shape, layout)
192 else:
193 offsets = (0,) * len(local_tensor.shape)
195 sizes = local_tensor.shape
196 chunk = ChunkStorageMetadata(offsets=offsets, sizes=sizes)
197 # Get tensor properties
198 dtype_str = str(local_tensor.dtype) if hasattr(local_tensor, 'dtype') else 'unknown'
199 properties = TensorProperties(dtype=dtype_str)
200 # Create write item for this tensor
201 index = MetadataIndex(fqn=fqn, offset=offsets, index=None)
202 write_item = WriteItem(
203 index=index,
204 type=WriteItemType.TENSOR,
205 tensor_data={
206 'chunk': chunk,
207 'properties': properties,
208 'size': obj.shape,
209 }
210 )
211 items.append(write_item)
212 elif isinstance(obj, Tensor):
213 # Create write item for platform.Tensor: build single chunk with tensor's own size
214 dtype_str = str(obj.dtype) if hasattr(obj, 'dtype') else 'unknown'
215 properties = TensorProperties(dtype=dtype_str)
216 # handle Tensor with shard information
217 if hasattr(obj, CHUNK_INFO):
218 if not isinstance(getattr(obj, CHUNK_INFO), ChunkInfo):
219 raise ValueError("The attr CHUNK_INFO should be a ChunkInfo instance")
220 chunk = getattr(obj, CHUNK_INFO).chunk
221 # Single chunk covering the whole tensor (offsets=0, sizes=shape)
222 else:
223 chunk = ChunkStorageMetadata(
224 offsets=(0,) * len(obj.shape),
225 sizes=obj.shape,
226 )
227 index = MetadataIndex(fqn=fqn, offset=chunk.offsets, index=None)
228 write_item = WriteItem(
229 index=index,
230 type=WriteItemType.TENSOR,
231 tensor_data={
232 'chunk': chunk,
233 'properties': properties,
234 'size': getattr(obj, CHUNK_INFO).global_shape if hasattr(obj, CHUNK_INFO) else obj.shape,
235 }
236 )
237 items.append(write_item)
238 else:
239 # Handle non-tensor types (bytes, etc.)
240 index = MetadataIndex(fqn=fqn)
241 write_item = WriteItem(
242 index=index,
243 type=WriteItemType.BYTE_IO,
244 bytes_io_data=None
245 )
246 items.append(write_item)
248 plan = SavePlan(items=items)
249 if self.flatten_state_dict:
250 plan.planner_data = self.name_mapping
251 return plan
253 def build_global_plan(self, all_plans: list[SavePlan]) -> tuple[list[SavePlan], Metadata]:
254 """
255 Build global plan from all local plans.
257 Collects chunks from all ranks, validates consistency, and creates metadata for the checkpoint.
259 Args:
260 all_plans (list[SavePlan]): List of local plans from all ranks.
262 Returns:
263 tuple[list[SavePlan], Metadata]: Updated plans and checkpoint metadata.
264 """
265 # Deduplicate plans if redundancy removal is enabled
266 if self.remove_redundancy and len(all_plans) > 1:
267 all_plans = remove_redundant_plans(all_plans, save_to_minimum_rank=self.save_to_minimum_rank)
269 # Collect all write items by FQN
270 fqn_to_chunks: dict[str, list[ChunkStorageMetadata]] = {}
271 fqn_to_properties: dict[str, TensorProperties] = {}
272 fqn_to_size: dict[str, tuple] = {}
273 state_dict_metadata: dict[str, Union[TensorStorageMetadata, BytesStorageMetadata]] = {}
275 final_global_plans: list[SavePlan] = []
276 for plan in all_plans:
277 with_index_items = []
278 for item in plan.items:
279 if item.type == WriteItemType.TENSOR and item.tensor_data:
280 fqn = item.index.fqn
281 chunk = item.tensor_data['chunk']
282 properties = item.tensor_data['properties']
283 size = item.tensor_data['size']
285 # Validate consistency across ranks
286 if fqn in fqn_to_chunks and (fqn_to_properties[fqn] != properties or fqn_to_size[fqn] != size):
287 raise ValueError(f"The {fqn} in different rank has different properties and size.")
289 # Initialize FQN entry if not exists
290 if fqn not in fqn_to_chunks:
291 fqn_to_properties[fqn] = properties
292 fqn_to_size[fqn] = size
293 fqn_to_chunks[fqn] = []
295 # Append chunk and set index (platform.Tensor has exactly one chunk)
296 new_index = dataclasses.replace(item.index, index=len(fqn_to_chunks[fqn]))
297 with_index_item = dataclasses.replace(item, index=new_index)
298 with_index_items.append(with_index_item)
299 fqn_to_chunks[fqn].append(chunk)
301 elif item.type == WriteItemType.BYTE_IO:
302 with_index_items.append(item)
303 state_dict_metadata[item.index.fqn] = BytesStorageMetadata()
304 else:
305 raise ValueError(f"Unsupported write item type: {item.type}")
307 final_global_plans.append(dataclasses.replace(plan, items=with_index_items))
309 # Create metadata for all tensors
310 for fqn, chunks in fqn_to_chunks.items():
311 state_dict_metadata[fqn] = TensorStorageMetadata(
312 properties=fqn_to_properties[fqn],
313 size=fqn_to_size[fqn],
314 chunks=chunks
315 )
317 metadata = Metadata(state_dict_metadata=state_dict_metadata)
318 if self.flatten_state_dict:
319 merged_mapping = {}
320 for p in all_plans:
321 merged_mapping.update(p.planner_data)
322 metadata.planner_data = merged_mapping
323 return final_global_plans, metadata
325 def finalize_plan(self, plan: SavePlan) -> SavePlan:
326 """
327 Finalize the plan.
329 Args:
330 plan (SavePlan): Plan to finalize.
332 Returns:
333 SavePlan: Finalized plan.
334 """
335 return plan
337 def get_cached(self) -> Optional[CachedSaveResult]:
338 """Return cached finalized plan and metadata when plan caching is enabled."""
339 if (
340 not self._enable_plan_caching
341 or self._cached_plans_key not in StandardSavePlanner._cached_save_result
342 ):
343 return None
344 return StandardSavePlanner._cached_save_result[self._cached_plans_key]
346 def cache_result(self, final_plan: SavePlan, metadata: Metadata) -> None:
347 """Store finalized plan and metadata in the class-level planner cache."""
348 if not self._enable_plan_caching:
349 return
350 StandardSavePlanner._cached_save_result[self._cached_plans_key] = CachedSaveResult(
351 final_plan=final_plan,
352 metadata=metadata,
353 )
355 def get_data(self, item: WriteItem) -> Any:
356 """
357 Get current runtime data from state_dict for a write item.
359 Args:
360 item (WriteItem): Write item describing what to write.
362 Returns:
363 Any: Runtime object to be written.
364 """
365 if self.state_dict is None:
366 raise RuntimeError("Planner not set up")
367 fqn = item.index.fqn
368 if fqn not in self.state_dict:
369 raise KeyError(f"Key {fqn} not found in state_dict")
370 obj = self.state_dict[fqn]
371 if item.type == WriteItemType.TENSOR:
372 if isinstance(obj, DTensor):
373 if obj.layout is not None and obj.layout.ragged_shard is not None:
374 return get_ragged_box_tensor(obj, item.index).detach().cpu()
375 return obj.to_local().detach().cpu()
376 if isinstance(obj, Tensor):
377 return obj.detach().cpu()
378 raise TypeError(f"Write item {fqn} expected tensor-like object, got {type(obj)}")
379 if item.type == WriteItemType.BYTE_IO:
380 return obj
381 raise TypeError(f"Unsupported write item type: {item.type}")
384def create_read_items_for_chunk_list(
385 fqn: str,
386 checkpoint_md: TensorStorageMetadata,
387 local_chunks: list[ChunkStorageMetadata],
388) -> list[ReadItem]:
389 """
390 Create ReadItems by matching local chunks (what this rank needs) with
391 saved chunks (checkpoint_md.chunks), including resharding overlaps.
393 Mirrors torch create_read_items_for_chunk_list behavior.
395 Args:
396 fqn (str): Fully qualified name of the tensor.
397 checkpoint_md (TensorStorageMetadata): Tensor storage metadata from checkpoint.
398 local_chunks (list[ChunkStorageMetadata]): List of local chunks needed by this rank.
400 Returns:
401 list[ReadItem]: List of ReadItems for loading the required data.
402 """
403 read_items: list[ReadItem] = []
404 saved_chunks = checkpoint_md.chunks
405 if not local_chunks or not saved_chunks:
406 return read_items
408 for local_idx, local_chunk in enumerate(local_chunks):
409 local_area = chunk_to_area(local_chunk)
410 for storage_idx, storage_chunk in enumerate(saved_chunks):
411 saved_area = chunk_to_area(storage_chunk)
412 overlap = infer_intersection(local_area, saved_area)
413 if overlap is None:
414 continue
416 dest_offsets = tuple(overlap[i][0] - local_chunk.offsets[i] for i in range(len(overlap)))
417 storage_offsets = tuple(overlap[i][0] - storage_chunk.offsets[i] for i in range(len(overlap)))
418 lengths = tuple(overlap[i][1] - overlap[i][0] for i in range(len(overlap)))
420 read_items.append(
421 ReadItem(
422 type=LoadItemType.TENSOR,
423 dest_index=MetadataIndex(fqn=fqn, offset=local_chunk.offsets, index=local_idx),
424 dest_offsets=dest_offsets,
425 storage_index=MetadataIndex(fqn=fqn, offset=storage_chunk.offsets, index=storage_idx),
426 storage_offsets=storage_offsets,
427 lengths=lengths,
428 )
429 )
430 return read_items
433class StandardLoadPlanner(LoadPlanner):
434 """
435 Standard implementation of LoadPlanner.
437 Iterate state_dict and creates load plans via chunk list for resharding support.
438 """
440 def __init__(self, allow_partial_load: bool = False):
441 """
442 Args:
443 allow_partial_load (bool): If True, allow loading when checkpoint has fewer keys than state_dict.
444 Default False.
445 """
446 self.state_dict: Optional[dict[str, Any]] = None
447 self.metadata: Optional[Metadata] = None
448 self.is_coordinator: bool = False
449 self.rank: int = 0
450 self.allow_partial_load = allow_partial_load
451 self.flatten_state_dict: bool = True
453 def configure_planner(self, state_dict: dict[str, Any], metadata: Metadata, **kwargs) -> None:
454 """
455 Configure planner with state dict and metadata.
457 Args:
458 state_dict (dict[str, Any]): The state_dict to load into (modified in-place).
459 metadata (Metadata): Checkpoint metadata.
460 **kwargs: Additional keyword arguments (e.g., is_coordinator, rank).
461 """
462 self.state_dict = state_dict
463 self.metadata = metadata
464 self.is_coordinator = kwargs.get("is_coordinator", False)
465 self.rank = kwargs.get("rank", 0)
466 self.flatten_state_dict = kwargs.get("flatten_state_dict", True)
467 self.original_state_dict = state_dict
468 if self.flatten_state_dict:
469 state_dict, self.name_mapping = flatten_state_dict(state_dict)
470 self.state_dict = state_dict
472 def build_local_plan(self) -> LoadPlan:
473 """
474 Build local load plan.
476 Iterate state_dict and creates load plans via chunk list for resharding support.
478 Returns:
479 LoadPlan: Local load plan containing ReadItems for this rank.
480 """
481 if self.state_dict is None or self.metadata is None:
482 raise RuntimeError("Planner not configured")
484 requests: list[ReadItem] = []
485 strict = not self.allow_partial_load
486 for fqn, obj in self.state_dict.items():
487 if fqn not in self.metadata.state_dict_metadata:
488 if fqn.endswith(('matched_adamw_rms', 'step')):
489 continue
490 if strict:
491 raise RuntimeError(f"Missing key in checkpoint state_dict: {fqn}.")
492 continue
493 md = self.metadata.state_dict_metadata[fqn]
494 if isinstance(md, TensorStorageMetadata):
495 obj_size = getattr(obj, CHUNK_INFO).global_shape if hasattr(obj, CHUNK_INFO) \
496 else getattr(obj, "shape", None)
497 if obj_size is None or md.size != tuple(obj_size):
498 raise ValueError(
499 f"Size mismatch between saved {md.size} and current: {obj_size} for {fqn}",
500 )
501 if isinstance(obj, DTensor):
502 layout = getattr(obj, "layout", None)
503 rank_list = getattr(layout, "rank_list", None) if layout else None
504 if rank_list is None and layout is not None:
505 rank_list = getattr(layout, "_rank_list", None)
506 if layout is not None and rank_list is not None:
507 if get_platform().get_rank() not in rank_list:
508 continue
509 # Both DTensor and platform.Tensor: create local chunks and read items
510 local_chunks = create_chunk_list_for_tensor(obj)
511 requests += create_read_items_for_chunk_list(fqn, md, local_chunks)
512 else:
513 requests.append(
514 ReadItem(
515 type=LoadItemType.BYTE_IO,
516 dest_index=MetadataIndex(fqn=fqn),
517 dest_offsets=(0,),
518 storage_index=MetadataIndex(fqn=fqn),
519 storage_offsets=(0,),
520 lengths=(0,),
521 )
522 )
523 return LoadPlan(items=requests)
525 def build_global_plan(self, all_plans: list[LoadPlan]) -> list[LoadPlan]:
526 """
527 Build global plan from all local plans.
529 For now, returns plans as-is. In a more sophisticated implementation, you might need to coordinate across ranks.
531 Args:
532 all_plans (list[LoadPlan]): List of local plans from all ranks.
534 Returns:
535 list[LoadPlan]: Global plans (currently returns plans as-is).
536 """
537 return all_plans
539 def finalize_plan(self, plan: LoadPlan) -> LoadPlan:
540 """
541 Finalize the plan (no-op for default implementation).
543 Args:
544 plan (LoadPlan): Plan to finalize.
546 Returns:
547 LoadPlan: Finalized plan.
548 """
549 return plan
551 def acquire_tensor(self, read_item: ReadItem) -> Any:
552 """
553 Acquire the destination slice (narrow view) for this read_item.
555 StorageReader uses this to copy loaded data into the correct region.
556 Torch-aligned behavior.
558 Args:
559 read_item (ReadItem): The read item specifying what to load.
561 Returns:
562 Any: The destination tensor slice where data should be written
563 (tensor-like object).
564 """
565 if self.state_dict is None:
566 raise RuntimeError("Planner not configured")
568 fqn = read_item.dest_index.fqn
569 if fqn not in self.state_dict:
570 raise KeyError(f"Key {fqn} not found in state_dict")
572 target = self.state_dict[fqn]
573 if (
574 isinstance(target, DTensor)
575 and target.layout is not None
576 and target.layout.ragged_shard is not None
577 ):
578 box_tensor = get_ragged_box_tensor(target, read_item.dest_index)
579 return narrow_tensor_by_index(
580 box_tensor,
581 read_item.dest_offsets,
582 read_item.lengths,
583 )
585 local_tensor = target.to_local().detach() if isinstance(target, DTensor) else target.detach()
586 return narrow_tensor_by_index(
587 local_tensor,
588 read_item.dest_offsets,
589 read_item.lengths,
590 )
592 def apply_tensor(self, read_item: ReadItem, tensor: Any) -> None:
593 """
594 Apply tensor after reading.
596 After read_data copies into the slice, this is no-op when tensor is the
597 same slice. When the backend has no copy_ (e.g. mindspore), read_data
598 passes the loaded slice here; we copy it into the destination slice.
600 Args:
601 read_item (ReadItem): The read item that was processed.
602 tensor (Any): The tensor data to apply (tensor-like object).
603 """
604 if tensor is None:
605 return
606 dest_slice = self.acquire_tensor(read_item)
607 if dest_slice is tensor:
608 return
609 if hasattr(dest_slice, "copy_"):
610 dest_slice.copy_(tensor)
611 else:
612 # Fallback: assign into state_dict if supported
613 dest_slice[...] = tensor
615 def apply_bytes(self, read_item: ReadItem, value: bytes) -> None:
616 """
617 Load bytes data into state_dict.
619 Args:
620 read_item (ReadItem): The read item specifying the destination.
621 value (bytes): The bytes data to deserialize and load.
622 """
623 if self.state_dict is None:
624 raise RuntimeError("Planner not set up")
626 fqn = read_item.dest_index.fqn
627 # Deserialize bytes
628 obj = pickle.loads(value)
629 self.state_dict[fqn] = obj
630 if self.flatten_state_dict:
631 set_element(self.original_state_dict, self.name_mapping[fqn], obj)
635class _DcpMergeLoadPlanner(StandardLoadPlanner):
636 """Load planner that builds distributed checkpoint from dcp into fully ``state_dict`` (in-place)."""
638 def __init__(self) -> None:
639 super().__init__()
641 def configure_planner(self, state_dict: dict[str, Any], metadata: Metadata, **kwargs) -> None:
642 if len(state_dict) > 0:
643 raise ValueError(
644 "state_dict must be empty for _DcpMergeLoadPlanner; "
645 "it is populated in-place from checkpoint metadata."
646 )
648 if metadata is None:
649 raise ValueError("metadata must not be None for _DcpMergeLoadPlanner.")
651 self.is_coordinator = kwargs.get("is_coordinator", False)
652 for k, v in metadata.state_dict_metadata.items():
653 if isinstance(v, TensorStorageMetadata):
654 v = platform.empty(
655 platform.list_to_size(v.size),
656 dtype=platform.str_to_dtype(v.properties.dtype),
657 )
659 state_dict[k] = v
660 if metadata.planner_data is not None and k in metadata.planner_data:
661 set_element(state_dict, metadata.planner_data[k], v)
663 super().configure_planner(
664 state_dict,
665 metadata,
666 is_coordinator=self.is_coordinator,
667 flatten_state_dict=True,
668 )