Simplify type reflection implementation #7539

jcrist · 2025-11-26T03:10:49Z

This is a fairly substantial rewrite/refactor of our existing type reflection system. There should be no user-visible changes from this refactor (though in the future we do want to make changes). For now this is just trying to simplify the internals so the existing system is easier to understand, reason about, and modify.

Moved (almost) all logic to a single file cuml.internals.outputs instead of being strewn around 5+ files.
Removed all the contextmanagers in api_context_managers.py in favor a simpler, more readable mechanisms
Removed all the decorators in api_decorators.py in favor of a single reflect decorator with sane defaults and only a few configurable knobs
Removed set_api_output_type; this feature was unnecessary, the reflect decorator can handle everything without an escape hatch.
Reduced state management for type reflection decisions down to 3 places (a combination of GlobalSettings().output_type, Base.output_type, and an array input type, depending on the call). The decision around what output type to return is now entirely in one location, and the conversion is also encompassed within a single function. This should hopefully be much easier to understand.
Removed the auto-decorating Base metaclass in favor of explicit decorators. This was done by logging the original auto-decorated versions, then inspecting each one when adding explicit versions to ensure they were accurate. Not everything that was decorated before needed to be decorated.
Removed decorators on functions that don't need them. This is mostly functions that return non-arrays and don't make any nested calls requiring a CumlArray output
Fixed a few decorators that weren't applied properly (e.g. LinearRegression.predict). These are bugfixes.

