Skip to content

tests: fix datafusion_memory_mb wiring and parameterize on workers/hosts #6754

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
blp merged 2 commits into from
Jul 29, 2026
Merged

tests: fix datafusion_memory_mb wiring and parameterize on workers/hosts #6754

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cb2b703
tests: fix datafusion_memory_mb wiring and parameterize on workers/hosts
blp Jul 29, 2026
ef7b488
[dbsp] Avoid race updating checkpoint catalog file.
blp Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 24 additions & 25 deletions crates/dbsp/src/circuit/checkpointer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<dyn StorageBackend>) -> Result<Self, Error> {
let checkpoint_list = Self::read_checkpoints(&*backend)?;
let checkpoint_list = Self::read_checkpoints_at_startup(&*backend)?;

let this = Checkpointer {
backend,
Expand Down Expand Up @@ -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<VecDeque<CheckpointMetadata>, 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<String> = Vec::new();
Expand All @@ -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::<CheckpointMetadata>::new())?;

Comment thread
swanandx marked this conversation as resolved.
Ok(VecDeque::new())
}
Err(error) => Err(error)?,
}
}

/// Reads the list of checkpoints available through `backend`.
pub fn read_checkpoints(
backend: &dyn StorageBackend,
) -> Result<VecDeque<CheckpointMetadata>, 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
Expand Down Expand Up @@ -736,7 +743,6 @@ impl<T: Default> Checkpoint for EmptyCheckpoint<T> {
#[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};
Expand All @@ -755,7 +761,6 @@ mod test {
struct CatalogFailingBackend {
inner: Arc<dyn StorageBackend>,
fail_on: StoragePath,
count_down: AtomicIsize,
}

impl feldera_storage::StorageBackend for CatalogFailingBackend {
Expand All @@ -764,7 +769,7 @@ mod test {
name: &StoragePath,
) -> Result<Box<dyn feldera_storage::FileWriter>, 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",
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -1134,12 +1139,6 @@ ERROR dbsp::circuit::checkpointer: 1 checkpoint(s) need missing file: w0-aaaaaaa
let backend: Arc<dyn StorageBackend> = 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();

Expand Down
4 changes: 4 additions & 0 deletions python/feldera/runtime_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
28 changes: 28 additions & 0 deletions python/feldera/testutils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
"Utility functions for writing tests against a Feldera instance."

import math
import os
import platform
import re
Expand Down Expand Up @@ -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."
Expand Down Expand Up @@ -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)

Expand All @@ -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()

Expand Down
19 changes: 14 additions & 5 deletions python/tests/runtime/test_udp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -97,17 +101,22 @@ 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"),
sql=sql,
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,
),
Expand Down
23 changes: 16 additions & 7 deletions python/tests/workloads/test_now.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -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()
Expand Down
Loading