Fix numba JIT configuration failing on airflow runner#1117
Merged
LucaMarconato merged 21 commits intomainfrom May 6, 2026
Merged
Fix numba JIT configuration failing on airflow runner#1117LucaMarconato merged 21 commits intomainfrom
LucaMarconato merged 21 commits intomainfrom
Conversation
Cache `blobs(256, 300, 3)` and `BlobsDataset()._labels_blobs()` once per session via private session-scoped fixtures, then deepcopy into each function-scoped fixture. Cuts fixture setup from 44.8s to 35.0s (-9.8s) and total suite from 186s to 180s. All 1332 tests still pass. Benchmark CSVs committed for reference (pytest_*.csv). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
to_circles() on labels scales linearly with pixel count. Dropping from 512×512 to 128×128 cuts test_labels_2d_to_circles from ~3.9s to ~1.0s per parametrized variant (−4.7s across the file). Updated hardcoded coordinate/radius assertions to match the new size. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
… deepcopy Two orthogonal wins: 1. _elements.py: `get_model()` already calls `schema.validate()` internally; the explicit second validate() + get_axes_names() call in every __setitem__ was redundant. Removing it halves the DataTree (_to_dataset_view) overhead per element insertion — directly speeds up fixture setup. 2. conftest.py: introduce `_fast_deepcopy_sdata` (copy.deepcopy + manual attrs restoration for DaskDataFrame/#503 and GeoDataFrame/#286) that is ~13x faster than sd_deepcopy (7ms vs 93ms for full_sdata). Session-scope full_sdata, images, labels and the 'full' sdata parametrized fixture; each test gets a fresh 7ms copy instead of an 87ms full reconstruction. Also switch sdata_blobs from sd_deepcopy to fast_deepcopy_sdata (2ms vs 25ms). Full suite: 186s → 163s (~12% reduction, ~23s saved). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…ion comment - get_model() now accepts validate=False to skip schema.validate() when the caller only needs to infer the element type without re-running validation - add comment to Elements.__setitem__ noting that subclass overrides call get_model() which performs the validation - drop Python 3.13 CI matrix entries (superseded by 3.14) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
… elements Add skip_element_validation() context manager (backed by a ContextVar) that makes __setitem__ call get_model(validate=False) — type inference only, no schema.validate(). Use it in every code path that constructs a SpatialData from elements that originated from an existing SpatialData and were never externally mutated: bounding_box_query, polygon_query, query_by_coordinate_system, transform_to_coordinate_system, subset, and init_from_elements. test_query_spatial_data: 0.77s → 0.64s (the remaining time is the query work itself — filtering, shapely ops, raster cropping). Also inline a minimal 2-image SpatialData in test_transformations_between_coordinate_systems instead of relying on the full 8-element images fixture; the test only ever uses image2d and image2d_multiscale, so writing the other 6 to disk was pure waste. test_transformations_between_coordinate_systems: 0.61s → 0.44s. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…dy-valid elements" This reverts commit acb25b1.
Replace per-slice O(H+W) approach (512 dask compute() calls for a 256×256 array) with a single array materialization + np.bincount O(n_pixels) pass. This speeds up get_centroids() on labels from ~1.5s to ~50ms, cutting to_circles(labels) from ~1.6s to ~53ms. Affects test_validation dataloader variants (~2.5s → ~0.2s each, saving ~9s), test_labels_2d_to_circles, and any production call to get_centroids or to_circles on label arrays. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Add blobs_factory fixture that returns fast deepcopies of the session-scoped blobs dataset, and update test_concatenate_* and test_no_shared_transformations to use it instead of calling blobs() per test. Also update test_get_attrs.py sdata_attrs fixture and test_empty_attrs to use sdata_blobs. The concatenate parametrized tests (6 variants) drop from ~0.55s to ~0.02s each (~3s total savings); test_get_attrs drops from ~0.28s to ~0.02s per test. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…ned_dataframe Drop from 200 to 20 for both N_PARTITIONS and the number of gene categories. The test's assertions (round-trip values, category-order mismatch between pandas and dask, dtype) are all preserved at the smaller scale; the 200-partition stress test was over-specified for what's actually being verified. ~1.56s → ~0.95s (39% faster) for that test. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…seed Drop from 20 to 10 for both N_PARTITIONS and the number of gene categories, and replace the module-level RNG with a local default_rng(seed=0) so the test is independent of other tests' RNG consumption. With seed=0 and N=10, partition 0 is deterministically missing gene_4, which is sufficient to trigger the dask category-order mismatch being tested. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Add pytest-xdist to test dependencies and pass -n auto --dist worksteal to pytest in CI. worksteal distributes tests dynamically across workers so slow IO tests don't block fast unit tests on a single worker. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…dexing - Remove reference to the old per-slice implementation from the docstring - Add inline comment explaining why indexing="ij" is correct for any number of spatial dimensions (2D and 3D labels) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…fixtures to session scope - test_delete_element_from_disk: subset full_sdata to [element_name, points_0_1] before writing, cutting the initial write from 19 elements to 2 (2–13× speedup per parametrize case) - conftest: set NUMBA_DISABLE_JIT=1 to avoid ~1.4s JIT overhead per worker on first datashader/rasterize call - test_partial_read: promote all module-scoped fixtures to session scope so the corrupted zarr stores are built once per session instead of once per module - test_transform: use small_translation=True to avoid out-of-bounds raster operations that trigger large dask computations - test_spatialdata_operations: reduce target_width 1000→100 in test_transform_to_data_extent to shrink the output raster Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
… napari plugin Napari registers a pytest11 entry point that loads (and breaks) numba before conftest.py runs. Fix: block it in pyproject.toml addopts with -p no:napari. datashader's @jit(cache=True) raises "no locator available" on Python 3.13 + numba. With napari blocked, conftest.py now runs before any plugin imports numba, so NUMBA_DISABLE_JIT=1 consistently disables all @jit decorators (datashader, xrspatial) and avoids the mixed JIT/non-JIT crash. Promote all module-scoped fixtures in test_partial_read to session scope so corrupted zarr stores are built once per session. For tests that write a full SpatialData object but only need a single element (or a small subset) to trigger the condition under test, subset the fixture before writing. The initial full write dominates test time (19 elements including 3-D multiscale rasters), so even subsetting to 1-6 elements gives 2-13× speedups per parametrize case: - test_delete_element_from_disk: subset to [element_name, points_0_1] - test_incremental_io_on_disk: subset to the 7 elements the loop accesses - test_overwrite_fails_when_zarr_store_present: use empty SpatialData() - test_element_already_on_disk_different_type: subset to [element_name] - test_self_contained: subset to [image2d, labels2d, points_0, circles] - test_change_path_of_subset: subset to the 5 elements + points_0_1 (needed so only_on_disk > 0 assertion passes) - test_validate_can_write_metadata_on_element: subset to [element_name] - test_save_transformations_incremental: subset to [element_name, image2d] (image2d anchors the non-self-contained assertion for the circles case) - test_consolidated_metadata: subset to one element per type - test_channel_names_raster_images_v1_to_v2_to_v3: subset to [image2d, image2d_multiscale] Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Two issues caused collection-time failures when tests ran via the Airflow
runner (ome_sdc conda env):
1. `os.environ.setdefault("NUMBA_DISABLE_JIT", "1")` is a no-op if the
runner environment already exports NUMBA_DISABLE_JIT=0. Switched to a
force-set assignment, and added a sys.modules patch for the case where a
pytest plugin (fast-array-utils, npe2, napari-plugin-engine) has already
imported numba — numba reads its config once at import time, so the env
var alone is too late if that happened first.
2. Module-level `blobs()` calls in four test files executed during pytest
collection, before conftest.py had a chance to take effect. Replaced each
with a `@pytest.fixture(scope="module")`, which fires during test
execution where the env is already configured. No performance change —
blobs() is still called once per file.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
The sub-conftest files in core/operations/ and dataloader/ duplicated the same NUMBA_DISABLE_JIT logic already present in the root tests/conftest.py, which applies to all tests. Remove the duplicates. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1117 +/- ##
=======================================
Coverage 91.97% 91.97%
=======================================
Files 51 51
Lines 7750 7750
=======================================
Hits 7128 7128
Misses 622 622 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Disabling the numba JIT in tests was not occurring in the airflow runner due to the use of
setdefault()not overriding an existing state. This PR addresses this. It also refactor the use of some globalsdatavariables as fixtures withmodulescope.