inspect-ai 0.3.90__py3-none-any.whl → 0.3.92__py3-none-any.whl

inspect_ai/_eval/task/run.py

@@ -2,7 +2,7 @@ import contextlib
 import functools
 import sys
 import time
-from copy import deepcopy
+from copy import copy, deepcopy
 from dataclasses import dataclass, field
 from datetime import datetime
 from logging import getLogger
@@ -307,6 +307,7 @@ async def task_run(options: TaskRunOptions) -> EvalLog:
                         functools.partial(
                             task_run_sample,
                             task_name=task.name,
+                            log_location=profile.log_location,
                             sample=sample,
                             state=state,
                             sandbox=sandbox,
@@ -325,6 +326,8 @@ async def task_run(options: TaskRunOptions) -> EvalLog:
                                 config.fail_on_error is None
                                 or config.fail_on_error is True
                             ),
+                            retry_on_error=config.retry_on_error or 0,
+                            error_retries=[],
                             time_limit=config.time_limit,
                             working_limit=config.working_limit,
                             semaphore=sample_semaphore,
@@ -484,7 +487,9 @@ def update_metrics_display_fn(
 
 
 async def task_run_sample(
+    *,
     task_name: str,
+    log_location: str,
     sample: Sample,
     state: TaskState,
     sandbox: SandboxEnvironmentSpec | None,
@@ -500,6 +505,8 @@ async def task_run_sample(
     sample_error: SampleErrorHandler,
     sample_complete: Callable[[dict[str, SampleScore]], None],
     fails_on_error: bool,
+    retry_on_error: int,
+    error_retries: list[EvalError],
     time_limit: int | None,
     working_limit: int | None,
     semaphore: anyio.Semaphore | None,
@@ -531,6 +538,9 @@ async def task_run_sample(
             sample_complete(sample_scores)
             return sample_scores
 
+    # copy variables that we may pass back to ourselves on a retry
+    initial_state = deepcopy(state)
+
     # use semaphore if provided
     semaphore_cm: anyio.Semaphore | contextlib.AbstractAsyncContextManager[None] = (
         semaphore if semaphore else contextlib.nullcontext()
@@ -561,20 +571,31 @@ async def task_run_sample(
 
     # helper to handle exceptions (will throw if we've exceeded the limit)
     def handle_error(ex: BaseException) -> tuple[EvalError, BaseException | None]:
-        err = sample_error(ex)
-        # if we aren't raising the error then print a warning
-        if err[1] is None:
-            py_logger.warning(
-                f"Sample error (id: {sample.id}, epoch: {state.epoch}): {exception_message(ex)})"
-            )
-        transcript()._event(ErrorEvent(error=err[0]))
-        return err
+        # helper to log sample error
+        def log_sample_error() -> None:
+            msg = f"Sample error (id: {sample.id}, epoch: {state.epoch}): {exception_message(ex)})"
+            if retry_on_error > 0:
+                msg = f"{msg}. Sample will be retried."
+            py_logger.warning(msg)
+
+        # if we have retries left then return EvalError
+        if retry_on_error > 0:
+            log_sample_error()
+            return eval_error(ex, type(ex), ex, ex.__traceback__), None
+        else:
+            err = sample_error(ex)
+            # if we aren't raising the error then print a warning
+            if err[1] is None:
+                log_sample_error()
+            transcript()._event(ErrorEvent(error=err[0]))
+            return err
 
     # solver loop
     async with (
         semaphore_cm,
         active_sample(
             task=task_name,
+            log_location=log_location,
             model=str(state.model),
             sample=sample,
             epoch=state.epoch,
@@ -582,7 +603,7 @@ async def task_run_sample(
             token_limit=state.token_limit,
             time_limit=time_limit,
             working_limit=working_limit,
-            fails_on_error=fails_on_error,
+            fails_on_error=fails_on_error or (retry_on_error > 0),
             transcript=sample_transcript,
         ) as active,
     ):
@@ -606,7 +627,7 @@ async def task_run_sample(
 
             async with sandboxenv_cm:
                 timeout_cm: (
-                    contextlib._GeneratorContextManager[anyio.CancelScope, None, None]
+                    contextlib._GeneratorContextManager[anyio.CancelScope]
                     | contextlib.nullcontext[None]
                 ) = contextlib.nullcontext()
                 try:
@@ -791,40 +812,84 @@ async def task_run_sample(
         except Exception as ex:
             error, raise_error = handle_error(ex)
 
-        # complete the sample
-        progress(SAMPLE_TOTAL_PROGRESS_UNITS)
+        # complete the sample if there is no error or if there is no retry_on_error in play
+        if not error or (retry_on_error == 0):
+            progress(SAMPLE_TOTAL_PROGRESS_UNITS)
+
+            # log it
+            if logger is not None:
+                # if we are logging images then be sure to base64 images injected by solvers
+                if log_images:
+                    state = (await states_with_base64_content([state]))[0]
+
+                # otherwise ensure there are no base64 images in sample or messages
+                else:
+                    sample = sample_without_base64_content(sample)
+                    state = state_without_base64_content(state)
+
+                # log the sample
+                await log_sample(
+                    start_time=start_time,
+                    logger=logger,
+                    sample=sample,
+                    state=state,
+                    scores=results,
+                    error=error,
+                    error_retries=error_retries,
+                    log_images=log_images,
+                )
 
-        # log it
+    # error that should be retried (we do this outside of the above scope so that we can
+    # retry outside of the original semaphore -- our retry will therefore go to the back
+    # of the sample queue)
+    if error and retry_on_error > 0:
+        # remove any buffered sample events
         if logger is not None:
-            # if we are logging images then be sure to base64 images injected by solvers
-            if log_images:
-                state = (await states_with_base64_content([state]))[0]
+            logger.remove_sample(state.sample_id, state.epoch)
 
-            # otherwise ensure there are no base64 images in sample or messages
-            else:
-                sample = sample_without_base64_content(sample)
-                state = state_without_base64_content(state)
-
-            # log the sample
-            await log_sample(
-                start_time=start_time,
-                logger=logger,
-                sample=sample,
-                state=state,
-                scores=results,
-                error=error,
-                log_images=log_images,
-            )
+        # recurse w/ tick down of retry_on_error and append of error to error_retries
+        return await task_run_sample(
+            task_name=task_name,
+            log_location=log_location,
+            sample=sample,
+            # state was deep copied at the outset
+            state=initial_state,
+            sandbox=sandbox,
+            max_sandboxes=max_sandboxes,
+            sandbox_cleanup=sandbox_cleanup,
+            plan=plan,
+            scorers=scorers,
+            generate=generate,
+            progress=progress,
+            logger=logger,
+            log_images=log_images,
+            sample_source=sample_source,
+            sample_error=sample_error,
+            sample_complete=sample_complete,
+            fails_on_error=fails_on_error,
+            # tick retry count down
+            retry_on_error=retry_on_error - 1,
+            # forward on error that caused retry
+            error_retries=copy(error_retries) + [error],
+            time_limit=time_limit,
+            working_limit=working_limit,
+            semaphore=semaphore,
+        )
 
-        # return
-        if error is None:
-            if results is not None:
-                sample_complete(results)
-            return results
-        elif raise_error:
-            raise raise_error
-        else:
-            return None
+    # no error
+    elif error is None:
+        # call sample_complete callback if we have score results
+        if results is not None:
+            sample_complete(results)
+        return results
+
+    # we have an error and should raise it
+    elif raise_error is not None:
+        raise raise_error
+
+    # we have an error and should not raise it
+    else:
+        return None
 
 
 async def log_sample(
@@ -834,6 +899,7 @@ async def log_sample(
     state: TaskState,
     scores: dict[str, SampleScore],
     error: EvalError | None,
+    error_retries: list[EvalError],
     log_images: bool,
 ) -> None:
     # sample must have id to be logged
@@ -879,6 +945,7 @@ async def log_sample(
         if total_time is not None
         else None,
         error=error,
+        error_retries=error_retries,
         limit=limit,
     )
 

inspect_ai/_eval/task/sandbox.py

@@ -17,6 +17,7 @@ from inspect_ai._eval.task.task import Task
 from inspect_ai._eval.task.util import task_run_dir
 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.path import chdir
 from inspect_ai._util.registry import registry_unqualified_name
 from inspect_ai._util.url import data_uri_to_base64, is_data_uri, is_http_url
 from inspect_ai.dataset import Sample
@@ -29,6 +30,7 @@ from inspect_ai.util._sandbox.environment import (
     SandboxEnvironment,
     SandboxEnvironmentConfigType,
     SandboxEnvironmentSpec,
+    TaskInitEnvironment,
 )
 from inspect_ai.util._sandbox.registry import registry_find_sandboxenv
 
@@ -42,7 +44,7 @@ async def sandboxenv_context(
     sample: Sample,
 ) -> AsyncGenerator[None, None]:
     # resolve sandbox
-    sandbox = resolve_sandbox(sandbox, sample)
+    sandbox = await resolve_sandbox(sandbox, sample)
     if not sandbox:
         raise ValueError("sandboxenv_context called with no sandbox specified")
 
@@ -143,22 +145,34 @@ async def read_sandboxenv_file(contents: str) -> bytes:
 class TaskSandboxEnvironment(NamedTuple):
     sandbox: SandboxEnvironmentSpec
     run_dir: str
+    env: tuple[tuple[str, str], ...]
 
 
-def resolve_sandbox_for_task(
+async def resolve_sandbox_for_task_and_sample(
     eval_sandbox: SandboxEnvironmentSpec | None,
     task: Task,
     sample: Sample,
 ) -> TaskSandboxEnvironment | None:
     # eval_sandbox overrides task or sample sandbox
-    sandbox = eval_sandbox or resolve_sandbox(task.sandbox, sample)
+    sandbox = eval_sandbox or await resolve_sandbox(task.sandbox, sample)
     if sandbox is not None:
-        return TaskSandboxEnvironment(sandbox, task_run_dir(task))
+        # see if there are environment variables required for init of this sample
+        run_dir = task_run_dir(task)
+        with chdir(run_dir):
+            sandboxenv_type = registry_find_sandboxenv(sandbox.type)
+            task_init_environment = cast(
+                TaskInitEnvironment, getattr(sandboxenv_type, "task_init_environment")
+            )
+            env = await task_init_environment(sandbox.config, sample.metadata or {})
+
+        return TaskSandboxEnvironment(
+            sandbox=sandbox, run_dir=run_dir, env=tuple(sorted(env.items()))
+        )
     else:
         return None
 
 
-def resolve_sandbox(
+async def resolve_sandbox(
     sandbox: SandboxEnvironmentSpec | None,
     sample: Sample,
 ) -> SandboxEnvironmentSpec | None:

inspect_ai/_util/_async.py

@@ -69,7 +69,7 @@ async def tg_collect(
         if exception_group:
             raise
         else:
-            raise ex.exceptions[0]
+            raise ex.exceptions[0] from None
 
 
 async def coro_print_exceptions(

inspect_ai/_util/constants.py

@@ -26,6 +26,7 @@ ALL_LOG_LEVELS = [
 DEFAULT_LOG_LEVEL = "warning"
 DEFAULT_LOG_LEVEL_TRANSCRIPT = "info"
 DEFAULT_LOG_SHARED = 10
+DEFAULT_RETRY_ON_ERROR = 1
 ALL_LOG_FORMATS = ["eval", "json"]
 DEFAULT_LOG_FORMAT: Literal["eval", "json"] = "eval"
 JSON_LOG_FORMAT = "json"

inspect_ai/_util/environ.py

@@ -24,3 +24,35 @@ def environ_var(name: str, value: str) -> Iterator[None]:
             os.environ.pop(name, None)
         else:
             os.environ[name] = previous_value
+
+
+@contextmanager
+def environ_vars(env_vars: dict[str, str]) -> Iterator[None]:
+    """
+    Temporarily set multiple environment variables within a context.
+
+    Args:
+        env_vars: Dictionary mapping environment variable names to values
+
+    Yields:
+        None
+    """
+    # save previous values
+    previous_values = {}
+    for name in env_vars:
+        previous_values[name] = os.environ.get(name)
+
+    # set new values
+    for name, value in env_vars.items():
+        os.environ[name] = value
+
+    try:
+        yield
+    finally:
+        # Restore previous environment
+        for name in env_vars:
+            previous_value = previous_values[name]
+            if previous_value is None:
+                os.environ.pop(name, None)
+            else:
+                os.environ[name] = previous_value

inspect_ai/_util/file.py

@@ -271,8 +271,15 @@ class FileSystem:
         if "mtime" not in file.keys() and file["type"] == "file":
             file["mtime"] = self.fs.created(file).timestamp()
 
+        # adjust mtime to be milliseconds
         if "mtime" in file.keys():
-            file["mtime"] = file["mtime"] * 1000
+            mtime = file["mtime"]
+            if isinstance(mtime, datetime.datetime):
+                file["mtime"] = mtime.timestamp() * 1000
+            elif isinstance(mtime, int | float):
+                file["mtime"] = mtime * 1000
+            else:
+                raise ValueError(f"Unexpected type for mtime ({type(mtime)}): {mtime}")
         else:
             file["mtime"] = None
 

inspect_ai/_util/httpx.py

@@ -1,10 +1,13 @@
 import logging
 from typing import Callable
 
-from httpx import ConnectError, ConnectTimeout, HTTPStatusError, ReadTimeout
+import httpcore
+import httpx
+from httpx import HTTPStatusError
 from tenacity import RetryCallState
 
 from inspect_ai._util.constants import HTTP
+from inspect_ai._util.http import is_retryable_http_status
 
 logger = logging.getLogger(__name__)
 
@@ -20,25 +23,10 @@ def httpx_should_retry(ex: BaseException) -> bool:
     Returns:
       True if a retry should occur
     """
-    # httpx status exception
     if isinstance(ex, HTTPStatusError):
-        # request timeout
-        if ex.response.status_code == 408:
-            return True
-        # lock timeout
-        elif ex.response.status_code == 409:
-            return True
-        # rate limit
-        elif ex.response.status_code == 429:
-            return True
-        # internal errors
-        elif ex.response.status_code >= 500:
-            return True
-        else:
-            return False
-
-    # connection error
-    elif is_httpx_connection_error(ex):
+        return is_retryable_http_status(ex.response.status_code)
+
+    elif httpx_should_retry_no_status_code(ex):
         return True
 
     # don't retry
@@ -50,11 +38,106 @@ def log_httpx_retry_attempt(context: str) -> Callable[[RetryCallState], None]:
     def log_attempt(retry_state: RetryCallState) -> None:
         logger.log(
             HTTP,
-            f"{context} connection retry {retry_state.attempt_number} after waiting for {retry_state.idle_for}",
+            f"{context} connection retry {retry_state.attempt_number} (retrying in {retry_state.upcoming_sleep:,.0f} seconds)",
         )
 
     return log_attempt
 
 
-def is_httpx_connection_error(ex: BaseException) -> bool:
-    return isinstance(ex, ConnectTimeout | ConnectError | ConnectionError | ReadTimeout)
+def httpx_should_retry_no_status_code(ex: BaseException) -> bool:
+    """
+    Check whether an exception (without an HTTP status code) should be retried.
+
+    To understand this function, it may be helpful to look at the exception hierarchies for
+    httpx and httpcore, which are reproduced below.
+
+
+    # HTTPX Exception Hierarchy
+    Exception (Python built-in)
+    |
+    +-- HTTPError
+    |   |
+    |   +-- RequestError
+    |   |   |
+    |   |   +-- TransportError
+    |   |   |   |
+    |   |   |   +-- TimeoutException
+    |   |   |   |   |
+    |   |   |   |   +-- ConnectTimeout
+    |   |   |   |   +-- ReadTimeout
+    |   |   |   |   +-- WriteTimeout
+    |   |   |   |   +-- PoolTimeout
+    |   |   |   |
+    |   |   |   +-- NetworkError
+    |   |   |   |   |
+    |   |   |   |   +-- ConnectError
+    |   |   |   |   +-- ReadError
+    |   |   |   |   +-- WriteError
+    |   |   |   |   +-- CloseError
+    |   |   |   |
+    |   |   |   +-- ProtocolError
+    |   |   |   |   |
+    |   |   |   |   +-- LocalProtocolError
+    |   |   |   |   +-- RemoteProtocolError
+    |   |   |   |
+    |   |   |   +-- ProxyError
+    |   |   |   +-- UnsupportedProtocol
+    |   |   |
+    |   |   +-- DecodingError
+    |   |   +-- TooManyRedirects
+    |   |
+    |   +-- HTTPStatusError
+    |
+    +-- InvalidURL
+    +-- CookieConflict
+    +-- RuntimeError (Python built-in)
+        |
+        +-- StreamError
+            |
+            +-- StreamConsumed
+            +-- StreamClosed
+            +-- ResponseNotRead
+            +-- RequestNotRead
+
+
+    # HTTPCore Exception Hierarchy
+    Exception (Python built-in)
+    |
+    +-- ConnectionNotAvailable
+    +-- ProxyError
+    +-- UnsupportedProtocol
+    +-- ProtocolError
+    |   |
+    |   +-- RemoteProtocolError
+    |   +-- LocalProtocolError
+    |
+    +-- TimeoutException
+    |   |
+    |   +-- PoolTimeout
+    |   +-- ConnectTimeout
+    |   +-- ReadTimeout
+    |   +-- WriteTimeout
+    |
+    +-- NetworkError
+        |
+        +-- ConnectError
+        +-- ReadError
+        +-- WriteError
+    """
+    # Base class for all exceptions that occur at the level of the Transport API.
+    is_transport_error = isinstance(ex, httpx.TransportError)
+
+    # Sometimes exceptions are raised directly by httpcore, the lower-level library that httpx uses
+    is_httpcore_network_error = isinstance(ex, httpcore.NetworkError)
+    is_httpcore_timeout_error = isinstance(ex, httpcore.TimeoutException)
+    is_httpcore_protocol_error = isinstance(ex, httpcore.ProtocolError)
+
+    # extensible in case we notice other cases
+    return any(
+        [
+            is_transport_error,
+            is_httpcore_network_error,
+            is_httpcore_timeout_error,
+            is_httpcore_protocol_error,
+        ]
+    )

inspect_ai/_util/registry.py

@@ -1,6 +1,17 @@
+from __future__ import annotations
+
 import inspect
 from inspect import get_annotations, isclass
-from typing import Any, Callable, Literal, TypedDict, TypeGuard, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Literal,
+    TypedDict,
+    TypeGuard,
+    cast,
+    overload,
+)
 
 from pydantic import BaseModel, Field
 from pydantic_core import to_jsonable_python
@@ -11,19 +22,30 @@ from inspect_ai._util.package import get_installed_package_name
 from .constants import PKG_NAME
 from .entrypoints import ensure_entry_points
 
+if TYPE_CHECKING:
+    from inspect_ai import Task
+    from inspect_ai.agent import Agent
+    from inspect_ai.approval import Approver
+    from inspect_ai.model import ModelAPI
+    from inspect_ai.scorer import Metric, Scorer, ScoreReducer
+    from inspect_ai.solver import Plan, Solver
+    from inspect_ai.tool import Tool
+    from inspect_ai.util import SandboxEnvironment
+
 obj_type = type
 
 RegistryType = Literal[
-    "task",
-    "solver",
     "agent",
-    "tool",
-    "scorer",
+    "approver",
     "metric",
-    "score_reducer",
     "modelapi",
+    "plan",
     "sandboxenv",
-    "approver",
+    "score_reducer",
+    "scorer",
+    "solver",
+    "task",
+    "tool",
 ]
 """Enumeration of registry object types.
 
@@ -184,7 +206,59 @@ def registry_find(predicate: Callable[[RegistryInfo], bool]) -> list[object]:
         return o
 
 
-def registry_create(type: RegistryType, name: str, **kwargs: Any) -> object:
+@overload
+def registry_create(type: Literal["agent"], name: str, **kwargs: Any) -> Agent: ...
+
+
+@overload
+def registry_create(
+    type: Literal["approver"], name: str, **kwargs: Any
+) -> Approver: ...
+
+
+@overload
+def registry_create(type: Literal["metric"], name: str, **kwargs: Any) -> Metric: ...
+
+
+@overload
+def registry_create(
+    type: Literal["modelapi"], name: str, **kwargs: Any
+) -> ModelAPI: ...
+
+
+@overload
+def registry_create(type: Literal["plan"], name: str, **kwargs: Any) -> Plan: ...
+
+
+@overload
+def registry_create(
+    type: Literal["sandboxenv"], name: str, **kwargs: Any
+) -> SandboxEnvironment: ...
+
+
+@overload
+def registry_create(type: Literal["scorer"], name: str, **kwargs: Any) -> Scorer: ...
+
+
+@overload
+def registry_create(
+    type: Literal["score_reducer"], name: str, **kwargs: Any
+) -> ScoreReducer: ...
+
+
+@overload
+def registry_create(type: Literal["solver"], name: str, **kwargs: Any) -> Solver: ...
+
+
+@overload
+def registry_create(type: Literal["task"], name: str, **kwargs: Any) -> Task: ...
+
+
+@overload
+def registry_create(type: Literal["tool"], name: str, **kwargs: Any) -> Tool: ...
+
+
+def registry_create(type: RegistryType, name: str, **kwargs: Any) -> object:  # type: ignore[return]
     r"""Create a registry object.
 
     Creates objects registered via decorator (e.g. `@task`, `@solver`). Note
@@ -230,7 +304,7 @@ def registry_create(type: RegistryType, name: str, **kwargs: Any) -> object:
     if isclass(obj):
         return with_registry_info(obj(**kwargs))
     elif callable(obj):
-        return_type = get_annotations(obj).get("return")
+        return_type = get_annotations(obj, eval_str=True).get("return")
         # Until we remove the MetricDeprecated symbol we need this extra
         # bit to map the Metric union back to Metric
         if "_metric.Metric" in str(return_type):