Repository: UKGovernmentBEIS/inspect_ai
Issue: `SandboxEnvironment.sample_init` is run sequentially
Root cause: Algorithmic Error

sample_init calls are not made in parallel, even though it appears they should be from the code. I'm not sure yet why this is the case, but actively looking into it. 

Here is a test that should reproduce the issue (test should fail):

```python
import tempfile
from pathlib import Path
from typing import Literal, Union, cast, overload

import aiofiles
from typing_extensions import override
import asyncio

from inspect_ai.util._sandbox.environment import SandboxEnvironment
from inspect_ai.util._sandbox.registry import sandboxenv
from inspect_ai.util._subprocess import ExecResult, subprocess

from inspect_ai import Task, eval
from inspect_ai.dataset import Sample
from inspect_ai.model import ModelOutput, get_model
from inspect_ai.scorer import includes
from inspect_ai.solver import generate
import itertools
from typing import Generic, Iterable, Iterator, TypeVar

from inspect_ai.model import ModelOutput
import logging
import time

T = TypeVar("T")
logger = logging.getLogger(__name__)

class _OneThenFixedIterator(Generic[T]):
    def __init__(self, first_item: T, fixed_item: T):
        self.first_item = first_item
        self.fixed_item = fixed_item
        self.first_returned = False

    def __iter__(self) -> Iterator[T]:
        return self

    def __next__(self) -> T:
        if not self.first_returned:
            self.first_returned = True
            return self.first_item
        return self.fixed_item


def _mock_submit(submission_value: str = "value of submission") -> ModelOutput:
    return ModelOutput.for_tool_call(
        model="mockllm/model",
        tool_name="submit_answer",
        tool_arguments={"submission": submission_value},
    )


def mock_submit_then_spin(
    submission_value: str = "value of submission",
) -> Iterable[ModelOutput]:
    return _OneThenFixedIterator(
        _mock_submit(submission_value),
        ModelOutput.from_content(
            model="mockllm/model",
            content="I already submitted the answer why are you calling generate on me again?",
        ),
    )
    
@sandboxenv(name="mylocal")
class MyLocalSandboxEnvironment(SandboxEnvironment):
    
    start_times: list[float] = []
    end_times: list[float] = []
    
    @override
    @classmethod
    async def sample_init(
        cls, task_name: str, config: str | None, metadata: dict[str, str]
    ) -> dict[str, SandboxEnvironment]:
        logger.info(f"Initializing {cls} with task {task_name}")
        # print(f"STARTING {cls} with task {task_name}")
        start = time.time()
        cls.start_times.append(start)
        await asyncio.sleep(15.0)
        end = time.time()
        cls.end_times.append(end)
        # print(f"STARTED {cls} with task {task_name}")

        logger.info(f"Initalized {cls} with task {task_name}")
        return {"default": MyLocalSandboxEnvironment()}

    @override
    @classmethod
    async def sample_cleanup(
        cls,
        task_name: str,
        config: str | None,
        environments: dict[str, SandboxEnvironment],
        interrupted: bool,
    ) -> None:
        for environment in environments.values():
            env = cast(MyLocalSandboxEnvironment, environment)
            env.directory.cleanup()

    def __init__(self) -> None:
        self.directory = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)

    @override
    async def exec(
        self,
        cmd: list[str],
        input: str | bytes | None = None,
        cwd: str | None = None,
        env: dict[str, str] = {},
        timeout: int | None = None,
    ) -> ExecResult[str]:
        return await subprocess(
            args=cmd,
            input=input,
            cwd=cwd if cwd else self.directory.name,
            env=env,
            timeout=timeout,
        )

    @override
    async def write_file(self, file: str, contents: str | bytes) -> None:
        # resolve file and ensure the parent dir exists
        file = self._resolve_file(file)
        Path(file).parent.mkdir(parents=True, exist_ok=True)

        if isinstance(contents, str):
            async with aiofiles.open(file, "w", encoding="utf-8") as f:
                await f.write(contents)
        else:
            async with aiofiles.open(file, "wb") as f:
                await f.write(contents)

    @overload
    async def read_file(self, file: str, text: Literal[True] = True) -> str: ...

    @overload
    async def read_file(self, file: str, text: Literal[False]) -> bytes: ...

    @override
    async def read_file(self, file: str, text: bool = True) -> Union[str | bytes]:
        file = self._resolve_file(file)
        if text:
            async with aiofiles.open(file, "r", encoding="utf-8") as f:
                return await f.read()
        else:
            async with aiofiles.open(file, "rb") as f:
                return await f.read()

    def _resolve_file(self, file: str) -> str:
        path = Path(file)
        if path.is_absolute():
            return file
        else:
            return (Path(self.directory.name) / Path(file)).as_posix()

def test_sandbox_env_parallel_startup():
    dataset = [
        Sample(
            input="What are the contents of file foo.txt?",
            target="unused",
            files={"foo.txt": "contents_of_foo.txt"},
        )
    ]
    tasks = [Task(
        dataset=dataset,
        plan=[generate()],
        scorer=includes(),
        sandbox="mylocal",
    ), Task(
        dataset=dataset,
        plan=[generate()],
        scorer=includes(),
        sandbox="mylocal",
    )]
    eval(
        tasks,
        model=get_model(
        "mockllm/model",
        custom_outputs=itertools.chain(
            [
                ModelOutput.from_content(
                    model="mockllm/model",
                    content="Sorry, but your prompt is too long",
                ),
            ],
            mock_submit_then_spin(),
        ),
        max_samples=10,
        max_tasks=10,
    ),
        max_messages=5,  # otherwise we can get into an infinite loop if the tools error
    )[0]
    
    # if the environments are started in parallel, the start time of the second environment should be less than the end time of the first environment
    assert MyLocalSandboxEnvironment.start_times[-1] < MyLocalSandboxEnvironment.end_times[0]

test_sandbox_env_parallel_startup()
```

