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