Once this is in, we should have an easier time making behavior changes and deprecating features (#7426) since the new implementation is simpler and has fewer moving pieces.

Fixes #5022.

jcrist · 2025-11-26T03:16:26Z

This still needs:

Docstring cleanups in all public functions/methods. The existing docstrings from before had bitrotted, and I want to improve them generally as developer docs. Same goes for Base
Another pass through for readability/comments where needed.
New tests to better specify the behavior. The existing tests still pass, but aren't good or sufficient, I plan to replace them with more thorough tests.
Probably some bugs have leaked through, pushed this up mostly to get a full CI run.

I'll ping when ready for review.

I recognize this is a large PR; most of the meat of it occurs before d6d8d21. The latter commits touch many many files, but are mostly mechanical changes applying explicit decorators where previously they were magically applied by BaseMetaClass.

If you pick through the changes in order they start off with fairly reasonable and small simplifications as I picked apart the existing complexity. Once things got simple enough to understand there are then a few more meaty commits that redo the implementation without the need for such complex contextmanagers/decorators.

These had no functional difference with `api_return_array`/`api_base_return_array`. 2 fewer decorators to understand.

- Simplify `ProcessReturn` from 3 classes into one - Remove parametrizing the context manager by return processing - Remove `__class_getitem__` usage entirely in favor of defining `ProcessEnter_Type` explicitly on subclasses At this point, return handling for reflection is the same everywhere, and the only switch is whether to enable it or not.

No need for second class at all at this point.

Some further trivial simplifications

These are now redundant and were aliases to their array counterparts. Removing lets us simplify further.

Simplify, just move all logic to InternalContextAPIBase. Terrible name, but we now have fewer moving pieces. All call logic is still the same as before. Some of the branches don't make sense to me, but we're not changing the logic, just the plumbing.

These should only be used on fits, but any fit should do all of them.

This can be easily inferred. Further simplifying.

Old implementation was scattered between too many files. Also deletes some dead code.

Further simplifications - you can actually follow how output type is inferred and propagated now! Now backdoor state.

jcrist

Annotated the diff for anything interesting/worth calling out.

Files without comments should contain uninteresting mechanical changes like swapping out decorator names (unless I missed something).

jcrist · 2025-11-26T22:49:21Z

python/cuml/cuml/_thirdparty/sklearn/preprocessing/_discretization.py


    """
-
-    bin_edges_internal_ = CumlArrayDescriptor()


Before bin_edges_ existed but would reflect incorrectly since it contains a numpy array of cupy arrays (upstream sklearn does a numpy array of arrays 🤷, the port here made minimal changes so that's what you get). We now move it to be a non-reflected attribute. I consider the previous behavior to be a bug since it wasn't working correctly.

jcrist · 2025-11-26T22:50:26Z

python/cuml/cuml/cluster/hdbscan/hdbscan.pyx

        raise ValueError("batch_size must be > 0")

-    # Reflect the output type from global settings or the clusterer
-    cuml.internals.set_api_output_type(clusterer._get_output_type())


The reflect decorator now handles cases like this automatically.

jcrist · 2025-11-26T22:50:44Z

python/cuml/cuml/datasets/arima.pyx

-    # Set the default output type to "cupy". This will be ignored if the user
-    # has set `cuml.global_settings.output_type`. Only necessary for array
-    # generation methods that do not take an array as input
-    cuml.internals.set_api_output_type("cupy")


The reflect decorator now handles cases like this automatically.

jcrist · 2025-11-26T22:52:18Z

python/cuml/cuml/internals/array.py

        return array_to_memory_order(arr) is not None


+def cuda_ptr(X):


This was in memory_utils.py, moved it here since that file went away.

jcrist · 2025-11-26T22:56:03Z

python/cuml/cuml/internals/global_settings.py

        self.shared_state = {
-            "root_cm": None,
            "_output_type": None,
+            "_external_output_type": False,


State on GlobalSettings is:

_output_type: None or a valid output type string reflecting the current output type. Can also be "mirror" (this was term in the old implementation), which effectively means "return CumlArray values except for CumlArrayDescriptor where instead the original value set on the descriptor is returned". This could/should probably be merged with output_type="cuml" except there are too many places (mostly in _thirdparty where a non-CumlArray value is set. If we decide to simplify this and merge the terms this will have to be done in a followup.

_external_output_type: The originally configured output type from outside the internal API context, or False if not running in an internal API.

jcrist · 2025-11-26T23:08:39Z

python/cuml/tests/test_array.py

 }


+def determine_array_memtype(X):


This used to be in cuml.internals.memory_utils, but was only used in this test file. Moved it here instead.

Same with the other change below - that was inlining a utility function from cuml.internals.memory_utils that was only used in that single test.

jcrist · 2025-11-26T23:09:55Z

python/cuml/tests/test_cuml_descr_decor.py

@@ -1,317 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.


The tests in this file passed fine in the new system (if you rollback a few commits you can run them).

I migrated the intentions (but not necessarily the implementation) to a more complete set of tests in test_reflection.py.

jcrist · 2025-11-26T23:10:36Z

python/cuml/tests/test_module_config.py

@@ -1,130 +0,0 @@
-#


The tests in this file passed fine in the new system (if you rollback a few commits you can run them).

I migrated the intentions (but not necessarily the implementation) to a more complete set of tests in test_reflection.py.

jcrist · 2025-11-26T23:11:54Z

docs/source/api.rst

- .. autofunction:: cuml.internals.memory_utils.using_output_type
+.. autofunction:: cuml.set_global_output_type
+
+.. autofunction:: cuml.using_output_type


These functions were only ever demonstrated from the top-level namespace, but did indeed contain the old filename in the full path.

If necessary we can add a shim cuml.internals.memory_utils file to keep around the old import paths for a deprecation cycle. I doubt it's needed, but 🤷.

jcrist · 2025-11-26T23:12:45Z

python/cuml/tests/test_reflection.py

@@ -0,0 +1,383 @@
+# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.


This file contains the bulk of the reflection tests. It's also worth reviewing.

jcrist · 2025-11-26T23:16:05Z

Gah, everything was passing last night, but then I rebased before tossing in a new test file. I assume something snuck in in the rebase (or maybe an upstream C++ dep issue, looking at the failures), I'll take another look next week after the holiday.

This was now dead code.

This was fully dead code

jcrist self-assigned this Nov 26, 2025

jcrist requested a review from a team as a code owner November 26, 2025 03:10

jcrist requested a review from divyegala November 26, 2025 03:10

jcrist added Tech Debt Issues related to debt improvement Improvement / enhancement to an existing function non-breaking Non-breaking change labels Nov 26, 2025

github-actions bot added the Cython / Python Cython or Python issue label Nov 26, 2025

jcrist requested review from csadorf and removed request for divyegala November 26, 2025 03:11

jcrist added 20 commits November 26, 2025 14:43

Remove lingering cupy allocator setting

a004b9c

Remove in_internal_api (move forward)

d87b398

Remove api_return_sparse_array/api_base_return_sparse_array

1c8521b

These had no functional difference with `api_return_array`/`api_base_return_array`. 2 fewer decorators to understand.

Drop some trivial contextmanager subclasses

db07390

Merge ProcessReturn into InternalAPIContextBase

3940964

No need for second class at all at this point.

Remove duplicate *ReturnAny, remove mixin class

f14f616

Some further trivial simplifications

Remove *generic decorators

1697a47

These are now redundant and were aliases to their array counterparts. Removing lets us simplify further.

Rip out ProcessEnter classes

d59849a

Simplify, just move all logic to InternalContextAPIBase. Terrible name, but we now have fewer moving pieces. All call logic is still the same as before. Some of the branches don't make sense to me, but we're not changing the logic, just the plumbing.

Couple set_n_features_in to set_output_type

ed547db

These should only be used on fits, but any fit should do all of them.

Remove needs_self

dd1b90a

This can be easily inferred. Further simplifying.

Add reflect, simplify frontend

cd48af1

Prepare to rip out set_api_output_type

11df678

Rip out last external ref to root_cm

f7b2c43

Reorg output handling into cuml.internals.outputs

c6ddca8

Old implementation was scattered between too many files. Also deletes some dead code.

Simplify and unify output_type validation

ddfb630

Remove root_cm, api_context_managers

4fc6645

Further simplifications - you can actually follow how output type is inferred and propagated now! Now backdoor state.

Simplify skip=True case

1353aeb

Explicit decorators

f3f968e

Explicit decorators for MG models

65df8e4

jcrist added 5 commits November 26, 2025 14:43

Remove base metaclass and old decorator names

2b3929c

Update doc refs

df7799f

Improve docstrings

b935d76

Rename test_module_config.py -> test_reflection.py

d58ac80

Consolidate and expand reflection tests

a44e6df

jcrist force-pushed the simplify-reflection branch from d07722d to a44e6df Compare November 26, 2025 20:43

Fixup for SpectralClustering change that snuck in on rebase

fd6bddd

jcrist commented Nov 26, 2025

View reviewed changes

jcrist added 4 commits November 26, 2025 17:53

Merge branch 'main' into simplify-reflection

a823909

Merge branch 'main' into simplify-reflection

ecda0c8

Remove type_utils

6e78370

This was now dead code.

Remove base_return_types

2e13e1d

This was fully dead code

		return array_to_memory_order(arr) is not None


		def cuda_ptr(X):

		@@ -1,317 +0,0 @@
		# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.

		@@ -0,0 +1,383 @@
		# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.

Simplify type reflection implementation #7539

Are you sure you want to change the base?

Simplify type reflection implementation #7539

Uh oh!

Conversation

jcrist commented Nov 26, 2025

Uh oh!

jcrist commented Nov 26, 2025

Uh oh!

jcrist left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jcrist commented Nov 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

jcrist commented Nov 26, 2025 •

edited

Loading