Function install_context_stack

axlearn/common/module.py:574–598

Installs the given context stack. `install_context_stack` should be called in every child thread to ensure that each thread uses its own context stack: ``` def my_thread(context_stack): install_context_stack(context_stack) ... t = Thread(target=my_thread, kwarg

(stack: list[InvocationContext])

def clone_context_stack() -> list[InvocationContext]:  # pytype: disable=invalid-annotation
    """Returns a copy of the current InvocationContext stack.

    This is often used together with `install_context_stack` to ensure that different threads
    do not interfere with each other. See `install_context_stack`.
    """
    return list(_global_context_stack.stack)


def install_context_stack(stack: list[InvocationContext]):  # pytype: disable=invalid-annotation
    """Installs the given context stack.

    `install_context_stack` should be called in every child thread to ensure that each thread
    uses its own context stack:

    ```
    def my_thread(context_stack):
        install_context_stack(context_stack)
        ...

    t = Thread(target=my_thread, kwargs=dict(context_stack=clone_context_stack()))
    ```

    Args:
        stack: The context stack to install. This is usually returned by `clone_context_stack()`.

    Raises:
        ValueError: if the current thread id does not match the global context stack thread id.