def await_ml_tasks()

in example-apps/chatbot-rag-app/data/index_data.py [0:0]


def await_ml_tasks(max_timeout=1200, interval=5):
    """
    Waits for all machine learning tasks to complete within a specified timeout period.

    Parameters:
        max_timeout (int): Maximum time to wait for tasks to complete, in seconds.
        interval (int): Time to wait between status checks, in seconds.

    Raises:
        TimeoutError: If the timeout is reached and machine learning tasks are still running.
    """
    start_time = time.time()
    ml_tasks = get_ml_tasks()
    if not ml_tasks:
        return  # likely a lost race on tasks

    print(f"Awaiting {len(ml_tasks)} ML tasks")

    while time.time() - start_time < max_timeout:
        ml_tasks = get_ml_tasks()
        if not ml_tasks:
            return
        time.sleep(interval)

    raise TimeoutError(
        f"Timeout reached. ML tasks are still running: {', '.join(ml_tasks)}"
    )