def pipeline_stage(fn: Callable[[], Dataset[T]],
context: DatasetContext) -> Dataset[T]:
DatasetContext._set_current(context)
try:
prev = set_progress_bars(False)
return fn()
finally:
set_progress_bars(prev)
def run(self, fn: Callable[[], Dataset[T]],
context: DatasetContext) -> Dataset[T]:
DatasetContext._set_current(context)
try:
prev = set_progress_bars(False)
# Force eager evaluation of all blocks in the pipeline stage. This
# prevents resource deadlocks due to overlapping stage execution
# (e.g., task -> actor stage).
return fn().force_reads()
finally:
set_progress_bars(prev)
def pipeline_stage(fn: Callable[[], Dataset[T]]) -> Dataset[T]:
try:
prev = set_progress_bars(False)
return fn()
finally:
set_progress_bars(prev)