inspect-ai 0.3.99__py3-none-any.whl → 0.3.101__py3-none-any.whl

inspect_ai/util/__init__.py

@@ -6,7 +6,9 @@ from inspect_ai.util._limit import (
     LimitScope,
     apply_limits,
     message_limit,
+    time_limit,
     token_limit,
+    working_limit,
 )
 
 from ._collect import collect
@@ -81,6 +83,8 @@ __all__ = [
     "subtask",
     "throttle",
     "token_limit",
+    "time_limit",
+    "working_limit",
     "trace_action",
     "trace_message",
     "RegistryType",

inspect_ai/util/_json.py

@@ -3,6 +3,7 @@ import typing
 from copy import deepcopy
 from dataclasses import is_dataclass
 from datetime import date, datetime, time
+from enum import EnumMeta
 from typing import (
     Any,
     Dict,
@@ -101,6 +102,8 @@ def json_schema(t: Type[Any]) -> JSONSchema:
             or (isinstance(t, type) and issubclass(t, BaseModel))
         ):
             return cls_json_schema(t)
+        elif isinstance(t, EnumMeta):
+            return JSONSchema(enum=[item.value for item in t])
         elif t is type(None):
             return JSONSchema(type="null")
         else:

inspect_ai/util/_limit.py

@@ -7,6 +7,7 @@ from contextvars import ContextVar
 from types import TracebackType
 from typing import TYPE_CHECKING, Generic, Iterator, Literal, TypeVar
 
+import anyio
 from typing_extensions import Self
 
 from inspect_ai._util.logger import warn_once
@@ -33,22 +34,23 @@ class LimitExceededError(Exception):
        value: Value compared to.
        limit: Limit applied.
        message (str | None): Optional. Human readable message.
-       source (Limit | None): Optional. The `Limit` instance which was responsible for
-         raising this error.
+       source (Limit | None): Optional. The `Limit` instance which was responsible for raising this error.
     """
 
     def __init__(
         self,
         type: Literal["message", "time", "working", "token", "operator", "custom"],
         *,
-        value: int,
-        limit: int,
+        value: float,
+        limit: float,
         message: str | None = None,
         source: Limit | None = None,
     ) -> None:
         self.type = type
         self.value = value
+        self.value_str = self._format_float_or_int(value)
         self.limit = limit
+        self.limit_str = self._format_float_or_int(limit)
         self.message = f"Exceeded {type} limit: {limit:,}"
         self.source = source
         super().__init__(message)
@@ -60,6 +62,12 @@ class LimitExceededError(Exception):
         )
         return self
 
+    def _format_float_or_int(self, value: float | int) -> str:
+        if isinstance(value, int):
+            return f"{value:,}"
+        else:
+            return f"{value:,.2f}"
+
 
 class Limit(abc.ABC):
     """Base class for all limit context managers."""
@@ -80,6 +88,12 @@ class Limit(abc.ABC):
     ) -> None:
         pass
 
+    @property
+    @abc.abstractmethod
+    def usage(self) -> float:
+        """The current usage of the resource being limited."""
+        pass
+
     def _check_reuse(self) -> None:
         if self._entered:
             raise RuntimeError(
@@ -112,18 +126,20 @@ def apply_limits(
         False, all `LimitExceededError` exceptions will be allowed to propagate.
     """
     limit_scope = LimitScope()
-    with ExitStack() as stack:
-        for limit in limits:
-            stack.enter_context(limit)
-        try:
+    # Try scope is outside the `with ExitStack()` so that we can catch any errors raised
+    # when exiting it (which will be where time_limit() would raise LimitExceededError).
+    try:
+        with ExitStack() as stack:
+            for limit in limits:
+                stack.enter_context(limit)
             yield limit_scope
-        except LimitExceededError as e:
-            # If it was not one of the limits we applied.
-            if e.source is None or e.source not in limits:
-                raise
-            limit_scope.limit_error = e
-            if not catch_errors:
-                raise
+    except LimitExceededError as e:
+        # If it was not one of the limits we applied.
+        if e.source is None or e.source not in limits:
+            raise
+        limit_scope.limit_error = e
+        if not catch_errors:
+            raise
 
 
 class LimitScope:
@@ -140,8 +156,6 @@ def token_limit(limit: int | None) -> _TokenLimit:
     """Limits the total number of tokens which can be used.
 
     The counter starts when the context manager is opened and ends when it is closed.
-    The context manager can be opened multiple times, even in different execution
-    contexts.
 
     These limits can be stacked.
 
@@ -186,8 +200,7 @@ def message_limit(limit: int | None) -> _MessageLimit:
     """Limits the number of messages in a conversation.
 
     The total number of messages in the conversation are compared to the limit (not just
-    "new" messages). The context manager can be opened multiple times, even in different
-    execution contexts.
+    "new" messages).
 
     These limits can be stacked.
 
