Repository: UKGovernmentBEIS/inspect_ai
Issue: (SandboxEnvironment - potential bug): Task_init is not actually ran concurrently!
Root cause: Algorithmic Error

## Problem

**`task_init()` on SandboxEnvironment is an asynchronous function, but it gets called in a sequential loop in the actual run.py**

In `inspect_ai/_eval/run.py`, the `startup_sandbox_environments()` function is:

```python
for sandboxenv in sandboxenvs:
    await task_init("startup", sandboxenv.sandbox.config)  # Blocks!
```

Even with `max_tasks > 1`, `each task_init()` waits for its predecessor to complete before starting.

## Fix

In `inspect_ai/_eval/run.py:517-531`:

**Before:**

```python
for sandboxenv in sandboxenvs:
    # find type
    sandboxenv_type = registry_find_sandboxenv(sandboxenv.sandbox.type)

    # run startup
    task_init = cast(TaskInit, getattr(sandboxenv_type, "task_init"))
    with chdir(sandboxenv.run_dir), environ_vars(dict(sandboxenv.env)):
        await task_init("startup", sandboxenv.sandbox.config)

    # append cleanup method
    task_cleanup = cast(TaskCleanup, getattr(sandboxenv_type, "task_cleanup"))
    cleanups.append(
        (task_cleanup, sandboxenv.sandbox.config, sandboxenv.run_dir)
    )
```

**After:**

```python
async with anyio.create_task_group() as tg:
    for sandboxenv in sandboxenvs:
        # find type
        sandboxenv_type = registry_find_sandboxenv(sandboxenv.sandbox.type)

        # run startup
        task_init = cast(TaskInit, getattr(sandboxenv_type, "task_init"))
        task_cleanup = cast(TaskCleanup, getattr(sandboxenv_type, "task_cleanup"))

        # Closure to capture loop variables for concurrent execution
        async def init_env(env=sandboxenv, init_fn=task_init, cleanup_fn=task_cleanup):
            with chdir(env.run_dir), environ_vars(dict(env.env)):
                await init_fn("startup", env.sandbox.config)
            cleanups.append((cleanup_fn, env.sandbox.config, env.run_dir))

        tg.start_soon(init_env)
```



## Reproduction

Running this mock Sandbox:

```python
   @classmethod
    async def task_init(
        cls, task_name: str, config: SandboxEnvironmentConfigType | None
    ) -> None:
        """Simulate slow initialization (like building Docker images or pulling)."""
        task_id = config.task_id  # type: ignore

        log(f"Task {task_id}: Starting task_init")

        # Simulate slow setup (docker build, pip install, etc)
        await asyncio.sleep(20)

        log(f"Task {task_id}: Completed task_init")
```

With:

```python
 eval(
        tasks,
        max_tasks=2,
        max_sandboxes=2,
        max_samples=1,
    )
```

We get this output normally:

```
12:00:00.000 - Task task_1: Starting task_init
12:00:20.000 - Task task_1: Completed task_init
12:00:20.001 - Task task_2: Starting task_init  # waits 20s
12:00:40.001 - Task task_2: Completed task_init
```


With our fix we get:

```
12:00:00.000 - Task task_1: Starting task_init
12:00:00.001 - Task task_2: Starting task_init  # starts immediately
12:00:20.000 - Task task_1: Completed task_init
12:00:20.001 - Task task_2: Completed task_init
```

## Appendix
_complete minimal reproduction:_


[sandbox_logging.py](https://github.com/user-attachments/files/22803628/sandbox_logging.py)
[sandbox.py](https://github.com/user-attachments/files/22803627/sandbox.py)
[task_eval.py](https://github.com/user-attachments/files/22803629/task_eval.py)


We do this serially so that the user can clearly see the Docker build/pull progress and details (as there may often be errors they need to read and diagnose). We prefer this behavior as I think if we interleaved all of these it would be super difficult to debug container build/pull errors.

Thanks for the response! Not sure I agree with this:
1. This assumes direct docker access in task_init - but presumably custom SandboxEnvironments should support a variety of different backend setups. 
2. If someone's making a custom sandbox they can handle setting up logs that respect the async nature of task init.

Options (imo):
1. Control this with a flag - have the option to run task_init asynchronously. Default to current behavior.
2. Make this run async & update the existing docker sandbox s.t. it has clear logs (timestamped? Artificially serialized?)
3. Make task_init synchronous to be explicit about how it's actually used. 

Appreciate the ideation here, the thing we are relying on for Docker is the ability of Docker to show it's full Docker pull/build UI (which often has critical information about failure modes -- when we originally allowed this to be async people couldn't successfully diagnose failures due to interleaved output). We could serialize this another way but we'd need to turn off ansi and actually wait until all the pull/build steps are done to show anything (which can often take minutes). Based on all this experience I am pretty confident that sync is the way to go for Docker. It's trading off init time for transparency/usability, and since the init time is re-used across invocations (once you've pulled or built you don't need to do it again) it's not too penal a tradeoff.

My fear is that any sandbox that has non-trivial initialization will have the same concerns (users need to be appraised of progress and errors). Are you using or implementing another sandbox type where this would be an unambiguous win? Totally open to making "real" async possible but definitely want to understand the use case before I do.