diff --git a/bindings/distrdf/python/DistRDF/Backends/Dask/Backend.py b/bindings/distrdf/python/DistRDF/Backends/Dask/Backend.py index 84a9a0bbd4abc..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,24 @@ 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) + + 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: """ @@ -115,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. @@ -141,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. @@ -183,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 @@ -213,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: @@ -237,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) @@ -273,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. @@ -295,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 @@ -327,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 010c96c299417..929a364c30d3a 100644 --- a/roottest/python/distrdf/backends/check_backend.py +++ b/roottest/python/distrdf/backends/check_backend.py @@ -62,6 +62,49 @@ 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"""