@@ -220,6 +233,62 @@ def check_message_limit(count: int, raise_for_equal: bool) -> None:
     node.check(count, raise_for_equal)
 
 
+def time_limit(limit: float | None) -> _TimeLimit:
+    """Limits the wall clock time which can elapse.
+
+    The timer starts when the context manager is opened and stops when it is closed.
+
+    These limits can be stacked.
+
+    When a limit is exceeded, the code block is cancelled and a `LimitExceededError` is
+    raised.
+
+    Uses anyio's cancellation scopes meaning that the operations within the context
+    manager block are cancelled if the limit is exceeded. The `LimitExceededError` is
+    therefore raised at the level that the `time_limit()` context manager was opened,
+    not at the level of the operation which caused the limit to be exceeded (e.g. a call
+    to `generate()`). Ensure you handle `LimitExceededError` at the level of opening the context manager.
+
+    Args:
+      limit: The maximum number of seconds that can pass while the context manager is
+        open. A value of None means unlimited time.
+    """
+    return _TimeLimit(limit)
+
+
+def working_limit(limit: float | None) -> _WorkingLimit:
+    """Limits the working time which can elapse.
+
+    Working time is the wall clock time minus any waiting time e.g. waiting before
+    retrying in response to rate limits or waiting on a semaphore.
+
+    The timer starts when the context manager is opened and stops when it is closed.
+
+    These limits can be stacked.
+
+    When a limit is exceeded, a `LimitExceededError` is raised.
+
+    Args:
+      limit: The maximum number of seconds of working that can pass while the context
+        manager is open. A value of None means unlimited time.
+    """
+    return _WorkingLimit(limit)
+
+
+def record_waiting_time(waiting_time: float) -> None:
+    node = working_limit_tree.get()
+    if node is None:
+        return
+    node.record_waiting_time(waiting_time)
+
+
+def check_working_limit() -> None:
+    node = working_limit_tree.get()
+    if node is None:
+        return
+    node.check()
+
+
 class _Tree(Generic[TNode]):
     """A tree data structure of limit nodes.
 
@@ -253,6 +322,7 @@ token_limit_tree: _Tree[_TokenLimit] = _Tree("token_limit_tree")
 # Store the message limit leaf node so that we know which limit to check in
 # check_message_limit().
 message_limit_tree: _Tree[_MessageLimit] = _Tree("message_limit_tree")
+working_limit_tree: _Tree[_WorkingLimit] = _Tree("working_limit_tree")
 
 
 class _Node:
@@ -296,6 +366,10 @@ class _TokenLimit(Limit, _Node):
     ) -> None:
         self._pop_and_check_identity(token_limit_tree)
 
+    @property
+    def usage(self) -> float:
+        return self._usage.total_tokens
+
     @property
     def limit(self) -> int | None:
         """Get the configured token limit value."""
@@ -312,7 +386,7 @@ class _TokenLimit(Limit, _Node):
         self._limit = value
 
     def record(self, usage: ModelUsage) -> None:
-        """Record model usage for this node and its parent nodes."""
+        """Record model usage for this node and its ancestor nodes."""
         if self.parent is not None:
             self.parent.record(usage)
         self._usage += usage
@@ -369,6 +443,13 @@ class _MessageLimit(Limit, _Node):
     ) -> None:
         self._pop_and_check_identity(message_limit_tree)
 
+    @property
+    def usage(self) -> float:
+        raise NotImplementedError(
+            "Retrieving the message count from a limit is not supported. Please query "
+            "the messages property on the task or agent state instead."
+        )
+
     @property
     def limit(self) -> int | None:
         """Get the configured message limit value."""
@@ -414,3 +495,132 @@ class _MessageLimit(Limit, _Node):
             raise ValueError(
                 f"Message limit value must be a non-negative integer or None: {value}"
             )
+
+
+class _TimeLimit(Limit):
+    def __init__(self, limit: float | None) -> None:
+        super().__init__()
+        _validate_time_limit("Time", limit)
+        self._limit = limit
+        self._start_time: float | None = None
+        self._end_time: float | None = None
+
+    def __enter__(self) -> Limit:
+        super()._check_reuse()
+        # Unlike the other limits, this one is not stored in a tree. Anyio handles all
+        # of the state.
+        self._cancel_scope = anyio.move_on_after(self._limit)
+        self._cancel_scope.__enter__()
+        self._start_time = anyio.current_time()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        from inspect_ai.log._transcript import SampleLimitEvent, transcript
+
+        self._cancel_scope.__exit__(exc_type, exc_val, exc_tb)
+        self._end_time = anyio.current_time()
+        if self._cancel_scope.cancel_called and self._limit is not None:
+            message = f"Time limit exceeded. limit: {self._limit} seconds"
+            assert self._start_time is not None
+            # Note we've measured the elapsed time independently of anyio's cancel scope
+            # so this is an approximation.
+            time_elapsed = self._end_time - self._start_time
+            transcript()._event(
+                SampleLimitEvent(type="time", message=message, limit=self._limit)
+            )
+            raise LimitExceededError(
+                "time",
+                value=time_elapsed,
+                limit=self._limit,
+                message=message,
+                source=self,
+            ) from exc_val
+
+    @property
+    def usage(self) -> float:
+        if self._start_time is None:
+            return 0.0
+        if self._end_time is None:
+            return anyio.current_time() - self._start_time
+        return self._end_time - self._start_time
+
+
+class _WorkingLimit(Limit, _Node):
+    def __init__(self, limit: float | None) -> None:
+        super().__init__()
+        _validate_time_limit("Working time", limit)
+        self._limit = limit
+        self.parent: _WorkingLimit | None = None
+        self._start_time: float | None = None
+        self._end_time: float | None = None
+
+    def __enter__(self) -> Limit:
+        super()._check_reuse()
+        self._start_time = anyio.current_time()
+        self._waiting_time = 0.0
+        working_limit_tree.push(self)
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        self._end_time = anyio.current_time()
+        self._pop_and_check_identity(working_limit_tree)
+
+    @property
+    def usage(self) -> float:
+        if self._start_time is None:
+            return 0.0
+        if self._end_time is None:
+            return anyio.current_time() - self._start_time - self._waiting_time
+        return self._end_time - self._start_time - self._waiting_time
+
+    def record_waiting_time(self, waiting_time: float) -> None:
+        """Record waiting time for this node and its ancestor nodes."""
+        if self.parent is not None:
+            self.parent.record_waiting_time(waiting_time)
+        self._waiting_time += waiting_time
+
+    def check(self) -> None:
+        """Check if this working time limit or any ancestor limits have been exceeded.
+
+        The checks occur from root to leaf. This is so that if multiple limits are
+        simultaneously exceeded, the outermost (closest to root) one raises the error,
+        preventing certain sub-agent architectures from ending up in an infinite loop.
+        """
+        if self.parent is not None:
+            self.parent.check()
+        self._check_self()
+
+    def _check_self(self) -> None:
+        from inspect_ai.log._transcript import SampleLimitEvent, transcript
+
+        if self._limit is None:
+            return
+        if self.usage > self._limit:
+            message = f"Working time limit exceeded. limit: {self._limit} seconds"
+            transcript()._event(
+                SampleLimitEvent(type="working", message=message, limit=self._limit)
+            )
+            raise LimitExceededError(
+                "working",
+                value=self.usage,
+                limit=self._limit,
+                message=message,
+                source=self,
+            )
+
+
+def _validate_time_limit(name: str, value: float | None) -> None:
+    if value is not None and value < 0:
+        raise ValueError(
+            f"{name} limit value must be a non-negative float or None: {value}"
+        )

inspect_ai/util/_sandbox/docker/compose.py

@@ -10,6 +10,7 @@ from pydantic import BaseModel
 
 from inspect_ai._util.error import PrerequisiteError
 from inspect_ai._util.trace import trace_message
+from inspect_ai.util._concurrency import concurrency
 from inspect_ai.util._display import display_type
 from inspect_ai.util._subprocess import ExecResult, subprocess
 
@@ -303,18 +304,26 @@ async def compose_command(
     # build final command
     compose_command = compose_command + command
 
-    # function to run command
+    # set a concurrency limit for docker CLI invocations.
+    # this should help with running more containers in parallel while avoiding hangs on some systems
+    DEFAULT_CLI_CONCURRENCY = max((os.cpu_count() or 1) * 2, 4)
+    docker_cli_concurrency = int(
+        os.environ.get("INSPECT_DOCKER_CLI_CONCURRENCY", DEFAULT_CLI_CONCURRENCY)
+    )
+
+    # function to run command (wrapped in concurrency limiter)
     async def run_command(command_timeout: int | None) -> ExecResult[str]:
-        result = await subprocess(
-            compose_command,
-            input=input,
-            cwd=cwd,
-            env=env,
-            timeout=command_timeout,
-            capture_output=capture_output,
-            output_limit=output_limit,
-        )
-        return result
+        async with concurrency("docker-cli", docker_cli_concurrency):
+            result = await subprocess(
+                compose_command,
+                input=input,
+                cwd=cwd,
+                env=env,
+                timeout=command_timeout,
+                capture_output=capture_output,
+                output_limit=output_limit,
+            )
+            return result
 
     # we have observed underlying unreliability in docker compose in some linux
     # environments on EC2 -- this exhibits in very simple commands (e.g. compose config)

inspect_ai/util/_span.py

@@ -35,7 +35,7 @@ async def span(name: str, *, type: str | None = None) -> AsyncIterator[None]:
             SpanBeginEvent(
                 id=id,
                 parent_id=parent_id,
-                type=type,
+                type=type or name,
                 name=name,
             )
         )

{inspect_ai-0.3.99.dist-info → inspect_ai-0.3.101.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: inspect_ai
-Version: 0.3.99
+Version: 0.3.101
 Summary: Framework for large language model evaluations
 Author: UK AI Security Institute
 License: MIT License
@@ -40,7 +40,7 @@ Requires-Dist: nest_asyncio
 Requires-Dist: numpy
 Requires-Dist: platformdirs>=2.3.0
 Requires-Dist: psutil
-Requires-Dist: pydantic>=2
+Requires-Dist: pydantic>=2.11.4
 Requires-Dist: python-dotenv>=0.16.0
 Requires-Dist: pyyaml
 Requires-Dist: rich<14.0.0,>=13.3.3
@@ -66,7 +66,7 @@ Requires-Dist: markdown; extra == "dev"
 Requires-Dist: mcp; extra == "dev"
 Requires-Dist: mistralai; extra == "dev"
 Requires-Dist: moto[server]; extra == "dev"
-Requires-Dist: mypy; extra == "dev"
+Requires-Dist: mypy>=1.16.0; extra == "dev"
 Requires-Dist: nbformat; extra == "dev"
 Requires-Dist: openai; extra == "dev"
 Requires-Dist: pandas>=2.0.0; extra == "dev"