def can_fuse(self, prev: Stage): context = DatasetContext.get_current() # TODO(ekl) also support fusing shuffle stages to subsequent 1:1 stages. if not context.optimize_fuse_shuffle_stages: return False if not self.supports_block_udf: return False if not isinstance(prev, OneToOneStage): return False if not is_task_compute(prev.compute): return False if any(k not in INHERITABLE_REMOTE_ARGS for k in prev.ray_remote_args): return False return True
def can_fuse(self, prev: Stage): if not isinstance(prev, OneToOneStage): return False # Allow fusing tasks->actors if the resources are compatible (read->map), but # not the other way around. The latter will be used as the compute if fused. if is_task_compute(self.compute) and prev.compute != self.compute: return False if (isinstance(self.fn, CallableClass) and isinstance(prev.fn, CallableClass) and (prev.fn != self.fn or (prev.fn_constructor_args != self.fn_constructor_args or prev.fn_constructor_kwargs != self.fn_constructor_kwargs))): # Fusing callable classes is only supported if they are the same function # AND their construction arguments are the same. # TODO(Clark): Support multiple callable classes instantiating in the same # actor worker constructor. return False if not _are_remote_args_compatible(prev.ray_remote_args, self.ray_remote_args): return False return True