Skip to content

Commit

Permalink
[CHORE] Update execution config to turn on Ray tracing (#3431)
Browse files Browse the repository at this point in the history
Co-authored-by: Jay Chia <[email protected]@users.noreply.github.com>
  • Loading branch information
jaychia and Jay Chia authored Nov 26, 2024
1 parent b201c61 commit 3ece13a
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions daft/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()`)
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 3ece13a

Please sign in to comment.