Cubed talk at SciPy 2025

https://github.com/cubed-dev/cubed

Transcript

1. Scalable array processing with bounded memory in Python *Tom Nicholas Tom White *tom@earthmover.io @TomNicholas @tegnicholas.bsky.social
2. What I will talk about: - Vision for Science at Scale - What is Cubed? - Xarray integration - Initial results - Pros and Cons - Next steps
3. Tale of two Toms: Tom White Tom Nicholas - Xarray core dev - Background in Plasma Physics + Oceanography - Now engineer at Earthmover - Sgkit developer - Hadoop maintainer - Cubed main developer
4. Vision for Science at Scale (Tom’s 🎄 list 🎁 ) - My perfect parallel executor…
5. Vision for Science at Scale (Tom’s 🎄 list 🎁 ) - (1) - Expressive - Scale without rewriting - Perfect weak horizontal scaling - (1000x problem in 1x time with 1000x CPUs) - Predictable (no nasty RAM surprises) - Robust to small failures
6. Vision for Science at Scale (Tom’s 🎄 list 🎁 ) - (2) - … - Resumable - Forget about the Cluster - Fully open - Not locked in to any one service or platform
7. What is Cubed? - Idea: All array operations take Zarr -> Zarr
8. - Many simple array operations are “blockwise” - Read: embarrassingly parallel across chunks
9. - But some array operations require a “rechunk” (diagram from pangeo-data/rechunker package)
10. - Some operations are more complex
11. - Turns out blockwise + rechunk cover all numpy-like array operations!
12. - Chain operations into a high-level “Plan” - Represented at array-level, not chunk-level
13. Bounded-memory operations - Blockwise processes one chunk at a time - Rechunk can be constant memory if spilling to intermediate Zarr store - (see pangeo Rechunker package)
14. Bounded-memory operations - Can therefore predict memory usage before launching!
15. Serverless execution - Every op is (a series of) embarrassingly parallel tasks - Just launch them all simultaneously - Ideal fit for ✨Serverless✨ cloud services - e.g. AWS Lambda, Google Cloud Functions - (Means no cluster to manage!)
16. Range of Executors - Serverless Modal Apps Coiled Functions
17. Range of Executors - Cluster / HPC - Single multicore machine Dragon?
18. Start on your Laptop - Process hundreds of GB on your laptop using all available cores
19. Overall design - Bounded-Memory Serverless Array Processing
20. Benchmark: “Quadratic Means” problem - Reduction computation - Input: 1.5TB - Time: 1m 40s - Lithops on AWS Lambda with 1000 workers
21. Benchmark: Rechunking ERA5 dataset variable - All-to-all rechunk - Multi-stage algorithm - Input: 1.5TB - Time: 8m 48s - Lithops on AWS Lambda with 1000 workers
22. Benchmark: Rechunking ERA5 dataset variable - Actual memory usage is always less than allowed maximum - Room to optimize further
23. Xarray Integration - Xarray has been generalized to wrap any chunked array type - Install cubed & cubed-xarray - Then specify the allowed memory - (And the location for intermediate Zarr stores) from cubed import Spec spec = Spec(work_dir='tmp', allowed_mem='1GB')
24. Xarray Integration - Now you can directly open from disk as cubed.Array objects ds = open_dataset( 'data.zarr', chunked_array_type='cubed', from_array_kwargs={'spec': spec}) chunks={}, )
25. Xarray Integration - Now just .compute, with your chosen serverless Executor! from cubed.runtime.executors.lithops import LithopsExecutor ds.compute(executor=LithopsExecutor())
26. Xarray Integration - Xarray now wraps Cubed OR Dask OR [new things??] Other??
27. Vision for Science at Scale (Tom’s 🎄 list 🎁 ) - Expressive - No Cluster - Predictable RAM usage - Retry failures - Resumable - Horizontal scaling - Fully open
28. Disadvantages - I/O to Zarr is slow compared to ideal dask case of staying in RAM - Serverless more expensive per CPU-hour - Only array operations
29. Next steps - We want your use cases to test on! - Optimizations - Other array types (JAX on GPU) - Other storage layers (obstore, zarrs-python)
30. Read the blog post! xarray.dev/blog/cubed-xarray - Join the discussion! Thanks to Tom White for writing Cubed!