# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from collections import defaultdict from dataclasses import dataclass import logging import os import pyarrow as pa import asyncio import ray import time from .friendly import new_friendly_name from datafusion_ray._datafusion_ray_internal import ( DFRayContext as DFRayContextInternal, DFRayDataFrame as DFRayDataFrameInternal, prettify, ) def setup_logging(): import logging logging.addLevelName(5, "TRACE") log_level = os.environ.get("DATAFUSION_RAY_LOG_LEVEL", "WARN").upper() # this logger gets captured and routed to rust. See src/lib.rs logging.getLogger("core_py").setLevel(log_level) logging.basicConfig() setup_logging() _log_level = os.environ.get("DATAFUSION_RAY_LOG_LEVEL", "ERROR").upper() _rust_backtrace = os.environ.get("RUST_BACKTRACE", "0") df_ray_runtime_env = { "worker_process_setup_hook": setup_logging, "env_vars": { "DATAFUSION_RAY_LOG_LEVEL": _log_level, "RAY_worker_niceness": "0", "RUST_BACKTRACE": _rust_backtrace, }, } log = logging.getLogger("core_py") def call_sync(coro): """call a coroutine in the current event loop or run a new one, and synchronously return the result""" try: loop = asyncio.get_running_loop() except RuntimeError: return asyncio.run(coro) else: return loop.run_until_complete(coro) # work around for https://github.com/ray-project/ray/issues/31606 async def _ensure_coro(maybe_obj_ref): return await maybe_obj_ref async def wait_for(coros, name=""): """Wait for all coros to complete and return their results. Does not preserve ordering.""" return_values = [] # wrap the coro in a task to work with python 3.10 and 3.11+ where asyncio.wait semantics # changed to not accept any awaitable start = time.time() done, _ = await asyncio.wait( [asyncio.create_task(_ensure_coro(c)) for c in coros] ) end = time.time() log.info(f"waiting for {name} took {end - start}s") for d in done: e = d.exception() if e is not None: log.error(f"Exception waiting {name}: {e}") raise e else: return_values.append(d.result()) return return_values class DFRayProcessorPool: """A pool of DFRayProcessor actors that can be acquired and released""" # TODO: We can probably manage this set in a better way # This is not a threadsafe implementation, though the DFRayContextSupervisor accesses it # from a single thread # # This is simple though and will suffice for now def __init__(self, min_processors: int, max_processors: int): self.min_processors = min_processors self.max_processors = max_processors # a map of processor_key (a random identifier) to stage actor reference self.pool = {} # a map of processor_key to listening address self.addrs = {} # holds object references from the start_up method for each processor # we know all processors are listening when all of these refs have # been waited on. When they are ready we remove them from this set self.processors_started = set() # an event that is set when all processors are ready to serve self.processors_ready = asyncio.Event() # processors that are started but we need to get their address self.need_address = set() # processors that we have the address for but need to start serving self.need_serving = set() # processors in use self.acquired = set() # processors available self.available = set() for _ in range(min_processors): self._new_processor() log.info( f"created ray processor pool (min_processors: {min_processors}, max_processors: {max_processors})" ) async def start(self): if not self.processors_ready.is_set(): await self._wait_for_processors_started() await self._wait_for_get_addrs() await self._wait_for_serve() self.processors_ready.set() async def wait_for_ready(self): await self.processors_ready.wait() async def acquire(self, need=1): processor_keys = [] have = len(self.available) total = len(self.available) + len(self.acquired) can_make = self.max_processors - total need_to_make = need - have if need_to_make > can_make: raise Exception( f"Cannot allocate processors above {self.max_processors}" ) if need_to_make > 0: log.debug(f"creating {need_to_make} additional processors") for _ in range(need_to_make): self._new_processor() await wait_for([self.start()], "waiting for created processors") assert len(self.available) >= need for _ in range(need): processor_key = self.available.pop() self.acquired.add(processor_key) processor_keys.append(processor_key) processors = [self.pool[sk] for sk in processor_keys] addrs = [self.addrs[sk] for sk in processor_keys] return (processors, processor_keys, addrs) def release(self, processor_keys: list[str]): for processor_key in processor_keys: self.acquired.remove(processor_key) self.available.add(processor_key) def _new_processor(self): self.processors_ready.clear() processor_key = new_friendly_name() log.debug(f"starting processor: {processor_key}") processor = DFRayProcessor.options( name=f"Processor : {processor_key}" ).remote(processor_key) self.pool[processor_key] = processor self.processors_started.add(processor.start_up.remote()) self.available.add(processor_key) async def _wait_for_processors_started(self): log.info("waiting for processors to be ready") started_keys = await wait_for( self.processors_started, "processors to be started" ) # we need the addresses of these processors still self.need_address.update(set(started_keys)) # we've started all the processors we know about self.processors_started = set() log.info("processors are all listening") async def _wait_for_get_addrs(self): # get the addresses in a pipelined fashion refs = [] processor_keys = [] for processor_key in self.need_address: processor = self.pool[processor_key] refs.append(processor.addr.remote()) processor_keys.append(processor_key) self.need_serving.add(processor_key) addrs = await wait_for(refs, "processor addresses") for key, addr in addrs: self.addrs[key] = addr self.need_address = set() async def _wait_for_serve(self): log.info("running processors") try: for processor_key in self.need_serving: log.info(f"starting serving of processor {processor_key}") processor = self.pool[processor_key] processor.serve.remote() self.need_serving = set() except Exception as e: log.error(f"ProcessorPool: Uhandled Exception in serve: {e}") raise e async def all_done(self): log.info("calling processor all done") refs = [ processor.all_done.remote() for processor in self.pool.values() ] await wait_for(refs, "processors to be all done") log.info("all processors shutdown") @ray.remote(num_cpus=0.01, scheduling_strategy="SPREAD") class DFRayProcessor: def __init__(self, processor_key): self.processor_key = processor_key # import this here so ray doesn't try to serialize the rust extension from datafusion_ray._datafusion_ray_internal import ( DFRayProcessorService, ) self.processor_service = DFRayProcessorService(processor_key) async def start_up(self): # this method is sync self.processor_service.start_up() return self.processor_key async def all_done(self): await self.processor_service.all_done() async def addr(self): return (self.processor_key, self.processor_service.addr()) async def update_plan( self, stage_id: int, stage_addrs: dict[int, dict[int, list[str]]], partition_group: list[int], plan_bytes: bytes, ): await self.processor_service.update_plan( stage_id, stage_addrs, partition_group, plan_bytes, ) async def serve(self): log.info( f"[{self.processor_key}] serving on {self.processor_service.addr()}" ) await self.processor_service.serve() log.info(f"[{self.processor_key}] done serving") @dataclass class StageData: stage_id: int plan_bytes: bytes partition_group: list[int] child_stage_ids: list[int] num_output_partitions: int full_partitions: bool @dataclass class InternalStageData: stage_id: int plan_bytes: bytes partition_group: list[int] child_stage_ids: list[int] num_output_partitions: int full_partitions: bool remote_processor: ... # ray.actor.ActorHandle[DFRayProcessor] remote_addr: str def __str__(self): return f"""Stage: {self.stage_id}, pg: {self.partition_group}, child_stages:{self.child_stage_ids}, listening addr:{self.remote_addr}""" @ray.remote(num_cpus=0.01, scheduling_strategy="SPREAD") class DFRayContextSupervisor: def __init__( self, processor_pool_min: int, processor_pool_max: int, ) -> None: log.info( f"Creating DFRayContextSupervisor processor_pool_min: {processor_pool_min}" ) self.pool = DFRayProcessorPool(processor_pool_min, processor_pool_max) self.stages: dict[str, InternalStageData] = {} log.info("Created DFRayContextSupervisor") async def start(self): await self.pool.start() async def wait_for_ready(self): await self.pool.wait_for_ready() async def get_stage_addrs(self, stage_id: int): addrs = [ sd.remote_addr for sd in self.stages.values() if sd.stage_id == stage_id ] return addrs async def new_query( self, stage_datas: list[StageData], ): if len(self.stages) > 0: self.pool.release(list(self.stages.keys())) remote_processors, remote_processor_keys, remote_addrs = ( await self.pool.acquire(len(stage_datas)) ) self.stages = {} for i, sd in enumerate(stage_datas): remote_processor = remote_processors[i] remote_processor_key = remote_processor_keys[i] remote_addr = remote_addrs[i] self.stages[remote_processor_key] = InternalStageData( sd.stage_id, sd.plan_bytes, sd.partition_group, sd.child_stage_ids, sd.num_output_partitions, sd.full_partitions, remote_processor, remote_addr, ) # sort out the mess of who talks to whom and ensure we can supply the correct # addresses to each of them addrs_by_stage_key = await self.sort_out_addresses() if log.level <= logging.DEBUG: # TODO: string builder here out = "" for stage_key, stage in self.stages.items(): out += f"[{stage_key}]: {stage}\n" out += f"child addrs: {addrs_by_stage_key[stage_key]}\n" log.debug(out) refs = [] # now tell the stages what they are doing for this query for stage_key, isd in self.stages.items(): log.info(f"going to update plan for {stage_key}") kid = addrs_by_stage_key[stage_key] refs.append( isd.remote_processor.update_plan.remote( isd.stage_id, { stage_id: val["child_addrs"] for (stage_id, val) in kid.items() }, isd.partition_group, isd.plan_bytes, ) ) log.info("that's all of them") await wait_for(refs, "updating plans") async def sort_out_addresses(self): """Iterate through our stages and gather all of their listening addresses. Then, provide the addresses to of peer stages to each stage. """ addrs_by_stage_key = {} for stage_key, isd in self.stages.items(): stage_addrs = defaultdict(dict) # using "isd" as shorthand to denote InternalStageData as a reminder for child_stage_id in isd.child_stage_ids: addrs = defaultdict(list) child_stage_keys, child_stage_datas = zip( *filter( lambda x: x[1].stage_id == child_stage_id, self.stages.items(), ) ) output_partitions = [ c_isd.num_output_partitions for c_isd in child_stage_datas ] # sanity check assert all( [op == output_partitions[0] for op in output_partitions] ) output_partitions = output_partitions[0] for child_stage_isd in child_stage_datas: if child_stage_isd.full_partitions: for partition in range(output_partitions): # this stage is the definitive place to read this output partition addrs[partition] = [child_stage_isd.remote_addr] else: for partition in range(output_partitions): # this output partition must be gathered from all stages with this stage_id addrs[partition] = [ c.remote_addr for c in child_stage_datas ] stage_addrs[child_stage_id]["child_addrs"] = addrs # not necessary but useful for debug logs stage_addrs[child_stage_id]["stage_keys"] = child_stage_keys addrs_by_stage_key[stage_key] = stage_addrs return addrs_by_stage_key async def all_done(self): await self.pool.all_done() class DFRayDataFrame: def __init__( self, internal_df: DFRayDataFrameInternal, supervisor, # ray.actor.ActorHandle[DFRayContextSupervisor], batch_size=8192, partitions_per_processor: int | None = None, prefetch_buffer_size=0, ): self.df = internal_df self.supervisor = supervisor self._stages = None self._batches = None self.batch_size = batch_size self.partitions_per_processor = partitions_per_processor self.prefetch_buffer_size = prefetch_buffer_size def stages(self): # create our coordinator now, which we need to create stages if not self._stages: self._stages = self.df.stages( self.batch_size, self.prefetch_buffer_size, self.partitions_per_processor, ) return self._stages def schema(self): return self.df.schema() def execution_plan(self): return self.df.execution_plan() def logical_plan(self): return self.df.logical_plan() def optimized_logical_plan(self): return self.df.optimized_logical_plan() def collect(self) -> list[pa.RecordBatch]: if not self._batches: t1 = time.time() self.stages() t2 = time.time() log.debug(f"creating stages took {t2 - t1}s") last_stage_id = max([stage.stage_id for stage in self._stages]) log.debug(f"last stage is {last_stage_id}") self.create_ray_stages() last_stage_addrs = ray.get( self.supervisor.get_stage_addrs.remote(last_stage_id) ) log.debug(f"last stage addrs {last_stage_addrs}") reader = self.df.read_final_stage( last_stage_id, last_stage_addrs[0] ) log.debug("got reader") self._batches = list(reader) return self._batches def show(self) -> None: batches = self.collect() print(prettify(batches)) def create_ray_stages(self): stage_datas = [] # note, whereas the PyDataFrameStage object contained in self.stages() # holds information for a numbered stage, # when we tell the supervisor about our query, it wants a StageData # object per actor that will be created. Hence the loop over partition_groups for stage in self.stages(): for partition_group in stage.partition_groups: stage_datas.append( StageData( stage.stage_id, stage.plan_bytes(), partition_group, stage.child_stage_ids, stage.num_output_partitions, stage.full_partitions, ) ) ref = self.supervisor.new_query.remote(stage_datas) call_sync(wait_for([ref], "creating ray stages")) class DFRayContext: def __init__( self, batch_size: int = 8192, prefetch_buffer_size: int = 0, partitions_per_processor: int | None = None, processor_pool_min: int = 1, processor_pool_max: int = 100, ) -> None: self.ctx = DFRayContextInternal() self.batch_size = batch_size self.partitions_per_processor = partitions_per_processor self.prefetch_buffer_size = prefetch_buffer_size self.supervisor = DFRayContextSupervisor.options( name="RayContextSupersisor", ).remote( processor_pool_min, processor_pool_max, ) # start up our super visor and don't check in on it until its # time to query, then we will await this ref start_ref = self.supervisor.start.remote() # ensure we are ready s = time.time() call_sync(wait_for([start_ref], "RayContextSupervisor start")) e = time.time() log.info( f"RayContext::__init__ waiting for supervisor to be ready took {e - s}s" ) def register_parquet(self, name: str, path: str): """ Register a Parquet file with the given name and path. The path can be a local filesystem path, absolute filesystem path, or a url. If the path is a object store url, the appropriate object store will be registered. Configuration of the object store will be gathered from the environment. For example for s3:// urls, credentials will be looked for by the AWS SDK, which will check environment variables, credential files, etc Parameters: path (str): The file path to the Parquet file. name (str): The name to register the Parquet file under. """ self.ctx.register_parquet(name, path) def register_csv(self, name: str, path: str): """ Register a csvfile with the given name and path. The path can be a local filesystem path, absolute filesystem path, or a url. If the path is a object store url, the appropriate object store will be registered. Configuration of the object store will be gathered from the environment. For example for s3:// urls, credentials will be looked for by the AWS SDK, which will check environment variables, credential files, etc Parameters: path (str): The file path to the csv file. name (str): The name to register the Parquet file under. """ self.ctx.register_csv(name, path) def register_listing_table( self, name: str, path: str, file_extention="parquet" ): """ Register a directory of parquet files with the given name. The path can be a local filesystem path, absolute filesystem path, or a url. If the path is a object store url, the appropriate object store will be registered. Configuration of the object store will be gathered from the environment. For example for s3:// urls, credentials will be looked for by the AWS SDK, which will check environment variables, credential files, etc Parameters: path (str): The file path to the Parquet file directory name (str): The name to register the Parquet file under. """ self.ctx.register_listing_table(name, path, file_extention) def sql(self, query: str) -> DFRayDataFrame: df = self.ctx.sql(query) return DFRayDataFrame( df, self.supervisor, self.batch_size, self.partitions_per_processor, self.prefetch_buffer_size, ) def set(self, option: str, value: str) -> None: self.ctx.set(option, value) def __del__(self): log.info("DFRayContext, cleaning up remote resources") ref = self.supervisor.all_done.remote() call_sync(wait_for([ref], "DFRayContextSupervisor all done"))