From 0da1421c095f7628abbebcf626db14d349f9031b Mon Sep 17 00:00:00 2001 From: JAGANNATHANJP Date: Sat, 23 May 2026 19:59:16 +0530 Subject: [PATCH 1/2] [df] Disallow Dask workers with more than one thread Using Dask workers with more than one thread does not provide any advantage, seeing how the RDataFrame computation graph runs in C++ and the Python threads are limited by the GIL anyway. Explicitly disallow using distributed RDataFrame with Dask when the workers have been configured to use more than one thread. Multithreading in RDataFrame can already be achieved on a single node via `ROOT.EnableImplicitMT`. Distributed processing runs with multiple processes instead. --- .../python/DistRDF/Backends/Dask/Backend.py | 13 ++++++ .../python/distrdf/backends/check_backend.py | 42 +++++++++++++++++++ 2 files changed, 55 insertions(+) diff --git a/bindings/distrdf/python/DistRDF/Backends/Dask/Backend.py b/bindings/distrdf/python/DistRDF/Backends/Dask/Backend.py index 84a9a0bbd4abc..fe7d65e4e369a 100644 --- a/bindings/distrdf/python/DistRDF/Backends/Dask/Backend.py +++ b/bindings/distrdf/python/DistRDF/Backends/Dask/Backend.py @@ -105,6 +105,19 @@ def __init__(self, daskclient: Optional[Client] = None): self.client = (daskclient if daskclient is not None else Client(LocalCluster(n_workers=os.cpu_count(), threads_per_worker=1, processes=True))) + workers = self.client.scheduler_info().get("workers", None) + + if workers is not None: + for worker in workers.values(): + threads = worker.get("nthreads", 1) + + if threads > 1: + raise RuntimeError( + "RDataFrame: running in distributed mode with Dask workers using more than one thread is " + "not supported. Please make sure that your Dask cluster was created with the appropriate " + "options, e.g. `processes=True` and `threads_per_worker=1`." + ) + def optimize_npartitions(self) -> int: """ Attempts to compute a clever number of partitions for the current diff --git a/roottest/python/distrdf/backends/check_backend.py b/roottest/python/distrdf/backends/check_backend.py index 010c96c299417..13af0a1b9a1fd 100644 --- a/roottest/python/distrdf/backends/check_backend.py +++ b/roottest/python/distrdf/backends/check_backend.py @@ -62,6 +62,48 @@ def test_optimize_npartitions(self, payload): backend = Backend.SparkBackend(sparkcontext=connection) assert backend.optimize_npartitions() == 2 + def test_dask_backend_handles_missing_workers(self, payload): + """ + Check that DaskBackend initialization succeeds when scheduler_info + does not provide worker information. + """ + connection, backend = payload + + if backend != "dask": + return + + from ROOT._distrdf.Backends.Dask import Backend + + original_scheduler_info = connection.scheduler_info + + try: + connection.scheduler_info = lambda: {} + + backend = Backend.DaskBackend(daskclient=connection) + assert backend.client is connection + + df = ROOT.RDataFrame(10, executor=connection) + assert df.Count().GetValue() == 10 + + finally: + connection.scheduler_info = original_scheduler_info + + def test_dask_backend_rejects_threaded_workers(self): + """ + Check that DaskBackend rejects threaded workers. + """ + from dask.distributed import Client, LocalCluster + from ROOT._distrdf.Backends.Dask import Backend + + with ( + LocalCluster(n_workers=1, threads_per_worker=2, processes=False, dashboard_address=":0") as cluster, + Client(cluster) as client, + pytest.raises( + RuntimeError, + match="running in distributed mode with Dask workers using more than one thread is not supported", + ) + ): + Backend.DaskBackend(daskclient=client) class TestInitialization: """Check initialization method in the Dask backend""" From 30314f02a4906ce51890964108880090bea897c6 Mon Sep 17 00:00:00 2001 From: Vincenzo Eduardo Padulano Date: Fri, 31 Jul 2026 11:25:01 +0200 Subject: [PATCH 2/2] [df] Format changed files with ruff --- .../python/DistRDF/Backends/Dask/Backend.py | 187 ++++++++++-------- .../python/distrdf/backends/check_backend.py | 3 +- 2 files changed, 104 insertions(+), 86 deletions(-) diff --git a/bindings/distrdf/python/DistRDF/Backends/Dask/Backend.py b/bindings/distrdf/python/DistRDF/Backends/Dask/Backend.py index fe7d65e4e369a..b9f0df9ef03e4 100644 --- a/bindings/distrdf/python/DistRDF/Backends/Dask/Backend.py +++ b/bindings/distrdf/python/DistRDF/Backends/Dask/Backend.py @@ -72,8 +72,10 @@ def get_total_cores_jobqueuecluster(cluster: JobQueueCluster) -> int: # 'cores' key for any type of dask-jobqueue cluster. return sum(spec["options"]["cores"] for spec in workers_spec.values()) except KeyError as e: - raise RuntimeError("Could not retrieve the provided worker specification from the Dask cluster object. " - "Please report this as a bug.") from e + raise RuntimeError( + "Could not retrieve the provided worker specification from the Dask cluster object. " + "Please report this as a bug." + ) from e def get_total_cores(client: Client) -> int: @@ -84,6 +86,7 @@ def get_total_cores(client: Client) -> int: # It may happen that the user is connected to a batch system. We try # to import the 'dask_jobqueue' module lazily to avoid a dependency. from dask_jobqueue import JobQueueCluster + if isinstance(client.cluster, JobQueueCluster): return get_total_cores_jobqueuecluster(client.cluster) except ModuleNotFoundError: @@ -102,8 +105,11 @@ def __init__(self, daskclient: Optional[Client] = None): # `daskclient` will be `None`. In this case, we create a default Dask # client connected to a cluster instance with N worker processes, where # N is the number of cores on the local machine. - self.client = (daskclient if daskclient is not None else - Client(LocalCluster(n_workers=os.cpu_count(), threads_per_worker=1, processes=True))) + self.client = ( + daskclient + if daskclient is not None + else Client(LocalCluster(n_workers=os.cpu_count(), threads_per_worker=1, processes=True)) + ) workers = self.client.scheduler_info().get("workers", None) @@ -128,12 +134,14 @@ def optimize_npartitions(self) -> int: return get_total_cores(self.client) @staticmethod - def dask_mapper(current_range: Tuple, - headers: List[str], - shared_libraries: List[str], - pcms: List[str], - files: List[str], - mapper: Callable) -> Callable: + def dask_mapper( + current_range: Tuple, + headers: List[str], + shared_libraries: List[str], + pcms: List[str], + files: List[str], + mapper: Callable, + ) -> Callable: """ Gets the paths to the file(s) in the current executor, then declares the headers found. @@ -154,34 +162,32 @@ def dask_mapper(current_range: Tuple, """ # Retrieve the current worker local directory localdir = get_worker().local_directory - - #Get and declare headers on each worker - headers_on_executor = [ - os.path.join(localdir, os.path.basename(filepath)) - for filepath in headers - ] + + # Get and declare headers on each worker + headers_on_executor = [os.path.join(localdir, os.path.basename(filepath)) for filepath in headers] Utils.distribute_headers(headers_on_executor) # Get and declare shared libraries on each worker - shared_libs_on_ex = [ - os.path.join(localdir, os.path.basename(filepath)) - for filepath in shared_libraries - ] - + shared_libs_on_ex = [os.path.join(localdir, os.path.basename(filepath)) for filepath in shared_libraries] + Utils.distribute_shared_libraries(shared_libs_on_ex) return mapper(current_range) - def ProcessAndMerge(self, - ranges: List[Any], - mapper: Callable[[Ranges.DataRange, - Callable[[Union[Ranges.EmptySourceRange, Ranges.TreeRangePerc]], - Base.TaskObjects], - Callable[[ROOT.RDF.RNode, int], List], - Callable], - Base.TaskResult], - reducer: Callable[[Base.TaskResult, Base.TaskResult], Base.TaskResult], - ) -> Base.TaskResult: + def ProcessAndMerge( + self, + ranges: List[Any], + mapper: Callable[ + [ + Ranges.DataRange, + Callable[[Union[Ranges.EmptySourceRange, Ranges.TreeRangePerc]], Base.TaskObjects], + Callable[[ROOT.RDF.RNode, int], List], + Callable, + ], + Base.TaskResult, + ], + reducer: Callable[[Base.TaskResult, Base.TaskResult], Base.TaskResult], + ) -> Base.TaskResult: """ Performs map-reduce using Dask framework. @@ -196,21 +202,21 @@ def ProcessAndMerge(self, Returns: list: A list representing the values of action nodes returned after computation (Map-Reduce). - """ - self.distribute_unique_paths(self.headers) + """ + self.distribute_unique_paths(self.headers) self.distribute_unique_paths(self.shared_libraries) self.distribute_unique_paths(self.pcms) self.distribute_unique_paths(self.files) - - + dmapper = dask.delayed(DaskBackend.dask_mapper) dreducer = dask.delayed(reducer) - mergeables_lists = [dmapper(range, self.headers, self.shared_libraries, self.pcms, self.files, mapper) for range in ranges] - + mergeables_lists = [ + dmapper(range, self.headers, self.shared_libraries, self.pcms, self.files, mapper) for range in ranges + ] + while len(mergeables_lists) > 1: - mergeables_lists.append( - dreducer(mergeables_lists.pop(0), mergeables_lists.pop(0))) + mergeables_lists.append(dreducer(mergeables_lists.pop(0), mergeables_lists.pop(0))) # Here we start the progressbar for the current RDF computation graph # running on the Dask client. This expects a future object, so we need @@ -226,19 +232,23 @@ def ProcessAndMerge(self, return final_results.compute() - def ProcessAndMergeLive(self, - ranges: List[Any], - mapper: Callable[[Ranges.DataRange, - Callable[[Union[Ranges.EmptySourceRange, Ranges.TreeRangePerc]], - Base.TaskObjects], - Callable[[ROOT.RDF.RNode, int], List], - Callable], - Base.TaskResult], - reducer: Callable[[Base.TaskResult, Base.TaskResult], Base.TaskResult], - drawables_info_dict: Dict[int, Tuple[List[Optional[Callable]], int, str]], - ) -> Base.TaskResult: + def ProcessAndMergeLive( + self, + ranges: List[Any], + mapper: Callable[ + [ + Ranges.DataRange, + Callable[[Union[Ranges.EmptySourceRange, Ranges.TreeRangePerc]], Base.TaskObjects], + Callable[[ROOT.RDF.RNode, int], List], + Callable, + ], + Base.TaskResult, + ], + reducer: Callable[[Base.TaskResult, Base.TaskResult], Base.TaskResult], + drawables_info_dict: Dict[int, Tuple[List[Optional[Callable]], int, str]], + ) -> Base.TaskResult: """ - Performs real-time map-reduce using Dask framework, retrieving the partial results + Performs real-time map-reduce using Dask framework, retrieving the partial results as soon as they are available, allowing real-time data representation. Args: @@ -250,23 +260,24 @@ def ProcessAndMergeLive(self, reducer (function): A function that merges two lists that were returned by the mapper. - drawables_info_dict (dict): A dictionary where keys are plot object IDs - and values are tuples containing optional callback functions, + drawables_info_dict (dict): A dictionary where keys are plot object IDs + and values are tuples containing optional callback functions, index of the plot object, and operation name. Returns: merged_results (TaskResult): The merged result of the computation. """ - - self.distribute_unique_paths(self.headers) + + self.distribute_unique_paths(self.headers) self.distribute_unique_paths(self.shared_libraries) self.distribute_unique_paths(self.pcms) self.distribute_unique_paths(self.files) - - + # Set up Dask mapper dmapper = dask.delayed(DaskBackend.dask_mapper) - mergeables_lists = [dmapper(range, self.headers, self.shared_libraries, self.pcms, self.files, mapper) for range in ranges] + mergeables_lists = [ + dmapper(range, self.headers, self.shared_libraries, self.pcms, self.files, mapper) for range in ranges + ] # Compute the delayed tasks to get Dask futures that can be passed to the as_completed method future_tasks = self.client.compute(mergeables_lists) @@ -286,7 +297,7 @@ def ProcessAndMergeLive(self, backend_pad.__destruct__() return merged_results - + def _setup_canvas(self, num_plots: int) -> ROOT.TCanvas: """ Set up a TCanvas for live visualization with divided pads based on the number of plots. @@ -308,29 +319,31 @@ def _setup_canvas(self, num_plots: int) -> ROOT.TCanvas: return c - def _process_partial_results(self, - canvas: ROOT.TCanvas, - drawables_info_dict: Dict[int, Tuple[List[Optional[Callable]], int, str]], - reducer: Callable[[Base.TaskResult, Base.TaskResult], Base.TaskResult], - future_tasks: List[dask.Future]) -> Base.TaskResult: + def _process_partial_results( + self, + canvas: ROOT.TCanvas, + drawables_info_dict: Dict[int, Tuple[List[Optional[Callable]], int, str]], + reducer: Callable[[Base.TaskResult, Base.TaskResult], Base.TaskResult], + future_tasks: List[dask.Future], + ) -> Base.TaskResult: """ Process partial results and display plots on the provided canvas. Args: canvas: The TCanvas object for displaying plots. - - drawables_info_dict (dict): A dictionary where keys are plot object IDs - and values are tuples containing optional callback functions, + + drawables_info_dict (dict): A dictionary where keys are plot object IDs + and values are tuples containing optional callback functions, index of the plot object, and operation name. - + reducer (function): A function for reducing partial results. - + future_tasks: Dask future tasks representing partial results. Returns: merged_results (TaskResult): The merged result of the computation. """ - merged_results: Base.TaskResult = None + merged_results: Base.TaskResult = None cumulative_plots: Dict[int, Any] = {} # Collect all futures in batches that had arrived since the last iteration @@ -340,38 +353,42 @@ def _process_partial_results(self, merged_results = reducer(merged_results, result) if merged_results else result mergeables = merged_results.mergeables - - for pad_num, (drawable_id, (callbacks_list, index, operation_name)) in enumerate(drawables_info_dict.items(), start=1): + + for pad_num, (drawable_id, (callbacks_list, index, operation_name)) in enumerate( + drawables_info_dict.items(), start=1 + ): cumulative_plots[index] = mergeables[index].GetValue() pad = canvas.cd(pad_num) self._apply_callbacks_and_draw(pad, cumulative_plots, operation_name, index, callbacks_list) - + return merged_results - def _apply_callbacks_and_draw(self, - pad: ROOT.TPad, - cumulative_plots: Dict[int, Any], - operation_name: str, - index: int, - callbacks_list: List[Optional[Callable]]) -> None: + def _apply_callbacks_and_draw( + self, + pad: ROOT.TPad, + cumulative_plots: Dict[int, Any], + operation_name: str, + index: int, + callbacks_list: List[Optional[Callable]], + ) -> None: """ Apply callbacks and draw plots on the provided pad. Args: pad: The TPad object for drawing plots. - + cumulative_plots: A dictionary of the current merged partial results. - + callbacks_list: A list of callback functions to be applied. - + operation_name (str): Name of the operation associated with the plot. - + index (int): Index of the plot in cumulative_plots dictionary. """ for callback in callbacks_list: - if callback is not None: - callback(cumulative_plots[index]) + if callback is not None: + callback(cumulative_plots[index]) if operation_name in ["Graph", "GraphAsymmErrors"]: cumulative_plots[index].Draw("AP") diff --git a/roottest/python/distrdf/backends/check_backend.py b/roottest/python/distrdf/backends/check_backend.py index 13af0a1b9a1fd..929a364c30d3a 100644 --- a/roottest/python/distrdf/backends/check_backend.py +++ b/roottest/python/distrdf/backends/check_backend.py @@ -101,10 +101,11 @@ def test_dask_backend_rejects_threaded_workers(self): pytest.raises( RuntimeError, match="running in distributed mode with Dask workers using more than one thread is not supported", - ) + ), ): Backend.DaskBackend(daskclient=client) + class TestInitialization: """Check initialization method in the Dask backend"""