diff --git a/crates/dbsp/src/circuit/checkpointer.rs b/crates/dbsp/src/circuit/checkpointer.rs index 41b5ba33f1d..5fa0dce0f86 100644 --- a/crates/dbsp/src/circuit/checkpointer.rs +++ b/crates/dbsp/src/circuit/checkpointer.rs @@ -51,7 +51,7 @@ impl Checkpointer { /// Creates a new checkpointer for directory `storage_path`. Deletes any /// unreferenced files in the directory. pub fn new(backend: Arc) -> Result { - let checkpoint_list = Self::read_checkpoints(&*backend)?; + let checkpoint_list = Self::read_checkpoints_at_startup(&*backend)?; let this = Checkpointer { backend, @@ -432,18 +432,17 @@ impl Checkpointer { Ok(self.checkpoint_list.clone().into()) } - /// Reads the list of checkpoints available through `backend`. - /// - /// A missing `checkpoints.feldera` is treated as "no checkpoints yet" - /// only when the storage directory holds no UUID-shaped subdirectories. - /// If UUID directories exist, the catalog has been lost while the - /// checkpoints themselves are likely still on disk; proceeding would - /// let `gc_startup` recursively delete them. Refuse to start instead. - pub fn read_checkpoints( + /// Reads the list of checkpoints available through `backend`. This is like + /// [Self::read_checkpoints] except that, if `checkpoints.feldera` is + /// missing, we check whether the storage directory hold any UUID-shaped + /// subdirectories. If UUID directories do exist, the catalog has been lost + /// while the checkpoints themselves are likely still on disk, and + /// proceeding would let `gc_startup` recursively delete them, so we refuse + /// to start instead. + fn read_checkpoints_at_startup( backend: &dyn StorageBackend, ) -> Result, Error> { - let file_name = StoragePath::from(CHECKPOINT_FILE_NAME); - match backend.read_json(&file_name) { + match backend.read_json(&StoragePath::from(CHECKPOINT_FILE_NAME)) { Ok(checkpoints) => Ok(checkpoints), Err(error) if error.kind() == ErrorKind::NotFound => { let mut orphan_uuid_dirs: Vec = Vec::new(); @@ -466,16 +465,24 @@ impl Checkpointer { })); } - // Write an empty checkpoint file to save the cost of listing - // all the files next time. - backend.write_json(&file_name, &VecDeque::::new())?; - Ok(VecDeque::new()) } Err(error) => Err(error)?, } } + /// Reads the list of checkpoints available through `backend`. + pub fn read_checkpoints( + backend: &dyn StorageBackend, + ) -> Result, Error> { + backend + .read_json(&StoragePath::from(CHECKPOINT_FILE_NAME)) + .or_else(|error| match error.kind() { + ErrorKind::NotFound => Ok(VecDeque::new()), + _ => Err(error.into()), + }) + } + fn update_checkpoint_file(&self) -> Result<(), Error> { Ok(self .backend @@ -736,7 +743,6 @@ impl Checkpoint for EmptyCheckpoint { #[cfg(test)] mod test { use std::sync::Arc; - use std::sync::atomic::{AtomicIsize, Ordering}; use feldera_storage::{DirEntry, StorageBackend, StoragePath}; use feldera_types::config::{FileBackendConfig, StorageCacheConfig}; @@ -755,7 +761,6 @@ mod test { struct CatalogFailingBackend { inner: Arc, fail_on: StoragePath, - count_down: AtomicIsize, } impl feldera_storage::StorageBackend for CatalogFailingBackend { @@ -764,7 +769,7 @@ mod test { name: &StoragePath, ) -> Result, feldera_storage::error::StorageError> { - if name == &self.fail_on && self.count_down.fetch_sub(1, Ordering::Relaxed) <= 0 { + if name == &self.fail_on { return Err(feldera_storage::error::StorageError::StdIo { kind: std::io::ErrorKind::PermissionDenied, operation: "injected catalog write failure", @@ -820,7 +825,7 @@ mod test { LogCapture::new(|| Checkpointer::new(backend.clone()).unwrap()).into_parts(); assert_eq!( log, - " INFO dbsp::circuit::checkpointer: GC kept 1/0/0 expected files/directories/other, kept 0/0/0 unexpected, and deleted 0/0/0 unused; 0 error(s) reading directory entries\n" + " INFO dbsp::circuit::checkpointer: GC kept 0/0/0 expected files/directories/other, kept 0/0/0 unexpected, and deleted 0/0/0 unused; 0 error(s) reading directory entries\n" ); let uuid = uuid::Uuid::now_v7(); @@ -1134,12 +1139,6 @@ ERROR dbsp::circuit::checkpointer: 1 checkpoint(s) need missing file: w0-aaaaaaa let backend: Arc = Arc::new(CatalogFailingBackend { inner: posix, fail_on: CHECKPOINT_FILE_NAME.into(), - count_down: { - // `Checkpointer::new` will create an empty checkpoints catalog. - // Let that succeed. Then fail the second attempt, the one - // intended to add to it. - AtomicIsize::new(1) - }, }); let mut checkpointer = Checkpointer::new(backend).unwrap(); diff --git a/python/feldera/runtime_config.py b/python/feldera/runtime_config.py index d898c626f6f..2d02a1d3930 100644 --- a/python/feldera/runtime_config.py +++ b/python/feldera/runtime_config.py @@ -85,9 +85,13 @@ def __init__( dev_tweaks: Optional[dict] = None, env: Optional[dict[str, str]] = None, logging: Optional[str] = None, + datafusion_memory_mb: Optional[int] = None, + max_rss_mb: Optional[int] = None, ): self.workers = workers self.hosts = hosts + self.datafusion_memory_mb = datafusion_memory_mb + self.max_rss_mb = max_rss_mb self.tracing = tracing self.tracing_endpoint_jaeger = tracing_endpoint_jaeger self.cpu_profiler = cpu_profiler diff --git a/python/feldera/testutils.py b/python/feldera/testutils.py index 98ba218eac5..6f877b14792 100644 --- a/python/feldera/testutils.py +++ b/python/feldera/testutils.py @@ -1,5 +1,6 @@ "Utility functions for writing tests against a Feldera instance." +import math import os import platform import re @@ -44,6 +45,31 @@ def _get_effective_api_key(): FELDERA_TEST_NUM_WORKERS = int(os.environ.get("FELDERA_TEST_NUM_WORKERS", "8")) FELDERA_TEST_NUM_HOSTS = int(os.environ.get("FELDERA_TEST_NUM_HOSTS", "1")) +# Ad-hoc queries reserve 64 MiB per worker for sort spill before sorting a +# single row (SORT_SPILL_RESERVATION_BYTES in +# adapterlib/src/utils/datafusion.rs), so a `datafusion_memory_mb` pool +# smaller than `workers * 64 MiB` fails any ORDER BY/EXCEPT/hash-join ad-hoc +# query with "Resources exhausted", independent of data size. This holds +# regardless of `hosts`: the reservation is sized from the *total* configured +# `workers`, whether they all run on one host or are split across several +# (the coordinator's own ad-hoc engine also uses the total worker count as +# its partition count -- see crates/coord/src/adhoc.rs). +_ADHOC_SORT_RESERVATION_MB_PER_WORKER = (64 * 1024 * 1024) / 1_000_000 + + +def min_datafusion_memory_mb( + workers: int = FELDERA_TEST_NUM_WORKERS, + hosts: int = FELDERA_TEST_NUM_HOSTS, + headroom_mb: int = 512, +) -> int: + """Minimum `datafusion_memory_mb` for a sort-heavy ad-hoc query to avoid + "Resources exhausted" with the given number of workers, plus + `headroom_mb` for the query's actual data. `hosts` doesn't change the + result (see module comment above); it's accepted so callers can pass + both test parameters without worrying about which one matters. + """ + return math.ceil(workers * _ADHOC_SORT_RESERVATION_MB_PER_WORKER) + headroom_mb + class _LazyClient: "Construct the FelderaClient only when accessed as opposed to when imported." @@ -291,6 +317,7 @@ def build_pipeline( views: List[ViewSpec], resources: Optional[Resources] = None, dev_tweaks: Optional[dict] = None, + datafusion_memory_mb: Optional[int] = None, ) -> Pipeline: sql = generate_program(tables, views) @@ -307,6 +334,7 @@ def build_pipeline( workers=FELDERA_TEST_NUM_WORKERS, hosts=FELDERA_TEST_NUM_HOSTS, dev_tweaks=dev_tweaks, + datafusion_memory_mb=datafusion_memory_mb, ), ).create_or_replace() diff --git a/python/tests/runtime/test_udp.py b/python/tests/runtime/test_udp.py index c3690661935..4edb94f0fb1 100644 --- a/python/tests/runtime/test_udp.py +++ b/python/tests/runtime/test_udp.py @@ -4,7 +4,11 @@ from tests import TEST_CLIENT from tests.platform.helper import PipelineTestCase from feldera.runtime_config import Resources, RuntimeConfig -from feldera.testutils import FELDERA_TEST_NUM_WORKERS, FELDERA_TEST_NUM_HOSTS +from feldera.testutils import ( + FELDERA_TEST_NUM_WORKERS, + FELDERA_TEST_NUM_HOSTS, + min_datafusion_memory_mb, +) # Test user-defined preprocessor @@ -97,6 +101,13 @@ def test_local(self): tracing = { version = "0.1.40" } """ + # Scales with FELDERA_TEST_NUM_WORKERS/_HOSTS so the ad-hoc ORDER BY + # query below doesn't fail with "Resources exhausted" if either + # changes. See min_datafusion_memory_mb. + datafusion_memory_mb = min_datafusion_memory_mb( + FELDERA_TEST_NUM_WORKERS, FELDERA_TEST_NUM_HOSTS + ) + pipeline = PipelineBuilder( TEST_CLIENT, name=self.register_for_cleanup("test_udps"), @@ -104,10 +115,8 @@ def test_local(self): udf_rust=udfs, udf_toml=toml, runtime_config=RuntimeConfig( - resources=Resources( - memory_mb_min=1024, - config={"datafusion_memory_mb": 512}, - ), + datafusion_memory_mb=datafusion_memory_mb, + resources=Resources(memory_mb_min=datafusion_memory_mb + 512), workers=FELDERA_TEST_NUM_WORKERS, hosts=FELDERA_TEST_NUM_HOSTS, ), diff --git a/python/tests/workloads/test_now.py b/python/tests/workloads/test_now.py index 5aca91bd0da..2e98dcf31e8 100644 --- a/python/tests/workloads/test_now.py +++ b/python/tests/workloads/test_now.py @@ -5,9 +5,12 @@ from feldera.pipeline import Pipeline from feldera.runtime_config import Resources from feldera.testutils import ( + FELDERA_TEST_NUM_HOSTS, + FELDERA_TEST_NUM_WORKERS, ViewSpec, build_pipeline, log, + min_datafusion_memory_mb, validate_outputs, unique_pipeline_name, ) @@ -96,21 +99,27 @@ def test_now(self): ), ] + # 12288 MB covers this test's data volume (>2GB ad-hoc query storage, + # peaks above 5GB memory on arm64); take the max with the per-worker + # sort-reservation floor (see min_datafusion_memory_mb) so a future + # bump to FELDERA_TEST_NUM_WORKERS can't undersize it. + datafusion_memory_mb = max( + min_datafusion_memory_mb(FELDERA_TEST_NUM_WORKERS, FELDERA_TEST_NUM_HOSTS), + 12288, + ) pipeline = build_pipeline( unique_pipeline_name("now-test"), tables, views, - # This test uses >2GB of storage in the ad hoc query and peaks - # above 5 GB of memory on arm64; an honest request keeps k8s - # from scheduling the pipeline onto a node where that overshoot - # gets it OOM-killed. No memory_mb_max: it would cap the - # DataFusion pool at 5% and the big ad-hoc validation queries - # exhaust that. + # An honest memory request keeps k8s from scheduling the pipeline + # onto a node where that overshoot gets it OOM-killed. No + # memory_mb_max: it would cap the DataFusion pool at 5% and the + # big ad-hoc validation queries exhaust that. resources=Resources( storage_mb_max=16384, memory_mb_min=16384, - config={"datafusion_memory_mb": 12288}, ), + datafusion_memory_mb=datafusion_memory_mb, ) pipeline.start()