feat: advanced pytree support (#40)

Author: mrapplexz

d9d/core/pytree/__init__.py

@@ -0,0 +1,21 @@
+"""Recursive traversal of nested tensor structures ("pytrees")."""
+
+from .ops import (
+    PyTreeSpec,
+    tree_flatten,
+    tree_leaves,
+    tree_leaves_with_path,
+    tree_map,
+    tree_map_only,
+    tree_unflatten,
+)
+
+__all__ = [
+    "PyTreeSpec",
+    "tree_flatten",
+    "tree_leaves",
+    "tree_leaves_with_path",
+    "tree_map",
+    "tree_map_only",
+    "tree_unflatten",
+]

d9d/core/pytree/flatten.py

@@ -0,0 +1,90 @@
+import dataclasses
+from collections.abc import Callable
+from typing import Any
+
+import optree
+import optree.dataclasses
+
+_NAMESPACE = "d9d"
+
+IsLeaf = Callable[[Any], bool]
+
+
+class PyTreeFlattener:
+    """Flattens pytrees into their leaves, descending into dataclasses as internal nodes."""
+
+    def __init__(self):
+        """Constructs a PyTreeFlattener."""
+        self._registered: set[type] = set()
+        self._is_dataclass_memo: dict[type, bool] = {}
+
+    def flatten(self, tree: Any, is_leaf: IsLeaf | None = None) -> tuple[list[Any], optree.PyTreeSpec]:
+        """Flattens ``tree``, lazily registering any dataclass types discovered as leaves.
+
+        Args:
+            tree: The nested structure to flatten.
+            is_leaf: Optional predicate; when it returns ``True`` for a node, that node is treated as
+                a leaf and not traversed further (even if it is a dataclass or container).
+
+        Returns:
+            A tuple of the leaf list and the ``optree`` structure specification.
+        """
+        while True:
+            leaves, treespec = optree.tree_flatten(tree, is_leaf=is_leaf, namespace=_NAMESPACE)
+            if self._register_unknown_nodes(leaves, is_leaf):
+                return leaves, treespec
+
+    def flatten_with_path(
+        self, tree: Any, is_leaf: IsLeaf | None = None
+    ) -> tuple[list[tuple[Any, ...]], list[Any], optree.PyTreeSpec]:
+        """Flattens ``tree`` alongside the path to each leaf, lazily registering dataclass types.
+
+        Args:
+            tree: The nested structure to flatten.
+            is_leaf: Optional predicate; when it returns ``True`` for a node, that node is treated as
+                a leaf and not traversed further (even if it is a dataclass or container).
+
+        Returns:
+            A tuple of the per-leaf paths, the leaf list, and the ``optree`` structure specification.
+        """
+        while True:
+            paths, leaves, treespec = optree.tree_flatten_with_path(tree, is_leaf=is_leaf, namespace=_NAMESPACE)
+            if self._register_unknown_nodes(leaves, is_leaf):
+                return paths, leaves, treespec
+
+    def _register_unknown_nodes(self, leaves: list[Any], is_leaf: IsLeaf | None) -> bool:
+        """Registers any not-yet-registered dataclass types among ``leaves``.
+
+        Args:
+            leaves: The leaves produced by a flatten pass.
+            is_leaf: The predicate passed to the flatten call, if any.
+
+        Returns:
+            ``True`` if all leaves were already fully expanded (nothing to register), else ``False``.
+        """
+        unregistered = {
+            type(leaf)
+            for leaf in leaves
+            if type(leaf) not in self._registered
+            and self._is_dataclass_type(type(leaf))
+            and not (is_leaf is not None and is_leaf(leaf))
+        }
+        if not unregistered:
+            return True
+
+        for cls in unregistered:
+            self._register(cls)
+        return False
+
+    def _is_dataclass_type(self, tp: type) -> bool:
+        cached = self._is_dataclass_memo.get(tp)
+        if cached is None:
+            cached = dataclasses.is_dataclass(tp)
+            self._is_dataclass_memo[tp] = cached
+        return cached
+
+    def _register(self, cls: type) -> None:
+        if cls in self._registered:
+            return
+        optree.dataclasses.register_node(cls, namespace=_NAMESPACE)
+        self._registered.add(cls)

d9d/core/pytree/ops.py

@@ -0,0 +1,116 @@
+from collections.abc import Callable
+from typing import Any, TypeVar, cast
+
+import optree
+
+from d9d.core.types import PyTree
+
+from .flatten import IsLeaf, PyTreeFlattener
+
+TLeaf = TypeVar("TLeaf")
+TMapped = TypeVar("TMapped")
+TTree = TypeVar("TTree", bound=PyTree)
+
+PyTreeSpec = optree.PyTreeSpec
+
+
+_flattener = PyTreeFlattener()
+
+
+def tree_flatten(tree: PyTree[TLeaf], is_leaf: IsLeaf | None = None) -> tuple[list[TLeaf], PyTreeSpec]:
+    """Flattens a pytree into its leaves and a structure specification.
+
+    Args:
+        tree: The nested structure to flatten.
+        is_leaf: Optional predicate; when it returns ``True`` for a node, that node is kept as a leaf
+            and not traversed further.
+
+    Returns:
+        A tuple of the leaf list and a ``PyTreeSpec`` that can rebuild the structure via
+        `tree_unflatten`.
+    """
+    return _flattener.flatten(tree, is_leaf)
+
+
+def tree_unflatten(treespec: PyTreeSpec, leaves: list[TLeaf]) -> PyTree[TLeaf]:
+    """Reconstructs a pytree from leaves and a structure specification.
+
+    Args:
+        treespec: A specification produced by `tree_flatten`.
+        leaves: The leaves to place into the structure, in flatten order.
+
+    Returns:
+        The reconstructed nested structure.
+    """
+    return optree.tree_unflatten(treespec, leaves)
+
+
+def tree_leaves(tree: PyTree[TLeaf], is_leaf: IsLeaf | None = None) -> list[TLeaf]:
+    """Returns the leaves of a pytree in deterministic (sorted-key) order.
+
+    Args:
+        tree: The nested structure to flatten.
+        is_leaf: Optional predicate; when it returns ``True`` for a node, that node is kept as a leaf
+            and not traversed further.
+
+    Returns:
+        The list of leaves.
+    """
+    return _flattener.flatten(tree, is_leaf)[0]
+
+
+def tree_map(func: Callable[[TLeaf], TMapped], tree: PyTree[TLeaf]) -> PyTree[TMapped]:
+    """Applies ``func`` to every leaf of a pytree, returning a structurally-identical tree.
+
+    Args:
+        func: The function to apply to each leaf.
+        tree: The nested structure to map over.
+
+    Returns:
+        A new tree with ``func`` applied to each leaf.
+    """
+    leaves, treespec = _flattener.flatten(tree)
+    return optree.tree_unflatten(treespec, [func(leaf) for leaf in leaves])
+
+
+def tree_map_only(
+    filters: type | tuple[type, ...],
+    func: Callable[[Any], Any],
+    tree: TTree,
+) -> TTree:
+    """Applies ``func`` only to leaves that are instances of ``type_or_types``.
+
+    Leaves of any other type are returned unchanged. This is the common case for tensor
+    operations over trees that also carry non-tensor bookkeeping (e.g. moving only tensors to a
+    device while leaving strings and ints alone). The returned tree preserves the structure and
+    leaf types of the input.
+
+    Args:
+        filters: The leaf type(s) that ``func`` should be applied to.
+        func: The function to apply to matching leaves.
+        tree: The nested structure to map over.
+
+    Returns:
+        A new tree with ``func`` applied to matching leaves only.
+    """
+    leaves, treespec = _flattener.flatten(tree)
+    mapped = [func(leaf) if isinstance(leaf, filters) else leaf for leaf in leaves]
+    return cast(TTree, optree.tree_unflatten(treespec, mapped))
+
+
+def tree_leaves_with_path(tree: PyTree[TLeaf], is_leaf: IsLeaf | None = None) -> list[tuple[tuple[Any, ...], TLeaf]]:
+    """Returns ``(path, leaf)`` pairs for every leaf of a pytree.
+
+    Each path is a tuple of keys and indices reaching the leaf from the root: ``str`` for dict keys
+    and dataclass field names, ``int`` for sequence indices.
+
+    Args:
+        tree: The nested structure to flatten.
+        is_leaf: Optional predicate; when it returns ``True`` for a node, that node is kept as a leaf
+            and not traversed further.
+
+    Returns:
+        A list of ``(path, leaf)`` tuples in deterministic (sorted-key) order.
+    """
+    paths, leaves, _ = _flattener.flatten_with_path(cast(Any, tree), is_leaf)
+    return list(zip(paths, leaves, strict=True))

d9d/internals/metric_collector/collector.py

@@ -1,7 +1,7 @@
 import torch
-import torch.utils._pytree as pytree  # noqa: PLC2701
 from torch.profiler import record_function
 
+from d9d.core import pytree
 from d9d.core.dist_context import DistributedContext
 from d9d.core.types import PyTree
 from d9d.metric import Metric

d9d/loop/component/job_logger.py

@@ -3,9 +3,9 @@
 from typing import Any
 
 import torch
-import torch.utils._pytree as pytree  # noqa: PLC2701
 from torch.distributed.checkpoint.stateful import Stateful
 
+from d9d.core import pytree
 from d9d.core.dist_context import DistributedContext
 from d9d.core.types import PyTree, ScalarTree
 from d9d.internals.metric_collector import AsyncMetricCollector
@@ -22,20 +22,7 @@ def _flatten_pytree_for_metrics(tree: PyTree[float]) -> dict[str, float]:
     flat_dict = {}
 
     for path_tuple, value in pytree.tree_leaves_with_path(tree):
-        path_segments = []
-
-        for key in path_tuple:
-            match key:
-                case pytree.MappingKey(k):
-                    path_segments.append(str(k))
-                case pytree.SequenceKey(idx):
-                    path_segments.append(str(idx))
-                case pytree.GetAttrKey(name):
-                    path_segments.append(name)
-                case _:
-                    path_segments.append(str(key))
-
-        flat_key = "/".join(path_segments)
+        flat_key = "/".join(str(segment) for segment in path_tuple)
         flat_dict[flat_key] = value
 
     return flat_dict

d9d/loop/component/pipeline_state.py

@@ -1,30 +1,14 @@
 from collections.abc import Iterator
 from contextlib import contextmanager
-from typing import Generic, TypeVar, cast
+from typing import Generic, TypeVar
 
 import torch
-import torch.utils._pytree as pytree  # noqa: PLC2701
 
+from d9d.core import pytree
 from d9d.core.types import PyTree
 
 TState = TypeVar("TState", bound=PyTree)
 
-TMap = TypeVar("TMap")
-
-
-def _detach_leaf(x: TMap) -> TMap:
-    """Detaches a tensor from the computation graph if the input is a tensor.
-
-    Args:
-        x: The input object.
-
-    Returns:
-        The detached tensor or original object.
-    """
-    if isinstance(x, torch.Tensor):
-        return cast(TMap, x.detach())
-    return x
-
 
 class PipelineStateHandler(Generic[TState]):
     """Holds the transient per-microbatch side-data of one step.
@@ -46,7 +30,7 @@ def store(self, microbatch_idx: int, state: TState):
             microbatch_idx: The index of the microbatch within the current pack.
             state: The side-data PyTree to store; every tensor leaf is detached.
         """
-        self._state[microbatch_idx] = pytree.tree_map(_detach_leaf, state)
+        self._state[microbatch_idx] = pytree.tree_map_only(torch.Tensor, lambda x: x.detach(), state)
 
     @contextmanager
     def scope(self, microbatch_idx: int) -> Iterator[TState]:
@@ -66,7 +50,7 @@ def scope(self, microbatch_idx: int) -> Iterator[TState]:
         try:
             yield state
         finally:
-            self._state[microbatch_idx] = pytree.tree_map(_detach_leaf, state)
+            self._state[microbatch_idx] = pytree.tree_map_only(torch.Tensor, lambda x: x.detach(), state)
 
     def reset(self):
         """Resets the underlying storage, clearing all state."""

d9d/loop/run/_device.py

@@ -1,6 +1,6 @@
 import torch
-import torch.utils._pytree as pytree  # noqa: PLC2701
 
+from d9d.core import pytree
 from d9d.core.types import MicrobatchPack
 
 
@@ -14,6 +14,4 @@ def move_pack_to_device(pack: MicrobatchPack, device: torch.types.Device) -> Mic
     Returns:
         A new pack with all tensors moved to the device.
     """
-    return [
-        pytree.tree_map(lambda x: x.to(device) if isinstance(x, torch.Tensor) else x, microbatch) for microbatch in pack
-    ]
+    return [pytree.tree_map_only(torch.Tensor, lambda x: x.to(device), microbatch) for microbatch in pack]