From 1d733bb75ebd7245d1830f4cc0d5326f41e45d86 Mon Sep 17 00:00:00 2001 From: Rostan Tabet Date: Tue, 24 Feb 2026 08:55:49 +0000 Subject: [PATCH 1/3] Replace dm-tree with optree Signed-off-by: Rostan Tabet --- conda/dali_python_bindings/recipe/meta.yaml | 4 ++-- dali/python/nvidia/dali/_conditionals.py | 24 +++++++++---------- dali/python/nvidia/dali/ops/__init__.py | 4 ++-- dali/python/setup.py.in | 2 +- .../test/python/operator_2/test_enum_types.py | 4 ++-- docker/Dockerfile | 2 +- qa/TL1_custom_src_pattern_build/test.sh | 2 +- 7 files changed, 21 insertions(+), 21 deletions(-) diff --git a/conda/dali_python_bindings/recipe/meta.yaml b/conda/dali_python_bindings/recipe/meta.yaml index 6cacc967034..cb2cf8e669e 100644 --- a/conda/dali_python_bindings/recipe/meta.yaml +++ b/conda/dali_python_bindings/recipe/meta.yaml @@ -84,7 +84,7 @@ requirements: - future - astunparse >=1.6.0 - gast >=0.3.3 - - dm-tree >=0.1.8 + - optree - packaging - nvtx - makefun @@ -100,7 +100,7 @@ requirements: - future - astunparse >=1.6.0 - gast >=0.3.3 - - dm-tree >=0.1.8 + - optree - packaging - nvtx - makefun diff --git a/dali/python/nvidia/dali/_conditionals.py b/dali/python/nvidia/dali/_conditionals.py index d3300e5b81f..57cacb1d2f2 100644 --- a/dali/python/nvidia/dali/_conditionals.py +++ b/dali/python/nvidia/dali/_conditionals.py @@ -43,7 +43,7 @@ from enum import Enum -import tree +import optree def _data_node_repr(data_node): @@ -51,11 +51,11 @@ def _data_node_repr(data_node): def _map_structure(func, *structures, **kwargs): - """Custom wrapper over tree.map_structure that filters it out from the user-visible stack trace + """Custom wrapper over optree.tree_map that filters it out from the user-visible stack trace for error reporting purposes. """ - with _autograph.CustomModuleFilter(tree): - return tree.map_structure(func, *structures, **kwargs) + with _autograph.CustomModuleFilter(optree): + return optree.tree_map(func, *structures, **kwargs) class _Branch(Enum): @@ -610,14 +610,14 @@ def if_stmt(self, cond, body, orelse, get_state, set_state, symbol_names, nouts) " same set of keys, the values may be different.\n" ) - try: - tree.assert_same_structure(body_outputs, orelse_outputs, check_types=True) - except ValueError as e: - # Suppress the original exception, add DALI explanation at the beginning, - # raise the full error message. - raise ValueError(err_msg + str(e)) from None - except TypeError as e: - raise TypeError(err_msg + str(e)) from None + body_structure = optree.tree_structure(body_outputs) + orelse_structure = optree.tree_structure(orelse_outputs) + if body_structure != orelse_structure: + raise ValueError( + f"{err_msg}\n" + f"'If' output structure: {optree.tree_map(lambda _: '*', body_outputs)}\n" + f"'Else' output structure: {optree.tree_map(lambda _: '*', orelse_outputs)}" + ) def merge_branches(new_body_val, new_orelse_val): logging.log( diff --git a/dali/python/nvidia/dali/ops/__init__.py b/dali/python/nvidia/dali/ops/__init__.py index 0cd4882f05e..31ed86faa3e 100644 --- a/dali/python/nvidia/dali/ops/__init__.py +++ b/dali/python/nvidia/dali/ops/__init__.py @@ -15,7 +15,7 @@ # pylint: disable=no-member import sys import threading -import tree +import optree import warnings import weakref from itertools import count @@ -853,7 +853,7 @@ def _promote_scalar_constant(value, input_device): dev = get_input_device(schema, idx) # Process the single ScalarConstant or list possibly containing ScalarConstants # and promote each of them into a DataNode - inp = tree.map_structure(lambda val: _promote_scalar_constant(val, dev), inp) + inp = optree.tree_map(lambda val: _promote_scalar_constant(val, dev), inp) inputs[idx] = inp return inputs diff --git a/dali/python/setup.py.in b/dali/python/setup.py.in index 17f3d71d032..ff67251bb1a 100644 --- a/dali/python/setup.py.in +++ b/dali/python/setup.py.in @@ -89,7 +89,7 @@ For more details please check the # the latest astunparse (1.6.3) doesn't work with any other six than # 1.16 or later on python 3.12 due to import six.moves 'six >= 1.16, <= 1.17', - 'dm-tree <= 0.1.9; python_version>="3.10"', + 'optree', 'packaging <= 25.0', 'numpy', 'nvtx', diff --git a/dali/test/python/operator_2/test_enum_types.py b/dali/test/python/operator_2/test_enum_types.py index d86683975c8..15937ee0574 100644 --- a/dali/test/python/operator_2/test_enum_types.py +++ b/dali/test/python/operator_2/test_enum_types.py @@ -15,7 +15,7 @@ from nvidia.dali import fn, pipeline_def, types import numpy as np -import tree +import optree from nose_utils import assert_raises from nose2.tools import params @@ -31,7 +31,7 @@ lambda value, dtype: types.Constant(value=value), # Explicit type when passed the underlying numeric value of the enum lambda value, dtype: types.Constant( - value=tree.map_structure(lambda v: v.value, value), dtype=dtype + value=optree.tree_map(lambda v: v.value, value), dtype=dtype ), ] ) diff --git a/docker/Dockerfile b/docker/Dockerfile index ada6b5db333..e8cc930cef8 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -50,7 +50,7 @@ ${LIBRARY_PATH} RUN ln -s /opt/python/cp${PYV}* /opt/python/v # install Python bindings and patch it to use the clang we have here -RUN pip install future setuptools wheel clang==14.0 flake8 bandit astunparse gast dm-tree "black[jupyter]"==25.12.0 nvtx makefun && \ +RUN pip install future setuptools wheel clang==14.0 flake8 bandit astunparse gast optree "black[jupyter]"==25.12.0 nvtx makefun && \ PY_CLANG_PATH=$(echo $(pip show clang) | sed 's/.*Location: \(.*\) Requires.*/\1/')/clang/cindex.py && \ LIBCLANG_PATH=/usr/lib64/libclang.so && \ sed -i "s|library_file = None|library_file = \"${LIBCLANG_PATH}\"|" ${PY_CLANG_PATH} && \ diff --git a/qa/TL1_custom_src_pattern_build/test.sh b/qa/TL1_custom_src_pattern_build/test.sh index 4fac4a04e62..4668a9dda6f 100644 --- a/qa/TL1_custom_src_pattern_build/test.sh +++ b/qa/TL1_custom_src_pattern_build/test.sh @@ -1,6 +1,6 @@ #!/bin/bash -e -pip_packages='astunparse gast dm-tree black nvtx makefun' +pip_packages='astunparse gast optree black nvtx makefun' build_and_check() { make -j From 2b3ed39295978c53848e777d32810c68ac0411f9 Mon Sep 17 00:00:00 2001 From: Rostan Tabet Date: Tue, 24 Feb 2026 11:48:34 +0000 Subject: [PATCH 2/3] Install NumPy in the Dockerfile. It was previously transitively included by dm-tree. Signed-off-by: Rostan Tabet --- docker/Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index e8cc930cef8..90484779135 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -50,7 +50,7 @@ ${LIBRARY_PATH} RUN ln -s /opt/python/cp${PYV}* /opt/python/v # install Python bindings and patch it to use the clang we have here -RUN pip install future setuptools wheel clang==14.0 flake8 bandit astunparse gast optree "black[jupyter]"==25.12.0 nvtx makefun && \ +RUN pip install future setuptools wheel clang==14.0 flake8 bandit astunparse gast optree numpy "black[jupyter]"==25.12.0 nvtx makefun && \ PY_CLANG_PATH=$(echo $(pip show clang) | sed 's/.*Location: \(.*\) Requires.*/\1/')/clang/cindex.py && \ LIBCLANG_PATH=/usr/lib64/libclang.so && \ sed -i "s|library_file = None|library_file = \"${LIBCLANG_PATH}\"|" ${PY_CLANG_PATH} && \ From 8aaeb811dcfb028f63d527fbc17a2810a23f3f0a Mon Sep 17 00:00:00 2001 From: Rostan Tabet Date: Wed, 25 Feb 2026 09:39:40 +0000 Subject: [PATCH 3/3] Update error message checks in conditionals error test Signed-off-by: Rostan Tabet --- dali/test/python/conditionals/test_nests.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/dali/test/python/conditionals/test_nests.py b/dali/test/python/conditionals/test_nests.py index 1bf5c7c9ec1..f28ff0fd645 100644 --- a/dali/test/python/conditionals/test_nests.py +++ b/dali/test/python/conditionals/test_nests.py @@ -146,12 +146,9 @@ def pipeline(): "*Divergent data found in different branches of `if/else` control" " flow statement. Variables in all code paths are merged into common" " output batches. The values assigned to a given variable need to" - " have the same nesting structure in every code path" - " (both `if` branches).*" - "*The two structures don't have the same nested structure*" - "*The two dictionaries don't have the same set of keys." - " First structure has keys type=list str=*'out', 'mismatched'*," - " while second structure has keys type=list str=*'out'*" + " have the same nesting structure in every code path (both `if` branches).*" + "'If' output structure:*" + "'Else' output structure:*" ), ): _ = pipeline()