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-21 04:29 +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_slice_area_by_rank, 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
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_rank(
171 mesh_shape=dtensor_layout.mesh_shape,
172 tensor_map=dtensor_layout.tensor_map,
173 rank_id=inner_rank_id,
174 full_shape=global_shape
175 )
176 # Extract offsets (start values) from slice_area
177 return tuple(start for start, _ in slice_area)
179 items = []
180 for fqn, obj in self.state_dict.items():
181 # Check if it's a DTensor
182 if isinstance(obj, DTensor):
183 if obj.layout is not None and obj.layout.ragged_shard is not None:
184 items.extend(create_ragged_write_items(fqn, obj))
185 continue
186 # Create write item for DTensor
187 local_tensor = obj.to_local()
188 layout = obj.layout
190 # Get chunk metadata with offsets
191 if layout:
192 offsets = compute_global_offsets(obj.shape, layout)
193 else:
194 offsets = (0,) * len(local_tensor.shape)
196 sizes = local_tensor.shape
197 chunk = ChunkStorageMetadata(offsets=offsets, sizes=sizes)
198 # Get tensor properties
199 dtype_str = str(local_tensor.dtype) if hasattr(local_tensor, 'dtype') else 'unknown'
200 properties = TensorProperties(dtype=dtype_str)
201 # Create write item for this tensor
202 index = MetadataIndex(fqn=fqn, offset=offsets, index=None)
203 write_item = WriteItem(
204 index=index,
205 type=WriteItemType.TENSOR,
206 tensor_data={
207 'chunk': chunk,
208 'properties': properties,
209 'size': obj.shape,
210 }
211 )
212 items.append(write_item)
213 elif isinstance(obj, Tensor):
214 # Create write item for platform.Tensor: build single chunk with tensor's own size
215 dtype_str = str(obj.dtype) if hasattr(obj, 'dtype') else 'unknown'
216 properties = TensorProperties(dtype=dtype_str)
217 # handle Tensor with shard information
218 if hasattr(obj, CHUNK_INFO):
219 if not isinstance(getattr(obj, CHUNK_INFO), ChunkInfo):
220 raise ValueError("The attr CHUNK_INFO should be a ChunkInfo instance")
221 chunk = getattr(obj, CHUNK_INFO).chunk
222 # Single chunk covering the whole tensor (offsets=0, sizes=shape)
223 else:
224 chunk = ChunkStorageMetadata(
225 offsets=(0,) * len(obj.shape),
226 sizes=obj.shape,
227 )
228 index = MetadataIndex(fqn=fqn, offset=chunk.offsets, index=None)
229 write_item = WriteItem(
230 index=index,
231 type=WriteItemType.TENSOR,
232 tensor_data={
233 'chunk': chunk,
234 'properties': properties,
235 'size': getattr(obj, CHUNK_INFO).global_shape if hasattr(obj, CHUNK_INFO) else obj.shape,
236 }
237 )
238 items.append(write_item)
239 else:
240 # Handle non-tensor types (bytes, etc.)
241 index = MetadataIndex(fqn=fqn)
242 write_item = WriteItem(
243 index=index,
244 type=WriteItemType.BYTE_IO,
245 bytes_io_data=None
246 )
247 items.append(write_item)
249 plan = SavePlan(items=items)
250 if self.flatten_state_dict:
251 plan.planner_data = self.name_mapping
252 return plan
254 def build_global_plan(self, all_plans: list[SavePlan]) -> tuple[list[SavePlan], Metadata]:
255 """
256 Build global plan from all local plans.
258 Collects chunks from all ranks, validates consistency, and creates metadata for the checkpoint.
260 Args:
261 all_plans (list[SavePlan]): List of local plans from all ranks.
263 Returns:
264 tuple[list[SavePlan], Metadata]: Updated plans and checkpoint metadata.
265 """
266 # Deduplicate plans if redundancy removal is enabled
267 if self.remove_redundancy and len(all_plans) > 1:
268 all_plans = remove_redundant_plans(all_plans, save_to_minimum_rank=self.save_to_minimum_rank)
270 # Collect all write items by FQN
271 fqn_to_chunks: dict[str, list[ChunkStorageMetadata]] = {}
272 fqn_to_properties: dict[str, TensorProperties] = {}
273 fqn_to_size: dict[str, tuple] = {}
274 state_dict_metadata: dict[str, Union[TensorStorageMetadata, BytesStorageMetadata]] = {}
276 final_global_plans: list[SavePlan] = []
277 for plan in all_plans:
278 with_index_items = []
279 for item in plan.items:
280 if item.type == WriteItemType.TENSOR and item.tensor_data:
281 fqn = item.index.fqn
282 chunk = item.tensor_data['chunk']
283 properties = item.tensor_data['properties']
284 size = item.tensor_data['size']
286 # Validate consistency across ranks
287 if fqn in fqn_to_chunks and (fqn_to_properties[fqn] != properties or fqn_to_size[fqn] != size):
288 raise ValueError(f"The {fqn} in different rank has different properties and size.")
290 # Initialize FQN entry if not exists
291 if fqn not in fqn_to_chunks:
292 fqn_to_properties[fqn] = properties
293 fqn_to_size[fqn] = size
294 fqn_to_chunks[fqn] = []
296 # Append chunk and set index (platform.Tensor has exactly one chunk)
297 new_index = dataclasses.replace(item.index, index=len(fqn_to_chunks[fqn]))
298 with_index_item = dataclasses.replace(item, index=new_index)
299 with_index_items.append(with_index_item)
300 fqn_to_chunks[fqn].append(chunk)
302 elif item.type == WriteItemType.BYTE_IO:
303 with_index_items.append(item)
304 state_dict_metadata[item.index.fqn] = BytesStorageMetadata()
305 else:
306 raise ValueError(f"Unsupported write item type: {item.type}")
308 final_global_plans.append(dataclasses.replace(plan, items=with_index_items))
310 # Create metadata for all tensors
311 for fqn, chunks in fqn_to_chunks.items():
312 state_dict_metadata[fqn] = TensorStorageMetadata(
313 properties=fqn_to_properties[fqn],
314 size=fqn_to_size[fqn],
315 chunks=chunks
316 )
318 metadata = Metadata(state_dict_metadata=state_dict_metadata)
319 if self.flatten_state_dict:
320 merged_mapping = {}
321 for p in all_plans:
322 merged_mapping.update(p.planner_data)
323 metadata.planner_data = merged_mapping
324 return final_global_plans, metadata
326 def finalize_plan(self, plan: SavePlan) -> SavePlan:
327 """
328 Finalize the plan.
330 Args:
331 plan (SavePlan): Plan to finalize.
333 Returns:
334 SavePlan: Finalized plan.
335 """
336 return plan
338 def get_cached(self) -> Optional[CachedSaveResult]:
339 """Return cached finalized plan and metadata when plan caching is enabled."""
340 if (
341 not self._enable_plan_caching
342 or self._cached_plans_key not in StandardSavePlanner._cached_save_result
343 ):
344 return None
345 return StandardSavePlanner._cached_save_result[self._cached_plans_key]
347 def cache_result(self, final_plan: SavePlan, metadata: Metadata) -> None:
348 """Store finalized plan and metadata in the class-level planner cache."""
349 if not self._enable_plan_caching:
350 return
351 StandardSavePlanner._cached_save_result[self._cached_plans_key] = CachedSaveResult(
352 final_plan=final_plan,
353 metadata=metadata,
354 )
356 def get_data(self, item: WriteItem) -> Any:
357 """
358 Get current runtime data from state_dict for a write item.
360 Args:
361 item (WriteItem): Write item describing what to write.
363 Returns:
364 Any: Runtime object to be written.
365 """
366 if self.state_dict is None:
367 raise RuntimeError("Planner not set up")
368 fqn = item.index.fqn
369 if fqn not in self.state_dict:
370 raise KeyError(f"Key {fqn} not found in state_dict")
371 obj = self.state_dict[fqn]
372 if item.type == WriteItemType.TENSOR:
373 if isinstance(obj, DTensor):
374 if obj.layout is not None and obj.layout.ragged_shard is not None:
375 return get_ragged_box_tensor(obj, item.index).detach().cpu()
376 return obj.to_local().detach().cpu()
377 if isinstance(obj, Tensor):
378 return obj.detach().cpu()
379 raise TypeError(f"Write item {fqn} expected tensor-like object, got {type(obj)}")
380 if item.type == WriteItemType.BYTE_IO:
381 return obj
382 raise TypeError(f"Unsupported write item type: {item.type}")
385def create_read_items_for_chunk_list(
386 fqn: str,
387 checkpoint_md: TensorStorageMetadata,
388 local_chunks: list[ChunkStorageMetadata],
389) -> list[ReadItem]:
390 """
391 Create ReadItems by matching local chunks (what this rank needs) with
392 saved chunks (checkpoint_md.chunks), including resharding overlaps.
394 Mirrors torch create_read_items_for_chunk_list behavior.
396 Args:
397 fqn (str): Fully qualified name of the tensor.
398 checkpoint_md (TensorStorageMetadata): Tensor storage metadata from checkpoint.
399 local_chunks (list[ChunkStorageMetadata]): List of local chunks needed by this rank.
401 Returns:
402 list[ReadItem]: List of ReadItems for loading the required data.
403 """
404 read_items: list[ReadItem] = []
405 saved_chunks = checkpoint_md.chunks
406 if not local_chunks or not saved_chunks:
407 return read_items
409 for local_idx, local_chunk in enumerate(local_chunks):
410 local_area = chunk_to_area(local_chunk)
411 for storage_idx, storage_chunk in enumerate(saved_chunks):
412 saved_area = chunk_to_area(storage_chunk)
413 overlap = infer_intersection(local_area, saved_area)
414 if overlap is None:
415 continue
417 dest_offsets = tuple(overlap[i][0] - local_chunk.offsets[i] for i in range(len(overlap)))
418 storage_offsets = tuple(overlap[i][0] - storage_chunk.offsets[i] for i in range(len(overlap)))
419 lengths = tuple(overlap[i][1] - overlap[i][0] for i in range(len(overlap)))
421 read_items.append(
422 ReadItem(
423 type=LoadItemType.TENSOR,
424 dest_index=MetadataIndex(fqn=fqn, offset=local_chunk.offsets, index=local_idx),
425 dest_offsets=dest_offsets,
426 storage_index=MetadataIndex(fqn=fqn, offset=storage_chunk.offsets, index=storage_idx),
427 storage_offsets=storage_offsets,
428 lengths=lengths,
429 )
430 )
431 return read_items
434class StandardLoadPlanner(LoadPlanner):
435 """
436 Standard implementation of LoadPlanner.
438 Iterate state_dict and creates load plans via chunk list for resharding support.
439 """
441 def __init__(self, allow_partial_load: bool = False):
442 """
443 Args:
444 allow_partial_load (bool): If True, allow loading when checkpoint has fewer keys than state_dict.
445 Default False.
446 """
447 self.state_dict: Optional[dict[str, Any]] = None
448 self.metadata: Optional[Metadata] = None
449 self.is_coordinator: bool = False
450 self.rank: int = 0
451 self.allow_partial_load = allow_partial_load
452 self.flatten_state_dict: bool = True
454 def configure_planner(self, state_dict: dict[str, Any], metadata: Metadata, **kwargs) -> None:
455 """
456 Configure planner with state dict and metadata.
458 Args:
459 state_dict (dict[str, Any]): The state_dict to load into (modified in-place).
460 metadata (Metadata): Checkpoint metadata.
461 **kwargs: Additional keyword arguments (e.g., is_coordinator, rank).
462 """
463 self.state_dict = state_dict
464 self.metadata = metadata
465 self.is_coordinator = kwargs.get("is_coordinator", False)
466 self.rank = kwargs.get("rank", 0)
467 self.flatten_state_dict = kwargs.get("flatten_state_dict", True)
468 self.original_state_dict = state_dict
469 if self.flatten_state_dict:
470 state_dict, self.name_mapping = flatten_state_dict(state_dict)
471 self.state_dict = state_dict
473 def build_local_plan(self) -> LoadPlan:
474 """
475 Build local load plan.
477 Iterate state_dict and creates load plans via chunk list for resharding support.
479 Returns:
480 LoadPlan: Local load plan containing ReadItems for this rank.
481 """
482 if self.state_dict is None or self.metadata is None:
483 raise RuntimeError("Planner not configured")
485 requests: list[ReadItem] = []
486 strict = not self.allow_partial_load
487 for fqn, obj in self.state_dict.items():
488 if fqn not in self.metadata.state_dict_metadata:
489 if fqn.endswith(('matched_adamw_rms', 'step')):
490 continue
491 if strict:
492 raise RuntimeError(f"Missing key in checkpoint state_dict: {fqn}.")
493 continue
494 md = self.metadata.state_dict_metadata[fqn]
495 if isinstance(md, TensorStorageMetadata):
496 obj_size = getattr(obj, CHUNK_INFO).global_shape if hasattr(obj, CHUNK_INFO) \
497 else getattr(obj, "shape", None)
498 if obj_size is None or md.size != tuple(obj_size):
499 raise ValueError(
500 f"Size mismatch between saved {md.size} and current: {obj_size} for {fqn}",
501 )
502 if isinstance(obj, DTensor):
503 layout = getattr(obj, "layout", None)
504 rank_list = getattr(layout, "rank_list", None) if layout else None
505 if rank_list is None and layout is not None:
506 rank_list = getattr(layout, "_rank_list", None)
507 if layout is not None and rank_list is not None:
508 if get_platform().get_rank() not in rank_list:
509 continue
510 # Both DTensor and platform.Tensor: create local chunks and read items
511 local_chunks = create_chunk_list_for_tensor(obj)
512 requests += create_read_items_for_chunk_list(fqn, md, local_chunks)
513 else:
514 requests.append(
515 ReadItem(
516 type=LoadItemType.BYTE_IO,
517 dest_index=MetadataIndex(fqn=fqn),
518 dest_offsets=(0,),
519 storage_index=MetadataIndex(fqn=fqn),
520 storage_offsets=(0,),
521 lengths=(0,),
522 )
523 )
524 return LoadPlan(items=requests)
526 def build_global_plan(self, all_plans: list[LoadPlan]) -> list[LoadPlan]:
527 """
528 Build global plan from all local plans.
530 For now, returns plans as-is. In a more sophisticated implementation, you might need to coordinate across ranks.
532 Args:
533 all_plans (list[LoadPlan]): List of local plans from all ranks.
535 Returns:
536 list[LoadPlan]: Global plans (currently returns plans as-is).
537 """
538 return all_plans
540 def finalize_plan(self, plan: LoadPlan) -> LoadPlan:
541 """
542 Finalize the plan (no-op for default implementation).
544 Args:
545 plan (LoadPlan): Plan to finalize.
547 Returns:
548 LoadPlan: Finalized plan.
549 """
550 return plan
552 def acquire_tensor(self, read_item: ReadItem) -> Any:
553 """
554 Acquire the destination slice (narrow view) for this read_item.
556 StorageReader uses this to copy loaded data into the correct region.
557 Torch-aligned behavior.
559 Args:
560 read_item (ReadItem): The read item specifying what to load.
562 Returns:
563 Any: The destination tensor slice where data should be written
564 (tensor-like object).
565 """
566 if self.state_dict is None:
567 raise RuntimeError("Planner not configured")
569 fqn = read_item.dest_index.fqn
570 if fqn not in self.state_dict:
571 raise KeyError(f"Key {fqn} not found in state_dict")
573 target = self.state_dict[fqn]
574 if (
575 isinstance(target, DTensor)
576 and target.layout is not None
577 and target.layout.ragged_shard is not None
578 ):
579 box_tensor = get_ragged_box_tensor(target, read_item.dest_index)
580 return narrow_tensor_by_index(
581 box_tensor,
582 read_item.dest_offsets,
583 read_item.lengths,
584 )
586 local_tensor = target.to_local().detach() if isinstance(target, DTensor) else target.detach()
587 return narrow_tensor_by_index(
588 local_tensor,
589 read_item.dest_offsets,
590 read_item.lengths,
591 )
593 def apply_tensor(self, read_item: ReadItem, tensor: Any) -> None:
594 """
595 Apply tensor after reading.
597 After read_data copies into the slice, this is no-op when tensor is the
598 same slice. When the backend has no copy_ (e.g. mindspore), read_data
599 passes the loaded slice here; we copy it into the destination slice.
601 Args:
602 read_item (ReadItem): The read item that was processed.
603 tensor (Any): The tensor data to apply (tensor-like object).
604 """
605 if tensor is None:
606 return
607 dest_slice = self.acquire_tensor(read_item)
608 if dest_slice is tensor:
609 return
610 if hasattr(dest_slice, "copy_"):
611 dest_slice.copy_(tensor)
612 else:
613 # Fallback: assign into state_dict if supported
614 dest_slice[...] = tensor
616 def apply_bytes(self, read_item: ReadItem, value: bytes) -> None:
617 """
618 Load bytes data into state_dict.
620 Args:
621 read_item (ReadItem): The read item specifying the destination.
622 value (bytes): The bytes data to deserialize and load.
623 """
624 if self.state_dict is None:
625 raise RuntimeError("Planner not set up")
627 fqn = read_item.dest_index.fqn
628 # Deserialize bytes
629 obj = pickle.loads(value)
630 self.state_dict[fqn] = obj
631 if self.flatten_state_dict:
632 set_element(self.original_state_dict, self.name_mapping[fqn], obj)
636class _DcpMergeLoadPlanner(StandardLoadPlanner):
637 """Load planner that builds distributed checkpoint from dcp into fully ``state_dict`` (in-place)."""
639 def __init__(self) -> None:
640 super().__init__()
642 def configure_planner(self, state_dict: dict[str, Any], metadata: Metadata, **kwargs) -> None:
643 if len(state_dict) > 0:
644 raise ValueError(
645 "state_dict must be empty for _DcpMergeLoadPlanner; "
646 "it is populated in-place from checkpoint metadata."
647 )
649 if metadata is None:
650 raise ValueError("metadata must not be None for _DcpMergeLoadPlanner.")
652 self.is_coordinator = kwargs.get("is_coordinator", False)
653 for k, v in metadata.state_dict_metadata.items():
654 if isinstance(v, TensorStorageMetadata):
655 v = platform.empty(
656 platform.list_to_size(v.size),
657 dtype=platform.str_to_dtype(v.properties.dtype),
658 )
660 state_dict[k] = v
661 if metadata.planner_data is not None and k in metadata.planner_data:
662 set_element(state_dict, metadata.planner_data[k], v)
664 super().configure_planner(
665 state_dict,
666 metadata,
667 is_coordinator=self.is_coordinator,
668 flatten_state_dict=True,
669 )