Skip to content

Commit

Permalink
style
Browse files Browse the repository at this point in the history
  • Loading branch information
Colin Ho authored and Colin Ho committed Dec 10, 2024
1 parent cb75de2 commit c1957e3
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions daft/runners/ray_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
# The ray runner is not a top-level module, so we don't need to lazily import pyarrow to minimize
# import times. If this changes, we first need to make the daft.lazy_import.LazyImport class
# serializable before importing pa from daft.dependencies.
import pyarrow as pa
import pyarrow as pa # noqa: TID253
import ray.experimental # noqa: TID253

from daft.arrow_utils import ensure_array
Expand Down Expand Up @@ -1018,7 +1018,9 @@ def _build_partitions(
else reduce_pipeline
)
if task.node_id is not None:
ray_options["scheduling_strategy"] = ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(task.node_id, soft=True)
ray_options["scheduling_strategy"] = ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
task.node_id, soft=True
)
else:
ray_options["scheduling_strategy"] = "SPREAD"
build_remote = build_remote.options(**ray_options).with_tracing(runner_tracer, task)
Expand Down

0 comments on commit c1957e3

Please sign in to comment.