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

inspect_ai/util/_sandbox/docker/docker.py

@@ -2,10 +2,11 @@ import base64
 import errno
 import json
 import os
+import shlex
 import tempfile
 from logging import getLogger
 from pathlib import Path, PurePosixPath
-from typing import Literal, Union, overload
+from typing import Literal, NamedTuple, Union, overload
 
 from typing_extensions import override
 
@@ -121,6 +122,37 @@ class DockerSandboxEnvironment(SandboxEnvironment):
             await project_cleanup_shutdown(True)
             raise ex
 
+    @override
+    @classmethod
+    async def task_init_environment(
+        cls, config: SandboxEnvironmentConfigType | None, metadata: dict[str, str]
+    ) -> dict[str, str]:
+        # get interpolated environment variables and underlying config path and text
+        resolved = resolve_config_environment(config, metadata)
+
+        # don't even consider sample-specific environment if there are no sample metadata refs
+        if resolved and len(resolved.env) > 0:
+            # resolve images using our env vars
+            result = await subprocess(
+                ["docker", "compose", "-f", resolved.config_file, "config", "--images"],
+                env=resolved.env,
+            )
+            if result.success:
+                # look through the images, if one of them doesn't apper in the the
+                # config text then this compose file requires its own sample specific
+                # environment for resolution
+                images = result.stdout.strip().splitlines()
+                for image in images:
+                    if image not in resolved.config_text:
+                        return resolved.env
+            else:
+                logger.warning(
+                    f"Unexpected error reading compose file '{resolved.config_file}': {result.stderr}"
+                )
+
+        # no per-sample environment required
+        return {}
+
     @override
     @classmethod
     async def sample_init(
@@ -130,17 +162,8 @@ class DockerSandboxEnvironment(SandboxEnvironment):
         metadata: dict[str, str],
     ) -> dict[str, SandboxEnvironment]:
         # create environment variables for sample metadata
-        env: dict[str, str] = {}
-        if isinstance(config, str) and Path(config).exists():
-            # read the config file
-            with open(config, "r") as f:
-                config_text = f.read()
-
-            # only add metadata files if the key is in the file
-            for key, value in metadata.items():
-                key = f"SAMPLE_METADATA_{key.replace(' ', '_').upper()}"
-                if key in config_text:
-                    env[key] = str(value)
+        resolved = resolve_config_environment(config, metadata)
+        env = resolved.env if resolved is not None else {}
 
         # create project
         from inspect_ai.log._samples import sample_active
@@ -154,6 +177,9 @@ class DockerSandboxEnvironment(SandboxEnvironment):
             env=env,
         )
 
+        # note that the project is running
+        project_startup(project)
+
         try:
             # enumerate the services that will be created
             services = await compose_services(project)
@@ -171,9 +197,6 @@ class DockerSandboxEnvironment(SandboxEnvironment):
                     f"No services started.\nCompose up stderr: {result.stderr}"
                 )
 
-            # note that the project is running
-            project_startup(project)
-
             # create sandbox environments for all running services
             default_service: str | None = None
             environments: dict[str, SandboxEnvironment] = {}
@@ -407,7 +430,7 @@ class DockerSandboxEnvironment(SandboxEnvironment):
                     return f.read()
 
     @override
-    async def connection(self) -> SandboxConnection:
+    async def connection(self, *, user: str | None = None) -> SandboxConnection:
         # find container for service
         services = await compose_ps(project=self._project)
         container = next(
@@ -419,15 +442,33 @@ class DockerSandboxEnvironment(SandboxEnvironment):
             None,
         )
 
+        # vscode doesn't support attaching to a container as a specific user,
+        # so don't include the vscode command if a user is specified
+        vscode_command = (
+            [
+                "remote-containers.attachToRunningContainer",
+                container,
+            ]
+            if user is None
+            else None
+        )
+
         # return container connection
         if container:
             return SandboxConnection(
                 type="docker",
-                command=f"docker exec -it {container} bash -l",
-                vscode_command=[
-                    "remote-containers.attachToRunningContainer",
-                    container,
-                ],
+                command=shlex.join(
+                    [
+                        "docker",
+                        "exec",
+                        "-it",
+                        *(["--user", user] if user else []),
+                        container,
+                        "bash",
+                        "-l",
+                    ]
+                ),
+                vscode_command=vscode_command,
                 ports=await get_ports_info(container),
                 container=container,
             )
@@ -517,3 +558,33 @@ def parse_docker_inspect_ports(json_str: str) -> list[PortMapping] | None:
         )
         port_mappings.append(port_mapping)
     return port_mappings if port_mappings else None
+
+
+class ConfigEnvironment(NamedTuple):
+    config_file: str
+    config_text: str
+    env: dict[str, str]
+
+
+def resolve_config_environment(
+    config: SandboxEnvironmentConfigType | None,
+    metadata: dict[str, str],
+) -> ConfigEnvironment | None:
+    # create environment variables for sample metadata
+    if isinstance(config, str) and Path(config).exists():
+        # read the config file
+        config_file = config
+        with open(config, "r") as f:
+            config_text = f.read()
+
+        # only add metadata files if the key is in the file
+        env: dict[str, str] = {}
+        for key, value in metadata.items():
+            key = f"SAMPLE_METADATA_{key.replace(' ', '_').upper()}"
+            if key in config_text:
+                env[key] = str(value)
+
+        # return resolved
+        return ConfigEnvironment(config_file, config_text, env)
+    else:
+        return None

inspect_ai/util/_sandbox/environment.py

@@ -27,6 +27,10 @@ logger = logging.getLogger(__name__)
 ST = TypeVar("ST", bound="SandboxEnvironment")
 
 TaskInit = Callable[[str, Union["SandboxEnvironmentConfigType", None]], Awaitable[None]]