oh wait hang on - the above is initialising two *tasks* separately, rather than two samples within a task. If you change it to the below, then the test no longer fails:

```
import tempfile
from pathlib import Path
from typing import Literal, Union, cast, overload

import aiofiles
from typing_extensions import override
import asyncio

from inspect_ai.util._sandbox.environment import SandboxEnvironment
from inspect_ai.util._sandbox.registry import sandboxenv
from inspect_ai.util._subprocess import ExecResult, subprocess

from inspect_ai import Task, eval
from inspect_ai.dataset import Sample
from inspect_ai.model import ModelOutput, get_model
from inspect_ai.scorer import includes
from inspect_ai.solver import generate
import itertools
from typing import Generic, Iterable, Iterator, TypeVar

from inspect_ai.model import ModelOutput
import logging
import time

T = TypeVar("T")
logger = logging.getLogger(__name__)

class _OneThenFixedIterator(Generic[T]):
    def __init__(self, first_item: T, fixed_item: T):
        self.first_item = first_item
        self.fixed_item = fixed_item
        self.first_returned = False

    def __iter__(self) -> Iterator[T]:
        return self

    def __next__(self) -> T:
        if not self.first_returned:
            self.first_returned = True
            return self.first_item
        return self.fixed_item


def _mock_submit(submission_value: str = "value of submission") -> ModelOutput:
    return ModelOutput.for_tool_call(
        model="mockllm/model",
        tool_name="submit_answer",
        tool_arguments={"submission": submission_value},
    )


def mock_submit_then_spin(
    submission_value: str = "value of submission",
) -> Iterable[ModelOutput]:
    return _OneThenFixedIterator(
        _mock_submit(submission_value),
        ModelOutput.from_content(
            model="mockllm/model",
            content="I already submitted the answer why are you calling generate on me again?",
        ),
    )
    
@sandboxenv(name="mylocal")
class MyLocalSandboxEnvironment(SandboxEnvironment):
    
    start_times: list[float] = []
    end_times: list[float] = []
    
    @override
    @classmethod
    async def sample_init(
        cls, task_name: str, config: str | None, metadata: dict[str, str]
    ) -> dict[str, SandboxEnvironment]:
        logger.info(f"Initializing {cls} with task {task_name}")
        # print(f"STARTING {cls} with task {task_name}")
        start = time.time()
        cls.start_times.append(start)
        await asyncio.sleep(15.0)
        end = time.time()
        cls.end_times.append(end)
        # print(f"STARTED {cls} with task {task_name}")

        logger.info(f"Initalized {cls} with task {task_name}")
        return {"default": MyLocalSandboxEnvironment()}

    @override
    @classmethod
    async def sample_cleanup(
        cls,
        task_name: str,
        config: str | None,
        environments: dict[str, SandboxEnvironment],
        interrupted: bool,
    ) -> None:
        for environment in environments.values():
            env = cast(MyLocalSandboxEnvironment, environment)
            env.directory.cleanup()

    def __init__(self) -> None:
        self.directory = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)

    @override
    async def exec(
        self,
        cmd: list[str],
        input: str | bytes | None = None,
        cwd: str | None = None,
        env: dict[str, str] = {},
        timeout: int | None = None,
    ) -> ExecResult[str]:
        return await subprocess(
            args=cmd,
            input=input,
            cwd=cwd if cwd else self.directory.name,
            env=env,
            timeout=timeout,
        )

    @override
    async def write_file(self, file: str, contents: str | bytes) -> None:
        # resolve file and ensure the parent dir exists
        file = self._resolve_file(file)
        Path(file).parent.mkdir(parents=True, exist_ok=True)

        if isinstance(contents, str):
            async with aiofiles.open(file, "w", encoding="utf-8") as f:
                await f.write(contents)
        else:
            async with aiofiles.open(file, "wb") as f:
                await f.write(contents)

    @overload
    async def read_file(self, file: str, text: Literal[True] = True) -> str: ...

    @overload
    async def read_file(self, file: str, text: Literal[False]) -> bytes: ...

    @override
    async def read_file(self, file: str, text: bool = True) -> Union[str | bytes]:
        file = self._resolve_file(file)
        if text:
            async with aiofiles.open(file, "r", encoding="utf-8") as f:
                return await f.read()
        else:
            async with aiofiles.open(file, "rb") as f:
                return await f.read()

    def _resolve_file(self, file: str) -> str:
        path = Path(file)
        if path.is_absolute():
            return file
        else:
            return (Path(self.directory.name) / Path(file)).as_posix()

def test_sandbox_env_parallel_startup():
    dataset = [
        Sample(
            input="What are the contents of file foo.txt?",
            target="unused",
            files={"foo.txt": "contents_of_foo.txt"},
        ),
        Sample(
            input="What are the contents of file foo.txt?",
            target="unused",
            files={"foo.txt": "contents_of_foo.txt"},
        ),
    ]
    tasks = [Task(
        dataset=dataset,
        plan=[generate()],
        scorer=includes(),
        sandbox="mylocal",
    )]
    eval(
        tasks,
        model=get_model(
        "mockllm/model",
        custom_outputs=itertools.chain(
            [
                ModelOutput.from_content(
                    model="mockllm/model",
                    content="Sorry, but your prompt is too long",
                ),
            ],
            mock_submit_then_spin(),
        ),
        max_samples=10,
        max_tasks=10,
    ),
        max_messages=5,  # otherwise we can get into an infinite loop if the tools error
    )[0]
    
    # if the environments are started in parallel, the start time of the second environment should be less than the end time of the first environment
    assert MyLocalSandboxEnvironment.start_times[-1] < MyLocalSandboxEnvironment.end_times[0]

test_sandbox_env_parallel_startup()
```

So it might be more accurate to say "sample_init runs sequentially iff there's multiple tasks"? I think? I'm not sure if this is expected behaviour, but it feels like it shouldn't be with task parallelism