import time
import ray
import numpy as np

ray.init()

def memory_hungry_fn(batch):
    """Allocate memory to trigger OOM."""
    # Allocate ~500MB per batch to stress memory
    _ = np.zeros((500, 1024, 1024), dtype=np.uint8)
    time.sleep(0.1)
    return batch

# Create dataset with enough data to keep actors busy
ds = ray.data.range(1000)

# Use actor pool with small memory allocation to trigger OOM
ds = ds.map_batches(
    memory_hungry_fn,
    batch_size=10,
    concurrency=4,  # Multiple actors
)

# This will hang after actor OOM kill
try:
    for batch in ds.iter_batches(batch_size=100):
        print(f"Got batch with {len(batch['id'])} rows")
except Exception as e:
    print(f"Error: {e}")

print("Done")  # Never reached if bug is present

Alternative minimal reproduction (unit test level):

import time
import ray
from ray.data._internal.compute import ActorPoolStrategy
from ray.data._internal.execution.operators.map_operator import MapOperator
from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
from ray.data._internal.execution.operators.map_transformer import (
    BlockMapTransformFn, MapTransformer,
)
from ray.data._internal.execution.streaming_executor_state import (
    build_streaming_topology, process_completed_tasks,
)
from ray.data._internal.execution.interfaces import ExecutionOptions
from ray.data._internal.execution.util import make_ref_bundles
from ray.data.context import DataContext

ray.init()

@ray.remote(num_cpus=0)
class BlockingSignal:
    def wait_forever(self):
        time.sleep(999999)

signal = BlockingSignal.remote()
ctx = DataContext.get_current()

def block_forever_then_passthrough(block_iter):
    ray.get(signal.wait_forever.remote())  # Block forever
    yield from block_iter

transformer = MapTransformer([BlockMapTransformFn(block_forever_then_passthrough)])
input_op = InputDataBuffer(ctx, make_ref_bundles([[0]]))

op = MapOperator.create(
    transformer,
    input_op=input_op,
    data_context=ctx,
    name="TestActorDiesCleanup",
    compute_strategy=ActorPoolStrategy(size=1),
    ray_remote_args={"max_restarts": 0},
)
topology = build_streaming_topology(op, ExecutionOptions())

# Wait for actor to be ready
while not op.should_add_input():
    process_completed_tasks(topology, [], max_errored_blocks=-1)

# Submit task
op.add_input(input_op.get_next(), 0)
print(f"Before kill: num_tasks_in_flight={op._actor_pool.num_tasks_in_flight()}")

# Kill actor mid-task (simulates OOM kill)
actor = next(iter(op._actor_pool.running_actors().keys()))
ray.kill(actor, no_restart=True)

# Process the failure
deadline = time.time() + 10
while True:
    process_completed_tasks(topology, [], max_errored_blocks=-1)
    tasks_in_flight = op._actor_pool.num_tasks_in_flight()
    print(f"After kill: num_tasks_in_flight={tasks_in_flight}")

    if tasks_in_flight == 0:
        print("SUCCESS: Task cleanup happened correctly!")
        break
    if time.time() > deadline:
        print("BUG: num_tasks_in_flight stuck at", tasks_in_flight)
        print("This causes pipeline hang - no new tasks can be scheduled")
        break
    time.sleep(0.5)