Add enriched dataset metadata and storage config registration - #396
Add enriched dataset metadata and storage config registration#396misiugodfrey wants to merge 8 commits into
Conversation
Enriches metadata.json written by generate_data_files.py to capture all
generation parameters (benchmark_type, generator, generator_version,
parquet_version, codec_definitions, convert_decimals_to_floats,
max_rows_per_file, approx_row_group_bytes, data_dir_path) so a dataset
can be fully replicated from metadata alone.
Adds register_storage_config.py, a new tool that POSTs a dataset to
/api/storage-configs/ and writes the resulting storage_configuration_name
back into metadata.json. Storage system type and GDS are auto-detected
from the filesystem (findmnt/lsblk/lsmod) with CLI overrides; the config
name is auto-generated as {machine}-{storage_system}-{benchmark_type}-sf{N}.
A pre-flight GET check prevents duplicate names. A --register flag group
on generate_data_files.py calls this step immediately after generation for
the common case; datasets that need modification before registration can
call register_storage_config.py separately.
Updates post_results.py to make --storage-configuration-name optional:
it is auto-detected from storage_configuration_name in metadata.json
found via data_dir in benchmark_result.json context. A mismatch between
an explicit flag value and the metadata value produces a warning rather
than an error, preserving intentional overrides.
findmnt without --target requires an exact mount point; using --target correctly traverses up to find the filesystem for any given path.
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Moves the constant to the top of the file so it serves as the single source of truth for both the actual tpchgen-cli invocation and the metadata.json record, preventing them from silently diverging.
httpx is already in benchmark_data_tools/requirements.txt and the script is imported as a module, so the PEP 723 block serves no purpose.
When --storage-configuration-name is explicitly supplied but the dataset has no registered name in metadata.json, print a warning so the user knows the provided name was not checked against the actual dataset.
| rc = register_storage_config( | ||
| data_dir=Path(args.data_dir_path), | ||
| machine=args.register_machine, | ||
| name=args.register_name, |
There was a problem hiding this comment.
Most of these should be automatically determined, but they can be manually overwritten via flag if desired.
| default=False, | ||
| help="Register the dataset as a storage configuration after generation.", | ||
| ) | ||
| reg.add_argument( |
There was a problem hiding this comment.
It's debatable if we want to include these overrides; they should usually not be used.
| return 1 | ||
| storage_configuration_name = resolved_storage_config | ||
| elif resolved_storage_config and resolved_storage_config != storage_configuration_name: | ||
| print( |
There was a problem hiding this comment.
Right now we give some deference to the user when there is a mismatch - assuming the user wants to submit using an older storage configuration and they know what they are doing. We could make this a proper error, but right now I don't want to block older workflows.
Perhaps we strengthen this once we have generated new datasets and hardened their usage?
| def _detect_gds_enabled() -> bool: | ||
| """Return True if the nvidia-fs kernel module (GPU Direct Storage) is loaded.""" | ||
| try: | ||
| result = subprocess.run(["lsmod"], capture_output=True, text=True) |
There was a problem hiding this comment.
Just a note that I think this is reflecting an issue with the database's data model: whether some dataset is over accessible over GDS is less about the dataset itself and more about the compute node and its path to the storage system. This check doesn't touch the dataset or its storage system: just the system that lsmod is running on.
Nothing to do for now, but LMK if this felt awkward to you too and maybe we can find a better way to model that.
There was a problem hiding this comment.
That's a fair point. I think for now there should not be a difference since the system should always be running on the same nodes that access the data (this is true even on our slurm systems); but it's something I'll confirm in testing.
Motivation
Closes #389, closes #388, partially addresses #164.
When running benchmarks and posting results via
post_results.py, the--storage-configuration-nameflag has to be supplied manually by the user. There is no mechanism to verify that the storage configuration matches the actual dataset the benchmark ran on — a mismatch is silent and corrupts the database record. Similarly,metadata.jsonwritten bygenerate_data_files.pyonly recordedscale_factorandapprox_row_group_bytes, making it impossible to fully reproduce a dataset or trace what configuration produced it.Summary
Enriched
metadata.json:generate_data_files.pynow writes a complete record of how a dataset was generated — benchmark type, generator (tpchgen/duckdb) and version, scale factor, parquet version, codec definitions, row group settings, and the absolute data path. The schema is self-contained enough to fully replicate a dataset from metadata alone.New
register_storage_config.py: Standalone CLI that registers a dataset as a storage configuration in the benchmarking database (POST /api/storage-configs/). Auto-detects storage system type (viafindmnt+lsblk) and GDS status (vialsmod); derives compression from codec definitions; auto-generates the config name as{machine}-{storage_system}-{benchmark_type}-sf{N}. Does a pre-flightGETto prevent duplicate names. Writesstorage_configuration_nameback intometadata.jsonafter successful registration. Also callable inline via a--registerflag group ongenerate_data_files.pyfor simple generate-and-register workflows — datasets that need modification before registration can call the script separately.Auto-detection and validation in
post_results.py:--storage-configuration-nameis now optional. The script readsstorage_configuration_namefrommetadata.jsonin the dataset directory (viadata_dirinbenchmark_result.jsoncontext). When a name is provided explicitly, it is validated against the registered name inmetadata.json:register_storage_config.pyTest plan
metadata.jsonbenchmark_typeand labels are correctregister_storage_config.py --dry-runon tpchgen dataset — verify payload shape, auto-detected fields, auto-generated name, andcodec_definitionsembedded intable_metadataregister_storage_config.py --dry-runon DuckDB dataset — verify tpchgen-only fields absent fromtable_metadatapost_results.py --dry-runwithout--storage-configuration-name— verified auto-detection frommetadata.jsonviadata_dirin benchmark context_resolve_storage_configurationreturnsNonewith helpful message whenstorage_configuration_nameabsent frommetadata.json_resolve_storage_configurationreturnsNonewith helpful message when nometadata.jsonfoundBenchmarkMetadata.data_dircorrectly populated from benchmark result contextregister_storage_config.pywith a live API key — verify pre-flight collision check andstorage_configuration_namewritten back tometadata.jsonpost_results.pywith mismatched--storage-configuration-name— verify warning is printedpost_results.pywith manually provided--storage-configuration-nameand no registered name inmetadata.json— verify warning is printedpost_results.pywith nometadata.jsonand no flag — verify helpful error message--registerinline flag ongenerate_data_files.py--convert-decimals-to-floatsreflected correctly inmetadata.json--nameoverride onregister_storage_config.py