You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Stream validation and test eval outputs instead of one in-memory frame - #159
This PR lowers host memory in the validation and test stages by streaming the evaluation output table and its files instead of building them as one in-memory DataFrame. Outputs, metrics and the calibrator are unchanged.
Problem
Validation and test stop after max_val_steps / max_test_steps batches. The largest allowed setting is 8000 batches of 2048 rows (16.4M rows); the default is 4000 × 1024 (4.1M rows). The stages are bounded by those caps, but at the cap the output handling was expensive:
Output table._eval built the whole eval table as one DataFrame. Multiclass and multilabel tables have one row per entity and top-k class (up to 20). With 100 classes at the default test cap, that is an 82M-row frame and an 11.3 GiB peak.
EvalOutput.save(). It wrote that frame to Parquet in a BytesIO, copied the bytes into an in-memory zip, and then copied the zip again to write it.
save_df(). It ran memory_usage(deep=True) over every row to size the 1 GiB holdout file.
Validation inside fit(). Every epoch kept each step's keys, logits and targets, joined them and computed probabilities. Nothing reads them while fitting, and the last epoch's copy stayed alive into test().
Change
In-fit validation keeps no step outputs. Standalone validate() and test() are unchanged.
Chunked formatting.DiskgraphNestedTrainer._eval formats 65,536 entities at a time and appends each formatted chunk to a new EvalFrames. That object keeps the chunks on local disk, and each chunk carries the index the whole frame would have had.
Streaming EvalOutput.
save() streams the chunks as Parquet straight into the zip entry of {stage}_output.pt, using full 1,048,576-row row groups like DataFrame.to_parquet.
save_df() writes the holdout from the leading chunks and measures size only up to the cap.
The regression histogram reads one column. The heatmap draws the same rows as DataFrame.sample.
EvalOutput(metrics, df, embedding) still works, and .df loads the chunks on first access.
Results
Peak RSS and time of the test-stage tail are measured from synthetic step outputs to written files. The run covers metrics, calibration, formatting, save() and save_df(), on CPU.
Case
Test rows
Peak before
Peak after
Time before
Time after
Binary, default cap
4.1M
2.02 GiB
1.82 GiB
13 s
16 s
Binary, largest cap
16.4M
4.56 GiB
3.96 GiB
55 s
55 s
100 classes
1.0M
3.60 GiB
2.68 GiB
45 s
49 s
100 classes, default cap
4.1M
11.25 GiB
7.59 GiB
128 s
116 s
After the change, the peak is the module's end-of-epoch join of logits and the calibration step. Neither is changed by this PR. During fit(), each validation epoch at the default cap (2.05M rows) no longer holds about 3.6 GiB for 100 classes or about 0.7 GiB for binary.
Back-compat
Situation
Behavior
{stage}_output.pt read by EvalOutput.load
Same format version and entries. The frame reads back equal to before. The zip now uses ZIP64 data descriptors, which Python's zipfile reads.
Holdout test_df.parquet, not truncated
Same rows and index.
Holdout truncated at the 1 GiB cap
Keeps the leading rows whose measured size fits the cap. Before, it took rows × cap / total rows, which assumes every row is the average size. See the table below.
test_df_meta.json
Same keys and meaning.
EvalOutput(...) built by callers
Same arguments, plus a keyword-only frames. It is no longer a dataclass, so dataclass == and repr are gone.
pl_model.outputs after validation inside fit()
{}. Standalone validate() and test() are unchanged.
Local disk
Chunks are spooled under the system temp directory, using about the in-memory size of the output table. They are removed when the EvalOutput is released.
Truncated holdout row counts on a 100k-row fixture with a string column:
Row sizes
Cap
Rows before
Size before
Rows after
Size after
Uniform
10%
9,999
1.00× cap
10,000
1.00× cap
Small rows first
10%
9,999
0.43× cap
20,121
1.00× cap
Large rows first
10%
9,999
1.57× cap
6,303
1.00× cap
Testing Done
Identity at fixture level. For binary, multiclass, mixed-type classes, multilabel classification, multilabel ranking, regression and forecasting, these all match the whole-frame result exactly: the chunked .df, the {stage}_output.pt frame, the holdout file and the heatmap sample. The comparison uses 2-entity chunks.
Identity end to end. Each kumo.engine fit (binary, regression, multiclass, multilabel classification, multilabel ranking) also wrote its val and test outputs through main's EvalOutput, from the same whole frame. With 64-entity chunks, all files read back identical.
Regression extras. The regression histogram JSON, heatmap sample and holdout are byte-identical between chunked and whole frames.
Memory. Counting tracemalloc and the Arrow memory pool, save()/save_df() peak memory is flat from 4 to 40 chunks, and _build_eval_frames is flat from 4k to 40k rows.
Suites.test/train/test_output.py, test_diskgraph_nested_trainer.py, test_base_node_module.py and test_trainer.py pass (192). packages/kumo-engine/test/test_holdout_df.py passes (11, real fits). flake8 is clean.
Validation inside fit() only needs metrics, yet every epoch kept each
step's pkey, time, logits and target, joined them and computed
probabilities. Nothing reads those outputs while fitting, and the last
epoch's copy stayed alive into test(). Keep them only for standalone
validate() and test().
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
EvalOutput can now hold its rows as EvalFrames: formatted chunks kept
on local disk. save() streams them as Parquet into the zip entry, which
is written straight to its destination, so neither the Parquet bytes
nor the zip are buffered in memory. save_df() writes the holdout from
the leading chunks and measures size only up to the cap. The regression
histogram reads one column and the heatmap draws the same rows as
DataFrame.sample.
The files read back to the same frames as before, and EvalOutput(df=...)
keeps working. One change: the holdout keeps the leading rows whose
measured size fits the cap. It no longer estimates the count from the
average row size, which could overshoot the cap or stop well short of
it when row sizes vary.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
_eval built the whole output table in memory. For a multiclass task
that is one row per entity and top-k class, so 4.1M test rows with 100
classes became an 82M-row frame. Format 65,536 entities at a time and
spool each chunk to EvalFrames, with the same index as the whole
frame. Peak memory for that case falls from 11.3 GiB to 7.6 GiB.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.
No actionable comments were generated in the recent review.
Walkthrough
Evaluation output now supports ordered chunks that can be sampled and written to Parquet without first materializing one full DataFrame. The nested trainer builds indexed chunks, and EvalOutput uses them for saving and prediction summaries. During fitting without explainers, evaluation processing returns superclass outputs before validation and prediction conversion. Validation and test output collection also share a tensor-to-CPU helper.
Sequence Diagram(s)
sequenceDiagram
participant DiskGraphNestedTrainer
participant EvalFrames
participant EvalOutput
DiskGraphNestedTrainer->>EvalFrames: Build and collect evaluation chunks
DiskGraphNestedTrainer->>EvalOutput: Pass frames for evaluation output
EvalOutput->>EvalFrames: Read chunks for saving and summaries
Loading
Merge Risk: Low · previous review 1c9e6
The change is broadly mergeable, but capped evaluation saves do not strictly honor the configured memory limit for variable-width rows.
Pre-merge checks: 2 · 1
Failed checks: 1 inconclusive
Check name
Status
Explanation
Resolution
Change Map Design Review
The required PR conversation Change map is unavailable. The supplied conversation contains summaries and review findings, but no comment beginning with <!-- kumo-pr-change-map -->; repository inspec…
Provide the complete marked Change map comment, including the linked problem, observed failure, per-change rationale, abstraction choices, focused verification, and stopping condition. Re-run the design review with that comment.
Passed checks: 2 passed
Check name
Status
Explanation
Pr Size
The pull request changes 9 files. The limit is 40 files, so the check passes.
Public Resource Access
The pull request adds no external URL, repository location, or package source location. The changed files contain only Python implementation and test changes, with local temporary paths such as `kumo-…
Full details: Public Resource Access
Explanation
The pull request adds no external URL, repository location, or package source location. The changed files contain only Python implementation and test changes, with local temporary paths such as kumo-eval-frames- and .pkl files. No public-access check is required.
Full details: Change Map Design Review
Explanation
The required PR conversation Change map is unavailable. The supplied conversation contains summaries and review findings, but no comment beginning with <!-- kumo-pr-change-map -->; repository inspection found only the template and review guidance marker. The diff is substantial (675 additions and 78 deletions across nine files), but without the author's proposed design, rationale, verification, and stopping condition, this check cannot assess necessity, abstraction layers, or omitted behavior.
Resolution
Provide the complete marked Change map comment, including the linked problem, observed failure, per-change rationale, abstraction choices, focused verification, and stopping condition. Re-run the design review with that comment.
Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Reviewing files that changed from the base of the PR and between d4c7a8b and 1c9e652.
Files selected for processing: 2
packages/kumo-core/src/kumo_core/train/output.py
test/train/test_output.py
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 8 remain after this review.
The reason will be displayed to describe this comment to others. Learn more.
Follow-up Comments
Consider making the spool directory configurable.
236-238: 🧹 Nitpick | 🔵 Trivial
Consider making the spool directory configurable.
tempfile.TemporaryDirectory(prefix=...) uses the system temp directory. In many container images, /tmp is a memory-backed tmpfs. Spooled chunks then count toward host memory, and this PR is meant to reduce host memory. Add an optional dir parameter or an environment override. This lets operators point the spool at real local disk.
tracemalloc does not measure the Arrow allocations that this test targets.
tracemalloc tracks Python and NumPy allocations. It does not track Arrow's memory pool. The Arrow tables from pa.Table.from_pandas, pending, and _write_row_groups therefore do not appear in peak. If the writer buffered all tables, test_save_memory_does_not_grow_with_rows would still pass. Track the Arrow pool too. For example, set pa.proxy_memory_pool(pa.default_memory_pool()) as the default pool for the measurement. Then add its max_memory() to the peak, and restore the previous pool afterward. The import pyarrow as pa statement is required. Verify the pool API against the pinned pyarrow version (>=15).
The reason will be displayed to describe this comment to others. Learn more.
Follow-up Comments
Handle zero-row evaluations in the heatmap path.
268-281: Minor · Stability & Availability · Quick win
Handle zero-row evaluations in the heatmap path.
_build_eval_frames() deliberately creates one empty chunk when the evaluation has zero rows. save_prediction_heatmap() then calls sample(0), which raises from pd.concat([]) before sampled_df.parquet is written. The broad exception handler suppresses the error, so the sampled-data and heatmap artifacts documented by save_all() are missing. Add an explicit zero-row policy at the heatmap boundary. The policy must write the defined empty sampled-data artifact and avoid heatmap calculations that require numeric min/max values.
Preserve a valid Parquet file for a zero-row holdout selection.
283-337: Minor · Data Integrity & Integration · Quick win
Preserve a valid Parquet file for a zero-row holdout selection.
When the first chunk exceeds the holdout byte cap, head_chunks() can round the selected prefix down to zero rows and return no chunks. save_df() then opens *_df.parquet and calls write_parquet() with an empty list. No writer can emit a footer, so the file can remain unreadable. The metadata write is also not guaranteed because it occurs after write_parquet() returns.
Retain a zero-row chunk with its schema, or handle the empty selection in write_parquet() by emitting a valid schema-preserving empty Parquet file.
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
The reason will be displayed to describe this comment to others. Learn more.
Re-tested 7f60f77. The attached ColumnAnalysis epoch-hook regression now passes with a real ColumnAnalysis instance. The zero-row holdout finding still reproduces: EvalOutput({}, pd.DataFrame({"ENTITY": ["large row"], "SCORE": [1.]})).save_df(tmpdir, TrainingStage.TEST, max_size=1) raises ArrowInvalid: Must provide at least one schema to unify and leaves no valid holdout artifact/metadata. Please preserve the input schema and write a typed empty Parquet file when truncation selects zero rows, matching the previous behavior, with this focused regression. Withholding approval pending that correction; the byte-cap example is intentionally minimal, not a claim that the normal default cap usually triggers it.
The reason will be displayed to describe this comment to others. Learn more.
Follow-up review of 7f60f77fb:
Before approval:
Fix the previously reported zero-row holdout selection: retain the source schema, write readable empty Parquet, and emit metadata with saved_rows=0 / truncated=true. Add the minimal byte-cap regression from my earlier comment.
Preserve empty sampling behavior too. Independently reproduced: EvalFrames.of(empty_typed_frame).sample(0) raises ValueError: No objects to concatenate, whereas the previous pandas frame.sample(0) returns the typed empty frame. This stops save_prediction_heatmap() before it writes sampled_df.parquet. Return a schema-preserving empty sample and explicitly skip numeric heatmap calculations when there are no rows; test that boundary.
Reuse opportunity:EvalFrames.write_parquet() / _write_row_groups() introduce another Parquet accumulator alongside BP's StreamingFrameWriter in train/io.py. Sharing its low-level byte-bounded buffering/flush mechanics would reduce duplication and avoid relying solely on a row-count buffer limit. Preserve evaluation's single-file ZIP member, index and mixed-CLASS schema handling; the BP writer is not a drop-in replacement and also skips empty batches. Keep this small or leave the consolidation as a follow-up rather than expanding the merge fix.
Verification/claims:tracemalloc alone does not establish total Arrow/native peak memory. Supplement it with Arrow-pool or process-memory measurement, or narrow the stated guarantee. Standalone evaluation still retains/concatenates split logits, so describe this as reducing formatting/serialization peaks, not end-to-end constant-memory evaluation.
The actual attached ColumnAnalysis epoch-hook regression now passes; that earlier finding is resolved.
258-259: Minor · Stability & Availability · Quick win
Keep EvalFrames alive for the whole iterator.
When a caller retains only the iterator, the current generator expression retains the chunk list but not the EvalFrames owner. The owner can be finalized, removing the spool directory before a later chunk.load() call.
Suggested fix
def __iter__(self) -> Iterator[pd.DataFrame]:
- return (chunk.load() for chunk in self._chunks)+ for chunk in self._chunks:+ yield chunk.load()
Skipped pred conversion gives in-fit ColumnAnalysis wrong quantile values.
When explain() raises after it attaches ColumnAnalysis, the explainer stays attached. A later non-refit fit() then runs explain_end during validation. For MultiQuantileLoss, evaluation_epoch_end returns early during fit and never builds pred. transform_prediction then flattens the (N, Q) logits and truncates them to num_nodes. The result can be written to column_analysis.json under the validation stage folder when root_dir is set.
Build pred with the same loss conversion as the non-fit path when explainers are attached:
Suggested fix
outputs = super().evaluation_epoch_end(stage, step_outputs)
if self.is_fitting:
+ if self.explainers and 'pred' not in outputs and 'logits' in outputs:+ outputs['pred'] = self.loss.to_prediction(outputs['logits'])
return outputs
pa.Table.from_pandas() allocates buffers for the object-string ENTITY column in the Arrow memory pool. Those native allocations are not reliably included in tracemalloc. If write_parquet() retains all converted tables, the 40-chunk case can use substantially more memory while this test still passes.
Measure the Arrow pool high-water mark in addition to tracemalloc. A per-run default_memory_pool().max_memory() baseline helps for these sequential runs, but it is not a reset when earlier tests have already raised the process-wide high-water mark. Use an isolated Arrow pool or run each measurement in a fresh subprocess.
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
Ensure the configured temporary filesystem has sufficient capacity.
EvalFrames already supports operator selection through Python's TMPDIR, TEMP, and TMP settings, so a new dir argument would duplicate existing behavior. The spool path still depends on that filesystem: frame.to_pickle(source) can raise ENOSPC when it fills, and a RAM-backed temporary directory can reduce the intended memory benefit. Set the temporary-directory environment before the evaluation worker starts and use disk-backed storage with sufficient capacity. Debug logging may improve diagnosis, but it does not prevent the failure.
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
Re-reviewed 638ac662f696. The zero-row holdout and empty sample/heatmap fixes now pass; 174 focused output/node-module/nested-trainer tests pass locally.
I independently reproduced the two newer bot findings:
Spool iterator lifetime (EvalFrames.__iter__): append a frame, retain it = iter(frames), drop frames, run gc.collect(), then consume it. The owner is finalized and its temp directory removed, so iteration raises FileNotFoundError. Make the iterator retain the owner (a generator method using yield is sufficient) and add this regression.
Attached regression explainer + multi-quantile loss (BaseNodeLightningModule.evaluation_epoch_end): the in-fit early return keeps logits but skips loss.to_prediction. The real ColumnAnalysis.transform_prediction(..., num_nodes=2) then flattens quantile logits and takes the first two values. With two rows of ordered quantiles starting at 1 and 101, it returns [1, 2] instead of median predictions [14, 114]. The same probe using the base method returns the correct medians. Preserve prediction conversion when explainers are attached, keeping the no-explainer fast path. This concerns the attached-explainer/reused-module path, not every normal post-fit explanation.
These are small correctness fixes; withholding approval pending them. Parquet-writer consolidation remains a follow-up rather than a merge prerequisite. Current CI is still running/queued.
The reason will be displayed to describe this comment to others. Learn more.
Caution
Some comments are outside the diff and can’t be posted inline due to GitHub limitations.
⚠️ Outside diff range comments (1)
Minor · Keep materialized EvalOutput.df changes in saved predictions. · output.py:341-375
packages/kumo-core/src/kumo_core/train/output.py:341-375 Minor · Data Integrity & Integration · Quick win
Keep materialized EvalOutput.df changes in saved predictions.
When a frame-backed output has multiple chunks, accessing and modifying output.df can leave self.frames unchanged. save() and save_df() then write the older chunks instead of the public dataframe.
Rebind self.frames when the property materializes df:
Suggested fix
def df(self) -> pd.DataFrame | None:
if self._df is None and self.frames is not None:
- self._df = self.frames.to_frame()+ df = self.frames.to_frame()+ self._df = df+ self.frames = EvalFrames.of(df)
return self._df
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.
Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.
Wrapping df once when it was first read cached its schema and row
counts, so a column added or rows dropped afterwards broke save() or
wrote stale counts. Build the frames from df on each read instead.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.
The reason will be displayed to describe this comment to others. Learn more.
Caution
Some comments are outside the diff and can’t be posted inline due to GitHub limitations.
⚠️ Outside diff range comments (2)
Minor · Measure the retained prefix against max_bytes. · output.py:294-296
packages/kumo-core/src/kumo_core/train/output.py:294-296 Minor · Functional Correctness · Quick win
Measure the retained prefix against max_bytes.
If row sizes vary within the crossing chunk, multiplying its row count by the remaining-byte fraction does not bound the size of frame.head(rows). A large first row can make save_df(max_size=...) write more than the requested cap. Check the retained prefix's measured size and reduce its row count until it fits. (raw.githubusercontent.com)
Minor · Detect a numeric CLASS chunk alongside a string CLASS chunk. · output.py:201-203
packages/kumo-core/src/kumo_core/train/output.py:201-203 Minor · Data Integrity & Integration · Quick win
Detect a numeric CLASS chunk alongside a string CLASS chunk.
EvalFrames.append accepts any DataFrame, so a caller can add one chunk with numeric CLASS values and one with string CLASS values. _Chunk.of records None for the numeric chunk and 'string' for the string chunk. _has_mixed_class ignores None, so it sees only {'string'} and returns False. write_parquet then passes the int64 and string schemas to pa.unify_schemas, which can fail before the mixed-class conversion runs. The same failure applies to EvalOutput(frames=...) when save runs.
Trainer-produced chunks share one CLASS dtype per evaluation, so this failure affects only frames built directly through the public API.
The fix flags a numeric chunk only when a string chunk is also present. Numeric-only chunks keep their current schema promotion.
Suggested fix
def _has_mixed_class(chunks: list[_Chunk]) -> bool:
dtypes = {c.class_dtype for c in chunks} - {None, 'empty'}
- return len(dtypes) > 1 or any(d.startswith('mixed') for d in dtypes)+ if len(dtypes) > 1 or any(d.startswith('mixed') for d in dtypes):+ return True+ return 'string' in dtypes and any(+ c.class_dtype is None and 'CLASS' in c.schema.names for c in chunks)
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 9 remain after this review.
A chunk with numeric CLASS values next to one with strings failed to
unify schemas on save. Convert both to strings, as for a whole frame.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.
The reason will be displayed to describe this comment to others. Learn more.
Re-reviewed 6354419. Earlier spool-lifetime, attached-explainer and DataFrame metadata regressions are fixed; 204 focused CPU tests passed. No remaining code blocker found. The sole failed check is the independently verified missing restricted BigQuery test identities; this approval does not waive that CI gate.
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
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.
GitHub Actions
Inactive
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Victsm
left a comment
This PR lowers host memory in the validation and test stages by streaming the evaluation output table and its files instead of building them as one in-memory DataFrame. Outputs, metrics and the calibrator are unchanged.
Problem
Validation and test stop after
max_val_steps/max_test_stepsbatches. The largest allowed setting is 8000 batches of 2048 rows (16.4M rows); the default is 4000 × 1024 (4.1M rows). The stages are bounded by those caps, but at the cap the output handling was expensive:_evalbuilt the whole eval table as one DataFrame. Multiclass and multilabel tables have one row per entity and top-k class (up to 20). With 100 classes at the default test cap, that is an 82M-row frame and an 11.3 GiB peak.EvalOutput.save(). It wrote that frame to Parquet in aBytesIO, copied the bytes into an in-memory zip, and then copied the zip again to write it.save_df(). It ranmemory_usage(deep=True)over every row to size the 1 GiB holdout file.fit(). Every epoch kept each step's keys, logits and targets, joined them and computed probabilities. Nothing reads them while fitting, and the last epoch's copy stayed alive intotest().Change
validate()andtest()are unchanged.DiskgraphNestedTrainer._evalformats 65,536 entities at a time and appends each formatted chunk to a newEvalFrames. That object keeps the chunks on local disk, and each chunk carries the index the whole frame would have had.EvalOutput.save()streams the chunks as Parquet straight into the zip entry of{stage}_output.pt, using full 1,048,576-row row groups likeDataFrame.to_parquet.save_df()writes the holdout from the leading chunks and measures size only up to the cap.DataFrame.sample.EvalOutput(metrics, df, embedding)still works, and.dfloads the chunks on first access.Results
Peak RSS and time of the test-stage tail are measured from synthetic step outputs to written files. The run covers metrics, calibration, formatting,
save()andsave_df(), on CPU.After the change, the peak is the module's end-of-epoch join of logits and the calibration step. Neither is changed by this PR. During
fit(), each validation epoch at the default cap (2.05M rows) no longer holds about 3.6 GiB for 100 classes or about 0.7 GiB for binary.Back-compat
{stage}_output.ptread byEvalOutput.loadzipfilereads.test_df.parquet, not truncatedrows × cap / totalrows, which assumes every row is the average size. See the table below.test_df_meta.jsonEvalOutput(...)built by callersframes. It is no longer a dataclass, so dataclass==andreprare gone.pl_model.outputsafter validation insidefit(){}. Standalonevalidate()andtest()are unchanged.EvalOutputis released.Truncated holdout row counts on a 100k-row fixture with a string column:
Testing Done
.df, the{stage}_output.ptframe, the holdout file and the heatmap sample. The comparison uses 2-entity chunks.kumo.enginefit (binary, regression, multiclass, multilabel classification, multilabel ranking) also wrote its val and test outputs through main'sEvalOutput, from the same whole frame. With 64-entity chunks, all files read back identical.save()/save_df()peak memory is flat from 4 to 40 chunks, and_build_eval_framesis flat from 4k to 40k rows.test/train/test_output.py,test_diskgraph_nested_trainer.py,test_base_node_module.pyandtest_trainer.pypass (192).packages/kumo-engine/test/test_holdout_df.pypasses (11, real fits).flake8is clean.🤖 Generated with Claude Code