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

inspect_ai/util/_subprocess.py

@@ -1,14 +1,20 @@
-import asyncio
+import functools
+import io
 import os
 import shlex
-import sys
-from asyncio.subprocess import Process
+from contextlib import aclosing
 from contextvars import ContextVar
 from dataclasses import dataclass
 from logging import getLogger
 from pathlib import Path
+from subprocess import DEVNULL, PIPE
 from typing import AsyncGenerator, Generic, Literal, TypeVar, Union, cast, overload
 
+import anyio
+from anyio import open_process
+from anyio.abc import ByteReceiveStream, Process
+
+from inspect_ai._util._async import tg_collect
 from inspect_ai._util.trace import trace_action
 
 from ._concurrency import concurrency
@@ -100,128 +106,103 @@ async def subprocess(
        TimeoutError: If the specified `timeout` expires.
     """
     # resolve input
-    input = input.encode() if isinstance(input, str) else input
+    input = (
+        input.encode()
+        if isinstance(input, str)
+        else bytes(input)
+        if input is not None
+        else None
+    )
 
-    # function to run command (we may or may not run it w/ concurrency)
     async def run_command() -> AsyncGenerator[
         Union[Process, ExecResult[str], ExecResult[bytes]], None
     ]:
-        if isinstance(args, str):
-            proc = await asyncio.create_subprocess_shell(
-                args,
-                stdin=asyncio.subprocess.PIPE,
-                stdout=asyncio.subprocess.PIPE if capture_output else None,
-                stderr=asyncio.subprocess.PIPE if capture_output else None,
-                cwd=cwd,
-                env={**os.environ, **env},
-            )
-        else:
-            proc = await asyncio.create_subprocess_exec(
-                args[0],
-                *args[1:],
-                stdin=asyncio.subprocess.PIPE,
-                stdout=asyncio.subprocess.PIPE if capture_output else None,
-                stderr=asyncio.subprocess.PIPE if capture_output else None,
-                cwd=cwd,
-                env={**os.environ, **env},
+        async with await open_process(
+            args,
+            stdin=PIPE if input else DEVNULL,
+            stdout=PIPE if capture_output else None,
+            stderr=PIPE if capture_output else None,
+            cwd=cwd,
+            env={**os.environ, **env},
+        ) as process:
+            # yield the process so the caller has a handle to it
+            yield process
+
+            # write to stdin (convert input to bytes)
+            if process.stdin and input:
+                await process.stdin.send(input)
+                await process.stdin.aclose()
+
+            # read streams incrementally so we can check output limits
+            async def read_stream(stream: ByteReceiveStream | None) -> bytes:
+                # return early for no stream
+                if stream is None:
+                    return bytes()
+
+                written = 0
+                buffer = io.BytesIO()
+                async for chunk in stream:
+                    buffer.write(chunk)
+                    written += len(chunk)
+                    if output_limit is not None and written > output_limit:
+                        process.kill()
+                        break
+
+                return buffer.getvalue()
+
+            stdout, stderr = await tg_collect(
+                [
+                    functools.partial(read_stream, process.stdout),
+                    functools.partial(read_stream, process.stderr),
+                ]
             )
 
-        # yield the proc
-        yield proc
-
-        # write stdin if specified
-        if proc.stdin is not None:
-            if input is not None:
-                proc.stdin.write(input)
-                await proc.stdin.drain()
-            proc.stdin.close()
-            await proc.stdin.wait_closed()
-
-        # read streams incrementally so we can check output limits
-        async def read_stream(stream: asyncio.StreamReader | None) -> bytes:
-            # return early for no stream
-            if stream is None:
-                return bytes()
-
-            # read 8k at a time
-            output = bytearray()
-            while True:
-                # read chunk and terminate if we are done
-                chunk = await stream.read(8192)
-                if not chunk:
-                    break
-
-                # append to output
-                output.extend(chunk)
-
-                # stop if we have a limit and we have exceeded it
-                if output_limit is not None and len(output) > output_limit:
-                    proc.kill()
-                    break
-
-            # return stream output
-            return bytes(output)
-
-        # wait for it to execute and yield result
-        stdout, stderr = await asyncio.gather(
-            read_stream(proc.stdout), read_stream(proc.stderr)
-        )
-        returncode = await proc.wait()
-        success = returncode == 0
-        if text:
-            yield ExecResult[str](
-                success=success,
-                returncode=returncode,
-                stdout=stdout.decode() if capture_output else "",
-                stderr=stderr.decode() if capture_output else "",
-            )
-        else:
-            yield ExecResult[bytes](
-                success=success,
-                returncode=returncode,
-                stdout=stdout if capture_output else bytes(),
-                stderr=stderr if capture_output else bytes(),
-            )
+            returncode = await process.wait()
+            success = returncode == 0
+            if text:
+                yield ExecResult[str](
+                    success=success,
+                    returncode=returncode,
+                    stdout=stdout.decode() if capture_output else "",
+                    stderr=stderr.decode() if capture_output else "",
+                )
+            else:
+                yield ExecResult[bytes](
+                    success=success,
+                    returncode=returncode,
+                    stdout=stdout if capture_output else bytes(),
+                    stderr=stderr if capture_output else bytes(),
+                )
 
     # wrapper for run command that implements timeout
     async def run_command_timeout() -> Union[ExecResult[str], ExecResult[bytes]]:
         # run the command and capture the process handle
-        rc = run_command()
-        proc = cast(Process, await anext(rc))
-
-        # await result wrapped in timeout handler if requested
-        if timeout:
-            try:
-                if sys.version_info >= (3, 11):
-                    async with asyncio.timeout(timeout):
+        async with aclosing(run_command()) as rc:
+            proc = cast(Process, await anext(rc))
+
+            # await result wrapped in timeout handler if requested
+            if timeout:
+                try:
+                    with anyio.fail_after(timeout):
                         result = await anext(rc)
                         return cast(Union[ExecResult[str], ExecResult[bytes]], result)
-                else:
-                    result = await asyncio.wait_for(anext(rc), timeout=timeout)
-                    return cast(Union[ExecResult[str], ExecResult[bytes]], result)
-            # wait_for raises asyncio.TimeoutError under Python 3.10, but TimeoutError
-            # under Python > 3.11! asynio.timeout (introduced in Python 3.11) always
-            # raises the standard TimeoutError
-            except (TimeoutError, asyncio.exceptions.TimeoutError):
-                # terminate timed out process -- try for graceful termination
-                # then be more forceful if requied
-                try:
-                    proc.terminate()
-                    await asyncio.sleep(2)
-                    if proc.returncode is None:
-                        proc.kill()
-                except Exception as ex:
-                    logger.warning(
-                        f"Unexpected error terminating timed out process '{args}': {ex}"
-                    )
-
-                # raise standard Python TimeoutError
-                raise TimeoutError
-
-        # await result without timeout
-        else:
-            result = await anext(rc)
-            return cast(Union[ExecResult[str], ExecResult[bytes]], result)
+                except TimeoutError:
+                    # terminate timed out process -- try for graceful termination
+                    # then be more forceful if requied
+                    with anyio.CancelScope(shield=True):
+                        try:
+                            proc.terminate()
+                            await anyio.sleep(2)
+                            if proc.returncode is None:
+                                proc.kill()
+                        except Exception:
+                            pass
+                    raise
+
+            # await result without timeout
+            else:
+                result = await anext(rc)
+                return cast(Union[ExecResult[str], ExecResult[bytes]], result)
 
     # run command
     async with concurrency("subprocesses", max_subprocesses_context_var.get()):

inspect_ai/util/_subtask.py

@@ -1,4 +1,3 @@
-import asyncio
 import inspect
 from datetime import datetime
 from functools import wraps
@@ -13,7 +12,7 @@ from typing import (
     runtime_checkable,
 )
 
-from inspect_ai._util._async import is_callable_coroutine
+from inspect_ai._util._async import is_callable_coroutine, tg_collect
 from inspect_ai._util.content import Content
 from inspect_ai._util.trace import trace_action
 from inspect_ai._util.working import sample_waiting_time
@@ -139,8 +138,7 @@ def subtask(
             transcript()._event(event)
 
             # create and run the task as a coroutine
-            asyncio_task = asyncio.create_task(run())
-            result, events = await asyncio_task
+            result, events = (await tg_collect([run]))[0]
 
             # time accounting
             completed = datetime.now()

{inspect_ai-0.3.72.dist-info → inspect_ai-0.3.73.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: inspect_ai
-Version: 0.3.72
+Version: 0.3.73
 Summary: Framework for large language model evaluations
 Author: UK AI Security Institute
 License: MIT License
@@ -21,11 +21,12 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: aiohttp>=3.9.0
-Requires-Dist: anyio>=4.4.0
+Requires-Dist: anyio>=4.8.0
 Requires-Dist: beautifulsoup4
 Requires-Dist: click>=8.1.3
 Requires-Dist: debugpy
 Requires-Dist: docstring-parser>=0.16
+Requires-Dist: exceptiongroup>=1.0.2; python_version < "3.11"
 Requires-Dist: fsspec<=2024.12.0,>=2023.1.0
 Requires-Dist: httpx
 Requires-Dist: ijson>=3.2.0
@@ -44,6 +45,7 @@ Requires-Dist: rich>=13.3.3
 Requires-Dist: s3fs>=2023
 Requires-Dist: semver>=3.0.0
 Requires-Dist: shortuuid
+Requires-Dist: sniffio
 Requires-Dist: tenacity
 Requires-Dist: textual>=0.86.2
 Requires-Dist: typing_extensions>=4.9.0
@@ -73,6 +75,7 @@ Requires-Dist: pytest-dotenv; extra == "dev"
 Requires-Dist: pytest-xdist; extra == "dev"
 Requires-Dist: ruff==0.9.6; extra == "dev"
 Requires-Dist: textual-dev>=0.86.2; extra == "dev"
+Requires-Dist: trio; extra == "dev"
 Requires-Dist: types-Markdown; extra == "dev"
 Requires-Dist: types-PyYAML; extra == "dev"
 Requires-Dist: types-beautifulsoup4; extra == "dev"
@@ -89,6 +92,7 @@ Requires-Dist: quarto-cli==1.5.57; extra == "doc"
 Requires-Dist: jupyter; extra == "doc"
 Requires-Dist: panflute; extra == "doc"
 Requires-Dist: markdown; extra == "doc"
+Requires-Dist: griffe; extra == "doc"
 Provides-Extra: dist
 Requires-Dist: twine; extra == "dist"
 Requires-Dist: build; extra == "dist"