From 3ece13aaf603e512f82b0005de2abf4d8e03250c Mon Sep 17 00:00:00 2001 From: Jay Chia <17691182+jaychia@users.noreply.github.com> Date: Tue, 26 Nov 2024 11:13:27 -0800 Subject: [PATCH] [CHORE] Update execution config to turn on Ray tracing (#3431) Co-authored-by: Jay Chia --- daft/context.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/daft/context.py b/daft/context.py index bb9aa3a0a5..fa1a61a637 100644 --- a/daft/context.py +++ b/daft/context.py @@ -350,6 +350,7 @@ def set_execution_config( default_morsel_size: int | None = None, shuffle_algorithm: str | None = None, pre_shuffle_merge_threshold: int | None = None, + enable_ray_tracing: bool | None = None, ) -> DaftContext: """Globally sets various configuration parameters which control various aspects of Daft execution. These configuration values are used when a Dataframe is executed (e.g. calls to `.write_*`, `.collect()` or `.show()`) @@ -388,6 +389,7 @@ def set_execution_config( default_morsel_size: Default size of morsels used for the new local executor. Defaults to 131072 rows. shuffle_algorithm: The shuffle algorithm to use. Defaults to "map_reduce". Other options are "pre_shuffle_merge". pre_shuffle_merge_threshold: Memory threshold in bytes for pre-shuffle merge. Defaults to 1GB + enable_ray_tracing: Enable tracing for Ray. Accessible in `/tmp/ray/session_latest/logs/daft` after the run completes. Defaults to False. """ # Replace values in the DaftExecutionConfig with user-specified overrides ctx = get_context() @@ -415,6 +417,7 @@ def set_execution_config( default_morsel_size=default_morsel_size, shuffle_algorithm=shuffle_algorithm, pre_shuffle_merge_threshold=pre_shuffle_merge_threshold, + enable_ray_tracing=enable_ray_tracing, ) ctx._daft_execution_config = new_daft_execution_config