Skip to content

Commit 996fa16

Browse files
Vedaant-Jclaude
andcommitted
Trim training-table lineage to essentials
The export sidecar path re-serialised the staged manifest behind a redundant directory check and a catch-all, and the tests restated fields and covered helpers owned by the snapshot-manifest PR. Copy the staged manifest bytes as the sidecar, inline the single-use source-lineage helper, shorten docstrings, and merge the lineage tests into one export/reload test plus best-effort and no-sidecar cases. Behaviour is unchanged. Verified: test_training_table_export, test_custom_source, test_custom_sql, test_predictive_query -k custom and kumo-core test_artifact_export pass. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent a35c8c5 commit 996fa16

6 files changed

Lines changed: 53 additions & 162 deletions

File tree

‎packages/kumo-core/src/kumo_core/artifact_export.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -205,8 +205,8 @@ def export_selected_files(files, output_path, *, fs, identity,
205205
206206
Filesystem/authentication and retry wrappers stay with the caller. The
207207
directory ownership, output names and CSV retry decisions are shared.
208-
``sidecars`` (name to bytes) are written into a directory destination
209-
after the data and before ``_SUCCESS``; ``fs`` then needs ``write_bytes``.
208+
``sidecars`` (name to bytes) land in a directory destination before
209+
``_SUCCESS``.
210210
"""
211211
import os.path as osp
212212
import tempfile

‎packages/kumo-core/src/kumo_core/object_store_sink.py‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -57,11 +57,7 @@ def write_files(self, connector, table_name, files,
5757
write_mode=WriteMode.OVERWRITE, project_id=None, *,
5858
export_id=None, extra_columns=None,
5959
sidecars: Mapping[str, bytes] = {}):
60-
"""Export selected Parquet without collecting their rows.
61-
62-
``sidecars`` are written into a directory destination before
63-
``_SUCCESS``.
64-
"""
60+
"""Export selected Parquet without collecting their rows."""
6561
del project_id
6662
if write_mode is not WriteMode.OVERWRITE:
6763
raise NotImplementedError(

‎packages/kumo-core/tests/test_artifact_export.py‎

Lines changed: 13 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -224,54 +224,33 @@ def interrupted(source, destination, **kwargs):
224224
assert (target / '_SUCCESS').exists()
225225

226226

227-
@pytest.mark.parametrize('directory', [True, False])
228-
def test_sidecars_are_written_before_success(tmp_path, monkeypatch,
229-
directory):
230-
import pyarrow.parquet as pq
231-
from kumo_core.artifact_export import export_selected_files
232-
from kumo_core.object_store_sink import _ArrowExportFiles
233-
from pyarrow.fs import LocalFileSystem
234-
235-
source = tmp_path / 'part-0.parquet'
236-
pq.write_table(pa.table({'id': [1]}), source)
237-
target = tmp_path / 'output'
238-
fs = _ArrowExportFiles(LocalFileSystem(), None)
239-
touched = fs.touch
240227

241-
def touch(path):
242-
assert (target / 'meta.json').read_bytes() == b'{}'
243-
touched(path)
244-
245-
monkeypatch.setattr(fs, 'touch', touch)
246-
export_selected_files([str(source)],
247-
str(target) + ('/' if directory else ''), fs=fs,
248-
identity='job-1', sidecars={'meta.json': b'{}'})
249-
if directory:
250-
assert (target / '_SUCCESS').exists()
251-
else:
252-
assert not (target / 'meta.json').exists()
253-
assert pq.read_table(tmp_path / 'output.parquet').num_rows == 1
254-
255-
256-
def test_failed_sidecar_write_still_publishes(tmp_path):
228+
def test_sidecars_are_best_effort_and_precede_success(tmp_path):
257229
import pyarrow.parquet as pq
258230
from kumo_core.artifact_export import export_selected_files
259231
from kumo_core.object_store_sink import _ArrowExportFiles
260232
from pyarrow.fs import LocalFileSystem
261233

262234
source = tmp_path / 'part-0.parquet'
263235
pq.write_table(pa.table({'id': [1]}), source)
236+
target = tmp_path / 'output'
264237
fs = _ArrowExportFiles(LocalFileSystem(), None)
265-
write_bytes = fs.write_bytes
238+
write_bytes, touch = fs.write_bytes, fs.touch
266239

267240
def denied(path, value):
268-
if path.endswith('meta.json'):
241+
if path.endswith('bad.json'):
269242
raise PermissionError('denied')
270243
write_bytes(path, value)
271244

272-
fs.write_bytes = denied
273-
target = tmp_path / 'output'
245+
def checked_touch(path):
246+
assert (target / 'meta.json').read_bytes() == b'{}'
247+
touch(path)
248+
249+
fs.write_bytes, fs.touch = denied, checked_touch
274250
with pytest.warns(UserWarning, match='denied'):
275251
export_selected_files([str(source)], f'{target}/', fs=fs,
276-
identity='job-1', sidecars={'meta.json': b'{}'})
252+
identity='job-1', sidecars={
253+
'meta.json': b'{}',
254+
'bad.json': b'{}'
255+
})
277256
assert (target / '_SUCCESS').exists()

‎packages/kumo-engine/kumo/engine/pquery/_custom_source.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,7 @@
1818

1919
def stage_source_table(source_table: SourceTable,
2020
destination: Path) -> Tuple[Path, Optional[str]]:
21-
"""Read ``source_table`` through its connector and retain its parquet
22-
parts; also returns the connector's version of the rows read."""
21+
"""Stage ``source_table``'s parquet parts; returns them and their version."""
2322
if not isinstance(source_table, SourceTable):
2423
raise TypeError("source_table must be a kumo.engine SourceTable")
2524
destination.mkdir(parents=True, exist_ok=False)

‎packages/kumo-engine/kumo/engine/pquery/training_table.py‎

Lines changed: 16 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -2,27 +2,24 @@
22
# All rights reserved.
33
# SPDX-License-Identifier: Apache-2.0
44

5-
import json
6-
import logging
75
import warnings
86
from dataclasses import dataclass, replace
97
from pathlib import Path
10-
from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence
8+
from typing import TYPE_CHECKING, Any, Dict, Optional
119

1210
from kumo_core.model_plan import TrainingTableGenerationPlan
1311
from kumoapi.model_plan import \
1412
TrainingTableGenerationPlan as PublicGenerationPlan
1513
from kumoapi.train import TrainingTableSpec
1614

15+
from kumo.engine import _provenance as provenance
1716
from kumo.engine._manifest import new_manifest, read_manifest, write_manifest
1817
from kumo.engine.connector.source_table import SourceTable
1918
from kumo.engine.graph import GraphSnapshot
2019

2120
if TYPE_CHECKING:
2221
from kumo.engine.trainer.output_config import TrainingTableOutputConfig
2322

24-
logger = logging.getLogger(__name__)
25-
2623
STAGING_MANIFEST_FILE = 'training-table-manifest.json'
2724
EXPORT_MANIFEST_FILE = '_kumo_manifest.json'
2825
_RANDOM_SPLITS = ('RandomSplit', 'RandomLinkSplit')
@@ -84,40 +81,22 @@ def export(self, output_config: "TrainingTableOutputConfig") -> str:
8481
root.mkdir(parents=True, exist_ok=True)
8582
with tempfile.TemporaryDirectory(dir=root) as output:
8683
files = export_train_table(str(self.staging_dir), output)
87-
try:
88-
payload = self._export_manifest(connector, output_config,
89-
files)
90-
except Exception as error:
91-
warnings.warn(f"Exporting without a training-table manifest: "
92-
f"{error}", UserWarning, stacklevel=2)
93-
payload = None
94-
extra = ({} if payload is None else {
95-
'sidecars': {
96-
EXPORT_MANIFEST_FILE: payload
97-
}
98-
})
9984
return write_files(connector, output_config.table_name, files,
10085
output_config.write_mode, export_id=export_id,
101-
**extra)
86+
**self._export_sidecars(connector,
87+
output_config))
10288