+TaskInitEnvironment = Callable[
+    [Union["SandboxEnvironmentConfigType", None], dict[str, str]],
+    Awaitable[dict[str, str]],
+]
 TaskCleanup = Callable[
     [str, Union["SandboxEnvironmentConfigType", None], bool], Awaitable[None]
 ]
@@ -187,11 +191,14 @@ class SandboxEnvironment(abc.ABC):
         """
         ...
 
-    async def connection(self) -> SandboxConnection:
+    async def connection(self, *, user: str | None = None) -> SandboxConnection:
         """Information required to connect to sandbox environment.
 
+        Args:
+          user: User to login as.
+
         Returns:
-           SandboxConnection: connection information
+           SandboxConnection: connection information.
 
         Raises:
            NotImplementedError: For sandboxes that don't provide connections
@@ -235,6 +242,30 @@ class SandboxEnvironment(abc.ABC):
         """
         pass
 
+    @classmethod
+    async def task_init_environment(
+        cls, config: SandboxEnvironmentConfigType | None, metadata: dict[str, str]
+    ) -> dict[str, str]:
+        """Called at task startup to identify environment variables required by task_init for a sample.
+
+        Return 1 or more environment variables to request a dedicated call to task_init
+        for samples that have exactly these environment variables (by default there is
+        only one call to task_init for all of the samples in a task if they share a
+        sandbox configuration).
+
+        This is useful for situations where config files are dynamic (e.g. through
+        sample metadata variable interpolation) and end up yielding different images
+        that need their own init (e.g. 'docker pull').
+
+        Args:
+            config: Implementation defined configuration (optional).
+            metadata: metadata: Sample `metadata` field
+
+        Returns:
+            Environment variables to set for call to task_init.
+        """
+        return {}
+
     @classmethod
     async def sample_init(
         cls,

inspect_ai/util/_sandbox/events.py

@@ -133,8 +133,8 @@ class SandboxEnvironmentProxy(SandboxEnvironment):
         return output
 
     @override
-    async def connection(self) -> SandboxConnection:
-        return await self._sandbox.connection()
+    async def connection(self, *, user: str | None = None) -> SandboxConnection:
+        return await self._sandbox.connection(user=user)
 
     @override
     def as_type(self, sandbox_cls: Type[ST]) -> ST:

inspect_ai/util/_sandbox/service.py

@@ -40,6 +40,7 @@ async def sandbox_service(
     methods: list[SandboxServiceMethod] | dict[str, SandboxServiceMethod],
     until: Callable[[], bool],
     sandbox: SandboxEnvironment,
+    user: str | None = None,
 ) -> None:
     """Run a service that is callable from within a sandbox.
 
@@ -49,9 +50,11 @@ async def sandbox_service(
         until (Callable[[], bool]): Function used to check whether
           the service should stop.
         sandbox (SandboxEnvironment): Sandbox to publish service to.
+        user (str | None): User to login as. Defaults to the sandbox environment's
+          default user.
     """
     # setup and start service
-    service = SandboxService(name, sandbox)
+    service = SandboxService(name, sandbox, user)
     if isinstance(methods, list):
         methods = {v.__name__: v for v in methods}
     for name, method in methods.items():
@@ -91,15 +94,20 @@ class SandboxService:
     ```
     """
 
-    def __init__(self, name: str, sandbox: SandboxEnvironment) -> None:
+    def __init__(
+        self, name: str, sandbox: SandboxEnvironment, user: str | None = None
+    ) -> None:
         """Create a SandboxService.
 
         Args:
             name (str): Service name
             sandbox (SandboxEnvironment): Sandbox to publish service to.
+            user (str | None): User to login as. Defaults to the sandbox environment's
+              default user.
         """
         self._name = name
         self._sandbox = sandbox
+        self._user = user
         self._service_dir = PurePosixPath(SERVICES_DIR, self._name)
         self._methods: dict[str, SandboxServiceMethod] = {}
         self._requests_dir: str = ""
@@ -250,7 +258,9 @@ class SandboxService:
 
     async def _exec(self, cmd: list[str], input: str | None = None) -> ExecResult[str]:
         try:
-            return await self._sandbox.exec(cmd, input=input, timeout=30)
+            return await self._sandbox.exec(
+                cmd, user=self._user, input=input, timeout=30
+            )
         except TimeoutError:
             raise RuntimeError(
                 f"Timed out executing command {' '.join(cmd)} in sandbox"

{inspect_ai-0.3.90.dist-info → inspect_ai-0.3.91.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: inspect_ai
-Version: 0.3.90
+Version: 0.3.91
 Summary: Framework for large language model evaluations
 Author: UK AI Security Institute
 License: MIT License
@@ -27,7 +27,7 @@ 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: fsspec<=2025.3.0,>=2023.1.0
 Requires-Dist: httpx
 Requires-Dist: ijson>=3.2.0
 Requires-Dist: jsonlines>=3.0.0
@@ -60,6 +60,7 @@ Requires-Dist: griffe; extra == "dev"
 Requires-Dist: groq; extra == "dev"
 Requires-Dist: ipython; extra == "dev"
 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"
@@ -88,6 +89,9 @@ Requires-Dist: types-jsonschema; extra == "dev"
 Requires-Dist: types-protobuf; extra == "dev"
 Requires-Dist: types-psutil; extra == "dev"
 Requires-Dist: types-python-dateutil; extra == "dev"
+Provides-Extra: dev-mcp-tests
+Requires-Dist: mcp-server-fetch; extra == "dev-mcp-tests"
+Requires-Dist: mcp_server_git; extra == "dev-mcp-tests"
 Provides-Extra: doc
 Requires-Dist: quarto-cli==1.5.57; extra == "doc"
 Requires-Dist: jupyter; extra == "doc"