inspect-ai 0.3.72__py3-none-any.whl → 0.3.73__py3-none-any.whl

inspect_ai/_eval/task/run.py

@@ -1,5 +1,5 @@
-import asyncio
 import contextlib
+import functools
 import sys
 import time
 from copy import deepcopy
@@ -9,6 +9,7 @@ from logging import getLogger
 from pathlib import PurePath
 from typing import Callable, Literal
 
+import anyio
 from typing_extensions import Unpack
 
 from inspect_ai._display import (
@@ -19,6 +20,7 @@ from inspect_ai._display import (
     display,
 )
 from inspect_ai._display.core.display import TaskDisplay, TaskDisplayMetric
+from inspect_ai._util._async import tg_collect
 from inspect_ai._util.constants import (
     DEFAULT_EPOCHS,
     DEFAULT_MAX_CONNECTIONS,
@@ -32,7 +34,6 @@ from inspect_ai._util.registry import (
     registry_log_name,
     registry_unqualified_name,
 )
-from inspect_ai._util.timeouts import Timeout, timeout
 from inspect_ai._util.working import (
     init_sample_working_limit,
     sample_waiting_time,
@@ -95,9 +96,9 @@ from .images import (
 )
 from .log import TaskLogger, collect_eval_data, log_start
 from .results import eval_results
-from .rundir import set_task_run_dir
+from .rundir import set_task_chdir
 from .sandbox import sandboxenv_context
-from .util import sample_messages, slice_dataset, task_run_dir
+from .util import sample_messages, slice_dataset
 
 py_logger = getLogger(__name__)
 
@@ -147,8 +148,8 @@ async def task_run(options: TaskRunOptions) -> EvalLog:
     # init task context
     init_task_context(model, options.task.approval, generate_config)
 
-    # establish run_dir for duration of execution
-    with set_task_run_dir(task_run_dir(task)):
+    # establish chdir for duration of execution (if a task has chdir=True)
+    with set_task_chdir(task):
         # track stats and error
         results: EvalResults | None = None
         reductions: list[EvalSampleReductions] | None = None
@@ -286,35 +287,6 @@ async def task_run(options: TaskRunOptions) -> EvalLog:
                             task.metrics,
                         )
 
-                    # create sample coroutines
-                    sample_coroutines = [
-                        task_run_sample(
-                            task_name=task.name,
-                            sample=sample,
-                            state=state,
-                            sandbox=sandbox,
-                            max_sandboxes=config.max_sandboxes,
-                            sandbox_cleanup=sandbox_cleanup,
-                            plan=plan,
-                            scorers=scorers,
-                            generate=generate,
-                            progress=progress,
-                            logger=logger if log_samples else None,
-                            log_images=log_images,
-                            sample_source=sample_source,
-                            sample_error=sample_error_handler,
-                            sample_complete=sample_complete,
-                            fails_on_error=(
-                                config.fail_on_error is None
-                                or config.fail_on_error is True
-                            ),
-                            time_limit=config.time_limit,
-                            working_limit=config.working_limit,
-                            semaphore=sample_semaphore,
-                        )
-                        for (sample, state) in zip(samples, states)
-                    ]
-
                     # initial progress
                     td.sample_complete(complete=0, total=len(samples))
 
@@ -327,7 +299,36 @@ async def task_run(options: TaskRunOptions) -> EvalLog:
                         task.metrics,
                     )
 
-                    sample_results = await asyncio.gather(*sample_coroutines)
+                    sample_results = await tg_collect(
+                        [
+                            functools.partial(
+                                task_run_sample,
+                                task_name=task.name,
+                                sample=sample,
+                                state=state,
+                                sandbox=sandbox,
+                                max_sandboxes=config.max_sandboxes,
+                                sandbox_cleanup=sandbox_cleanup,
+                                plan=plan,
+                                scorers=scorers,
+                                generate=generate,
+                                progress=progress,
+                                logger=logger if log_samples else None,
+                                log_images=log_images,
+                                sample_source=sample_source,
+                                sample_error=sample_error_handler,
+                                sample_complete=sample_complete,
+                                fails_on_error=(
+                                    config.fail_on_error is None
+                                    or config.fail_on_error is True
+                                ),
+                                time_limit=config.time_limit,
+                                working_limit=config.working_limit,
+                                semaphore=sample_semaphore,
+                            )
+                            for (sample, state) in zip(samples, states)
+                        ]
+                    )
 
                 # compute and record metrics if we have scores
                 completed_scores = [
@@ -362,17 +363,18 @@ async def task_run(options: TaskRunOptions) -> EvalLog:
                     )
                 )
 
