def apply(self, fn: Any, remote_args: dict, blocks: BlockList) -> BlockList: context = DatasetContext.get_current() # Handle empty datasets. if blocks.initial_num_blocks() == 0: return blocks blocks = list(blocks.iter_blocks_with_metadata()) map_bar = ProgressBar("Map Progress", total=len(blocks)) if context.block_splitting_enabled: map_block = cached_remote_fn(_map_block_split).options( **remote_args) refs = [map_block.remote(b, fn, m.input_files) for b, m in blocks] else: map_block = cached_remote_fn(_map_block_nosplit).options( **dict(remote_args, num_returns=2)) all_refs = [ map_block.remote(b, fn, m.input_files) for b, m in blocks ] data_refs = [r[0] for r in all_refs] refs = [r[1] for r in all_refs] # Common wait for non-data refs. try: results = map_bar.fetch_until_complete(refs) except (ray.exceptions.RayTaskError, KeyboardInterrupt) as e: # One or more mapper tasks failed, or we received a SIGINT signal # while waiting; either way, we cancel all map tasks. for ref in refs: ray.cancel(ref) # Wait until all tasks have failed or been cancelled. for ref in refs: try: ray.get(ref) except (ray.exceptions.RayTaskError, ray.exceptions.TaskCancelledError): pass # Reraise the original task failure exception. raise e from None new_blocks, new_metadata = [], [] if context.block_splitting_enabled: for result in results: for block, metadata in result: new_blocks.append(block) new_metadata.append(metadata) else: for block, metadata in zip(data_refs, results): new_blocks.append(block) new_metadata.append(metadata) return BlockList(list(new_blocks), list(new_metadata))
def apply(self, fn: Any, remote_args: dict, blocks: BlockList) -> BlockList: # Handle empty datasets. if blocks.initial_num_blocks() == 0: return blocks blocks = list(blocks.iter_blocks_with_metadata()) map_bar = ProgressBar("Map Progress", total=len(blocks)) map_block = cached_remote_fn(_map_block) refs = [ map_block.options(**remote_args).remote(b, fn, m.input_files) for b, m in blocks ] try: results = map_bar.fetch_until_complete(refs) except (ray.exceptions.RayTaskError, KeyboardInterrupt) as e: # One or more mapper tasks failed, or we received a SIGINT signal # while waiting; either way, we cancel all map tasks. for ref in refs: ray.cancel(ref) # Wait until all tasks have failed or been cancelled. for ref in refs: try: ray.get(ref) except (ray.exceptions.RayTaskError, ray.exceptions.TaskCancelledError): pass # Reraise the original task failure exception. raise e from None new_blocks, new_metadata = [], [] for result in results: for block, metadata in result: new_blocks.append(block) new_metadata.append(metadata) return BlockList(list(new_blocks), list(new_metadata))
def apply(self, fn: Any, remote_args: dict, blocks: BlockList) -> BlockList: blocks_in = list(blocks.iter_blocks_with_metadata()) orig_num_blocks = len(blocks_in) results = [] map_bar = ProgressBar("Map Progress", total=orig_num_blocks) class BlockWorker: def ready(self): return "ok" def process_block( self, block: Block, input_files: List[str] ) -> Iterable[Tuple[Block, BlockMetadata]]: output = [] for new_block in fn(block): accessor = BlockAccessor.for_block(new_block) new_metadata = BlockMetadata( num_rows=accessor.num_rows(), size_bytes=accessor.size_bytes(), schema=accessor.schema(), input_files=input_files) owner = DatasetContext.get_current().block_owner output.append((ray.put(new_block, _owner=owner), new_metadata)) return output if not remote_args: remote_args["num_cpus"] = 1 BlockWorker = ray.remote(**remote_args)(BlockWorker) self.workers = [BlockWorker.remote()] tasks = {w.ready.remote(): w for w in self.workers} ready_workers = set() while len(results) < orig_num_blocks: ready, _ = ray.wait(list(tasks), timeout=0.01, num_returns=1, fetch_local=False) if not ready: if len(ready_workers) / len(self.workers) > 0.8: w = BlockWorker.remote() self.workers.append(w) tasks[w.ready.remote()] = w map_bar.set_description( "Map Progress ({} actors {} pending)".format( len(ready_workers), len(self.workers) - len(ready_workers))) continue [obj_id] = ready worker = tasks[obj_id] del tasks[obj_id] # Process task result. if worker in ready_workers: results.append(obj_id) map_bar.update(1) else: ready_workers.add(worker) # Schedule a new task. if blocks_in: block, meta = blocks_in.pop() ref = worker.process_block.remote(block, meta.input_files) tasks[ref] = worker map_bar.close() new_blocks, new_metadata = [], [] for result in ray.get(results): for block, metadata in result: new_blocks.append(block) new_metadata.append(metadata) return BlockList(new_blocks, new_metadata)
def apply(self, fn: Any, remote_args: dict, blocks: BlockList) -> BlockList: context = DatasetContext.get_current() blocks_in = list(blocks.iter_blocks_with_metadata()) orig_num_blocks = len(blocks_in) results = [] map_bar = ProgressBar("Map Progress", total=orig_num_blocks) class BlockWorker: def ready(self): return "ok" def map_block_split(self, block: Block, input_files: List[str]) -> BlockPartition: return _map_block_split(block, fn, input_files) @ray.method(num_returns=2) def map_block_nosplit( self, block: Block, input_files: List[str]) -> Tuple[Block, BlockMetadata]: return _map_block_nosplit(block, fn, input_files) if not remote_args: remote_args["num_cpus"] = 1 BlockWorker = ray.remote(**remote_args)(BlockWorker) self.workers = [BlockWorker.remote()] tasks = {w.ready.remote(): w for w in self.workers} metadata_mapping = {} ready_workers = set() while len(results) < orig_num_blocks: ready, _ = ray.wait(list(tasks), timeout=0.01, num_returns=1, fetch_local=False) if not ready: if len(ready_workers) / len(self.workers) > 0.8: w = BlockWorker.remote() self.workers.append(w) tasks[w.ready.remote()] = w map_bar.set_description( "Map Progress ({} actors {} pending)".format( len(ready_workers), len(self.workers) - len(ready_workers))) continue [obj_id] = ready worker = tasks[obj_id] del tasks[obj_id] # Process task result. if worker in ready_workers: results.append(obj_id) map_bar.update(1) else: ready_workers.add(worker) # Schedule a new task. if blocks_in: block, meta = blocks_in.pop() if context.block_splitting_enabled: ref = worker.map_block_split.remote( block, meta.input_files) else: ref, meta_ref = worker.map_block_nosplit.remote( block, meta.input_files) metadata_mapping[ref] = meta_ref tasks[ref] = worker map_bar.close() new_blocks, new_metadata = [], [] if context.block_splitting_enabled: for result in ray.get(results): for block, metadata in result: new_blocks.append(block) new_metadata.append(metadata) else: for block in results: new_blocks.append(block) new_metadata.append(metadata_mapping[block]) return BlockList(new_blocks, new_metadata)
def apply(self, fn: Any, remote_args: dict, blocks: BlockList) -> BlockList: blocks_in = list(blocks.iter_blocks_with_metadata()) orig_num_blocks = len(blocks_in) blocks_out = [] map_bar = ProgressBar("Map Progress", total=orig_num_blocks) class BlockWorker: def ready(self): return "ok" @ray.method(num_returns=2) def process_block(self, block: Block, input_files: List[str] ) -> (Block, BlockMetadata): new_block = fn(block) accessor = BlockAccessor.for_block(new_block) new_metadata = BlockMetadata( num_rows=accessor.num_rows(), size_bytes=accessor.size_bytes(), schema=accessor.schema(), input_files=input_files) return new_block, new_metadata if not remote_args: remote_args["num_cpus"] = 1 BlockWorker = ray.remote(**remote_args)(BlockWorker) self.workers = [BlockWorker.remote()] metadata_mapping = {} tasks = {w.ready.remote(): w for w in self.workers} ready_workers = set() while len(blocks_out) < orig_num_blocks: ready, _ = ray.wait( list(tasks), timeout=0.01, num_returns=1, fetch_local=False) if not ready: if len(ready_workers) / len(self.workers) > 0.8: w = BlockWorker.remote() self.workers.append(w) tasks[w.ready.remote()] = w map_bar.set_description( "Map Progress ({} actors {} pending)".format( len(ready_workers), len(self.workers) - len(ready_workers))) continue [obj_id] = ready worker = tasks[obj_id] del tasks[obj_id] # Process task result. if worker in ready_workers: blocks_out.append(obj_id) map_bar.update(1) else: ready_workers.add(worker) # Schedule a new task. if blocks_in: block, meta = blocks_in.pop() block_ref, meta_ref = worker.process_block.remote( block, meta.input_files) metadata_mapping[block_ref] = meta_ref tasks[block_ref] = worker new_metadata = ray.get([metadata_mapping[b] for b in blocks_out]) map_bar.close() return BlockList(blocks_out, new_metadata)