Based on the research and issues I've found, I can provide some guidance on your issue with the RayTaskRunner and wait() behavior. There appear to be a few important points to consider: 1. First, make sure you're using the latest version of

prefect-ray

. The integration has had several updates to improve task execution behavior. 2. When using

wait()

with mapped tasks, you should ensure you're waiting on the entire FuturesList, not individual futures. Here's the recommended pattern:

python
from prefect import flow, task
from prefect_ray import RayTaskRunner

@task
def my_task(x):
    # your task code
    pass

@flow(task_runner=RayTaskRunner())
def my_flow(data):
    # Submit all tasks
    futures = my_task.map(data)
    
    # Wait for ALL futures to complete
    futures.wait()
    # Or if you need the results:
    results = futures.collect()
    
    return results

3. If you're still seeing tasks running after the flow completes, you might need to ensure the Ray cluster has enough resources to execute all tasks concurrently. If resources are constrained, some tasks might be queued but not yet executed when wait() returns. 4. You can also try setting explicit timeouts when waiting:

python
futures.wait(timeout=300)  # Wait up to 5 minutes

5. If you're using nested mapped tasks (maps within maps), make sure to wait for each level appropriately:

python
@flow(task_runner=RayTaskRunner())
def my_flow():
    outer_futures = outer_task.map(data1)
    for future in outer_futures:
        inner_futures = inner_task.map(future)
        inner_futures.wait()  # Wait for each inner map to complete
    outer_futures.wait()  # Wait for all outer tasks

If you're still experiencing issues after implementing these suggestions, it would be helpful if you could share: 1. Your current flow code structure 2. The version of

prefect-ray

you're using 3. Whether you're running Ray locally or in a distributed setup 4. The approximate number of tasks being mapped This would help identify if there might be other factors affecting the task completion behavior. Also, make sure you have proper error handling in your tasks, as unhandled exceptions might cause tasks to fail silently while the flow continues execution.