From 94e0e4cdcd7ef1ee6898f240ab9cddbbfef4a0d5 Mon Sep 17 00:00:00 2001 From: Raunak Bhagat Date: Thu, 21 Nov 2024 13:49:54 -0800 Subject: [PATCH] Change double-negative flag to single-negative flag --- benchmarking/tpch/__main__.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/benchmarking/tpch/__main__.py b/benchmarking/tpch/__main__.py index 125712ef92..ee7d2d191c 100644 --- a/benchmarking/tpch/__main__.py +++ b/benchmarking/tpch/__main__.py @@ -126,7 +126,7 @@ def run_all_benchmarks( csv_output_location: str | None, ray_job_dashboard_url: str | None = None, requirements: str | None = None, - no_pymodules: bool = False, + pickle_daft_module: bool = True, ): get_df = get_df_with_parquet_folder(parquet_folder) @@ -144,7 +144,7 @@ def run_all_benchmarks( tpch_qnum=i, working_dir=working_dir, entrypoint=entrypoint, - runtime_env=get_ray_runtime_env(requirements, no_pymodules=no_pymodules), + runtime_env=get_ray_runtime_env(requirements, pickle_daft_module=pickle_daft_module), ) # Run once as a warmup step @@ -203,10 +203,10 @@ def get_daft_benchmark_runner_name() -> Literal["ray"] | Literal["py"] | Literal return name -def get_ray_runtime_env(requirements: str | None, no_pymodules: bool = False) -> dict: +def get_ray_runtime_env(requirements: str | None, pickle_daft_module: bool = True) -> dict: daft_env_variables = dict(filter(lambda key_value: key_value[0].startswith("DAFT"), os.environ.items())) runtime_env = { - "py_modules": None if no_pymodules else [daft], + "py_modules": [daft] if pickle_daft_module else None, "eager_install": True, "env_vars": { **daft_env_variables, @@ -297,9 +297,10 @@ def warm_up_function(): help="Ray Dashboard URL to submit jobs instead of using Ray client, most useful when running on a remote cluster", ) parser.add_argument( - "--no_pymodules", - action="store_true", - help="Avoid pickling any modules in the ray-environment before initializing the Ray cluster; useful in CI", + "--pickle_daft_module", + type=lambda arg: arg.lower() == "true", + default=True, + help="Avoid pickling the daft module in the ray-environment before initializing the Ray cluster; useful to turn *off* in CI when wheels are being used instead", ) args = parser.parse_args() @@ -339,5 +340,5 @@ def warm_up_function(): csv_output_location=args.output_csv, ray_job_dashboard_url=args.ray_job_dashboard_url, requirements=args.requirements, - no_pymodules=True if args.no_pymodules else False, + pickle_daft_module=args.pickle_daft_module, )