Add pluggable backend dispatch system - #1151
Conversation
for more information, see https://pre-commit.ci
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1151 +/- ##
==========================================
+ Coverage 77.24% 77.42% +0.18%
==========================================
Files 63 64 +1
Lines 9378 9455 +77
Branches 1579 1585 +6
==========================================
+ Hits 7244 7321 +77
Misses 1532 1532
Partials 602 602
🚀 New features to boost your workflow:
|
| # Trusted but not installed | ||
| if canonical in TRUSTED_BACKENDS and get_backend(canonical) is None: | ||
| package = TRUSTED_BACKENDS[canonical]["package"] | ||
| raise ValueError( |
| # conformance test suite (squidpy.testing.backend_conformance). | ||
| TRUSTED_BACKENDS: dict[str, dict[str, Any]] = { | ||
| "rapids_singlecell": { | ||
| "aliases": ["rapids-singlecell", "rsc", "cuda", "gpu"], |
There was a problem hiding this comment.
why have multiple? maybe just have a error message helper that recognizes them and does “backend 'gpu' does not exist, did you mean 'rapids-singlecell'?”
There was a problem hiding this comment.
didn’t @selmanozleyen already build something like this? this looks like it’s similar code, so if the other version is merged, this should be unified with that.
There was a problem hiding this comment.
@selmanozleyen approach is similar but less extensible. It's still a draft PR like this one. This approach would immediately open the door for other backends. With 0 updates needed in squidpy.
| return wrapper | ||
|
|
||
|
|
||
| def _update_signatures() -> None: |
There was a problem hiding this comment.
This private function is not used in this file but it is exported. Therefore it shouldn't be private. The module _dispatch is already private so it is fine to remove the _ here
| def _get_param_sets( | ||
| func: Callable, | ||
| adapter_method: Callable, | ||
| func_name: str, |
There was a problem hiding this comment.
to be more clear func_name is not being used here
|
compared to #1093 I agree that renaming it to However, currently there is a blocker because the |
| Called once automatically after backend discovery so that ``help()`` / | ||
| IDE tooltips show the full parameter list (CPU + GPU + backend) with | ||
| documentation. |
There was a problem hiding this comment.
I am not sure about dynamically updating the signature and docstrings. I like the approach with just linking to the dispatched backend better in :#1093
@flying-sheep @ilan-gold wdyt? For example if I compiled the docs with gpu will it compile a mix of rsc and squidpy docs?
There was a problem hiding this comment.
I would ideally want to expose everything, even if dispatched to another function in the public function signature
| _update_signatures() | ||
|
|
||
|
|
||
| def _check_trusted(name: str) -> None: |
There was a problem hiding this comment.
again, should be check_trusted
I think But I think the right answer is to do what the rest of the pydata ecosystem does and just follow these patterns. I don't see any reason for us to stick out. Generally, I hope that a generalized version of this could be implemented into https://github.com/scverse/scverse-misc so that it would eventually be reused across all of our packages trivially. Edit: I just read on zulip:
Yeah SGTM but maybe |
But this way we can leave it to the user to either set the context with whatever backend of their choice. If we had jax backend for example we could do it either in cpu or gpu. |
| effective = local_backend or settings.backend | ||
|
|
||
| if effective == "cpu": | ||
| return func(*args, **kwargs) |
There was a problem hiding this comment.
This has a subtle bug if we want to generalize. We only bind *args by order and don't ever check if the order and the names match.
For example:
CPU: func(a, b)
GPU: func(b, a)
and both a and b are integers:
func(5, 10, backend="gpu")
becomes:
GPU sees b=5, a=10
If the backend computes something like a / b, you silently get 10 / 5 instead of 5 / 10. No exception, just incorrect output.
or imagine this
if you have
CPU: threshold, n_perms
GPU: n_perms, threshold
these won't error at the level they should. Maybe it will say n_perms should be an integer inside the function but it shouldn't have been dispatchable to begin with.
|
@Zethson We should also document the contract here. There are lots of assumptions hidden in this PR if we want to generalize. For example currently the dispatch is position based for args, see the review I did, if it's not addressed and generalized it will introduce subtle bugs. To expose these assumptions ideally we should have a For two functions
shared = [name for name in base_args if name in backend_arg_names]
backend_shared = [name for name in backend_args if name in base_arg_names]
# we want
shared == backend_sharedThese are the current terms for the contract. But instead of update-the-signature trick we can disallow |
4 participants
Summary
scverse-backends>=0.0.3for Squidpy backend dispatch.rapids-singlecell, withcuda,rapids, andrapids_singlecellaliases.spatial_autocorr,co_occurrence,ligrec, andcalculate_niche.backendusage forspatial_autocorr,co_occurrence, andligrecthrough explicit compatibility handling.scverse_backends.testing.run_conformancewith Squidpy-owned CPU-vs-backend checks.Backend Contract
Backends register an adapter under the
squidpy.backendsentry-point group. The companion RAPIDS SingleCell change supplies that adapter; this host PR keeps the dispatcher configuration and aliases identical to the Scanpy integration.Verification
uv run --group test pytest -q tests/test_backends.pyuv run --group test pytest -q tests/graph/test_ppatterns.pyuv run --group test pytest -q tests/graph/test_ligrec.py -k 'parallel_backend or legacy_parallel_backend'uv run ruff check ...