-            except asyncio.CancelledError:
-                # collect eval data
-                collect_eval_data(stats)
+            except anyio.get_cancelled_exc_class():
+                with anyio.CancelScope(shield=True):
+                    # collect eval data
+                    collect_eval_data(stats)
 
-                # finish w/ cancelled status
-                eval_log = await logger.log_finish(
-                    "cancelled", stats, results, reductions
-                )
+                    # finish w/ cancelled status
+                    eval_log = await logger.log_finish(
+                        "cancelled", stats, results, reductions
+                    )
 
-                # display task cancelled
-                td.complete(TaskCancelled(logger.samples_completed, stats))
+                    # display task cancelled
+                    td.complete(TaskCancelled(logger.samples_completed, stats))
 
             except BaseException as ex:
                 if options.debug_errors:
@@ -503,7 +505,7 @@ async def task_run_sample(
     fails_on_error: bool,
     time_limit: int | None,
     working_limit: int | None,
-    semaphore: asyncio.Semaphore | None,
+    semaphore: anyio.Semaphore | None,
 ) -> dict[str, SampleScore] | None:
     # if there is an existing sample then tick off its progress, log it, and return it
     if sample_source and sample.id is not None:
@@ -533,7 +535,7 @@ async def task_run_sample(
             return sample_scores
 
     # use semaphore if provided
-    semaphore_cm: asyncio.Semaphore | contextlib.AbstractAsyncContextManager[None] = (
+    semaphore_cm: anyio.Semaphore | contextlib.AbstractAsyncContextManager[None] = (
         semaphore if semaphore else contextlib.nullcontext()
     )
 
@@ -606,7 +608,7 @@ async def task_run_sample(
 
                     # initialise timeout context manager
                     timeout_cm = (
-                        timeout(time_limit)
+                        anyio.fail_after(time_limit)
                         if time_limit is not None
                         else contextlib.nullcontext()
                     )
@@ -616,7 +618,7 @@ async def task_run_sample(
                     init_sample_working_limit(start_time, working_limit)
 
                     # run sample w/ optional timeout
-                    async with timeout_cm:
+                    with timeout_cm:
                         # mark started
                         active.started = datetime.now().timestamp()
 
@@ -640,9 +642,9 @@ async def task_run_sample(
                     # capture most recent state for scoring
                     state = sample_state() or state
 
-                except asyncio.CancelledError as ex:
+                except anyio.get_cancelled_exc_class() as ex:
                     if active.interrupt_action:
-                        # record eve t
+                        # record event
                         transcript()._event(
                             SampleLimitEvent(
                                 type="operator",
@@ -660,6 +662,8 @@ async def task_run_sample(
                                 error, raise_error = handle_error(ex)
 
                     else:
+                        # task group provided by tg_collect will automatically
+                        # handle the cancel exception
                         raise
 
                 except SampleLimitExceededError as ex:
@@ -687,9 +691,8 @@ async def task_run_sample(
                 # the cause of the timeout is a hung container and scoring requires
                 # interacting with the container). as a middle ground we use half
                 # of the original timeout value for scoring.
-                if isinstance(timeout_cm, Timeout):
-                    assert time_limit
-                    timeout_cm = timeout(time_limit / 2)
+                if time_limit is not None:
+                    timeout_cm = anyio.fail_after(time_limit / 2)
 
                 # turn off message and token limits
                 state.message_limit = None
@@ -699,7 +702,7 @@ async def task_run_sample(
                 # scoring
                 try:
                     # timeout during scoring will result in an ordinary sample error
-                    async with timeout_cm:
+                    with timeout_cm:
                         if error is None:
                             for scorer in scorers or []:
                                 scorer_name = unique_scorer_name(
@@ -740,7 +743,7 @@ async def task_run_sample(
                             # propagate results into scores
                             state.scores = {k: v.score for k, v in results.items()}
 
-                except asyncio.CancelledError:
+                except anyio.get_cancelled_exc_class():
                     if active.interrupt_action:
                         transcript()._event(
                             SampleLimitEvent(
@@ -970,10 +973,10 @@ def create_sample_semaphore(
     config: EvalConfig,
     generate_config: GenerateConfig,
     modelapi: ModelAPI | None = None,
-) -> asyncio.Semaphore:
+) -> anyio.Semaphore:
     # if the user set max_samples then use that
     if config.max_samples is not None:
-        return asyncio.Semaphore(config.max_samples)
+        return anyio.Semaphore(config.max_samples)
 
     # use max_connections
     max_samples = (
@@ -985,4 +988,4 @@ def create_sample_semaphore(
     )
 
     # return the semaphore
-    return asyncio.Semaphore(max_samples)
+    return anyio.Semaphore(max_samples)

inspect_ai/_eval/task/rundir.py

@@ -6,9 +6,12 @@ from contextvars import ContextVar
 from functools import wraps
 from typing import Any, Callable, Iterator, TypeVar
 
+from inspect_ai._eval.task.task import Task
+from inspect_ai._eval.task.util import task_chdir
+
 TASK_DIRECTORY_ATTRIB = "task_directory"
 
-_task_run_dir = ContextVar[str]("task_run_dir")
+_task_chdir = ContextVar[str | None]("_task_chdir", default=None)
 
 T = TypeVar("T", bound="asyncio.BaseEventLoop")
 
@@ -46,12 +49,16 @@ def task_run_dir_switching() -> Iterator[None]:
 
 
 @contextmanager
-def set_task_run_dir(run_dir: str) -> Iterator[None]:
-    token = _task_run_dir.set(run_dir)
-    try:
+def set_task_chdir(task: Task) -> Iterator[None]:
+    chdir = task_chdir(task)
+    if chdir is not None:
+        token = _task_chdir.set(chdir)
+        try:
+            yield
+        finally:
+            _task_chdir.reset(token)
+    else:
         yield
-    finally:
-        _task_run_dir.reset(token)
 
 
 if sys.platform == "win32":
@@ -63,9 +70,9 @@ else:
 def _wrap_callback(callback: Callable[..., Any]) -> Callable[..., Any]:
     @wraps(callback)
     def wrapper(*args: Any, **kwargs: Any) -> Any:
-        run_dir = _task_run_dir.get(None)
-        if run_dir is not None and run_dir != os.getcwd():
-            os.chdir(run_dir)
+        chdir = _task_chdir.get(None)
+        if chdir is not None and chdir != os.getcwd():
+            os.chdir(chdir)
         return callback(*args, **kwargs)
 
     return wrapper

inspect_ai/_eval/task/sandbox.py

@@ -1,9 +1,9 @@
-import asyncio
 import base64
 import contextlib
 from random import random
 from typing import AsyncGenerator, Callable, NamedTuple, cast
 
+import anyio
 import httpx
 from tenacity import (
     retry,
@@ -15,10 +15,9 @@ from tenacity import (
 
 from inspect_ai._eval.task.task import Task
 from inspect_ai._eval.task.util import task_run_dir
-from inspect_ai._util.constants import DEFAULT_MAX_RETRIES, DEFAULT_TIMEOUT
 from inspect_ai._util.file import file, filesystem
+from inspect_ai._util.httpx import httpx_should_retry, log_httpx_retry_attempt
 from inspect_ai._util.registry import registry_unqualified_name
-from inspect_ai._util.retry import httpx_should_retry, log_retry_attempt
 from inspect_ai._util.url import data_uri_to_base64, is_data_uri, is_http_url
 from inspect_ai.dataset import Sample
 from inspect_ai.util._concurrency import concurrency
@@ -62,7 +61,7 @@ async def sandboxenv_context(
     # in and grab all of the sandboxes). Therefore, in this case we wait a random
     # delay so that all tasks/samples have an equal shot at getting scheduled.
     if max_sandboxes is not None:
-        await asyncio.sleep(random())
+        await anyio.sleep(random())
 
     # enforce concurrency if required
     sandboxes_cm = (
@@ -103,7 +102,7 @@ async def sandboxenv_context(
             # run sample
             yield
 
-        except asyncio.CancelledError as ex:
+        except anyio.get_cancelled_exc_class() as ex:
             interrupted = True
             raise ex
 
@@ -186,14 +185,14 @@ async def _retrying_httpx_get(
     url: str,
     client: httpx.AsyncClient = httpx.AsyncClient(),
     timeout: int = 30,  # per-attempt timeout
-    max_retries: int = DEFAULT_MAX_RETRIES,
-    total_timeout: int = DEFAULT_TIMEOUT,  #  timeout for the whole retry loop. not for an individual attempt
+    max_retries: int = 10,
+    total_timeout: int = 120,  #  timeout for the whole retry loop. not for an individual attempt
 ) -> bytes:
     @retry(
         wait=wait_exponential_jitter(),
         stop=(stop_after_attempt(max_retries) | stop_after_delay(total_timeout)),
         retry=retry_if_exception(httpx_should_retry),
-        before_sleep=log_retry_attempt(url),
+        before_sleep=log_httpx_retry_attempt(url),
     )
     async def do_get() -> bytes:
         response = await client.get(

inspect_ai/_eval/task/util.py

@@ -25,6 +25,13 @@ def task_run_dir(task: Task) -> str:
     return getattr(task, TASK_RUN_DIR_ATTR, os.getcwd())
 
 
+def task_chdir(task: Task) -> str | None:
+    if task.attribs.get("chdir", False) is True:
+        return task_run_dir(task)
+    else:
+        return None
+
+
 def task_file(task: Task, relative: bool = False) -> str | None:
     file = cast(str | None, getattr(task, TASK_FILE_ATTR, None))
     if file:

inspect_ai/_util/_async.py

@@ -1,20 +1,100 @@
 import asyncio
-from typing import Any, Coroutine, TypeVar
+import inspect
+import os
+import sys
+from logging import Logger
+from typing import Any, Awaitable, Callable, Coroutine, Literal, TypeVar, cast
 
+import anyio
 import nest_asyncio  # type: ignore
+import sniffio
+
+if sys.version_info >= (3, 11):
+    from typing import TypeVarTuple, Unpack
+else:
+    from exceptiongroup import ExceptionGroup
+    from typing_extensions import TypeVarTuple, Unpack
+
+
+PosArgsT = TypeVarTuple("PosArgsT")
 
 
 def is_callable_coroutine(func_or_cls: Any) -> bool:
-    if asyncio.iscoroutinefunction(func_or_cls):
+    if inspect.iscoroutinefunction(func_or_cls):
         return True
     elif callable(func_or_cls):
-        return asyncio.iscoroutinefunction(func_or_cls.__call__)
+        return inspect.iscoroutinefunction(func_or_cls.__call__)
     return False
 
 
 T = TypeVar("T")
 
 
+async def tg_collect(
+    funcs: list[Callable[[], Awaitable[T]]], exception_group: bool = False
+) -> list[T]:
+    """Runs all of the pased async functions and collects their results.
+
+    The results will be returned in the same order as the input `funcs`.
+
+    Args:
+       funcs: List of async functions.
+       exception_group: `True` to raise an ExceptionGroup or
+          `False` (the default) to raise only the first exception.
+
+    Returns:
+       List of results of type T.
+
+    Raises:
+       Exception: First exception occurring in failed tasks
+          (for `exception_group == False`, the default)
+       ExceptionGroup: Exceptions that occurred in failed tasks
+         (for `exception_group == True`)
+    """
+    try:
+        results: list[tuple[int, T]] = []
+
+        async with anyio.create_task_group() as tg:
+
+            async def run_task(index: int) -> None:
+                result = await funcs[index]()
+                results.append((index, result))
+
+            for i in range(0, len(funcs)):
+                tg.start_soon(run_task, i)
+
+        # sort results by original index and return just the values
+        return [r for _, r in sorted(results)]
+    except ExceptionGroup as ex:
+        if exception_group:
+            raise
+        else:
+            raise ex.exceptions[0]
+
+
+async def coro_print_exceptions(
+    context: str,
+    func: Callable[[Unpack[PosArgsT]], Awaitable[Any]],
+    *args: Unpack[PosArgsT],
+) -> None:
+    try:
+        await func(*args)
+    except Exception as ex:
+        print(f"Error {context}: {ex}")
+
+
+async def coro_log_exceptions(
+    logger: Logger,
+    context: str,
+    func: Callable[[Unpack[PosArgsT]], Awaitable[Any]],
+    *args: Unpack[PosArgsT],
+) -> None:
+    try:
+        await func(*args)
+    except Exception as ex:
+        logger.warning(f"Error {context}: {ex}")
+
+
 _initialised_nest_asyncio: bool = False
 
 
@@ -26,14 +106,42 @@ def init_nest_asyncio() -> None:
 
 
 def run_coroutine(coroutine: Coroutine[None, None, T]) -> T:
-    try:
-        # this will throw if there is no running loop
-        asyncio.get_running_loop()
+    from inspect_ai._util.platform import running_in_notebook
+
+    if current_async_backend() == "trio":
+        raise RuntimeError("run_coroutine cannot be used with trio")
 
-        # initialiase nest_asyncio then we are clear to run
+    if running_in_notebook():
         init_nest_asyncio()
         return asyncio.run(coroutine)
+    else:
+        try:
+            # this will throw if there is no running loop
+            asyncio.get_running_loop()
 
-    except RuntimeError:
-        # No running event loop so we are clear to run
-        return asyncio.run(coroutine)
+            # initialiase nest_asyncio then we are clear to run
+            init_nest_asyncio()
+            return asyncio.run(coroutine)
+
+        except RuntimeError:
+            # No running event loop so we are clear to run
+            return asyncio.run(coroutine)
+
+
+def current_async_backend() -> Literal["asyncio", "trio"] | None:
+    try:
+        return _validate_backend(sniffio.current_async_library().lower())
+    except sniffio.AsyncLibraryNotFoundError:
+        return None
+
+
+def configured_async_backend() -> Literal["asyncio", "trio"]:
+    backend = os.environ.get("INSPECT_ASYNC_BACKEND", "asyncio").lower()
+    return _validate_backend(backend)
+
+
+def _validate_backend(backend: str) -> Literal["asyncio", "trio"]:
+    if backend in ["asyncio", "trio"]:
+        return cast(Literal["asyncio", "trio"], backend)
+    else:
+        raise RuntimeError(f"Unknown async backend: {backend}")

inspect_ai/_util/constants.py

@@ -6,8 +6,6 @@ PKG_AUTHOR_DIR = "UK-AISI"
 PKG_NAME = Path(__file__).parent.parent.stem
 PKG_PATH = Path(__file__).parent.parent
 DEFAULT_EPOCHS = 1
-DEFAULT_MAX_RETRIES = 5
-DEFAULT_TIMEOUT = 120
 DEFAULT_MAX_CONNECTIONS = 10
 DEFAULT_MAX_TOKENS = 2048
 DEFAULT_VIEW_PORT = 7575

inspect_ai/_util/file.py

@@ -1,5 +1,3 @@
-import asyncio
-import contextlib
 import datetime
 import io
 import os
@@ -9,17 +7,19 @@ import unicodedata
 from contextlib import contextmanager
 from copy import deepcopy
 from pathlib import Path
-from typing import Any, AsyncIterator, BinaryIO, Iterator, Literal, cast, overload
+from typing import Any, BinaryIO, Iterator, Literal, cast, overload
 from urllib.parse import urlparse
 
 import fsspec  # type: ignore  # type: ignore
-from fsspec.asyn import AsyncFileSystem  # type: ignore
 from fsspec.core import split_protocol  # type: ignore  # type: ignore
 from fsspec.implementations.local import make_path_posix  # type: ignore
 from pydantic import BaseModel
 from s3fs import S3FileSystem  # type: ignore
 from shortuuid import uuid
 
+from inspect_ai._util._async import configured_async_backend, current_async_backend
+from inspect_ai._util.error import PrerequisiteError
+
 # https://filesystem-spec.readthedocs.io/en/latest/_modules/fsspec/spec.html#AbstractFileSystem
 # https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.generic.GenericFileSystem
 
@@ -298,30 +298,6 @@ def filesystem(path: str, fs_options: dict[str, Any] = {}) -> FileSystem:
     return FileSystem(fs)
 
 
-@contextlib.asynccontextmanager
-async def async_fileystem(
-    location: str, fs_options: dict[str, Any] = {}
-) -> AsyncIterator[AsyncFileSystem]:
-    # determine protocol
-    protocol, _ = split_protocol(location)
-    protocol = protocol or "file"
-
-    # build options
-    options = default_fs_options(location)
-    options.update(fs_options)
-
-    if protocol == "s3":
-        s3 = S3FileSystem(asynchronous=True, **options)
-        session = await s3.set_session()
-        try:
-            yield s3
-        finally:
-            await session.close()
-    else:
-        options.update({"asynchronous": True, "loop": asyncio.get_event_loop()})
-        yield fsspec.filesystem(protocol, **options)
-
-
 def absolute_file_path(file: str) -> str:
     # check for a relative dir, if we find one then resolve to absolute
     fs_scheme = urlparse(file).scheme
@@ -331,7 +307,17 @@ def absolute_file_path(file: str) -> str:
 
 
 def default_fs_options(file: str) -> dict[str, Any]:
-    options = deepcopy(DEFAULT_FS_OPTIONS.get(urlparse(file).scheme, {}))
+    scheme = urlparse(file).scheme
+    if (
+        scheme == "s3"
+        and configured_async_backend() == "trio"
+        and current_async_backend() == "trio"
+    ):
+        raise PrerequisiteError(
+            "ERROR: The s3 interface is not supported when running under the trio async backend."
+        )
+
+    options = deepcopy(DEFAULT_FS_OPTIONS.get(scheme, {}))
     # disable caching for all filesystems
     options.update(
         dict(

inspect_ai/_util/future.py

@@ -0,0 +1,37 @@
+from typing import Generic, TypeVar
+
+import anyio
+
+T = TypeVar("T")
+
+
+class Future(Generic[T]):
+    def __init__(self) -> None:
+        self._result: T | None = None
+        self._ex: Exception | None = None
+        self._event = anyio.Event()
+
+    def set_result(self, result: T) -> None:
+        self._result = result
+        self._event.set()
+
+    def set_exception(self, ex: Exception) -> None:
+        self._ex = ex
+        self._event.set()
+
+    async def result(self) -> T:
+        await self._event.wait()
+        if self._result is not None:
+            return self._result
+        elif self._ex is not None:
+            raise self._ex
+        else:
+            raise RuntimeError("Future completed without a result or error")
+
+    @staticmethod
+    def set_future_result(future: "Future[T]", result: T) -> None:
+        future.set_result(result)
+
+    @staticmethod
+    def set_future_exception(future: "Future[T]", error: Exception) -> None:
+        future.set_exception(error)