diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils_new_frontends.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils_new_frontends.py index b0d8ba0487f1..6305d0dda25f 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils_new_frontends.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils_new_frontends.py @@ -14,6 +14,7 @@ import logging import os import pprint +import shlex import sys import textwrap import time @@ -430,6 +431,8 @@ class RunConfig: timestamp: str = dataclasses.field( default_factory=lambda: datetime.now(UTC).isoformat() ) + command_line: str + capture_env_vars: str def __post_init__(self) -> None: # noqa: D105 if self.io_mode == "hot" and self.iterations < 2: @@ -438,6 +441,12 @@ def __post_init__(self) -> None: # noqa: D105 "iteration 0 warms the cache, iterations 1+ are the hot measurements." ) + # Update `extra_info.environment` with the captured environment variables. + self.extra_info.setdefault("environment", {}) + for var in self.capture_env_vars.split(","): + var_ = var.strip() + self.extra_info["environment"][var_] = os.environ.get(var_) + @classmethod def from_args(cls, args: argparse.Namespace) -> RunConfig: """Create a RunConfig from command line arguments.""" @@ -543,6 +552,8 @@ def from_args(cls, args: argparse.Namespace) -> RunConfig: duckdb_threads=args.duckdb_threads, duckdb_memory_limit=args.duckdb_memory_limit, duckdb_temp_dir=args.duckdb_temp_dir, + command_line=shlex.join(sys.argv), + capture_env_vars=args.capture_env_vars, ) def serialize(self, engine: StreamingEngine | None) -> dict: @@ -567,6 +578,7 @@ def serialize(self, engine: StreamingEngine | None) -> dict: "extra_info": self.extra_info, "run_id": str(self.run_id), "timestamp": self.timestamp, + "command_line": self.command_line, "streaming_options": { "rapidsmpf": opts.to_rapidsmpf_options().get_strings(), "executor": opts.to_executor_options(), @@ -1892,6 +1904,12 @@ def build_parser(num_queries: int = 22) -> argparse.ArgumentParser: default=None, help="Directory for DuckDB to spill temporary data to disk.", ) + parser.add_argument( + "--capture-env-vars", + type=str, + default="CUDF_POLARS_LOG_TRACES_MEMORY,CUDF_POLARS_LOG_TRACES,DASK_DISTRIBUTED__COMM__TIMEOUTS__CONNECT,DASK_DISTRIBUTED__COMM__UCX__CONNECT_TIMEOUT,KVIKIO_NTHREADS,LIBCUDF_NUM_HOST_WORKERS,OMP_NUM_THREADS,POLARS_MAX_THREADS,RAPIDSMPF_num_streaming_threads,UCX_MAX_RNDV_RAILS,UCX_PROTO_ENABLE,UCX_RNDV_FRAG_MEM_TYPES,UCX_RNDV_MTYPE_WORKER_FC_ENABLE,UCX_RNDV_MTYPE_WORKER_MAX_MEM,UCX_RNDV_PIPELINE_ERROR_HANDLING", + help="Comma-separated list of environment variables to capture. Written to ``extra_info.environment``.", + ) StreamingOptions._add_cli_args(parser)