103-
def _export_manifest(self, connector, output_config,
104-
files: Sequence[str]) -> Optional[bytes]:
89+
def _export_sidecars(self, connector, output_config) -> Dict[str, Any]:
10590
from kumo_core.object_store_sink import ObjectStoreSinkBackend
10691
from kumo_core.sources import layout
10792

108-
name = output_config.table_name
93+
manifest = self.staging_dir / STAGING_MANIFEST_FILE
94+
name = output_config.table_name.rstrip('/')
10995
if (not isinstance(connector.sink_backend, ObjectStoreSinkBackend)
110-
or not (name.endswith('/') or len(files) > 1)
111-
or layout.file_type(name.rstrip('/')) not in (
112-
None, layout.PARQUET)):
113-
logger.debug('Training-table manifests are only exported to '
114-
'Parquet directories on files or object stores.')
115-
return None
116-
manifest = self.manifest
117-
if manifest is None:
118-
return None
119-
return (json.dumps(manifest, indent=2, sort_keys=True) +
120-
'\n').encode('utf-8')
96+
or layout.file_type(name) == layout.CSV
97+
or not manifest.is_file()):
98+
return {}
99+
return {'sidecars': {EXPORT_MANIFEST_FILE: manifest.read_bytes()}}
121100

122101
@classmethod
123102
def from_source_table(
@@ -182,7 +161,6 @@ def from_source_table(
182161
run_dir = train_table_scratch_dir(run_id)
183162
rows_path, version = stage_source_table(source_table,
184163
run_dir / "source-table")
185-
lineage = _source_lineage(source_table, version)
186164
generation_plan = from_json_dict(
187165
to_json_dict(public_plan),
188166
TrainingTableGenerationPlan,
@@ -247,8 +225,11 @@ def from_source_table(
247225
required_stages=task_plan.task_dataset.required_stages,
248226
generation_plan=generation_plan,
249227
)
250-
_record_manifest(table, snapshot, task_plan.train_table, 'custom',
251-
source_table=lineage)
228+
_record_manifest(
229+
table, snapshot, task_plan.train_table, 'custom',
230+
source_table=provenance.source_table(source_table.connector,
231+
source_table.table_name,
232+
source_table.sql, version))
252233
return table
253234

254235

@@ -285,11 +266,3 @@ def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot,
285266
warnings.warn(
286267
f"Could not write the training-table manifest at {path}: "
287268
f"{error}", UserWarning, stacklevel=3)
288-
289-
290-
def _source_lineage(source_table: SourceTable,
291-
version: Optional[str]) -> Dict[str, Optional[str]]:
292-
from kumo.engine._provenance import source_table as lineage
293-
294-
return lineage(source_table.connector, source_table.table_name,
295-
source_table.sql, version)

‎packages/kumo-engine/test/test_training_table_export.py‎

Lines changed: 20 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -295,84 +295,38 @@ def test_custom_training_table_declares_integer_timestamp_unit(graph, unit,
295295
assert _rows(seeds[stage], columns) == _rows(original_seeds[stage], columns)
296296

297297

298-
def _reload(graph, name, pquery):
299-
_, snapshot, connector, _ = graph
300-
return kumo.TrainingTable.from_source_table(
301-
connector[name], TrainingTableSpec(split_col='SPLIT'), pq=pquery,
302-
snapshot=snapshot)
303-
304-
305298
def test_manifest_and_export_sidecar(graph, tmp_path):
306299
from kumo_core.training.export import export_train_table
307300

308301
_, snapshot, connector, out = graph
309302
query = 'PREDICT ents.LABEL FOR EACH ents.ID'
310303
pquery, table, _ = _export(graph, query, 'lineage')
311304
manifest = table.manifest
312-
assert manifest['kind'] == 'training_table'
313305
assert manifest['id'] == table.tt_id == _spec(table)['tt_id']
314306
assert manifest['parents'] == {'snapshot': snapshot.id}
315307
assert (manifest['query'], manifest['source']) == (query, 'pql')
316-
assert manifest['task_type'] == 'binary_classification'
317308
assert manifest['rows'] == _seeds(table)[1]
318309
assert manifest['split_seed'] == 0
319-
assert manifest['generation_plan']['split'] == 'inferred'
320310

321311
directory = out / 'lineage'
322-
parts = sorted(directory.glob('*.parquet'))
323312
baseline = tmp_path / 'baseline'
324-
expected = export_train_table(str(table.staging_dir), str(baseline))
325-
assert [p.name for p in parts] == [p.split('/')[-1] for p in expected]
326-
assert [p.read_bytes() for p in parts] == [
327-
(baseline / p.name).read_bytes() for p in parts
328-
]
329-
assert sorted(p.name for p in directory.iterdir()) == sorted(
330-
[p.name for p in parts] +
331-
['_KUMO_EXPORT', '_SUCCESS', '_kumo_manifest.json'])
332-
sidecar = (directory / '_kumo_manifest.json').read_bytes()
333-
assert json.loads(sidecar) == manifest
334-
335-
table.export(TrainingTableOutputConfig(connector, 'lineage'))
336-
assert (directory / '_kumo_manifest.json').read_bytes() == sidecar
337-
338-
339-
def test_custom_table_records_its_source(graph, recwarn):
340-
_, snapshot, connector, out = graph
341-
pquery, table, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID',
342-
'reloaded')
343-
(out / 'reloaded' / '_kumo_manifest.json').write_text('{not json')
344-
reloaded = _reload(graph, 'reloaded', pquery)
313+
export_train_table(str(table.staging_dir), str(baseline))
314+
for part in baseline.glob('*.parquet'):
315+
assert (directory / part.name).read_bytes() == part.read_bytes()
316+
assert json.loads(
317+
(directory / '_kumo_manifest.json').read_bytes()) == manifest
318+
319+
reloaded = kumo.TrainingTable.from_source_table(
320+
connector['lineage'], TrainingTableSpec(split_col='SPLIT'), pq=pquery,
321+
snapshot=snapshot)
345322
manifest = reloaded.manifest
346323
assert manifest['source'] == 'custom'
347-
assert manifest['id'] != table.tt_id
348-
assert manifest['parents'] == {
349-
'snapshot': snapshot.id,
350-
'source_table': {
351-
'connector': 'LOCAL_FILE',
352-
'location': str(out / 'reloaded'),
353-
'version': connector.source_backend.get_data_version(
354-
connector, 'reloaded'),
355-
},
324+
assert manifest['parents']['source_table'] == {
325+
'connector': 'LOCAL_FILE',
326+
'location': str(directory),
327+
'version': connector.source_backend.get_data_version(
328+
connector, 'lineage'),
356329
}
357-
assert manifest['parents']['source_table']['version']
358-
assert not [w for w in recwarn if 'manifest' in str(w.message)]
359-
assert _seeds(reloaded)[1] == _seeds(table)[1]
360-
361-
362-
def test_unrecordable_location_is_null(graph, monkeypatch):
363-
from kumo.engine import _provenance
364-
365-
pquery, table, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID',
366-
'unrecorded')
367-
368-
def fail(*args, **kwargs):
369-
raise OSError('listing denied')
370-
371-
monkeypatch.setattr(_provenance, 'source_location', fail)
372-
with pytest.warns(UserWarning, match='listing denied'):
373-
reloaded = _reload(graph, 'unrecorded', pquery)
374-
recorded = reloaded.manifest['parents']['source_table']
375-
assert recorded['location'] is None and recorded['version']
376330
assert _seeds(reloaded)[1] == _seeds(table)[1]
377331

378332

@@ -387,37 +341,27 @@ def fail(*args, **kwargs):
387341

388342
monkeypatch.setattr(training_table, 'write_manifest', fail)
389343
with pytest.warns(UserWarning, match='read-only scratch'):
390-
pquery, table, exported = _export(
344+
_, table, exported = _export(
391345
graph, 'PREDICT ents.LABEL FOR EACH ents.ID', 'no_manifest')
392346
assert table.manifest is None
393347
assert not (out / 'no_manifest' / '_kumo_manifest.json').exists()
394-
assert (out / 'no_manifest' / '_SUCCESS').exists()
395348
assert exported.num_rows == sum(_seeds(table)[1].values())
396349

397350

398-
def test_warehouse_sink_receives_no_sidecars(graph):
399-
graph_, snapshot, _, _ = graph
351+
def test_no_sidecar_for_warehouse_or_csv_exports(graph):
352+
graph_, snapshot, connector, out = graph
400353
table = kumo.PredictiveQuery(
401354
graph=graph_, query='PREDICT ents.LABEL FOR EACH ents.ID'
402355
).generate_training_table(snapshot)
403356
calls = []
404357

405358
def write_files(connector, table_name, files, write_mode, **kwargs):
406359
calls.append(kwargs)
407-
return table_name
408360

409-
connector = SimpleNamespace(
361+
warehouse = SimpleNamespace(
410362
source_type='SNOWFLAKE',
411363
sink_backend=SimpleNamespace(write_files=write_files))
412-
table.export(TrainingTableOutputConfig(connector, 'DB.SCHEMA.TT'))
364+
table.export(TrainingTableOutputConfig(warehouse, 'DB.SCHEMA.TT'))
413365
assert calls == [{'export_id': f'{table.tt_id}:tt-export-v1'}]
414-
415-
416-
def test_csv_directory_export_gets_no_sidecar(graph):
417-
graph_, snapshot, connector, out = graph
418-
table = kumo.PredictiveQuery(
419-
graph=graph_, query='PREDICT ents.LABEL FOR EACH ents.ID'
420-
).generate_training_table(snapshot)
421366
table.export(TrainingTableOutputConfig(connector, 'rows.csv/'))
422-
names = {p.name for p in (out / 'rows.csv').iterdir()}
423-
assert '_SUCCESS' in names and '_kumo_manifest.json' not in names
367+
assert not (out / 'rows.csv' / '_kumo_manifest.json').exists()

0 commit comments

Comments
 (0)