From 6d56b252dbcb52aa5a5ca995e2b5505fe553b6de Mon Sep 17 00:00:00 2001 From: WatchTree-19 <119982314+WatchTree-19@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:13:44 +0000 Subject: [PATCH 1/6] FEAT: program execution scorer in an isolated sandbox (scorer phase 10b) Adds ProgramBehaves, ProgramExecutionScorable, ProgramExecutionSource, the SandboxClient protocol with a DockerSandboxClient, and ProgramExecutionScorer. The scorer sends the source only the runtime profile and invocations; expected results stay with the scorer. Program failure is observed evidence, a sandbox failure is undetermined, and a stored execution can be judged again without running the program. --- doc/code/framework.md | 5 + doc/code/scoring/0_scoring.ipynb | 3 +- doc/code/scoring/0_scoring.py | 3 +- .../scoring/7_program_execution_scorer.ipynb | 310 ++++++++++ .../scoring/7_program_execution_scorer.py | 173 ++++++ doc/myst.yml | 1 + pyrit/models/__init__.py | 14 + pyrit/models/score/__init__.py | 13 + pyrit/models/score/condition.py | 68 ++- pyrit/models/score/observation.py | 105 +++- pyrit/models/score/program.py | 78 +++ pyrit/models/score/scorable.py | 49 +- pyrit/score/__init__.py | 11 + pyrit/score/observation/__init__.py | 9 + .../observation/docker_sandbox_client.py | 416 +++++++++++++ pyrit/score/observation/execution.py | 11 +- .../observation/program_execution_source.py | 222 +++++++ pyrit/score/observation/sandbox_client.py | 67 +++ .../true_false/program_execution_scorer.py | 293 +++++++++ tests/unit/models/test_program.py | 261 ++++++++ tests/unit/models/test_scorable.py | 8 + .../unit/score/test_docker_sandbox_client.py | 268 +++++++++ .../score/test_program_execution_scorer.py | 560 ++++++++++++++++++ 23 files changed, 2940 insertions(+), 8 deletions(-) create mode 100644 doc/code/scoring/7_program_execution_scorer.ipynb create mode 100644 doc/code/scoring/7_program_execution_scorer.py create mode 100644 pyrit/models/score/program.py create mode 100644 pyrit/score/observation/docker_sandbox_client.py create mode 100644 pyrit/score/observation/program_execution_source.py create mode 100644 pyrit/score/observation/sandbox_client.py create mode 100644 pyrit/score/true_false/program_execution_scorer.py create mode 100644 tests/unit/models/test_program.py create mode 100644 tests/unit/score/test_docker_sandbox_client.py create mode 100644 tests/unit/score/test_program_execution_scorer.py diff --git a/doc/code/framework.md b/doc/code/framework.md index dd67790ed2..3bbc794a57 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -289,6 +289,11 @@ If you are contributing to PyRIT, that work will most likely land in one of the - Surface sources read what a location holds for a `SurfaceScorable`, such as files under one root directory. `FileWriteScorer` evaluates its `ContentWritten` condition against that evidence. The caller owns workspace isolation; a snapshot does not identify its writer. +- Program-execution sources run a program for a `ProgramExecutionScorable` through an injected + `SandboxClient`, such as `DockerSandboxClient`. `ProgramExecutionScorer` splits its `ProgramBehaves` + condition so the source receives only the runtime profile and invocations, never the expected + results. A program that fails is observed evidence; a sandbox that cannot run it is an + acquisition error, so the verdict is undetermined rather than false. - Observation lifecycle is optional. A scorer that implements `AttemptObserver`, whether it is a root scorer or wrapped inside one, is asked by the attack to start observing after setup, once its attempt ID is allocated and before anything is sent, and to stop when the attempt ends. A diff --git a/doc/code/scoring/0_scoring.ipynb b/doc/code/scoring/0_scoring.ipynb index a58a705fd4..8c1f01439b 100644 --- a/doc/code/scoring/0_scoring.ipynb +++ b/doc/code/scoring/0_scoring.ipynb @@ -275,7 +275,8 @@ "the target need not be a language model. Its kind is `scorer_target_response`.\n", "Media observation capture remains deferred until its evidence can be snapshotted.\n", "Trace-backed tool observations are covered in [Tool-call scoring](5_tool_call_scorer.ipynb), and file-system\n", - "observations in [File-write scoring](6_file_write_scorer.ipynb).\n", + "observations in [File-write scoring](6_file_write_scorer.ipynb), and sandboxed program runs in\n", + "[Program-execution scoring](7_program_execution_scorer.ipynb).\n", "\n", "Replaying a judgment is different from evaluating a stored run against a new expectation.\n", "A retained target judgment answers the original expectation; changing that expectation\n", diff --git a/doc/code/scoring/0_scoring.py b/doc/code/scoring/0_scoring.py index 0dc5d1fa2b..006d7e7052 100644 --- a/doc/code/scoring/0_scoring.py +++ b/doc/code/scoring/0_scoring.py @@ -161,7 +161,8 @@ # the target need not be a language model. Its kind is `scorer_target_response`. # Media observation capture remains deferred until its evidence can be snapshotted. # Trace-backed tool observations are covered in [Tool-call scoring](5_tool_call_scorer.ipynb), and file-system -# observations in [File-write scoring](6_file_write_scorer.ipynb). +# observations in [File-write scoring](6_file_write_scorer.ipynb), and sandboxed program runs in +# [Program-execution scoring](7_program_execution_scorer.ipynb). # # Replaying a judgment is different from evaluating a stored run against a new expectation. # A retained target judgment answers the original expectation; changing that expectation diff --git a/doc/code/scoring/7_program_execution_scorer.ipynb b/doc/code/scoring/7_program_execution_scorer.ipynb new file mode 100644 index 0000000000..9089018cb3 --- /dev/null +++ b/doc/code/scoring/7_program_execution_scorer.ipynb @@ -0,0 +1,310 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Program-execution scoring\n", + "\n", + "`ProgramExecutionScorer` answers \"Does this program behave as intended for these inputs?\" It runs\n", + "the program in an isolated sandbox and judges what the program did, not what a response claims\n", + "it would do.\n", + "\n", + "The `ProgramBehaves` condition lists cases. Each `ProgramCase` pairs a `ProgramInvocation`\n", + "(arguments and standard input) with a `ProgramResult` (exit code, and optional `stdout` and\n", + "`stderr` matchers). The scorer splits them before anything runs: the runtime profile and the\n", + "invocations go to the source, and the expected results stay with the scorer. The code that runs\n", + "the program never sees which output passes.\n", + "\n", + "A `ProgramExecutionSource` asks an injected `SandboxClient` to run the program.\n", + "`DockerSandboxClient` runs each invocation in a fresh container with no network, a read-only\n", + "file system, no capabilities, an unprivileged user, and memory, CPU and process limits. A runtime\n", + "profile is a trusted name for a digest-pinned image and a command; the program cannot pick either.\n", + "Images are never pulled while scoring, so pull them first.\n", + "\n", + "A program that fails to compile, exits with the wrong code, prints the wrong output, times out, or\n", + "exceeds a resource limit is a `False` verdict: the sandbox observed it. A sandbox that cannot run\n", + "the program, or stops before every case ran, leaves the verdict undetermined, so broken\n", + "infrastructure never looks like a failing program.\n", + "\n", + "This walkthrough needs a Docker daemon and the pinned image below\n", + "(`docker pull python:3.12-slim`), and PyRIT's in-memory storage. It needs no model or credentials." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import httpx\n", + "\n", + "from pyrit.executor.attack import AttackScoringConfig, PromptSendingAttack\n", + "from pyrit.memory import CentralMemory\n", + "from pyrit.models import (\n", + " ContentScorable,\n", + " Equals,\n", + " ProgramBehaves,\n", + " ProgramCase,\n", + " ProgramInvocation,\n", + " ProgramResult,\n", + " ScoringExpectation,\n", + ")\n", + "from pyrit.prompt_target import HTTPTarget\n", + "from pyrit.score import DockerRuntimeProfile, DockerSandboxClient, ProgramExecutionScorer, ProgramExecutionSource\n", + "from pyrit.setup import IN_MEMORY, initialize_pyrit_async\n", + "\n", + "await initialize_pyrit_async( # type: ignore\n", + " memory_db_type=IN_MEMORY,\n", + " load_defaults=False,\n", + " env_files=[],\n", + " silent=True,\n", + ")\n", + "memory = CentralMemory.get_memory_instance()\n", + "\n", + "python_312 = DockerRuntimeProfile(\n", + " image=\"python@sha256:a6e34c598f2467ed0e9a8d349809fcd8b5c603269512df273a0bb1784edc11b1\",\n", + " command=(\"python3\", \"-I\"),\n", + ")\n", + "sandbox = DockerSandboxClient(profiles={\"python-3.12\": python_312}, timeout_seconds=5)\n", + "scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=sandbox))\n", + "\n", + "squares = ScoringExpectation(\n", + " conditions=(\n", + " ProgramBehaves(\n", + " runtime_profile=\"python-3.12\",\n", + " cases=(\n", + " ProgramCase(\n", + " invocation=ProgramInvocation(arguments=(\"3\",)), expected=ProgramResult(stdout=Equals(value=\"9\"))\n", + " ),\n", + " ProgramCase(\n", + " invocation=ProgramInvocation(stdin=\"12\\n\"), expected=ProgramResult(stdout=Equals(value=\"144\"))\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "## Score a program directly\n", + "\n", + "A `ContentScorable` holds the program text. The program is run verbatim, once per case." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "square: True - All 2 case(s) produced the expected result under python-3.12.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "doubles: False - The program did not produce the expected result under python-3.12: case 0 output did not match; case 1 output did not match.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "loops: False - The program did not produce the expected result under python-3.12: case 0 timed out; case 1 timed out.\n" + ] + } + ], + "source": [ + "square = \"import sys\\nn = int(sys.argv[1]) if len(sys.argv) > 1 else int(sys.stdin.read())\\nprint(n * n)\\n\"\n", + "doubles = \"import sys\\nn = int(sys.argv[1]) if len(sys.argv) > 1 else int(sys.stdin.read())\\nprint(n * 2)\\n\"\n", + "loops = \"while True:\\n pass\\n\"\n", + "\n", + "for name, program in ((\"square\", square), (\"doubles\", doubles), (\"loops\", loops)):\n", + " score = (await scorer.score_async(scorable=ContentScorable(value=program), expectation=squares))[0] # type: ignore\n", + " print(f\"{name}: {score.get_value()} - {score.score_rationale}\")" + ] + }, + { + "cell_type": "markdown", + "id": "4", + "metadata": {}, + "source": [ + "## A sandbox failure is not a program failure\n", + "\n", + "Here the profile names an image that is not on the machine, so no container starts. The\n", + "observation is an acquisition error and the score is undetermined, not `False`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Sandbox could not run program for execution d96c13b0-eb8f-4b1d-ad01-a7c2b160230d: the container did not start: docker: Error response from daemon: No such image: python@sha256:0000000000000000000000000000000000000000000000000000000000000000\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "undetermined: The program's behavior under python-3.12 could not be observed (sandbox error: the container did not start: docker: Error response from daemon: No such image: python@sha256:0000000000000000000000000000000000000000000000000000000000000000).\n" + ] + } + ], + "source": [ + "missing_image = DockerSandboxClient(\n", + " profiles={\"python-3.12\": DockerRuntimeProfile(image=\"python@sha256:\" + \"0\" * 64, command=(\"python3\", \"-I\"))}\n", + ")\n", + "unavailable = (\n", + " await ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=missing_image)).score_async( # type: ignore\n", + " scorable=ContentScorable(value=square), expectation=squares\n", + " )\n", + ")[0]\n", + "print(f\"{unavailable.status.value}: {unavailable.score_rationale}\")" + ] + }, + { + "cell_type": "markdown", + "id": "6", + "metadata": {}, + "source": [ + "## Re-judge saved evidence\n", + "\n", + "The observation keeps each invocation's exit state, a bounded copy of both streams, the resolved\n", + "command, and the image digest. Replay judges that evidence against new expected results without\n", + "running the program again. The invocations must match the ones that ran." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "case 0: ('python3', '-I', '/program/main.py', '3') exit 0, stdout '9\\n'\n", + "case 1: ('python3', '-I', '/program/main.py') exit 0, stdout '144\\n'\n", + "Replayed against a stricter expectation: False - The program did not produce the expected result under python-3.12: case 1 output did not match.\n" + ] + } + ], + "source": [ + "passing = (await scorer.score_async(scorable=ContentScorable(value=square), expectation=squares))[0] # type: ignore\n", + "saved = (await memory.get_observations_async(observation_ids=passing.observation_ids))[0]\n", + "for event in saved.payload.events:\n", + " print(f\"case {event.case_index}: {event.resolved_command} exit {event.exit_code}, stdout {event.stdout!r}\")\n", + "\n", + "stricter = ScoringExpectation(\n", + " conditions=(\n", + " ProgramBehaves(\n", + " runtime_profile=\"python-3.12\",\n", + " cases=(\n", + " ProgramCase(\n", + " invocation=ProgramInvocation(arguments=(\"3\",)), expected=ProgramResult(stdout=Equals(value=\"9\"))\n", + " ),\n", + " ProgramCase(\n", + " invocation=ProgramInvocation(stdin=\"12\\n\"), expected=ProgramResult(stdout=Equals(value=\"145\"))\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + ")\n", + "replayed = (await scorer.score_observation_async(observation=saved, expectation=stricter))[0] # type: ignore\n", + "print(f\"Replayed against a stricter expectation: {replayed.get_value()} - {replayed.score_rationale}\")" + ] + }, + { + "cell_type": "markdown", + "id": "8", + "metadata": { + "lines_to_next_cell": 2 + }, + "source": [ + "## Score a generated program through an attack\n", + "\n", + "The attack scores its response as a program. This local target returns code as its whole response.\n", + "A response that wraps code in explanation is run as it is and fails its cases; the scorer does not\n", + "guess where the code starts." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Target writes square -> outcome success\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Target writes doubles -> outcome failure\n" + ] + } + ], + "source": [ + "def coder(*, program: str) -> HTTPTarget:\n", + " \"\"\"Return a local target that answers every prompt with the same program.\"\"\"\n", + " return HTTPTarget(\n", + " http_request=\"POST / HTTP/1.1\\nHost: coder.test\\n\\n{PROMPT}\",\n", + " transport=httpx.MockTransport(lambda request: httpx.Response(200, text=program)),\n", + " )\n", + "\n", + "\n", + "for name, program in ((\"square\", square), (\"doubles\", doubles)):\n", + " attack = PromptSendingAttack(\n", + " objective_target=coder(program=program),\n", + " attack_scoring_config=AttackScoringConfig(objective_scorer=scorer),\n", + " max_attempts_on_failure=0,\n", + " )\n", + " result = await attack.execute_async(objective=\"write a program that squares its input\", expectation=squares) # type: ignore\n", + " print(f\"Target writes {name} -> outcome {result.outcome.value}\")" + ] + } + ], + "metadata": { + "jupytext": { + "cell_metadata_filter": "-all" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.16" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/doc/code/scoring/7_program_execution_scorer.py b/doc/code/scoring/7_program_execution_scorer.py new file mode 100644 index 0000000000..536aaf751b --- /dev/null +++ b/doc/code/scoring/7_program_execution_scorer.py @@ -0,0 +1,173 @@ +# --- +# jupyter: +# jupytext: +# cell_metadata_filter: -all +# text_representation: +# extension: .py +# format_name: percent +# format_version: '1.3' +# --- + +# %% [markdown] +# # Program-execution scoring +# +# `ProgramExecutionScorer` answers "Does this program behave as intended for these inputs?" It runs +# the program in an isolated sandbox and judges what the program did, not what a response claims +# it would do. +# +# The `ProgramBehaves` condition lists cases. Each `ProgramCase` pairs a `ProgramInvocation` +# (arguments and standard input) with a `ProgramResult` (exit code, and optional `stdout` and +# `stderr` matchers). The scorer splits them before anything runs: the runtime profile and the +# invocations go to the source, and the expected results stay with the scorer. The code that runs +# the program never sees which output passes. +# +# A `ProgramExecutionSource` asks an injected `SandboxClient` to run the program. +# `DockerSandboxClient` runs each invocation in a fresh container with no network, a read-only +# file system, no capabilities, an unprivileged user, and memory, CPU and process limits. A runtime +# profile is a trusted name for a digest-pinned image and a command; the program cannot pick either. +# Images are never pulled while scoring, so pull them first. +# +# A program that fails to compile, exits with the wrong code, prints the wrong output, times out, or +# exceeds a resource limit is a `False` verdict: the sandbox observed it. A sandbox that cannot run +# the program, or stops before every case ran, leaves the verdict undetermined, so broken +# infrastructure never looks like a failing program. +# +# This walkthrough needs a Docker daemon and the pinned image below +# (`docker pull python:3.12-slim`), and PyRIT's in-memory storage. It needs no model or credentials. + +# %% +import httpx + +from pyrit.executor.attack import AttackScoringConfig, PromptSendingAttack +from pyrit.memory import CentralMemory +from pyrit.models import ( + ContentScorable, + Equals, + ProgramBehaves, + ProgramCase, + ProgramInvocation, + ProgramResult, + ScoringExpectation, +) +from pyrit.prompt_target import HTTPTarget +from pyrit.score import DockerRuntimeProfile, DockerSandboxClient, ProgramExecutionScorer, ProgramExecutionSource +from pyrit.setup import IN_MEMORY, initialize_pyrit_async + +await initialize_pyrit_async( # type: ignore + memory_db_type=IN_MEMORY, + load_defaults=False, + env_files=[], + silent=True, +) +memory = CentralMemory.get_memory_instance() + +python_312 = DockerRuntimeProfile( + image="python@sha256:a6e34c598f2467ed0e9a8d349809fcd8b5c603269512df273a0bb1784edc11b1", + command=("python3", "-I"), +) +sandbox = DockerSandboxClient(profiles={"python-3.12": python_312}, timeout_seconds=5) +scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=sandbox)) + +squares = ScoringExpectation( + conditions=( + ProgramBehaves( + runtime_profile="python-3.12", + cases=( + ProgramCase( + invocation=ProgramInvocation(arguments=("3",)), expected=ProgramResult(stdout=Equals(value="9")) + ), + ProgramCase( + invocation=ProgramInvocation(stdin="12\n"), expected=ProgramResult(stdout=Equals(value="144")) + ), + ), + ), + ) +) + +# %% [markdown] +# ## Score a program directly +# +# A `ContentScorable` holds the program text. The program is run verbatim, once per case. + +# %% +square = "import sys\nn = int(sys.argv[1]) if len(sys.argv) > 1 else int(sys.stdin.read())\nprint(n * n)\n" +doubles = "import sys\nn = int(sys.argv[1]) if len(sys.argv) > 1 else int(sys.stdin.read())\nprint(n * 2)\n" +loops = "while True:\n pass\n" + +for name, program in (("square", square), ("doubles", doubles), ("loops", loops)): + score = (await scorer.score_async(scorable=ContentScorable(value=program), expectation=squares))[0] # type: ignore + print(f"{name}: {score.get_value()} - {score.score_rationale}") + +# %% [markdown] +# ## A sandbox failure is not a program failure +# +# Here the profile names an image that is not on the machine, so no container starts. The +# observation is an acquisition error and the score is undetermined, not `False`. + +# %% +missing_image = DockerSandboxClient( + profiles={"python-3.12": DockerRuntimeProfile(image="python@sha256:" + "0" * 64, command=("python3", "-I"))} +) +unavailable = ( + await ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=missing_image)).score_async( # type: ignore + scorable=ContentScorable(value=square), expectation=squares + ) +)[0] +print(f"{unavailable.status.value}: {unavailable.score_rationale}") + +# %% [markdown] +# ## Re-judge saved evidence +# +# The observation keeps each invocation's exit state, a bounded copy of both streams, the resolved +# command, and the image digest. Replay judges that evidence against new expected results without +# running the program again. The invocations must match the ones that ran. + +# %% +passing = (await scorer.score_async(scorable=ContentScorable(value=square), expectation=squares))[0] # type: ignore +saved = (await memory.get_observations_async(observation_ids=passing.observation_ids))[0] +for event in saved.payload.events: + print(f"case {event.case_index}: {event.resolved_command} exit {event.exit_code}, stdout {event.stdout!r}") + +stricter = ScoringExpectation( + conditions=( + ProgramBehaves( + runtime_profile="python-3.12", + cases=( + ProgramCase( + invocation=ProgramInvocation(arguments=("3",)), expected=ProgramResult(stdout=Equals(value="9")) + ), + ProgramCase( + invocation=ProgramInvocation(stdin="12\n"), expected=ProgramResult(stdout=Equals(value="145")) + ), + ), + ), + ) +) +replayed = (await scorer.score_observation_async(observation=saved, expectation=stricter))[0] # type: ignore +print(f"Replayed against a stricter expectation: {replayed.get_value()} - {replayed.score_rationale}") + +# %% [markdown] +# ## Score a generated program through an attack +# +# The attack scores its response as a program. This local target returns code as its whole response. +# A response that wraps code in explanation is run as it is and fails its cases; the scorer does not +# guess where the code starts. + + +# %% +def coder(*, program: str) -> HTTPTarget: + """Return a local target that answers every prompt with the same program.""" + return HTTPTarget( + http_request="POST / HTTP/1.1\nHost: coder.test\n\n{PROMPT}", + transport=httpx.MockTransport(lambda request: httpx.Response(200, text=program)), + ) + + +for name, program in (("square", square), ("doubles", doubles)): + attack = PromptSendingAttack( + objective_target=coder(program=program), + attack_scoring_config=AttackScoringConfig(objective_scorer=scorer), + max_attempts_on_failure=0, + ) + result = await attack.execute_async(objective="write a program that squares its input", expectation=squares) # type: ignore + print(f"Target writes {name} -> outcome {result.outcome.value}") diff --git a/doc/myst.yml b/doc/myst.yml index 6c1783749a..8f983dd9ac 100644 --- a/doc/myst.yml +++ b/doc/myst.yml @@ -162,6 +162,7 @@ project: - file: code/scoring/4_scorer_metrics.ipynb - file: code/scoring/5_tool_call_scorer.ipynb - file: code/scoring/6_file_write_scorer.ipynb + - file: code/scoring/7_program_execution_scorer.ipynb - file: code/memory/0_memory.md children: - file: code/memory/1_sqlite_memory.ipynb diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index f3d6d98464..9852c8f1b8 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -203,6 +203,13 @@ Observation, ObservationPayload, OutputMatches, + ProgramBehaves, + ProgramCase, + ProgramExecutionEvent, + ProgramExecutionObservationPayload, + ProgramExecutionScorable, + ProgramInvocation, + ProgramResult, Regex, Scorable, ScorableUnion, @@ -279,6 +286,13 @@ "Regex": "pyrit.models.score", "TextMatcher": "pyrit.models.score", "OutputMatches": "pyrit.models.score", + "ProgramBehaves": "pyrit.models.score", + "ProgramCase": "pyrit.models.score", + "ProgramExecutionEvent": "pyrit.models.score", + "ProgramExecutionObservationPayload": "pyrit.models.score", + "ProgramExecutionScorable": "pyrit.models.score", + "ProgramInvocation": "pyrit.models.score", + "ProgramResult": "pyrit.models.score", "ConversationScorable": "pyrit.models.score", "ConversationObservationPayload": "pyrit.models.score", "AllAvailableDatasetSize": "pyrit.models.scenario_dataset_size_estimate", diff --git a/pyrit/models/score/__init__.py b/pyrit/models/score/__init__.py index 8cda49e477..6d82a06631 100644 --- a/pyrit/models/score/__init__.py +++ b/pyrit/models/score/__init__.py @@ -22,6 +22,9 @@ DivergesFromRepetition, MatchesObjective, OutputMatches, + ProgramBehaves, + ProgramCase, + ProgramResult, ToolCallRequirement, ToolsCalled, ) @@ -34,15 +37,18 @@ ConversationObservationPayload, Observation, ObservationPayload, + ProgramExecutionObservationPayload, ScorerTargetResponsePayload, SurfaceObservationPayload, ToolEventsObservationPayload, ) + from pyrit.models.score.program import ProgramExecutionEvent, ProgramInvocation from pyrit.models.score.scorable import ( ContentEntryScorable, ContentScorable, ConversationScorable, MessageScorable, + ProgramExecutionScorable, Scorable, ScorableUnion, SurfaceScorable, @@ -69,6 +75,13 @@ ) _LAZY_EXPORTS: dict[str, str] = { + "ProgramBehaves": "pyrit.models.score.condition", + "ProgramCase": "pyrit.models.score.condition", + "ProgramExecutionEvent": "pyrit.models.score.program", + "ProgramExecutionObservationPayload": "pyrit.models.score.observation", + "ProgramExecutionScorable": "pyrit.models.score.scorable", + "ProgramInvocation": "pyrit.models.score.program", + "ProgramResult": "pyrit.models.score.condition", "ConversationScorable": "pyrit.models.score.scorable", "ConversationObservationPayload": "pyrit.models.score.observation", "Contains": "pyrit.models.score.text_matcher", diff --git a/pyrit/models/score/condition.py b/pyrit/models/score/condition.py index 8491522d21..196aa01e74 100644 --- a/pyrit/models/score/condition.py +++ b/pyrit/models/score/condition.py @@ -9,6 +9,10 @@ from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, SerializeAsAny, TypeAdapter, model_validator from pyrit.models.score._trace_validation import ToolName # noqa: TC001 (runtime-required by Pydantic) +from pyrit.models.score.program import ( # noqa: TC001 (runtime-required by Pydantic) + ProgramInvocation, + RuntimeProfileName, +) from pyrit.models.score.surface import SurfaceMatch # noqa: TC001 (runtime-required by Pydantic) from pyrit.models.score.text_matcher import TextMatcher # noqa: TC001 (runtime-required by Pydantic) @@ -240,6 +244,65 @@ class ContentWritten(Condition): require_content_change: bool = False +class ProgramResult(BaseModel): + """ + The expected result of one invocation. + + ``exit_code`` defaults to 0; ``None`` accepts any exit code from a program that exited on + its own. A stream with no matcher is not checked. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + exit_code: int | None = 0 + stdout: TextMatcher | None = None + stderr: TextMatcher | None = None + + +class ProgramCase(BaseModel): + """One invocation and the result it is expected to produce.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + invocation: ProgramInvocation = Field(default_factory=ProgramInvocation) + expected: ProgramResult = Field(default_factory=ProgramResult) + + +class ProgramBehaves(Condition): + """ + The program produces the expected result for every case. + + The condition carries test stimuli and assertions together, because both vary by dataset + row. A scorer splits them before acquisition: the runtime profile and each invocation go to + the source, and each expected result stays with the scorer's matcher. A run that times out + or exceeds a resource limit fails its case. + """ + + condition_type: Literal["program_behaves"] = "program_behaves" + runtime_profile: RuntimeProfileName + cases: tuple[ProgramCase, ...] + + @model_validator(mode="after") + def _validate_cases(self) -> ProgramBehaves: + """ + Require at least one case. + + Returns: + ProgramBehaves: The validated condition. + + Raises: + ValueError: If there are no cases. + """ + if not self.cases: + raise ValueError("ProgramBehaves requires at least one case.") + return self + + @property + def invocations(self) -> tuple[ProgramInvocation, ...]: + """The stimuli of every case, in case order, without their expected results.""" + return tuple(case.invocation for case in self.cases) + + def _parse_conditions(value: Any) -> Any: """ Rebuild discriminator-tagged conditions without losing subclass fields. @@ -263,8 +326,11 @@ class _ConditionTupleSchema: def __get_pydantic_json_schema__(self, core_schema: CoreSchema, handler: GetJsonSchemaHandler) -> JsonSchemaValue: """Return the tuple schema with a discriminated union of concrete conditions.""" schema = handler(core_schema) + # A condition that uses one nested model in several fields is emitted as a reference to its + # definition; resolve it so every entry describes its fields inline. condition_schemas = [ - handler(TypeAdapter(condition_class).core_schema) for _, condition_class in sorted(_CONDITION_TYPES.items()) + handler.resolve_ref_schema(handler(TypeAdapter(condition_class).core_schema)) + for _, condition_class in sorted(_CONDITION_TYPES.items()) ] for condition_schema in condition_schemas: required_fields = condition_schema.setdefault("required", []) diff --git a/pyrit/models/score/observation.py b/pyrit/models/score/observation.py index 857261fc69..0cde1a803e 100644 --- a/pyrit/models/score/observation.py +++ b/pyrit/models/score/observation.py @@ -16,11 +16,13 @@ ComponentIdentifier, # noqa: TC001 (runtime-required by Pydantic field annotations) ) from pyrit.models.literals import MEDIA_PATH_DATA_TYPES +from pyrit.models.score.program import ProgramExecutionEvent # noqa: TC001 (runtime-required by Pydantic) from pyrit.models.score.scorable import ( ContentEntryScorable, ContentScorable, ConversationScorable, MessageScorable, + ProgramExecutionScorable, ScorableUnion, # noqa: TC001 (runtime-required by Pydantic field annotations) SurfaceScorable, TraceScorable, @@ -419,11 +421,77 @@ def _validate_entries(self) -> SurfaceObservationPayload: return self +class ProgramExecutionObservationPayload(BaseModel): + """ + What one isolated execution attempt produced: one event per invocation that ran. + + ``program_sha256`` digests the exact code the sandbox received, so the evidence names the + program even if the message it came from is later edited. ``reasons`` explains why an + acquisition is not complete. A program that fails its invocations is still complete + evidence; only a sandbox that could not run it, or did not finish, leaves gaps. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + kind: Literal["program_execution"] = "program_execution" + schema_version: Literal[1] = 1 + program_sha256: str | None = Field(default=None, min_length=64, max_length=64, pattern=r"^[0-9a-f]{64}$") + events: tuple[ProgramExecutionEvent, ...] = () + reasons: tuple[str, ...] = () + + @field_validator("schema_version", mode="before") + @classmethod + def _validate_schema_version(cls, value: Any) -> Any: + """ + Require the exact supported schema version, without numeric coercion. + + Returns: + Any: The supported schema version. + + Raises: + ValueError: If the version is not exactly the supported integer. + """ + if type(value) is not int or value != 1: + raise ValueError("Unsupported program execution payload schema_version; expected 1.") + return value + + @model_validator(mode="after") + def _validate_events(self) -> ProgramExecutionObservationPayload: + """ + Keep one event per case, in case order. + + Returns: + ProgramExecutionObservationPayload: The validated payload. + + Raises: + ValueError: If a case repeats, events are out of order, or events lack a program digest. + """ + indices = [event.case_index for event in self.events] + if indices != sorted(set(indices)): + raise ValueError("Program execution events must name each case once, in case order.") + if self.events and self.program_sha256 is None: + raise ValueError("Program execution events require the digest of the program that ran.") + return self + + def event_for(self, case_index: int) -> ProgramExecutionEvent | None: + """ + Return the event for one case. + + Args: + case_index (int): The case's position in the condition. + + Returns: + ProgramExecutionEvent | None: The event, or None when that case did not run. + """ + return next((event for event in self.events if event.case_index == case_index), None) + + ObservationPayload = Annotated[ ScorerTargetResponsePayload | ToolEventsObservationPayload | ConversationObservationPayload - | SurfaceObservationPayload, + | SurfaceObservationPayload + | ProgramExecutionObservationPayload, Field(discriminator="kind"), ] @@ -528,6 +596,36 @@ def _validate_surface_acquisition(self) -> Observation: raise ValueError("Unavailable or failed surface acquisition cannot contain entries.") return self + @model_validator(mode="after") + def _validate_program_acquisition(self) -> Observation: + """ + Require acquisition status to agree with the execution events. + + Returns: + Observation: The validated observation. + + Raises: + ValueError: If the anchor, events, or acquisition are inconsistent. + """ + if not isinstance(self.payload, ProgramExecutionObservationPayload): + return self + if not isinstance(self.scorable, ProgramExecutionScorable): + raise ValueError("Program execution observations require a ProgramExecutionScorable.") + events = self.payload.events + if any(event.case_index >= len(self.scorable.invocations) for event in events): + raise ValueError("Program execution events must belong to the scorable's invocations.") + if any(event.runtime_profile != self.scorable.runtime_profile for event in events): + raise ValueError("Program execution events must use the scorable's runtime profile.") + if self.acquisition is Acquisition.COMPLETE and ( + len(events) != len(self.scorable.invocations) or self.payload.reasons + ): + raise ValueError("Complete program acquisition requires one event per invocation and no gaps.") + if self.acquisition is Acquisition.PARTIAL and len(events) == len(self.scorable.invocations): + raise ValueError("Partial program acquisition cannot hold an event for every invocation.") + if self.acquisition in (Acquisition.UNAVAILABLE, Acquisition.ERROR) and events: + raise ValueError("Unavailable or failed program acquisition cannot contain events.") + return self + @property def response_message_piece_ids(self) -> tuple[uuid.UUID, ...]: """The ordered message references retained by this payload.""" @@ -566,7 +664,10 @@ def validate_evidence( Raises: ValueError: If scored or response evidence is missing, modified, or unsupported. """ - if isinstance(self.payload, (ToolEventsObservationPayload, SurfaceObservationPayload)): + if isinstance( + self.payload, + (ToolEventsObservationPayload, SurfaceObservationPayload, ProgramExecutionObservationPayload), + ): return if isinstance(self.payload, ConversationObservationPayload): if not isinstance(self.scorable, ConversationScorable): diff --git a/pyrit/models/score/program.py b/pyrit/models/score/program.py new file mode 100644 index 0000000000..84248dd6f2 --- /dev/null +++ b/pyrit/models/score/program.py @@ -0,0 +1,78 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Backend-neutral evidence about how a program behaved when an isolated sandbox ran it.""" + +from __future__ import annotations + +from typing import Annotated + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +#: A trusted runtime profile name, such as ``python-3.12-restricted``. It names a sandbox +#: configuration the caller allowlisted; it is never an image name or a command. +RuntimeProfileName = Annotated[str, Field(min_length=1, max_length=64, pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*$")] + + +class ProgramInvocation(BaseModel): + """ + One bounded invocation of a program: its arguments and standard input. + + An invocation is a test stimulus. It carries no expected result, so a source that runs it + cannot see which output passes. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + arguments: tuple[Annotated[str, Field(pattern=r"^[^\x00]*$")], ...] = () + stdin: str | None = None + + +class ProgramExecutionEvent(BaseModel): + """ + What one invocation did when the sandbox ran it. + + The event records enough to audit the result and to judge it again. It does not make + execution reproducible: a program can read the clock, use randomness, or depend on its + runtime. Replay judges this stored evidence again without running the program. + + ``stdout`` and ``stderr`` hold at most the source's configured limit. Each stream has its + own truncation marker, so a matcher never treats a missing suffix as proof of absence. + ``exit_code`` is ``None`` when the program did not exit on its own, for example because the + sandbox stopped it at its time or resource limit. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + case_index: int = Field(ge=0) + runtime_profile: RuntimeProfileName + runtime_image_digest: str = Field(min_length=1) + resolved_command: tuple[str, ...] = Field(min_length=1) + exit_code: int | None + stdout: str = "" + stderr: str = "" + stdout_truncated: bool = False + stderr_truncated: bool = False + timed_out: bool = False + resource_limit_exceeded: bool = False + duration_ms: int = Field(ge=0) + + @model_validator(mode="after") + def _validate_exit_state(self) -> ProgramExecutionEvent: + """ + Keep the exit state consistent with how the run ended. + + Returns: + ProgramExecutionEvent: The validated event. + + Raises: + ValueError: If a run the sandbox stopped also reports an exit code of its own. + """ + if self.timed_out and self.exit_code is not None: + raise ValueError("A timed-out invocation did not exit on its own and has no exit code.") + return self + + @property + def output_truncated(self) -> bool: + """Whether either output stream was cut at the retention limit.""" + return self.stdout_truncated or self.stderr_truncated diff --git a/pyrit/models/score/scorable.py b/pyrit/models/score/scorable.py index e3de7f4474..cf744d4a6f 100644 --- a/pyrit/models/score/scorable.py +++ b/pyrit/models/score/scorable.py @@ -11,6 +11,10 @@ from pyrit.models.literals import PromptDataType # noqa: TC001 (runtime-required by Pydantic field annotations) from pyrit.models.score._trace_validation import TraceId # noqa: TC001 (runtime-required by Pydantic) +from pyrit.models.score.program import ( # noqa: TC001 (runtime-required by Pydantic) + ProgramInvocation, + RuntimeProfileName, +) from pyrit.models.score.surface import SurfaceMatch # noqa: TC001 (runtime-required by Pydantic) if TYPE_CHECKING: @@ -171,11 +175,54 @@ class SurfaceScorable(Scorable): match: SurfaceMatch = SurfaceMatch.EXACT +class ProgramExecutionScorable(Scorable): + """ + One isolated execution attempt of a program: the program, its runtime, and its invocations. + + This is an acquisition request, not a result: it holds no expected output, so the source + that runs it cannot see which output passes. ``program`` names the code to run, either a + stored message with one text piece or loose text, and is used verbatim; code is never + extracted from surrounding prose. + + ``execution_id`` identifies the attempt. Reusing a scorable reuses one observed execution, + and a caller that wants a fresh run mints a new id. + """ + + scorable_type: Literal["program_execution"] = "program_execution" + execution_id: uuid.UUID + program: Annotated[MessageScorable | ContentScorable, Field(discriminator="scorable_type")] + runtime_profile: RuntimeProfileName + invocations: tuple[ProgramInvocation, ...] + + @model_validator(mode="after") + def _validate_execution(self) -> ProgramExecutionScorable: + """ + Require text code and at least one invocation. + + Returns: + ProgramExecutionScorable: The validated scorable. + + Raises: + ValueError: If the program is not text, or there is nothing to run. + """ + if isinstance(self.program, ContentScorable) and self.program.data_type != "text": + raise ValueError("A program must be text content.") + if not self.invocations: + raise ValueError("A ProgramExecutionScorable requires at least one invocation.") + return self + + # Polymorphic union of scorables that can be stored on a Score. Every member declares a # ``scorable_type`` tag and Pydantic dispatches on it, so a new member is never mistaken for # an existing one and storage never depends on field shape. ScorableUnion = Annotated[ - MessageScorable | ContentScorable | ContentEntryScorable | ConversationScorable | TraceScorable | SurfaceScorable, + MessageScorable + | ContentScorable + | ContentEntryScorable + | ConversationScorable + | TraceScorable + | SurfaceScorable + | ProgramExecutionScorable, Field(discriminator="scorable_type"), ] diff --git a/pyrit/score/__init__.py b/pyrit/score/__init__.py index 2c6c0a47b4..121e615100 100644 --- a/pyrit/score/__init__.py +++ b/pyrit/score/__init__.py @@ -42,12 +42,15 @@ from pyrit.score.float_scale.video_float_scale_scorer import VideoFloatScaleScorer from pyrit.score.message_scorable_resolver import MessageScorableResolver from pyrit.score.message_scorer import MessageScorer + from pyrit.score.observation.docker_sandbox_client import DockerRuntimeProfile, DockerSandboxClient from pyrit.score.observation.execution import NonReplayableObservationError from pyrit.score.observation.lifecycle import AttemptObserver, SurfaceLifecycleSource from pyrit.score.observation.local_file_surface_source import LocalFileSurfaceSource from pyrit.score.observation.observation_source import ObservationSource from pyrit.score.observation.otel_span_exporter import InMemoryTraceExporter from pyrit.score.observation.otel_trace_source import OtelTraceSource + from pyrit.score.observation.program_execution_source import ProgramExecutionSource + from pyrit.score.observation.sandbox_client import SandboxClient, SandboxExecutionError, SandboxRun from pyrit.score.observation.trace_client import InMemoryTraceClient, TraceAcquisitionError, TraceClient from pyrit.score.response_handler import ( CallableResponseHandler, @@ -101,6 +104,7 @@ from pyrit.score.true_false.message_tool_call_scorer import MessageToolCallScorer from pyrit.score.true_false.otel_tool_call_scorer import OtelToolCallScorer from pyrit.score.true_false.output_matches_scorer import OutputMatchesScorer + from pyrit.score.true_false.program_execution_scorer import ProgramExecutionScorer from pyrit.score.true_false.prompt_shield_scorer import PromptShieldScorer from pyrit.score.true_false.question_answer_scorer import QuestionAnswerScorer from pyrit.score.true_false.regex.agent_threat_rules_scorer import AgentThreatRulesScorer @@ -161,6 +165,13 @@ from pyrit.score.true_false.wildguard_scorer import WildGuardScorer, render_wildguard_prompt _LAZY_EXPORTS: dict[str, str | tuple[str, str | None]] = { + "DockerRuntimeProfile": "pyrit.score.observation.docker_sandbox_client", + "DockerSandboxClient": "pyrit.score.observation.docker_sandbox_client", + "ProgramExecutionScorer": "pyrit.score.true_false.program_execution_scorer", + "ProgramExecutionSource": "pyrit.score.observation.program_execution_source", + "SandboxClient": "pyrit.score.observation.sandbox_client", + "SandboxExecutionError": "pyrit.score.observation.sandbox_client", + "SandboxRun": "pyrit.score.observation.sandbox_client", "OutputMatchesScorer": "pyrit.score.true_false.output_matches_scorer", "AnsiEscapeOutputScorer": "pyrit.score.true_false.regex.ansi_escape_output_scorer", "AnthraxKeywordScorer": "pyrit.score.true_false.regex.anthrax_keyword_scorer", diff --git a/pyrit/score/observation/__init__.py b/pyrit/score/observation/__init__.py index bbb4b94186..9d9ac31b08 100644 --- a/pyrit/score/observation/__init__.py +++ b/pyrit/score/observation/__init__.py @@ -10,15 +10,24 @@ if TYPE_CHECKING: from pyrit.score.observation.conversation_source import ConversationSource + from pyrit.score.observation.docker_sandbox_client import DockerRuntimeProfile, DockerSandboxClient from pyrit.score.observation.execution import NonReplayableObservationError from pyrit.score.observation.lifecycle import AttemptObserver, SurfaceLifecycleSource from pyrit.score.observation.local_file_surface_source import LocalFileSurfaceSource from pyrit.score.observation.observation_source import ObservationSource from pyrit.score.observation.otel_span_exporter import InMemoryTraceExporter from pyrit.score.observation.otel_trace_source import OtelTraceSource + from pyrit.score.observation.program_execution_source import ProgramExecutionSource + from pyrit.score.observation.sandbox_client import SandboxClient, SandboxExecutionError, SandboxRun from pyrit.score.observation.trace_client import InMemoryTraceClient, TraceAcquisitionError, TraceClient _LAZY_EXPORTS: dict[str, str] = { + "DockerRuntimeProfile": "pyrit.score.observation.docker_sandbox_client", + "DockerSandboxClient": "pyrit.score.observation.docker_sandbox_client", + "ProgramExecutionSource": "pyrit.score.observation.program_execution_source", + "SandboxClient": "pyrit.score.observation.sandbox_client", + "SandboxExecutionError": "pyrit.score.observation.sandbox_client", + "SandboxRun": "pyrit.score.observation.sandbox_client", "AttemptObserver": "pyrit.score.observation.lifecycle", "ConversationSource": "pyrit.score.observation.conversation_source", "InMemoryTraceClient": "pyrit.score.observation.trace_client", diff --git a/pyrit/score/observation/docker_sandbox_client.py b/pyrit/score/observation/docker_sandbox_client.py new file mode 100644 index 0000000000..fa16050f63 --- /dev/null +++ b/pyrit/score/observation/docker_sandbox_client.py @@ -0,0 +1,416 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Run programs in disposable, locked-down Docker containers.""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +import os +import re +import shutil +import tempfile +import time +import uuid +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +from pyrit.models import ComponentIdentifier, ProgramExecutionEvent +from pyrit.score.observation.sandbox_client import SandboxExecutionError, SandboxRun + +if TYPE_CHECKING: + from collections.abc import Mapping + + from pyrit.models import ProgramInvocation + +logger = logging.getLogger(__name__) + +_DIGEST_PATTERN = re.compile(r"@(sha256:[0-9a-f]{64})$") +_NEVER_STARTED = "0001-01-01T00:00:00Z" + + +@dataclass(frozen=True) +class DockerRuntimeProfile: + """ + One allowlisted runtime: a digest-pinned image and the command that runs a program in it. + + The container runs ``command``, then the program's path, then the invocation's arguments. + For example ``command=("python3", "-I")`` runs ``python3 -I /program/main.py ``. + ``image`` must be pinned by digest, so the profile always names the same runtime. + """ + + image: str + command: tuple[str, ...] + program_filename: str = "main.py" + + def __post_init__(self) -> None: + """ + Validate the image pin and the program file name. + + Raises: + ValueError: If the image is not pinned by digest, the command is empty, or the + file name is not a plain name. + """ + if _DIGEST_PATTERN.search(self.image) is None: + raise ValueError(f"Runtime image {self.image!r} must be pinned by digest (name@sha256:<64 hex>).") + if not self.command or any(not part for part in self.command): + raise ValueError("A runtime profile command must be a nonempty argument list.") + if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]*", self.program_filename): + raise ValueError(f"Program file name {self.program_filename!r} must be a plain file name.") + + @property + def image_digest(self) -> str: + """The ``sha256:`` digest the image is pinned to; ``__post_init__`` guarantees its presence.""" + return self.image.rsplit("@", 1)[1] + + +@dataclass(frozen=True) +class _ContainerExit: + exit_code: int | None + oom_killed: bool + + +class DockerSandboxClient: + """ + Run each invocation in its own disposable Docker container. + + Every container has no network, a read-only root file system, all capabilities dropped, + no privilege escalation, an unprivileged user, and caps on memory, CPU, and process count. + The program is mounted read-only; a small ``/tmp`` is the only writable location. Each + invocation gets a fresh container that is removed afterwards, so one run cannot leave state + for the next. + + Images are never pulled while scoring: a profile's image must already be present, or the + attempt fails as a sandbox error. Pull the pinned images before a run. + + A container that runs past ``timeout_seconds`` is killed and reported as timed out. One that + the kernel stops for exceeding its memory limit is reported as exceeding a resource limit. + Each stream keeps at most ``max_output_bytes`` bytes. The timeout includes container start-up. + + This client isolates submitted code with Docker's defaults plus the restrictions above. It + is not a hardened multi-tenant sandbox; run it on a host dedicated to evaluation. + """ + + _READ_CHUNK_BYTES = 1 << 16 + _CONTROL_TIMEOUT_SECONDS = 30.0 + + def __init__( + self, + *, + profiles: Mapping[str, DockerRuntimeProfile], + docker_executable: str = "docker", + timeout_seconds: float = 10.0, + memory: str = "256m", + cpus: float = 1.0, + pids_limit: int = 64, + max_output_bytes: int = 64_000, + ) -> None: + """ + Initialize with the allowlisted runtime profiles and per-container limits. + + Args: + profiles (Mapping[str, DockerRuntimeProfile]): Runtime profile names and their runtimes. + docker_executable (str): The Docker command line client to call. + timeout_seconds (float): Wall-clock limit for each invocation. + memory (str): Memory limit for each container, in Docker's notation, such as ``256m``. + cpus (float): CPU limit for each container. + pids_limit (int): Most processes a container may hold. + max_output_bytes (int): Most bytes of each stream kept per invocation. + + Raises: + ValueError: If no profile is given, a profile name is invalid, or a limit is not positive. + """ + if not profiles: + raise ValueError("DockerSandboxClient requires at least one runtime profile.") + invalid = [name for name in profiles if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]{0,63}", name)] + if invalid: + raise ValueError(f"Invalid runtime profile names: {invalid}.") + if timeout_seconds <= 0 or cpus <= 0 or pids_limit < 1 or max_output_bytes < 1: + raise ValueError("timeout_seconds, cpus, pids_limit and max_output_bytes must be positive.") + if not re.fullmatch(r"[1-9][0-9]*[bkmg]?", memory): + raise ValueError(f"Memory limit {memory!r} must look like 256m or 1g.") + self._profiles = dict(profiles) + self._docker = docker_executable + self._timeout_seconds = timeout_seconds + self._memory = memory + self._cpus = cpus + self._pids_limit = pids_limit + self._max_output_bytes = max_output_bytes + + def get_identifier(self) -> ComponentIdentifier: + """ + Return the profiles and limits. + + Returns: + ComponentIdentifier: The nonsecret client configuration. + """ + return ComponentIdentifier.of( + self, + params={ + "client_version": 1, + "profiles": { + name: { + "image": profile.image, + "command": list(profile.command), + "program_filename": profile.program_filename, + } + for name, profile in sorted(self._profiles.items()) + }, + "timeout_seconds": self._timeout_seconds, + "memory": self._memory, + "cpus": self._cpus, + "pids_limit": self._pids_limit, + "max_output_bytes": self._max_output_bytes, + }, + ) + + def get_runtime_profiles(self) -> frozenset[str]: + """ + Return the configured runtime profile names. + + Returns: + frozenset[str]: The profile names this client can run. + """ + return frozenset(self._profiles) + + async def execute_async( + self, *, program: str, runtime_profile: str, invocations: tuple[ProgramInvocation, ...] + ) -> SandboxRun: + """ + Run each invocation in a fresh container, in order. + + Args: + program (str): The program text, written to a read-only file in the container. + runtime_profile (str): The allowlisted profile to run under. + invocations (tuple[ProgramInvocation, ...]): The arguments and input of each run. + + Returns: + SandboxRun: One event per invocation that ran. If the sandbox fails after some + invocations ran, the run holds those events and the reason it stopped. + + Raises: + ValueError: If the runtime profile is not configured. + SandboxExecutionError: If no invocation could be run. + """ + profile = self._profiles.get(runtime_profile) + if profile is None: + raise ValueError(f"Runtime profile {runtime_profile!r} is not configured.") + program_dir = Path(tempfile.mkdtemp(prefix="pyrit-program-")) + try: + program_path = program_dir / profile.program_filename + program_path.write_text(program, encoding="utf-8") + program_path.chmod(0o444) + program_dir.chmod(0o755) + events: list[ProgramExecutionEvent] = [] + for index, invocation in enumerate(invocations): + try: + events.append( + await self._run_invocation_async( + case_index=index, + runtime_profile=runtime_profile, + profile=profile, + program_dir=program_dir, + invocation=invocation, + ) + ) + except SandboxExecutionError as error: + if not events: + raise + return SandboxRun(events=tuple(events), reasons=(f"sandbox error at case {index}: {error}",)) + return SandboxRun(events=tuple(events)) + finally: + shutil.rmtree(program_dir, ignore_errors=True) + + async def _run_invocation_async( + self, + *, + case_index: int, + runtime_profile: str, + profile: DockerRuntimeProfile, + program_dir: Path, + invocation: ProgramInvocation, + ) -> ProgramExecutionEvent: + """ + Run one invocation in a new container and remove the container afterwards. + + Returns: + ProgramExecutionEvent: What the invocation did. + + Raises: + SandboxExecutionError: If Docker cannot start or inspect the container. + """ + name = f"pyrit-sandbox-{uuid.uuid4().hex}" + resolved_command = (*profile.command, f"/program/{profile.program_filename}", *invocation.arguments) + argv = [ + self._docker, + "run", + "--name", + name, + "--interactive", + "--pull", + "never", + "--network", + "none", + "--read-only", + "--cap-drop", + "ALL", + "--security-opt", + "no-new-privileges", + "--user", + "65534:65534", + "--memory", + self._memory, + "--memory-swap", + self._memory, + "--cpus", + str(self._cpus), + "--pids-limit", + str(self._pids_limit), + "--tmpfs", + "/tmp:rw,nosuid,nodev,size=16m", + "--workdir", + "/tmp", + "--mount", + f"type=bind,source={os.fspath(program_dir)},target=/program,readonly", + profile.image, + *resolved_command, + ] + started = time.monotonic() + try: + process = await asyncio.create_subprocess_exec( + *argv, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + except OSError as error: + raise SandboxExecutionError(f"could not start {self._docker!r}: {error}") from error + timed_out = False + try: + readers = asyncio.gather( + self._read_capped_async(process.stdout), + self._read_capped_async(process.stderr), + self._feed_stdin_async(process.stdin, invocation.stdin), + ) + try: + (stdout, stdout_truncated), (stderr, stderr_truncated), _ = await asyncio.wait_for( + asyncio.shield(readers), timeout=self._timeout_seconds + ) + except TimeoutError: + timed_out = True + await self._control_async("kill", name) + try: + await asyncio.wait_for(asyncio.shield(readers), timeout=self._CONTROL_TIMEOUT_SECONDS) + except TimeoutError: + # The container would not stop; ending the attached client closes its streams. + with contextlib.suppress(ProcessLookupError): + process.kill() + (stdout, stdout_truncated), (stderr, stderr_truncated), _ = await readers + await process.wait() + duration_ms = int((time.monotonic() - started) * 1000) + container = await self._inspect_async(name=name, cli_stderr=stderr) + finally: + if process.returncode is None: + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() + await self._control_async("rm", "--force", name) + return ProgramExecutionEvent( + case_index=case_index, + runtime_profile=runtime_profile, + runtime_image_digest=profile.image_digest, + resolved_command=resolved_command, + exit_code=None if (timed_out or container.oom_killed) else container.exit_code, + stdout=stdout.decode("utf-8", errors="replace"), + stderr=stderr.decode("utf-8", errors="replace"), + stdout_truncated=stdout_truncated, + stderr_truncated=stderr_truncated, + timed_out=timed_out, + resource_limit_exceeded=container.oom_killed and not timed_out, + duration_ms=duration_ms, + ) + + async def _inspect_async(self, *, name: str, cli_stderr: bytes) -> _ContainerExit: + """ + Read how the container ended, distinguishing a program exit from a Docker failure. + + Returns: + _ContainerExit: The container's exit code and whether the kernel killed it for memory. + + Raises: + SandboxExecutionError: If the container does not exist or never started. + """ + returncode, output = await self._control_async( + "inspect", "--format", "{{.State.ExitCode}} {{.State.OOMKilled}} {{.State.StartedAt}}", name + ) + fields = output.split() + if returncode != 0 or len(fields) != 3 or fields[2] == _NEVER_STARTED: + lines = [line.strip() for line in cli_stderr.decode("utf-8", errors="replace").splitlines() if line.strip()] + detail = lines[0][:500] if lines else output.strip()[:500] + raise SandboxExecutionError(f"the container did not start: {detail or 'unknown Docker failure'}") + return _ContainerExit(exit_code=int(fields[0]), oom_killed=fields[1] == "true") + + async def _control_async(self, *arguments: str) -> tuple[int, str]: + """ + Run a short Docker management command, bounded in time. + + Returns: + tuple[int, str]: The exit status and combined output; status -1 when it failed to run. + """ + try: + process = await asyncio.create_subprocess_exec( + self._docker, + *arguments, + stdin=asyncio.subprocess.DEVNULL, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + ) + except OSError as error: + logger.warning("Docker %s failed to start: %s", arguments[0], error) + return -1, str(error) + try: + output, _ = await asyncio.wait_for(process.communicate(), timeout=self._CONTROL_TIMEOUT_SECONDS) + except TimeoutError: + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() + logger.warning("Docker %s timed out", arguments[0]) + return -1, "timed out" + return process.returncode if process.returncode is not None else -1, output.decode("utf-8", errors="replace") + + async def _read_capped_async(self, stream: asyncio.StreamReader | None) -> tuple[bytes, bool]: + """ + Read a stream to its end, keeping at most the configured number of bytes. + + Returns: + tuple[bytes, bool]: The retained bytes and whether more were discarded. + """ + if stream is None: + return b"", False + kept = bytearray() + truncated = False + while chunk := await stream.read(self._READ_CHUNK_BYTES): + room = self._max_output_bytes - len(kept) + if len(chunk) > room: + truncated = True + kept.extend(chunk[: max(room, 0)]) + return bytes(kept), truncated + + @staticmethod + async def _feed_stdin_async(stream: asyncio.StreamWriter | None, data: str | None) -> None: + """Write the invocation's input, then close standard input so the program sees its end.""" + if stream is None: + return + try: + if data: + stream.write(data.encode("utf-8")) + await stream.drain() + except (BrokenPipeError, ConnectionResetError): + pass + finally: + stream.close() + with contextlib.suppress(BrokenPipeError, ConnectionResetError): + await stream.wait_closed() diff --git a/pyrit/score/observation/execution.py b/pyrit/score/observation/execution.py index c7fd663ad3..b14b81ad97 100644 --- a/pyrit/score/observation/execution.py +++ b/pyrit/score/observation/execution.py @@ -15,6 +15,7 @@ MessagePiece, MessageScorable, Observation, + ProgramExecutionObservationPayload, Scorable, ScorableUnion, Score, @@ -38,7 +39,11 @@ class NonReplayableObservationError(ValueError): _ObservationEvidence: TypeAlias = ( - Message | ToolEventsObservationPayload | SurfaceObservationPayload | tuple[MessagePiece, ...] + Message + | ToolEventsObservationPayload + | SurfaceObservationPayload + | ProgramExecutionObservationPayload + | tuple[MessagePiece, ...] ) @@ -315,7 +320,9 @@ async def resolve_async(self, *, observation: Observation) -> _ObservationEviden NonReplayableObservationError: If referenced evidence is missing, modified, or unsupported. """ payload = observation.payload - if isinstance(payload, (ToolEventsObservationPayload, SurfaceObservationPayload)): + if isinstance( + payload, (ToolEventsObservationPayload, SurfaceObservationPayload, ProgramExecutionObservationPayload) + ): return payload pieces = await self._memory.get_message_pieces_async(prompt_ids=list(observation.evidence_message_piece_ids)) pieces_by_id = {piece.id: piece for piece in pieces} diff --git a/pyrit/score/observation/program_execution_source.py b/pyrit/score/observation/program_execution_source.py new file mode 100644 index 0000000000..4ea3911560 --- /dev/null +++ b/pyrit/score/observation/program_execution_source.py @@ -0,0 +1,222 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Acquire how a program behaves by asking an injected sandbox to run it.""" + +from __future__ import annotations + +import hashlib +import logging +from collections import OrderedDict +from typing import TYPE_CHECKING + +from pyrit.memory import CentralMemory +from pyrit.models import ( + Acquisition, + ComponentIdentifier, + ContentScorable, + MessageScorable, + Observation, + ProgramExecutionEvent, + ProgramExecutionObservationPayload, +) +from pyrit.score.message_scorable_resolver import MessageScorableResolver +from pyrit.score.observation.sandbox_client import SandboxExecutionError + +if TYPE_CHECKING: + from pyrit.models import ProgramExecutionScorable + from pyrit.score.observation.sandbox_client import SandboxClient + +logger = logging.getLogger(__name__) + + +class _NoProgramError(Exception): + """The named evidence holds nothing that can be submitted as a program.""" + + +class ProgramExecutionSource: + """ + Run the program a ``ProgramExecutionScorable`` names, once per invocation, in a sandbox. + + The source receives only the derived scorable: the program, the runtime profile, and the + invocations. It never sees expected results, so it can run the program but cannot tell + which output passes. An injected ``SandboxClient`` owns the disposable environment; the + source only submits the attempt and waits for it. + + The program is the named text, used verbatim. A message must name exactly one text piece; + code is never extracted from surrounding prose. A message that holds a blocked or failed + response has no program to run, so the evidence is unavailable. + + Program failure is not acquisition failure. A compile error, a nonzero exit, wrong output, + a timeout, or an enforced resource limit is complete evidence that the program misbehaved. + Only a sandbox that cannot run the program makes the acquisition an error, and only a + sandbox that stops early leaves it partial, so broken infrastructure never looks like a + program that failed its tests. + + Each stream keeps at most ``max_output_chars`` characters and records when it was cut. + A completed or partial observation is reused when the same scorable is acquired again, so + one execution attempt is observed once; a new ``execution_id`` runs the program again. + """ + + _MAX_REUSED_OBSERVATIONS = 256 + + def __init__(self, *, sandbox: SandboxClient, max_output_chars: int = 64_000) -> None: + """ + Initialize with the sandbox that runs submitted programs. + + Args: + sandbox (SandboxClient): Runs programs under its allowlisted runtime profiles. + max_output_chars (int): The most characters of each stream retained per invocation. + + Raises: + ValueError: If the output limit is not positive. + """ + if max_output_chars < 1: + raise ValueError("max_output_chars must be positive.") + self._sandbox = sandbox + self._max_output_chars = max_output_chars + self._observed: OrderedDict[ProgramExecutionScorable, Observation] = OrderedDict() + + def get_identifier(self) -> ComponentIdentifier: + """ + Return the acquisition version, output limit, and sandbox configuration. + + Returns: + ComponentIdentifier: The nonsecret source configuration. + """ + return ComponentIdentifier.of( + self, + params={"acquisition_version": 1, "max_output_chars": self._max_output_chars}, + children={"sandbox": self._sandbox.get_identifier()}, + ) + + async def acquire_async(self, *, scorable: ProgramExecutionScorable) -> Observation: + """ + Run every invocation once and return what the program did. + + Args: + scorable (ProgramExecutionScorable): The program, runtime profile, and invocations. + + Returns: + Observation: The execution evidence, with its acquisition state. + + Raises: + ValueError: If the runtime profile is not one the sandbox offers, the sandbox + returns events for invocations it was not given, or the program reference + names missing or several pieces. + """ + reused = self._observed.get(scorable) + if reused is not None: + self._observed.move_to_end(scorable) + return reused + if scorable.runtime_profile not in self._sandbox.get_runtime_profiles(): + raise ValueError( + f"Runtime profile {scorable.runtime_profile!r} is not offered by the sandbox; " + f"available profiles: {sorted(self._sandbox.get_runtime_profiles())}." + ) + try: + program = await self._resolve_program_async(scorable=scorable) + except _NoProgramError as error: + return self._observation(scorable=scorable, acquisition=Acquisition.UNAVAILABLE, reasons=(str(error),)) + try: + run = await self._sandbox.execute_async( + program=program, + runtime_profile=scorable.runtime_profile, + invocations=scorable.invocations, + ) + except SandboxExecutionError as error: + logger.warning("Sandbox could not run program for execution %s: %s", scorable.execution_id, error) + return self._observation( + scorable=scorable, acquisition=Acquisition.ERROR, reasons=(f"sandbox error: {error}",) + ) + events = tuple(self._bounded(event) for event in run.events) + if any(event.case_index >= len(scorable.invocations) for event in events): + raise ValueError("The sandbox returned an event for an invocation it was not given.") + program_sha256 = hashlib.sha256(program.encode("utf-8")).hexdigest() + if len(events) == len(scorable.invocations): + observation = self._observation( + scorable=scorable, acquisition=Acquisition.COMPLETE, events=events, program_sha256=program_sha256 + ) + else: + reasons = run.reasons or ("the sandbox stopped before running every invocation",) + acquisition = Acquisition.PARTIAL if events else Acquisition.ERROR + observation = self._observation( + scorable=scorable, + acquisition=acquisition, + events=events, + program_sha256=program_sha256 if events else None, + reasons=reasons, + ) + if observation.acquisition in (Acquisition.COMPLETE, Acquisition.PARTIAL): + self._remember(scorable=scorable, observation=observation) + return observation + + async def _resolve_program_async(self, *, scorable: ProgramExecutionScorable) -> str: + """ + Read the program text the scorable names. + + Returns: + str: The code to submit, verbatim. + + Raises: + _NoProgramError: If the named message holds a blocked or failed response. + ValueError: If the reference cannot be resolved to one text piece. + """ + reference = scorable.program + if isinstance(reference, ContentScorable): + return reference.value + if not isinstance(reference, MessageScorable) or len(reference.message_piece_ids) != 1: + raise ValueError("A program message reference must name exactly one message piece.") + message = await MessageScorableResolver().resolve_async( + scorable=reference, memory=CentralMemory.get_memory_instance() + ) + piece = message.get_piece() + if piece.has_error(): + raise _NoProgramError(f"the program message holds a {piece.response_error} response, not a program") + if piece.converted_value_data_type != "text": + raise ValueError(f"A program must be text, not {piece.converted_value_data_type}.") + return piece.converted_value + + def _bounded(self, event: ProgramExecutionEvent) -> ProgramExecutionEvent: + """ + Cut each stream to the retention limit, keeping any truncation the sandbox reported. + + Returns: + ProgramExecutionEvent: The event as it is retained. + """ + limit = self._max_output_chars + if len(event.stdout) <= limit and len(event.stderr) <= limit: + return event + return event.model_copy( + update={ + "stdout": event.stdout[:limit], + "stderr": event.stderr[:limit], + "stdout_truncated": event.stdout_truncated or len(event.stdout) > limit, + "stderr_truncated": event.stderr_truncated or len(event.stderr) > limit, + } + ) + + def _observation( + self, + *, + scorable: ProgramExecutionScorable, + acquisition: Acquisition, + events: tuple[ProgramExecutionEvent, ...] = (), + program_sha256: str | None = None, + reasons: tuple[str, ...] = (), + ) -> Observation: + return Observation( + source_identifier=self.get_identifier(), + acquisition=acquisition, + scorable=scorable, + payload=ProgramExecutionObservationPayload( + program_sha256=program_sha256, + events=events, + reasons=reasons, + ), + ) + + def _remember(self, *, scorable: ProgramExecutionScorable, observation: Observation) -> None: + self._observed[scorable] = observation + while len(self._observed) > self._MAX_REUSED_OBSERVATIONS: + self._observed.popitem(last=False) diff --git a/pyrit/score/observation/sandbox_client.py b/pyrit/score/observation/sandbox_client.py new file mode 100644 index 0000000000..c2a5a16b8f --- /dev/null +++ b/pyrit/score/observation/sandbox_client.py @@ -0,0 +1,67 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""The injected service that runs a program in a disposable, isolated environment.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +if TYPE_CHECKING: + from pyrit.models import ComponentIdentifier, ProgramExecutionEvent, ProgramInvocation + + +class SandboxExecutionError(Exception): + """ + The sandbox could not run the program. + + This is an infrastructure failure, such as an unreachable service, a failed + authentication, or an environment that could not be provisioned. It is never raised for a + program that fails, times out, or exceeds a limit: those are observed behavior. + """ + + +@dataclass(frozen=True) +class SandboxRun: + """ + The invocations a sandbox ran for one execution attempt. + + ``events`` holds one event for each invocation that ran, in case order. When the sandbox + stopped before running every invocation, ``reasons`` says why, so the missing cases stay + undetermined instead of failing. + """ + + events: tuple[ProgramExecutionEvent, ...] = () + reasons: tuple[str, ...] = () + + +@runtime_checkable +class SandboxClient(Protocol): + """ + Run a program under an allowlisted runtime profile in a disposable environment. + + The client owns the environment: it creates it for the attempt and destroys it afterwards. + It never retains an environment between attempts, builds runtime images, or starts a daemon. + A runtime profile is a trusted name the client maps to its own configuration; the program + cannot choose an image or a command. + """ + + def get_identifier(self) -> ComponentIdentifier: + """Return the nonsecret client configuration.""" + ... + + def get_runtime_profiles(self) -> frozenset[str]: + """Return the runtime profile names this client is configured to run.""" + ... + + async def execute_async( + self, *, program: str, runtime_profile: str, invocations: tuple[ProgramInvocation, ...] + ) -> SandboxRun: + """ + Run each invocation of the program once, in order, in a fresh isolated environment. + + Raises: + SandboxExecutionError: If the sandbox could not run the program at all. + """ + ... diff --git a/pyrit/score/true_false/program_execution_scorer.py b/pyrit/score/true_false/program_execution_scorer.py new file mode 100644 index 0000000000..009414c803 --- /dev/null +++ b/pyrit/score/true_false/program_execution_scorer.py @@ -0,0 +1,293 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Score whether a program behaves as expected, over evidence from an isolated execution.""" + +from __future__ import annotations + +import uuid +from typing import TYPE_CHECKING + +from pyrit.models import ( + Acquisition, + Contains, + ContentScorable, + MessageScorable, + ProgramBehaves, + ProgramExecutionObservationPayload, + ProgramExecutionScorable, + Score, + ScoreStatus, +) +from pyrit.score.observation.execution import NonReplayableObservationError, _collect_observation +from pyrit.score.text_matching import match_text +from pyrit.score.true_false.true_false_scorer import TrueFalseScorer + +if TYPE_CHECKING: + from pyrit.models import ( + ComponentIdentifier, + Observation, + ProgramExecutionEvent, + ProgramResult, + Scorable, + ScoringExpectation, + TextMatcher, + ) + from pyrit.score.observation.execution import _ObservationEvidence + from pyrit.score.observation.observation_source import ObservationSource + + +def _stream_matches(*, matcher: TextMatcher, text: str, truncated: bool) -> bool | None: + """ + Match one output stream, treating a cut stream as unknown unless a prefix proves a match. + + Returns: + bool | None: Whether the stream matches, or None when the retained text cannot tell. + """ + if truncated and not isinstance(matcher, Contains): + return None + if match_text(matcher=matcher, text=text): + return True + return None if truncated else False + + +def match_case(*, expected: ProgramResult, event: ProgramExecutionEvent) -> bool | None: + """ + Decide whether one invocation produced its expected result. + + A run the sandbox stopped at its time or resource limit fails, because the program did not + finish. Otherwise the exit code and each stream that has a matcher must all agree. + + Args: + expected (ProgramResult): The expected result. + event (ProgramExecutionEvent): What the invocation did. + + Returns: + bool | None: True when every check passes, False when one fails, None when a cut + stream leaves a check open. + """ + if event.timed_out or event.resource_limit_exceeded: + return False + if event.exit_code is None: + return None + if expected.exit_code is not None and event.exit_code != expected.exit_code: + return False + checks: list[bool | None] = [] + if expected.stdout is not None: + checks.append(_stream_matches(matcher=expected.stdout, text=event.stdout, truncated=event.stdout_truncated)) + if expected.stderr is not None: + checks.append(_stream_matches(matcher=expected.stderr, text=event.stderr, truncated=event.stderr_truncated)) + if False in checks: + return False + return None if None in checks else True + + +def match_program_behaves( + *, condition: ProgramBehaves, payload: ProgramExecutionObservationPayload, acquisition: Acquisition +) -> bool | None: + """ + Match every case against the execution evidence, without I/O. + + One failing case is enough for a false verdict, even when other cases did not run. A true + verdict needs every case to have run and passed. + + Args: + condition (ProgramBehaves): The cases and their expected results. + payload (ProgramExecutionObservationPayload): The acquired execution evidence. + acquisition (Acquisition): Whether every invocation ran. + + Returns: + bool | None: True when every case passes, False when one fails, otherwise None. + """ + if acquisition in (Acquisition.ERROR, Acquisition.UNAVAILABLE): + return None + unknown = False + for index, case in enumerate(condition.cases): + event = payload.event_for(index) + if event is None: + unknown = True + continue + result = match_case(expected=case.expected, event=event) + if result is False: + return False + unknown = unknown or result is None + return None if unknown else True + + +class ProgramExecutionScorer(TrueFalseScorer): + """ + Score whether a program behaves as expected, judged from what an isolated sandbox observed. + + The ``ProgramBehaves`` condition carries each case's stimulus and expected result. The scorer + splits them: it builds a ``ProgramExecutionScorable`` from the program it is given, the + runtime profile and the invocations, and sends only that to its source. Each expected result + stays with the scorer, so the code that runs the program never learns which output passes. + + The program is the scored content, run verbatim: loose text, or a stored message with one + text piece. Code is not extracted from surrounding prose, so a response that wraps code in + explanation fails its cases rather than being guessed at. + + The verdict is true when every case passes and false when any case fails. A compile error, + nonzero exit, wrong output, timeout, or resource limit is a failure the sandbox observed. A + sandbox that could not run the program leaves the verdict undetermined, as does a case that + did not run or a check that cut output cannot settle. + + Every scoring call is a new execution attempt. Scoring an explicit ``ProgramExecutionScorable`` + reuses that attempt's evidence when the source has already observed it, and the stored + observation can be judged again without running the program. + """ + + CONDITION_TYPE = ProgramBehaves + + def __init__(self, *, source: ObservationSource[ProgramExecutionScorable]) -> None: + """ + Initialize with a condition-independent, caller-configured execution source. + + Args: + source (ObservationSource[ProgramExecutionScorable]): Runs the programs scorables name. + """ + super().__init__() + self._source = source + + def _build_identifier(self) -> ComponentIdentifier: + return self._create_identifier( + params={"matching_version": 1}, + children={"source": self._source.get_identifier()}, + ) + + async def _score_scorable_async(self, *, scorable: Scorable, expectation: ScoringExpectation | None) -> list[Score]: + condition = self._get_required_condition(expectation=expectation, condition_type=ProgramBehaves) + if isinstance(scorable, ProgramExecutionScorable): + if (scorable.runtime_profile, scorable.invocations) != (condition.runtime_profile, condition.invocations): + raise ValueError( + "A ProgramExecutionScorable must run the runtime profile and invocations of the " + "ProgramBehaves condition." + ) + execution = scorable + elif isinstance(scorable, (MessageScorable, ContentScorable)): + execution = ProgramExecutionScorable( + execution_id=uuid.uuid4(), + program=scorable, + runtime_profile=condition.runtime_profile, + invocations=condition.invocations, + ) + else: + raise TypeError( + "ProgramExecutionScorer requires a MessageScorable, a ContentScorable, or a ProgramExecutionScorable." + ) + + observation = await self._source.acquire_async(scorable=execution) + if observation.scorable != execution: + raise ValueError("Program execution source changed the caller's evidence anchor.") + if not isinstance(observation.payload, ProgramExecutionObservationPayload): + raise ValueError("Program execution source returned incompatible evidence.") + # A reused attempt's evidence may already be stored; a new score then only links to it. + if execution is not scorable or not await self._is_stored_async(observation=observation): + _collect_observation(observation) + scores = self._score_observation(observation=observation, evidence=observation.payload, expectation=expectation) + for score in scores: + score.scorable = scorable + score.message_piece_id = self._piece_id_from_scorable(scorable) + return scores + + async def _is_stored_async(self, *, observation: Observation) -> bool: + """ + Check whether an observation the source reused is already stored unchanged. + + Returns: + bool: True when memory holds the same observation. + + Raises: + ValueError: If memory holds a different observation under the same id. + """ + stored = await self._memory.get_observations_async(observation_ids=[observation.id]) + if not stored: + return False + if stored[0] != observation: + raise ValueError(f"Observation {observation.id} does not match its stored evidence.") + return True + + def _score_observation( + self, + *, + observation: Observation, + evidence: _ObservationEvidence, + expectation: ScoringExpectation | None, + ) -> list[Score]: + if not isinstance(evidence, ProgramExecutionObservationPayload) or not isinstance( + observation.scorable, ProgramExecutionScorable + ): + raise NonReplayableObservationError("Program behavior scoring requires a stored program execution.") + condition = self._get_required_condition(expectation=expectation, condition_type=ProgramBehaves) + execution = observation.scorable + if (execution.runtime_profile, execution.invocations) != (condition.runtime_profile, condition.invocations): + raise NonReplayableObservationError( + "The stored execution ran different invocations or a different runtime profile than the " + "ProgramBehaves condition." + ) + value = match_program_behaves(condition=condition, payload=evidence, acquisition=observation.acquisition) + return [ + Score( + score_value=None if value is None else str(value).lower(), + status=ScoreStatus.UNDETERMINED if value is None else ScoreStatus.COMPLETE, + score_type="true_false", + score_rationale=self._rationale( + condition=condition, payload=evidence, acquisition=observation.acquisition, value=value + ), + score_value_description="The program produced the expected result for every case.", + scorer_class_identifier=self.get_identifier(), + scorable=observation.scorable, + message_piece_id=self._piece_id_from_scorable(observation.scorable), + observation_ids=[observation.id], + ) + ] + + @staticmethod + def _rationale( + *, + condition: ProgramBehaves, + payload: ProgramExecutionObservationPayload, + acquisition: Acquisition, + value: bool | None, + ) -> str: + """ + Explain the verdict in terms of the cases. + + Returns: + str: The rationale. + """ + total = len(condition.cases) + profile = condition.runtime_profile + if value is True: + return f"All {total} case(s) produced the expected result under {profile}." + if acquisition in (Acquisition.ERROR, Acquisition.UNAVAILABLE): + gaps = ", ".join(payload.reasons) or "no execution evidence" + return f"The program's behavior under {profile} could not be observed ({gaps})." + outcomes = [] + for index, case in enumerate(condition.cases): + event = payload.event_for(index) + if event is None: + outcomes.append(f"case {index} did not run") + continue + result = match_case(expected=case.expected, event=event) + if result is True: + continue + if event.timed_out: + detail = "timed out" + elif event.resource_limit_exceeded: + detail = "exceeded a resource limit" + elif event.exit_code is None: + detail = "did not report an exit code" + elif case.expected.exit_code is not None and event.exit_code != case.expected.exit_code: + detail = f"exited with {event.exit_code}, expected {case.expected.exit_code}" + elif result is False: + detail = "output did not match" + else: + detail = "output was cut before the check could be settled" + outcomes.append(f"case {index} {detail}") + summary = "; ".join(outcomes) + if value is False: + return f"The program did not produce the expected result under {profile}: {summary}." + gaps = ", ".join(payload.reasons) + suffix = f" ({gaps})" if gaps else "" + return f"The program's behavior under {profile} could not be fully established: {summary}{suffix}." diff --git a/tests/unit/models/test_program.py b/tests/unit/models/test_program.py new file mode 100644 index 0000000000..8cc9e826b0 --- /dev/null +++ b/tests/unit/models/test_program.py @@ -0,0 +1,261 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +import uuid + +import pytest +from pydantic import ValidationError + +from pyrit.models import ( + Acquisition, + ComponentIdentifier, + Condition, + ContentScorable, + Equals, + MessageScorable, + Observation, + ProgramBehaves, + ProgramCase, + ProgramExecutionEvent, + ProgramExecutionObservationPayload, + ProgramExecutionScorable, + ProgramInvocation, + ProgramResult, + Score, + ScoringExpectation, + SurfaceScorable, + scorable_from_dict, +) + +_SOURCE = ComponentIdentifier(class_name="FakeSource", class_module="tests.fake") +_DIGEST = "a" * 64 +_IMAGE_DIGEST = "sha256:" + "b" * 64 + + +def _event(*, case_index: int = 0, profile: str = "python-3.12", **updates: object) -> ProgramExecutionEvent: + values: dict[str, object] = { + "case_index": case_index, + "runtime_profile": profile, + "runtime_image_digest": _IMAGE_DIGEST, + "resolved_command": ("python3", "/program/main.py"), + "exit_code": 0, + "stdout": "ok", + "duration_ms": 5, + } + values.update(updates) + return ProgramExecutionEvent.model_validate(values) + + +def _execution(*, invocations: int = 2, profile: str = "python-3.12") -> ProgramExecutionScorable: + return ProgramExecutionScorable( + execution_id=uuid.uuid4(), + program=ContentScorable(value="print('ok')"), + runtime_profile=profile, + invocations=tuple(ProgramInvocation(arguments=(str(index),)) for index in range(invocations)), + ) + + +def _observation( + *, + scorable: ProgramExecutionScorable, + acquisition: Acquisition, + events: tuple[ProgramExecutionEvent, ...] = (), + reasons: tuple[str, ...] = (), +) -> Observation: + return Observation( + source_identifier=_SOURCE, + acquisition=acquisition, + scorable=scorable, + payload=ProgramExecutionObservationPayload( + program_sha256=_DIGEST if events else None, events=events, reasons=reasons + ), + ) + + +def test_program_behaves_round_trips_through_condition_registry() -> None: + condition = ProgramBehaves( + runtime_profile="python-3.12", + cases=( + ProgramCase( + invocation=ProgramInvocation(arguments=("3",), stdin="x"), + expected=ProgramResult(exit_code=None, stdout=Equals(value="9")), + ), + ), + ) + restored = Condition.model_validate(condition.model_dump(mode="json")) + assert restored == condition + assert ScoringExpectation(conditions=(condition,)).conditions == (condition,) + + +def test_program_behaves_invocations_drop_expected_results() -> None: + condition = ProgramBehaves( + runtime_profile="python-3.12", + cases=( + ProgramCase(invocation=ProgramInvocation(arguments=("1",)), expected=ProgramResult(exit_code=3)), + ProgramCase(invocation=ProgramInvocation(stdin="in")), + ), + ) + assert condition.invocations == (ProgramInvocation(arguments=("1",)), ProgramInvocation(stdin="in")) + + +def test_program_case_defaults_expect_a_clean_exit_with_no_arguments() -> None: + case = ProgramCase() + assert case.invocation == ProgramInvocation() + assert case.expected == ProgramResult(exit_code=0) + + +def test_program_behaves_requires_a_case() -> None: + with pytest.raises(ValidationError, match="at least one case"): + ProgramBehaves(runtime_profile="python-3.12", cases=()) + + +@pytest.mark.parametrize("profile", ["", "-leading-dash", "has space", "image:tag", "a/b", "x" * 65]) +def test_runtime_profile_must_be_a_plain_name(profile: str) -> None: + with pytest.raises(ValidationError): + ProgramBehaves(runtime_profile=profile, cases=(ProgramCase(),)) + + +def test_invocation_arguments_reject_nul() -> None: + with pytest.raises(ValidationError): + ProgramInvocation(arguments=("a\x00b",)) + + +def test_timed_out_event_cannot_report_an_exit_code() -> None: + with pytest.raises(ValidationError, match="timed-out"): + _event(timed_out=True, exit_code=0) + assert _event(timed_out=True, exit_code=None).output_truncated is False + assert _event(stderr_truncated=True).output_truncated is True + + +def test_execution_scorable_round_trips_as_a_stored_scorable() -> None: + message_program = ProgramExecutionScorable( + execution_id=uuid.uuid4(), + program=MessageScorable(message_piece_ids=(uuid.uuid4(),)), + runtime_profile="python-3.12", + invocations=(ProgramInvocation(),), + ) + for scorable in (message_program, _execution()): + assert scorable_from_dict(scorable.model_dump(mode="json")) == scorable + + +def test_execution_scorable_requires_text_and_an_invocation() -> None: + with pytest.raises(ValidationError, match="text"): + ProgramExecutionScorable( + execution_id=uuid.uuid4(), + program=ContentScorable(value="image.png", data_type="image_path"), + runtime_profile="python-3.12", + invocations=(ProgramInvocation(),), + ) + with pytest.raises(ValidationError, match="at least one invocation"): + ProgramExecutionScorable( + execution_id=uuid.uuid4(), + program=ContentScorable(value="print(1)"), + runtime_profile="python-3.12", + invocations=(), + ) + + +def test_execution_scorable_rejects_a_program_that_is_not_text_or_a_message() -> None: + with pytest.raises(ValidationError): + ProgramExecutionScorable.model_validate( + { + "execution_id": str(uuid.uuid4()), + "program": SurfaceScorable(uri="/x").model_dump(), + "runtime_profile": "python-3.12", + "invocations": [{}], + } + ) + + +def test_score_round_trips_an_execution_anchor() -> None: + scorable = _execution() + score = Score( + score_value="true", + score_type="true_false", + score_rationale="", + score_value_description="", + scorable=scorable, + scorer_class_identifier=_SOURCE, + ) + assert Score.model_validate(score.model_dump(mode="json")).scorable == scorable + + +@pytest.mark.parametrize( + "indices", + [(1, 0), (0, 0)], +) +def test_payload_requires_events_once_in_case_order(indices: tuple[int, ...]) -> None: + with pytest.raises(ValidationError, match="case order"): + ProgramExecutionObservationPayload( + program_sha256=_DIGEST, events=tuple(_event(case_index=index) for index in indices) + ) + + +def test_payload_events_require_the_program_digest() -> None: + with pytest.raises(ValidationError, match="digest"): + ProgramExecutionObservationPayload(events=(_event(),)) + + +def test_payload_rejects_a_coerced_schema_version() -> None: + with pytest.raises(ValidationError, match="schema_version"): + ProgramExecutionObservationPayload.model_validate({"schema_version": "1"}) + + +def test_payload_finds_events_by_case() -> None: + payload = ProgramExecutionObservationPayload(program_sha256=_DIGEST, events=(_event(case_index=1),)) + assert payload.event_for(1) is payload.events[0] + assert payload.event_for(0) is None + + +def test_complete_observation_holds_one_event_per_invocation() -> None: + scorable = _execution() + observation = _observation( + scorable=scorable, acquisition=Acquisition.COMPLETE, events=(_event(case_index=0), _event(case_index=1)) + ) + assert Observation.model_validate(observation.model_dump(mode="json")) == observation + with pytest.raises(ValidationError, match="Complete program acquisition"): + _observation(scorable=scorable, acquisition=Acquisition.COMPLETE, events=(_event(case_index=0),)) + with pytest.raises(ValidationError, match="Complete program acquisition"): + _observation( + scorable=scorable, + acquisition=Acquisition.COMPLETE, + events=(_event(case_index=0), _event(case_index=1)), + reasons=("gap",), + ) + + +def test_partial_observation_is_missing_an_invocation() -> None: + scorable = _execution() + _observation(scorable=scorable, acquisition=Acquisition.PARTIAL, events=(_event(case_index=0),), reasons=("x",)) + with pytest.raises(ValidationError, match="Partial program acquisition"): + _observation( + scorable=scorable, acquisition=Acquisition.PARTIAL, events=(_event(case_index=0), _event(case_index=1)) + ) + + +@pytest.mark.parametrize("acquisition", [Acquisition.ERROR, Acquisition.UNAVAILABLE]) +def test_failed_observation_holds_no_events(acquisition: Acquisition) -> None: + scorable = _execution() + _observation(scorable=scorable, acquisition=acquisition, reasons=("sandbox down",)) + with pytest.raises(ValidationError, match="cannot contain events"): + _observation(scorable=scorable, acquisition=acquisition, events=(_event(case_index=0),)) + + +def test_observation_events_must_belong_to_the_scorable() -> None: + scorable = _execution(invocations=1) + with pytest.raises(ValidationError, match="belong to the scorable"): + _observation(scorable=scorable, acquisition=Acquisition.PARTIAL, events=(_event(case_index=1),)) + with pytest.raises(ValidationError, match="runtime profile"): + _observation(scorable=scorable, acquisition=Acquisition.COMPLETE, events=(_event(profile="node-20"),)) + + +def test_program_payload_requires_an_execution_anchor() -> None: + with pytest.raises(ValidationError, match="ProgramExecutionScorable"): + Observation( + source_identifier=_SOURCE, + acquisition=Acquisition.ERROR, + scorable=ContentScorable(value="print(1)"), + payload=ProgramExecutionObservationPayload(reasons=("x",)), + ) diff --git a/tests/unit/models/test_scorable.py b/tests/unit/models/test_scorable.py index 79f9cd5798..2fb1b2980b 100644 --- a/tests/unit/models/test_scorable.py +++ b/tests/unit/models/test_scorable.py @@ -13,6 +13,8 @@ Message, MessagePiece, MessageScorable, + ProgramExecutionScorable, + ProgramInvocation, Scorable, Score, SurfaceScorable, @@ -155,6 +157,12 @@ def test_every_union_member_round_trips_to_its_own_type(): TraceScorable(trace_ids=("1" * 32,)), SurfaceScorable(uri="/data/out.txt"), ConversationScorable(conversation_id="whole-conversation"), + ProgramExecutionScorable( + execution_id=uuid.uuid4(), + program=ContentScorable(value="print(1)"), + runtime_profile="python-3.12", + invocations=(ProgramInvocation(),), + ), ] assert {type(case) for case in cases} == set(SCORABLE_TYPES) diff --git a/tests/unit/score/test_docker_sandbox_client.py b/tests/unit/score/test_docker_sandbox_client.py new file mode 100644 index 0000000000..74648b7a34 --- /dev/null +++ b/tests/unit/score/test_docker_sandbox_client.py @@ -0,0 +1,268 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Exercise DockerSandboxClient against a fake Docker command line client.""" + +from __future__ import annotations + +import json +import stat +import sys +from typing import TYPE_CHECKING + +import pytest + +from pyrit.models import ProgramInvocation +from pyrit.score import DockerRuntimeProfile, DockerSandboxClient, SandboxClient, SandboxExecutionError + +if TYPE_CHECKING: + from pathlib import Path + +_IMAGE = "python@sha256:" + "c" * 64 + +# The fake runs the "program" as a scenario keyword, and keeps container state in files so +# that inspect, kill and rm see what run did. +_FAKE_DOCKER = r""" +import json, os, signal, sys, time +from pathlib import Path + +state = Path(os.environ["FAKE_DOCKER_STATE"]) +args = sys.argv[1:] +command = args[0] + +def path(name, suffix): + return state / f"{name}.{suffix}" + +if command == "run": + name = args[args.index("--name") + 1] + path(name, "argv").write_text(json.dumps(args)) + mount = args[args.index("--mount") + 1] + source = dict(part.split("=", 1) for part in mount.split(",") if "=" in part)["source"] + image_index = next(i for i, a in enumerate(args) if "@sha256:" in a) + inner = args[image_index + 1:] + program_arg = next(a for a in inner if a.startswith("/program/")) + program = (Path(source) / program_arg.rsplit("/", 1)[1]).read_text() + user_args = inner[inner.index(program_arg) + 1:] + def finish(code, oom=False, started=True): + path(name, "state").write_text(json.dumps( + {"code": code, "oom": oom, "started": "2026-10-14T00:00:00Z" if started else "0001-01-01T00:00:00Z"})) + sys.stdout.flush() + sys.exit(code) + path(name, "pid").write_text(str(os.getpid())) + path(name, "state").write_text(json.dumps({"code": 0, "oom": False, "started": "2026-10-14T00:00:00Z"})) + if program == "ECHO": + data = sys.stdin.read() + print(" ".join(user_args)) + sys.stderr.write(data) + finish(0) + elif program.startswith("EXIT "): + finish(int(program.split()[1])) + elif program == "SLEEP": + time.sleep(60) + finish(0) + elif program == "OOM": + finish(137, oom=True) + elif program == "BIG": + sys.stdout.write("a" * 200_000) + finish(0) + elif program == "NOSTART": + sys.stderr.write("docker: Error response from daemon: No such image") + finish(125, started=False) + elif program == "FAILSECOND": + marker = state / "ran-once" + if marker.exists(): + sys.stderr.write("docker: daemon went away") + finish(125, started=False) + marker.write_text("x") + finish(0) + finish(0) +elif command == "inspect": + name = args[-1] + try: + info = json.loads(path(name, "state").read_text()) + except FileNotFoundError: + print("Error: No such object: " + name) + sys.exit(1) + print(info["code"], "true" if info["oom"] else "false", info["started"]) +elif command == "kill": + name = args[-1] + pid = int(path(name, "pid").read_text()) + os.kill(pid, signal.SIGKILL) + info = json.loads(path(name, "state").read_text()) + info["code"] = 137 + path(name, "state").write_text(json.dumps(info)) +elif command == "rm": + name = args[-1] + with (state / "removed.log").open("a") as log: + log.write(name + "\n") + path(name, "state").unlink(missing_ok=True) +""" + + +@pytest.fixture +def fake_docker(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + if sys.platform == "win32": + pytest.skip("The fake Docker client is a POSIX script.") + state = tmp_path / "state" + state.mkdir() + script = tmp_path / "docker" + script.write_text(f"#!{sys.executable}\n{_FAKE_DOCKER}") + script.chmod(script.stat().st_mode | stat.S_IXUSR) + monkeypatch.setenv("FAKE_DOCKER_STATE", str(state)) + return script + + +def _client(docker: Path, **overrides: object) -> DockerSandboxClient: + settings: dict[str, object] = { + "profiles": {"python-3.12": DockerRuntimeProfile(image=_IMAGE, command=("python3", "-I"))}, + "docker_executable": str(docker), + "timeout_seconds": 5.0, + } + settings.update(overrides) + return DockerSandboxClient(**settings) # type: ignore[arg-type] + + +async def _run(client: DockerSandboxClient, program: str, *invocations: ProgramInvocation): # type: ignore[no-untyped-def] + return await client.execute_async( + program=program, runtime_profile="python-3.12", invocations=invocations or (ProgramInvocation(),) + ) + + +def _removed(docker: Path) -> list[str]: + log = docker.parent / "state" / "removed.log" + return log.read_text().split() if log.exists() else [] + + +async def test_runs_each_invocation_with_its_arguments_and_input_async(fake_docker: Path) -> None: + run = await _run( + _client(fake_docker), + "ECHO", + ProgramInvocation(arguments=("a", "b c"), stdin="input one"), + ProgramInvocation(arguments=("z",)), + ) + + assert [event.case_index for event in run.events] == [0, 1] + first, second = run.events + assert (first.stdout, first.stderr, first.exit_code) == ("a b c\n", "input one", 0) + assert (second.stdout, second.stderr) == ("z\n", "") + assert first.resolved_command == ("python3", "-I", "/program/main.py", "a", "b c") + assert first.runtime_image_digest == "sha256:" + "c" * 64 + assert first.runtime_profile == "python-3.12" + assert run.reasons == () + + +async def test_containers_are_locked_down_and_removed_async(fake_docker: Path) -> None: + await _run(_client(fake_docker, memory="128m", cpus=0.5, pids_limit=8), "ECHO") + + argv_files = list((fake_docker.parent / "state").glob("*.argv")) + assert len(argv_files) == 1 + argv = json.loads(argv_files[0].read_text()) + for flag, value in [ + ("--network", "none"), + ("--pull", "never"), + ("--cap-drop", "ALL"), + ("--security-opt", "no-new-privileges"), + ("--user", "65534:65534"), + ("--memory", "128m"), + ("--memory-swap", "128m"), + ("--cpus", "0.5"), + ("--pids-limit", "8"), + ]: + assert argv[argv.index(flag) + 1] == value + assert "--read-only" in argv + assert "--privileged" not in argv + mount = argv[argv.index("--mount") + 1] + assert mount.endswith(",target=/program,readonly") + assert argv[argv.index(_IMAGE) + 1 :] == ["python3", "-I", "/program/main.py"] + assert _removed(fake_docker) == [argv_files[0].stem] + + +async def test_program_exit_codes_are_observed_not_errors_async(fake_docker: Path) -> None: + run = await _run(_client(fake_docker), "EXIT 125") + + assert run.events[0].exit_code == 125 + assert not run.events[0].timed_out + + +async def test_timeout_kills_the_container_async(fake_docker: Path) -> None: + run = await _run(_client(fake_docker, timeout_seconds=0.5), "SLEEP") + + event = run.events[0] + assert event.timed_out + assert event.exit_code is None + assert not event.resource_limit_exceeded + assert len(_removed(fake_docker)) == 1 + + +async def test_memory_kill_is_a_resource_limit_async(fake_docker: Path) -> None: + event = (await _run(_client(fake_docker), "OOM")).events[0] + + assert event.resource_limit_exceeded + assert event.exit_code is None + + +async def test_output_is_capped_per_stream_async(fake_docker: Path) -> None: + event = (await _run(_client(fake_docker, max_output_bytes=1000), "BIG")).events[0] + + assert event.stdout == "a" * 1000 + assert event.stdout_truncated + assert not event.stderr_truncated + + +async def test_a_container_that_never_starts_is_a_sandbox_error_async(fake_docker: Path) -> None: + with pytest.raises(SandboxExecutionError, match="No such image"): + await _run(_client(fake_docker), "NOSTART") + assert len(_removed(fake_docker)) == 1 + + +async def test_a_failure_after_some_runs_returns_what_ran_async(fake_docker: Path) -> None: + run = await _run(_client(fake_docker), "FAILSECOND", ProgramInvocation(), ProgramInvocation()) + + assert [event.case_index for event in run.events] == [0] + assert "case 1" in run.reasons[0] + + +async def test_a_missing_docker_client_is_a_sandbox_error_async(tmp_path: Path) -> None: + with pytest.raises(SandboxExecutionError, match="could not start"): + await _run(_client(tmp_path / "no-such-docker"), "ECHO") + + +async def test_unknown_profile_is_rejected_async(fake_docker: Path) -> None: + with pytest.raises(ValueError, match="not configured"): + await _client(fake_docker).execute_async( + program="ECHO", runtime_profile="node-20", invocations=(ProgramInvocation(),) + ) + + +def test_profile_image_must_be_pinned_by_digest() -> None: + with pytest.raises(ValueError, match="pinned by digest"): + DockerRuntimeProfile(image="python:3.12-slim", command=("python3",)) + with pytest.raises(ValueError, match="nonempty"): + DockerRuntimeProfile(image=_IMAGE, command=()) + with pytest.raises(ValueError, match="plain file name"): + DockerRuntimeProfile(image=_IMAGE, command=("python3",), program_filename="../main.py") + + +@pytest.mark.parametrize( + "overrides", + [ + {"profiles": {}}, + {"profiles": {"bad name": DockerRuntimeProfile(image=_IMAGE, command=("python3",))}}, + {"timeout_seconds": 0}, + {"pids_limit": 0}, + {"max_output_bytes": 0}, + {"memory": "lots"}, + ], +) +def test_client_rejects_invalid_settings(overrides: dict[str, object], tmp_path: Path) -> None: + with pytest.raises(ValueError): + _client(tmp_path / "docker", **overrides) + + +def test_client_satisfies_the_sandbox_protocol(tmp_path: Path) -> None: + client = _client(tmp_path / "docker") + assert isinstance(client, SandboxClient) + assert client.get_runtime_profiles() == frozenset({"python-3.12"}) + params = client.get_identifier().params + assert params["profiles"]["python-3.12"]["image"] == _IMAGE + assert params["timeout_seconds"] == 5.0 diff --git a/tests/unit/score/test_program_execution_scorer.py b/tests/unit/score/test_program_execution_scorer.py new file mode 100644 index 0000000000..59cd7f81ea --- /dev/null +++ b/tests/unit/score/test_program_execution_scorer.py @@ -0,0 +1,560 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +import hashlib +import uuid +from typing import TYPE_CHECKING + +import pytest + +from pyrit.models import ( + Acquisition, + ComponentIdentifier, + Contains, + ContentScorable, + ContentWritten, + Equals, + MessagePiece, + MessageScorable, + Observation, + ProgramBehaves, + ProgramCase, + ProgramExecutionEvent, + ProgramExecutionObservationPayload, + ProgramExecutionScorable, + ProgramInvocation, + ProgramResult, + Regex, + ScoreStatus, + ScoringExpectation, + SurfaceScorable, +) +from pyrit.score import ( + NonReplayableObservationError, + ProgramExecutionScorer, + ProgramExecutionSource, + SandboxClient, + SandboxExecutionError, + SandboxRun, + TrueFalseInverterScorer, +) +from pyrit.score.true_false.program_execution_scorer import match_case, match_program_behaves + +if TYPE_CHECKING: + from pyrit.memory import SQLiteMemory + from pyrit.models import TextMatcher + +pytestmark = pytest.mark.usefixtures("patch_central_database") + +_PROFILE = "python-3.12" +_IMAGE_DIGEST = "sha256:" + "b" * 64 + + +def _event(*, case_index: int = 0, **updates: object) -> ProgramExecutionEvent: + values: dict[str, object] = { + "case_index": case_index, + "runtime_profile": _PROFILE, + "runtime_image_digest": _IMAGE_DIGEST, + "resolved_command": ("python3", "/program/main.py"), + "exit_code": 0, + "stdout": "", + "duration_ms": 1, + } + values.update(updates) + return ProgramExecutionEvent.model_validate(values) + + +class _FakeSandbox: + """Answers each invocation from a script keyed by its first argument, and records calls.""" + + def __init__(self, *, outputs: dict[str, dict[str, object]] | None = None, profiles: frozenset[str] | None = None): + self.outputs = outputs or {} + self.profiles = profiles or frozenset({_PROFILE}) + self.calls: list[tuple[str, str, tuple[ProgramInvocation, ...]]] = [] + self.error: SandboxExecutionError | None = None + self.stop_after: int | None = None + + def get_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of(self, params={"fake": True}) + + def get_runtime_profiles(self) -> frozenset[str]: + return self.profiles + + async def execute_async( + self, *, program: str, runtime_profile: str, invocations: tuple[ProgramInvocation, ...] + ) -> SandboxRun: + self.calls.append((program, runtime_profile, invocations)) + if self.error is not None: + raise self.error + ran = invocations if self.stop_after is None else invocations[: self.stop_after] + events = tuple( + _event(case_index=index, **self.outputs.get(invocation.arguments[0] if invocation.arguments else "", {})) + for index, invocation in enumerate(ran) + ) + reasons = ("the sandbox lost its worker",) if len(ran) < len(invocations) else () + return SandboxRun(events=events, reasons=reasons) + + +def _case(argument: str, *, exit_code: int | None = 0, stdout: TextMatcher | None = None) -> ProgramCase: + return ProgramCase( + invocation=ProgramInvocation(arguments=(argument,)), + expected=ProgramResult(exit_code=exit_code, stdout=stdout), + ) + + +def _expectation(*cases: ProgramCase, profile: str = _PROFILE) -> ScoringExpectation: + return ScoringExpectation(conditions=(ProgramBehaves(runtime_profile=profile, cases=cases),)) + + +_SQUARES = (_case("3", stdout=Equals(value="9")), _case("4", stdout=Equals(value="16"))) + + +def _scorer(sandbox: _FakeSandbox, *, max_output_chars: int = 64_000) -> ProgramExecutionScorer: + return ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=sandbox, max_output_chars=max_output_chars)) + + +@pytest.mark.parametrize( + ("expected", "event", "result"), + [ + (ProgramResult(), {}, True), + (ProgramResult(), {"exit_code": 1}, False), + (ProgramResult(exit_code=None), {"exit_code": 1}, True), + (ProgramResult(exit_code=None), {"exit_code": None}, None), + (ProgramResult(exit_code=None), {"exit_code": None, "timed_out": True}, False), + (ProgramResult(exit_code=None), {"exit_code": None, "resource_limit_exceeded": True}, False), + (ProgramResult(stdout=Equals(value="9")), {"stdout": "9\n"}, True), + (ProgramResult(stdout=Equals(value="9")), {"stdout": "10"}, False), + (ProgramResult(stdout=Equals(value="9")), {"stdout": "9", "stdout_truncated": True}, None), + (ProgramResult(stdout=Contains(value="key")), {"stdout": "a key", "stdout_truncated": True}, True), + (ProgramResult(stdout=Contains(value="key")), {"stdout": "a lock", "stdout_truncated": True}, None), + (ProgramResult(stdout=Contains(value="key")), {"stdout": "a lock"}, False), + (ProgramResult(stdout=Regex(value=r"^\d+$")), {"stdout": "42"}, True), + (ProgramResult(stderr=Contains(value="Error")), {"stderr": "ValueError"}, True), + (ProgramResult(stderr=Contains(value="Error")), {"stderr": ""}, False), + (ProgramResult(stdout=Equals(value="9")), {"stdout": "9", "stderr_truncated": True}, True), + (ProgramResult(exit_code=1, stdout=Equals(value="9")), {"stdout": "9"}, False), + ( + ProgramResult(stdout=Equals(value="9"), stderr=Equals(value="x")), + {"stdout": "8", "stderr_truncated": True}, + False, + ), + ], +) +def test_match_case(*, expected: ProgramResult, event: dict[str, object], result: bool | None) -> None: + assert match_case(expected=expected, event=_event(**event)) is result + + +def test_match_program_behaves_needs_every_case() -> None: + condition = ProgramBehaves(runtime_profile=_PROFILE, cases=_SQUARES) + passing = (_event(case_index=0, stdout="9"), _event(case_index=1, stdout="16")) + payload = ProgramExecutionObservationPayload(program_sha256="a" * 64, events=passing) + assert match_program_behaves(condition=condition, payload=payload, acquisition=Acquisition.COMPLETE) is True + + one_ran = ProgramExecutionObservationPayload(program_sha256="a" * 64, events=passing[:1]) + assert match_program_behaves(condition=condition, payload=one_ran, acquisition=Acquisition.PARTIAL) is None + + failing_first = ProgramExecutionObservationPayload( + program_sha256="a" * 64, events=(_event(case_index=0, stdout="8"),) + ) + assert match_program_behaves(condition=condition, payload=failing_first, acquisition=Acquisition.PARTIAL) is False + + empty = ProgramExecutionObservationPayload(reasons=("down",)) + for acquisition in (Acquisition.ERROR, Acquisition.UNAVAILABLE): + assert match_program_behaves(condition=condition, payload=empty, acquisition=acquisition) is None + + +async def test_scores_true_when_every_case_passes_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "9"}, "4": {"stdout": "16"}}) + program = ContentScorable(value="import sys\nprint(int(sys.argv[1]) ** 2)\n") + + score = (await _scorer(sandbox).score_async(scorable=program, expectation=_expectation(*_SQUARES)))[0] + + assert score.score_value == "true" + assert score.status is ScoreStatus.COMPLETE + assert "All 2 case(s)" in score.score_rationale + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert isinstance(observation.scorable, ProgramExecutionScorable) + assert observation.acquisition is Acquisition.COMPLETE + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + assert observation.payload.program_sha256 == hashlib.sha256(program.value.encode()).hexdigest() + + +async def test_sandbox_receives_stimuli_but_not_expected_results_async() -> None: + sandbox = _FakeSandbox() + secret = "the-expected-answer" + expectation = _expectation( + ProgramCase( + invocation=ProgramInvocation(arguments=("3",), stdin="in"), + expected=ProgramResult(stdout=Equals(value=secret)), + ) + ) + + await _scorer(sandbox).score_async(scorable=ContentScorable(value="print(1)"), expectation=expectation) + + program, profile, invocations = sandbox.calls[0] + assert (program, profile) == ("print(1)", _PROFILE) + assert invocations == (ProgramInvocation(arguments=("3",), stdin="in"),) + assert secret not in repr(sandbox.calls) + + +async def test_scores_false_when_one_case_fails_async() -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "9"}, "4": {"stdout": "15"}}) + + score = ( + await _scorer(sandbox).score_async( + scorable=ContentScorable(value="print(1)"), expectation=_expectation(*_SQUARES) + ) + )[0] + + assert score.score_value == "false" + assert "case 1 output did not match" in score.score_rationale + + +@pytest.mark.parametrize( + ("outcome", "rationale"), + [ + ({"exit_code": 1}, "case 0 exited with 1, expected 0"), + ({"exit_code": None, "timed_out": True}, "case 0 timed out"), + ({"exit_code": None, "resource_limit_exceeded": True}, "case 0 exceeded a resource limit"), + ], +) +async def test_program_failure_is_a_false_verdict_not_undetermined_async( + outcome: dict[str, object], rationale: str +) -> None: + sandbox = _FakeSandbox(outputs={"1": outcome}) + + score = ( + await _scorer(sandbox).score_async(scorable=ContentScorable(value="x"), expectation=_expectation(_case("1"))) + )[0] + + assert (score.score_value, score.status) == ("false", ScoreStatus.COMPLETE) + assert rationale in score.score_rationale + + +async def test_sandbox_failure_is_undetermined_not_false_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox() + sandbox.error = SandboxExecutionError("service unavailable") + + score = ( + await _scorer(sandbox).score_async(scorable=ContentScorable(value="x"), expectation=_expectation(_case("1"))) + )[0] + + assert score.status is ScoreStatus.UNDETERMINED + assert score.score_value is None + assert "service unavailable" in score.score_rationale + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert observation.acquisition is Acquisition.ERROR + + +async def test_sandbox_stopping_early_leaves_passing_runs_undetermined_async() -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "9"}}) + sandbox.stop_after = 1 + + score = ( + await _scorer(sandbox).score_async(scorable=ContentScorable(value="x"), expectation=_expectation(*_SQUARES)) + )[0] + + assert score.status is ScoreStatus.UNDETERMINED + assert "case 1 did not run" in score.score_rationale + assert "lost its worker" in score.score_rationale + + +async def test_sandbox_stopping_early_still_fails_a_failed_case_async() -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "8"}}) + sandbox.stop_after = 1 + + score = ( + await _scorer(sandbox).score_async(scorable=ContentScorable(value="x"), expectation=_expectation(*_SQUARES)) + )[0] + + assert score.score_value == "false" + + +async def test_output_beyond_the_limit_is_cut_and_marked_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox(outputs={"1": {"stdout": "x" * 50 + "key", "stderr": "e" * 60}}) + scorer = _scorer(sandbox, max_output_chars=10) + + equals = ( + await scorer.score_async( + scorable=ContentScorable(value="x"), + expectation=_expectation(_case("1", stdout=Equals(value="x" * 50 + "key"))), + ) + )[0] + contains = ( + await scorer.score_async( + scorable=ContentScorable(value="x"), expectation=_expectation(_case("1", stdout=Contains(value="xxx"))) + ) + )[0] + + assert equals.status is ScoreStatus.UNDETERMINED + assert contains.score_value == "true" + observation = (await sqlite_instance.get_observations_async(observation_ids=equals.observation_ids))[0] + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + event = observation.payload.events[0] + assert (event.stdout, event.stdout_truncated, len(event.stderr), event.stderr_truncated) == ( + "x" * 10, + True, + 10, + True, + ) + + +async def test_replay_judges_stored_output_without_running_again_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "9"}, "4": {"stdout": "16"}}) + scorer = _scorer(sandbox) + score = (await scorer.score_async(scorable=ContentScorable(value="x"), expectation=_expectation(*_SQUARES)))[0] + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + + same = (await scorer.score_observation_async(observation=observation, expectation=_expectation(*_SQUARES)))[0] + stricter = ( + await scorer.score_observation_async( + observation=observation, + expectation=_expectation(_case("3", stdout=Equals(value="9")), _case("4", stdout=Equals(value="17"))), + ) + )[0] + + assert same.score_value == "true" + assert stricter.score_value == "false" + assert same.observation_ids == [observation.id] + assert len(sandbox.calls) == 1 + + +async def test_replay_rejects_different_stimuli_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "9"}, "4": {"stdout": "16"}}) + scorer = _scorer(sandbox) + score = (await scorer.score_async(scorable=ContentScorable(value="x"), expectation=_expectation(*_SQUARES)))[0] + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + + with pytest.raises(NonReplayableObservationError, match="different invocations"): + await scorer.score_observation_async( + observation=observation, expectation=_expectation(_case("5", stdout=Equals(value="25"))) + ) + with pytest.raises(NonReplayableObservationError, match="different invocations"): + await scorer.score_observation_async( + observation=observation, expectation=_expectation(*_SQUARES, profile="python-3.13") + ) + + +async def test_each_scoring_call_is_a_new_execution_async() -> None: + sandbox = _FakeSandbox() + scorer = _scorer(sandbox) + program = ContentScorable(value="x") + + first = (await scorer.score_async(scorable=program, expectation=_expectation(_case("1"))))[0] + second = (await scorer.score_async(scorable=program, expectation=_expectation(_case("1"))))[0] + + assert len(sandbox.calls) == 2 + assert first.observation_ids != second.observation_ids + + +async def test_explicit_execution_reuses_its_observed_run_async() -> None: + sandbox = _FakeSandbox() + scorer = _scorer(sandbox) + execution = ProgramExecutionScorable( + execution_id=uuid.uuid4(), + program=ContentScorable(value="x"), + runtime_profile=_PROFILE, + invocations=(ProgramInvocation(arguments=("1",)),), + ) + + first = (await scorer.score_async(scorable=execution, expectation=_expectation(_case("1"))))[0] + second = (await scorer.score_async(scorable=execution, expectation=_expectation(_case("1", exit_code=None))))[0] + + assert len(sandbox.calls) == 1 + assert first.observation_ids == second.observation_ids + assert first.scorable == execution + + +async def test_sandbox_errors_are_not_reused_async() -> None: + sandbox = _FakeSandbox() + sandbox.error = SandboxExecutionError("down") + source = ProgramExecutionSource(sandbox=sandbox) + execution = ProgramExecutionScorable( + execution_id=uuid.uuid4(), + program=ContentScorable(value="x"), + runtime_profile=_PROFILE, + invocations=(ProgramInvocation(),), + ) + + assert (await source.acquire_async(scorable=execution)).acquisition is Acquisition.ERROR + sandbox.error = None + assert (await source.acquire_async(scorable=execution)).acquisition is Acquisition.COMPLETE + assert len(sandbox.calls) == 2 + + +async def test_a_run_with_no_events_is_an_error_and_not_reused_async() -> None: + sandbox = _FakeSandbox() + sandbox.stop_after = 0 + source = ProgramExecutionSource(sandbox=sandbox) + execution = ProgramExecutionScorable( + execution_id=uuid.uuid4(), + program=ContentScorable(value="x"), + runtime_profile=_PROFILE, + invocations=(ProgramInvocation(),), + ) + + failed = await source.acquire_async(scorable=execution) + sandbox.stop_after = None + retried = await source.acquire_async(scorable=execution) + + assert failed.acquisition is Acquisition.ERROR + assert isinstance(failed.payload, ProgramExecutionObservationPayload) + assert failed.payload.reasons == ("the sandbox lost its worker",) + assert retried.acquisition is Acquisition.COMPLETE + + +async def test_explicit_execution_must_match_the_condition_async() -> None: + execution = ProgramExecutionScorable( + execution_id=uuid.uuid4(), + program=ContentScorable(value="x"), + runtime_profile=_PROFILE, + invocations=(ProgramInvocation(arguments=("2",)),), + ) + + with pytest.raises(RuntimeError, match="runtime profile and invocations"): + await _scorer(_FakeSandbox()).score_async(scorable=execution, expectation=_expectation(_case("1"))) + + +async def test_runtime_profile_must_be_offered_by_the_sandbox_async() -> None: + sandbox = _FakeSandbox() + + with pytest.raises(RuntimeError, match="not offered by the sandbox"): + await _scorer(sandbox).score_async( + scorable=ContentScorable(value="x"), expectation=_expectation(_case("1"), profile="node-20") + ) + assert sandbox.calls == [] + + +async def test_scores_a_stored_message_verbatim_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox() + piece = MessagePiece(role="assistant", original_value="print('hi')", conversation_id=str(uuid.uuid4())) + await sqlite_instance.add_message_to_memory_async(request=piece.to_message()) + scorable = MessageScorable(message_piece_ids=(piece.id,)) + + score = (await _scorer(sandbox).score_async(scorable=scorable, expectation=_expectation(_case("1"))))[0] + + assert score.score_value == "true" + assert sandbox.calls[0][0] == "print('hi')" + assert score.scorable == scorable + assert score.message_piece_id == piece.id + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert isinstance(observation.scorable, ProgramExecutionScorable) + assert observation.scorable.program == scorable + + +async def test_a_blocked_response_has_no_program_to_run_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox() + piece = MessagePiece( + role="assistant", + original_value="content filtered", + conversation_id=str(uuid.uuid4()), + response_error="blocked", + ) + await sqlite_instance.add_message_to_memory_async(request=piece.to_message()) + + score = ( + await _scorer(sandbox).score_async( + scorable=MessageScorable(message_piece_ids=(piece.id,)), expectation=_expectation(_case("1")) + ) + )[0] + + assert score.status is ScoreStatus.UNDETERMINED + assert "blocked response" in score.score_rationale + assert sandbox.calls == [] + + +async def test_a_message_reference_must_name_one_piece_async(sqlite_instance: SQLiteMemory) -> None: + conversation_id = str(uuid.uuid4()) + first = MessagePiece(role="assistant", original_value="a", conversation_id=conversation_id) + second = MessagePiece(role="assistant", original_value="b", conversation_id=conversation_id) + await sqlite_instance.add_message_to_memory_async(request=first.to_message()) + + with pytest.raises(RuntimeError, match="exactly one message piece"): + await _scorer(_FakeSandbox()).score_async( + scorable=MessageScorable(message_piece_ids=(first.id, second.id)), expectation=_expectation(_case("1")) + ) + + +async def test_rejects_other_scorables_async() -> None: + with pytest.raises(RuntimeError, match="requires a MessageScorable"): + await _scorer(_FakeSandbox()).score_async( + scorable=SurfaceScorable(uri="/out.txt"), expectation=_expectation(_case("1")) + ) + + +async def test_requires_the_program_behaves_condition_async() -> None: + with pytest.raises(ValueError, match="requires one ProgramBehaves condition"): + await _scorer(_FakeSandbox()).score_async( + scorable=ContentScorable(value="x"), + expectation=ScoringExpectation(conditions=(ContentWritten(uri="/out.txt"),)), + ) + + +async def test_inverter_keeps_undetermined_async() -> None: + sandbox = _FakeSandbox() + sandbox.error = SandboxExecutionError("down") + passing = _FakeSandbox() + + undetermined = ( + await TrueFalseInverterScorer(scorer=_scorer(sandbox)).score_async( + scorable=ContentScorable(value="x"), expectation=_expectation(_case("1")) + ) + )[0] + inverted = ( + await TrueFalseInverterScorer(scorer=_scorer(passing)).score_async( + scorable=ContentScorable(value="x"), expectation=_expectation(_case("1")) + ) + )[0] + + assert undetermined.status is ScoreStatus.UNDETERMINED + assert inverted.get_value() is False + + +def test_fake_sandbox_satisfies_the_protocol() -> None: + assert isinstance(_FakeSandbox(), SandboxClient) + + +def test_source_rejects_a_nonpositive_output_limit() -> None: + with pytest.raises(ValueError, match="positive"): + ProgramExecutionSource(sandbox=_FakeSandbox(), max_output_chars=0) + + +def test_identifier_names_the_source_and_sandbox() -> None: + identifier = _scorer(_FakeSandbox()).get_identifier() + source = identifier.children["source"] + assert isinstance(source, ComponentIdentifier) + assert source.class_name == "ProgramExecutionSource" + assert isinstance(source.children["sandbox"], ComponentIdentifier) + + +async def test_source_rejects_events_for_unknown_invocations_async() -> None: + class _Rogue(_FakeSandbox): + async def execute_async(self, *, program, runtime_profile, invocations): # type: ignore[no-untyped-def] + return SandboxRun(events=(_event(case_index=5),)) + + execution = ProgramExecutionScorable( + execution_id=uuid.uuid4(), + program=ContentScorable(value="x"), + runtime_profile=_PROFILE, + invocations=(ProgramInvocation(),), + ) + with pytest.raises(ValueError, match="not given"): + await ProgramExecutionSource(sandbox=_Rogue()).acquire_async(scorable=execution) + + +def test_observation_round_trip_keeps_the_program_reference() -> None: + execution = ProgramExecutionScorable( + execution_id=uuid.uuid4(), + program=ContentScorable(value="x"), + runtime_profile=_PROFILE, + invocations=(ProgramInvocation(),), + ) + observation = Observation( + source_identifier=ComponentIdentifier(class_name="S", class_module="m"), + acquisition=Acquisition.COMPLETE, + scorable=execution, + payload=ProgramExecutionObservationPayload(program_sha256="a" * 64, events=(_event(),)), + ) + assert Observation.model_validate(observation.model_dump(mode="json")) == observation From 715870b1abb6d07f3737f521c86ed16e85925441 Mon Sep 17 00:00:00 2001 From: WatchTree-19 <119982314+WatchTree-19@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:56:37 +0000 Subject: [PATCH 2/6] Address review: shared runs, container state, cleanup and stored content - ProgramExecutionSource shares one in-flight run between concurrent callers of the same attempt and cancels it only when every caller has gone. - DockerSandboxClient reports an exit code only for an exited container; a client that exits while the container is still running is a sandbox error. - Docker management commands are killed and reaped on cancellation as well as on timeout. - A failed container removal is reported as a warning with the container name, kept in the observation metadata, and stops further cases. - The read-only program file is made writable before its directory is removed, so cleanup also works on Windows. - ContentEntryScorable is accepted as a program reference, so a saved score anchor can be scored again. --- pyrit/models/score/scorable.py | 8 +- .../observation/docker_sandbox_client.py | 91 +++++++++++---- .../observation/program_execution_source.py | 50 +++++++- pyrit/score/observation/sandbox_client.py | 4 +- .../true_false/program_execution_scorer.py | 10 +- tests/unit/models/test_program.py | 18 +++ .../unit/score/test_docker_sandbox_client.py | 107 ++++++++++++++++- .../score/test_program_execution_scorer.py | 110 ++++++++++++++++++ 8 files changed, 356 insertions(+), 42 deletions(-) diff --git a/pyrit/models/score/scorable.py b/pyrit/models/score/scorable.py index cf744d4a6f..b711cd6112 100644 --- a/pyrit/models/score/scorable.py +++ b/pyrit/models/score/scorable.py @@ -180,8 +180,8 @@ class ProgramExecutionScorable(Scorable): One isolated execution attempt of a program: the program, its runtime, and its invocations. This is an acquisition request, not a result: it holds no expected output, so the source - that runs it cannot see which output passes. ``program`` names the code to run, either a - stored message with one text piece or loose text, and is used verbatim; code is never + that runs it cannot see which output passes. ``program`` names the code to run: a stored + message with one text piece, loose text, or stored loose text. It is used verbatim; code is never extracted from surrounding prose. ``execution_id`` identifies the attempt. Reusing a scorable reuses one observed execution, @@ -190,7 +190,7 @@ class ProgramExecutionScorable(Scorable): scorable_type: Literal["program_execution"] = "program_execution" execution_id: uuid.UUID - program: Annotated[MessageScorable | ContentScorable, Field(discriminator="scorable_type")] + program: Annotated[MessageScorable | ContentScorable | ContentEntryScorable, Field(discriminator="scorable_type")] runtime_profile: RuntimeProfileName invocations: tuple[ProgramInvocation, ...] @@ -205,7 +205,7 @@ def _validate_execution(self) -> ProgramExecutionScorable: Raises: ValueError: If the program is not text, or there is nothing to run. """ - if isinstance(self.program, ContentScorable) and self.program.data_type != "text": + if isinstance(self.program, (ContentScorable, ContentEntryScorable)) and self.program.data_type != "text": raise ValueError("A program must be text content.") if not self.invocations: raise ValueError("A ProgramExecutionScorable requires at least one invocation.") diff --git a/pyrit/score/observation/docker_sandbox_client.py b/pyrit/score/observation/docker_sandbox_client.py index fa16050f63..f1e90d038d 100644 --- a/pyrit/score/observation/docker_sandbox_client.py +++ b/pyrit/score/observation/docker_sandbox_client.py @@ -11,6 +11,7 @@ import os import re import shutil +import stat import tempfile import time import uuid @@ -199,30 +200,48 @@ async def execute_async( if profile is None: raise ValueError(f"Runtime profile {runtime_profile!r} is not configured.") program_dir = Path(tempfile.mkdtemp(prefix="pyrit-program-")) + program_path = program_dir / profile.program_filename try: - program_path = program_dir / profile.program_filename program_path.write_text(program, encoding="utf-8") program_path.chmod(0o444) program_dir.chmod(0o755) events: list[ProgramExecutionEvent] = [] for index, invocation in enumerate(invocations): try: - events.append( - await self._run_invocation_async( - case_index=index, - runtime_profile=runtime_profile, - profile=profile, - program_dir=program_dir, - invocation=invocation, - ) + event, cleanup_warning = await self._run_invocation_async( + case_index=index, + runtime_profile=runtime_profile, + profile=profile, + program_dir=program_dir, + invocation=invocation, ) except SandboxExecutionError as error: if not events: raise return SandboxRun(events=tuple(events), reasons=(f"sandbox error at case {index}: {error}",)) + events.append(event) + if cleanup_warning is not None: + # The environment may still be running, so no further case starts beside it. + remaining = len(invocations) - len(events) + reasons = (f"stopped after case {index}: {cleanup_warning}",) if remaining else () + return SandboxRun(events=tuple(events), reasons=reasons, warnings=(cleanup_warning,)) return SandboxRun(events=tuple(events)) finally: - shutil.rmtree(program_dir, ignore_errors=True) + self._remove_program_dir(program_dir=program_dir, program_path=program_path) + + @staticmethod + def _remove_program_dir(*, program_dir: Path, program_path: Path) -> None: + """ + Delete the submitted program, restoring write permission first. + + On Windows the read-only mode is a file attribute that blocks deletion, so the file + is made writable again before the directory is removed. + """ + with contextlib.suppress(OSError): + program_path.chmod(stat.S_IREAD | stat.S_IWRITE) + shutil.rmtree(program_dir, ignore_errors=True) + if program_dir.exists(): + logger.warning("Could not remove the program directory %s", program_dir) async def _run_invocation_async( self, @@ -232,15 +251,17 @@ async def _run_invocation_async( profile: DockerRuntimeProfile, program_dir: Path, invocation: ProgramInvocation, - ) -> ProgramExecutionEvent: + ) -> tuple[ProgramExecutionEvent, str | None]: """ Run one invocation in a new container and remove the container afterwards. Returns: - ProgramExecutionEvent: What the invocation did. + tuple[ProgramExecutionEvent, str | None]: What the invocation did, and a warning when + the container could not be removed. Raises: - SandboxExecutionError: If Docker cannot start or inspect the container. + SandboxExecutionError: If Docker cannot start or inspect the container, or the + container had not finished when the attached client exited. """ name = f"pyrit-sandbox-{uuid.uuid4().hex}" resolved_command = (*profile.command, f"/program/{profile.program_filename}", *invocation.arguments) @@ -289,6 +310,7 @@ async def _run_invocation_async( except OSError as error: raise SandboxExecutionError(f"could not start {self._docker!r}: {error}") from error timed_out = False + cleanup_warning: str | None = None try: readers = asyncio.gather( self._read_capped_async(process.stdout), @@ -311,14 +333,18 @@ async def _run_invocation_async( (stdout, stdout_truncated), (stderr, stderr_truncated), _ = await readers await process.wait() duration_ms = int((time.monotonic() - started) * 1000) - container = await self._inspect_async(name=name, cli_stderr=stderr) + container = await self._inspect_async(name=name, cli_stderr=stderr, timed_out=timed_out) finally: if process.returncode is None: with contextlib.suppress(ProcessLookupError): process.kill() await process.wait() - await self._control_async("rm", "--force", name) - return ProgramExecutionEvent( + removal_status, removal_output = await self._control_async("rm", "--force", name) + if removal_status != 0: + detail = removal_output.strip()[:300] or "unknown Docker failure" + cleanup_warning = f"could not remove container {name}: {detail}" + logger.warning("Sandbox cleanup failed: %s", cleanup_warning) + event = ProgramExecutionEvent( case_index=case_index, runtime_profile=runtime_profile, runtime_image_digest=profile.image_digest, @@ -332,26 +358,40 @@ async def _run_invocation_async( resource_limit_exceeded=container.oom_killed and not timed_out, duration_ms=duration_ms, ) + return event, cleanup_warning - async def _inspect_async(self, *, name: str, cli_stderr: bytes) -> _ContainerExit: + async def _inspect_async(self, *, name: str, cli_stderr: bytes, timed_out: bool) -> _ContainerExit: """ Read how the container ended, distinguishing a program exit from a Docker failure. + An exit code is only reported for a container that has finished. A container still + running when the attached client exited means the run was cut short by Docker, not by + the program, unless the sandbox stopped it at its time limit. + Returns: _ContainerExit: The container's exit code and whether the kernel killed it for memory. Raises: - SandboxExecutionError: If the container does not exist or never started. + SandboxExecutionError: If the container does not exist, never started, or had not + finished when the client exited. """ returncode, output = await self._control_async( - "inspect", "--format", "{{.State.ExitCode}} {{.State.OOMKilled}} {{.State.StartedAt}}", name + "inspect", + "--format", + "{{.State.Status}} {{.State.ExitCode}} {{.State.OOMKilled}} {{.State.StartedAt}}", + name, ) fields = output.split() - if returncode != 0 or len(fields) != 3 or fields[2] == _NEVER_STARTED: + if returncode != 0 or len(fields) != 4 or fields[3] == _NEVER_STARTED: lines = [line.strip() for line in cli_stderr.decode("utf-8", errors="replace").splitlines() if line.strip()] detail = lines[0][:500] if lines else output.strip()[:500] raise SandboxExecutionError(f"the container did not start: {detail or 'unknown Docker failure'}") - return _ContainerExit(exit_code=int(fields[0]), oom_killed=fields[1] == "true") + status = fields[0] + if status != "exited" and not timed_out: + raise SandboxExecutionError( + f"the Docker client exited while the container was still {status}, so the program did not finish" + ) + return _ContainerExit(exit_code=int(fields[1]), oom_killed=fields[2] == "true") async def _control_async(self, *arguments: str) -> tuple[int, str]: """ @@ -374,11 +414,14 @@ async def _control_async(self, *arguments: str) -> tuple[int, str]: try: output, _ = await asyncio.wait_for(process.communicate(), timeout=self._CONTROL_TIMEOUT_SECONDS) except TimeoutError: - with contextlib.suppress(ProcessLookupError): - process.kill() - await process.wait() logger.warning("Docker %s timed out", arguments[0]) return -1, "timed out" + finally: + # Runs on timeout and on cancellation alike, so no management command outlives its caller. + if process.returncode is None: + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() return process.returncode if process.returncode is not None else -1, output.decode("utf-8", errors="replace") async def _read_capped_async(self, stream: asyncio.StreamReader | None) -> tuple[bytes, bool]: diff --git a/pyrit/score/observation/program_execution_source.py b/pyrit/score/observation/program_execution_source.py index 4ea3911560..d034d6483f 100644 --- a/pyrit/score/observation/program_execution_source.py +++ b/pyrit/score/observation/program_execution_source.py @@ -5,9 +5,11 @@ from __future__ import annotations +import asyncio import hashlib import logging from collections import OrderedDict +from dataclasses import dataclass from typing import TYPE_CHECKING from pyrit.memory import CentralMemory @@ -34,6 +36,14 @@ class _NoProgramError(Exception): """The named evidence holds nothing that can be submitted as a program.""" +@dataclass +class _InFlight: + """One running acquisition and how many callers are waiting for it.""" + + task: asyncio.Future[Observation] + waiters: int = 0 + + class ProgramExecutionSource: """ Run the program a ``ProgramExecutionScorable`` names, once per invocation, in a sandbox. @@ -54,8 +64,11 @@ class ProgramExecutionSource: program that failed its tests. Each stream keeps at most ``max_output_chars`` characters and records when it was cut. - A completed or partial observation is reused when the same scorable is acquired again, so - one execution attempt is observed once; a new ``execution_id`` runs the program again. + A completed or partial observation is reused when the same scorable is acquired again, and + concurrent callers for one attempt share a single run, so one execution attempt is observed + once; a new ``execution_id`` runs the program again. Problems the sandbox reports after the + evidence was collected, such as a container it could not remove, are kept in the + observation's ``sandbox_warnings`` metadata. """ _MAX_REUSED_OBSERVATIONS = 256 @@ -76,6 +89,7 @@ def __init__(self, *, sandbox: SandboxClient, max_output_chars: int = 64_000) -> self._sandbox = sandbox self._max_output_chars = max_output_chars self._observed: OrderedDict[ProgramExecutionScorable, Observation] = OrderedDict() + self._in_flight: dict[ProgramExecutionScorable, _InFlight] = {} def get_identifier(self) -> ComponentIdentifier: """ @@ -109,6 +123,31 @@ async def acquire_async(self, *, scorable: ProgramExecutionScorable) -> Observat if reused is not None: self._observed.move_to_end(scorable) return reused + flight = self._in_flight.get(scorable) + if flight is None: + flight = _InFlight(task=asyncio.ensure_future(self._acquire_uncached_async(scorable=scorable))) + self._in_flight[scorable] = flight + flight.task.add_done_callback(lambda _task, key=scorable: self._in_flight.pop(key, None)) + flight.waiters += 1 + try: + return await asyncio.shield(flight.task) + finally: + flight.waiters -= 1 + if flight.waiters == 0 and not flight.task.done(): + # Every caller of this attempt was cancelled, so nobody needs its result. + flight.task.cancel() + + async def _acquire_uncached_async(self, *, scorable: ProgramExecutionScorable) -> Observation: + """ + Run the attempt once; ``acquire_async`` shares this work between concurrent callers. + + Returns: + Observation: The execution evidence, with its acquisition state. + + Raises: + ValueError: If the runtime profile is not offered, the program reference is + unusable, or the sandbox returns events for invocations it was not given. + """ if scorable.runtime_profile not in self._sandbox.get_runtime_profiles(): raise ValueError( f"Runtime profile {scorable.runtime_profile!r} is not offered by the sandbox; " @@ -147,6 +186,8 @@ async def acquire_async(self, *, scorable: ProgramExecutionScorable) -> Observat program_sha256=program_sha256 if events else None, reasons=reasons, ) + if run.warnings: + observation = observation.model_copy(update={"metadata": {"sandbox_warnings": "; ".join(run.warnings)}}) if observation.acquisition in (Acquisition.COMPLETE, Acquisition.PARTIAL): self._remember(scorable=scorable, observation=observation) return observation @@ -160,12 +201,13 @@ async def _resolve_program_async(self, *, scorable: ProgramExecutionScorable) -> Raises: _NoProgramError: If the named message holds a blocked or failed response. - ValueError: If the reference cannot be resolved to one text piece. + ValueError: If the reference cannot be resolved to one text piece, or stored + content is missing. """ reference = scorable.program if isinstance(reference, ContentScorable): return reference.value - if not isinstance(reference, MessageScorable) or len(reference.message_piece_ids) != 1: + if isinstance(reference, MessageScorable) and len(reference.message_piece_ids) != 1: raise ValueError("A program message reference must name exactly one message piece.") message = await MessageScorableResolver().resolve_async( scorable=reference, memory=CentralMemory.get_memory_instance() diff --git a/pyrit/score/observation/sandbox_client.py b/pyrit/score/observation/sandbox_client.py index c2a5a16b8f..bd66a62c81 100644 --- a/pyrit/score/observation/sandbox_client.py +++ b/pyrit/score/observation/sandbox_client.py @@ -29,11 +29,13 @@ class SandboxRun: ``events`` holds one event for each invocation that ran, in case order. When the sandbox stopped before running every invocation, ``reasons`` says why, so the missing cases stay - undetermined instead of failing. + undetermined instead of failing. ``warnings`` reports problems that do not invalidate the + evidence collected, such as an environment the sandbox could not destroy afterwards. """ events: tuple[ProgramExecutionEvent, ...] = () reasons: tuple[str, ...] = () + warnings: tuple[str, ...] = () @runtime_checkable diff --git a/pyrit/score/true_false/program_execution_scorer.py b/pyrit/score/true_false/program_execution_scorer.py index 009414c803..923aadba7e 100644 --- a/pyrit/score/true_false/program_execution_scorer.py +++ b/pyrit/score/true_false/program_execution_scorer.py @@ -11,6 +11,7 @@ from pyrit.models import ( Acquisition, Contains, + ContentEntryScorable, ContentScorable, MessageScorable, ProgramBehaves, @@ -123,8 +124,8 @@ class ProgramExecutionScorer(TrueFalseScorer): runtime profile and the invocations, and sends only that to its source. Each expected result stays with the scorer, so the code that runs the program never learns which output passes. - The program is the scored content, run verbatim: loose text, or a stored message with one - text piece. Code is not extracted from surrounding prose, so a response that wraps code in + The program is the scored content, run verbatim: loose text, stored loose text, or a stored + message with one text piece. Code is not extracted from surrounding prose, so a response that wraps code in explanation fails its cases rather than being guessed at. The verdict is true when every case passes and false when any case fails. A compile error, @@ -164,7 +165,7 @@ async def _score_scorable_async(self, *, scorable: Scorable, expectation: Scorin "ProgramBehaves condition." ) execution = scorable - elif isinstance(scorable, (MessageScorable, ContentScorable)): + elif isinstance(scorable, (MessageScorable, ContentScorable, ContentEntryScorable)): execution = ProgramExecutionScorable( execution_id=uuid.uuid4(), program=scorable, @@ -173,7 +174,8 @@ async def _score_scorable_async(self, *, scorable: Scorable, expectation: Scorin ) else: raise TypeError( - "ProgramExecutionScorer requires a MessageScorable, a ContentScorable, or a ProgramExecutionScorable." + "ProgramExecutionScorer requires a MessageScorable, a ContentScorable, a ContentEntryScorable, " + "or a ProgramExecutionScorable." ) observation = await self._source.acquire_async(scorable=execution) diff --git a/tests/unit/models/test_program.py b/tests/unit/models/test_program.py index 8cc9e826b0..1ea1e730c1 100644 --- a/tests/unit/models/test_program.py +++ b/tests/unit/models/test_program.py @@ -12,6 +12,7 @@ Acquisition, ComponentIdentifier, Condition, + ContentEntryScorable, ContentScorable, Equals, MessageScorable, @@ -140,6 +141,23 @@ def test_execution_scorable_round_trips_as_a_stored_scorable() -> None: assert scorable_from_dict(scorable.model_dump(mode="json")) == scorable +def test_execution_scorable_accepts_stored_text_content() -> None: + stored = ProgramExecutionScorable( + execution_id=uuid.uuid4(), + program=ContentEntryScorable(content_id=uuid.uuid4()), + runtime_profile="python-3.12", + invocations=(ProgramInvocation(),), + ) + assert scorable_from_dict(stored.model_dump(mode="json")) == stored + with pytest.raises(ValidationError, match="text"): + ProgramExecutionScorable( + execution_id=uuid.uuid4(), + program=ContentEntryScorable(content_id=uuid.uuid4(), data_type="image_path"), + runtime_profile="python-3.12", + invocations=(ProgramInvocation(),), + ) + + def test_execution_scorable_requires_text_and_an_invocation() -> None: with pytest.raises(ValidationError, match="text"): ProgramExecutionScorable( diff --git a/tests/unit/score/test_docker_sandbox_client.py b/tests/unit/score/test_docker_sandbox_client.py index 74648b7a34..4fc6e20aa1 100644 --- a/tests/unit/score/test_docker_sandbox_client.py +++ b/tests/unit/score/test_docker_sandbox_client.py @@ -5,14 +5,18 @@ from __future__ import annotations +import asyncio import json +import os +import shutil import stat import sys from typing import TYPE_CHECKING +from unittest.mock import patch import pytest -from pyrit.models import ProgramInvocation +from pyrit.models import ProgramExecutionEvent, ProgramInvocation from pyrit.score import DockerRuntimeProfile, DockerSandboxClient, SandboxClient, SandboxExecutionError if TYPE_CHECKING: @@ -43,13 +47,20 @@ def path(name, suffix): program_arg = next(a for a in inner if a.startswith("/program/")) program = (Path(source) / program_arg.rsplit("/", 1)[1]).read_text() user_args = inner[inner.index(program_arg) + 1:] + flags = {"rm_fails": program == "RMFAIL", "slow_inspect": program == "SLOWINSPECT"} + def save(**info): + path(name, "state").write_text(json.dumps({**flags, **info})) def finish(code, oom=False, started=True): - path(name, "state").write_text(json.dumps( - {"code": code, "oom": oom, "started": "2026-10-14T00:00:00Z" if started else "0001-01-01T00:00:00Z"})) + save(status="exited" if started else "created", code=code, oom=oom, + started="2026-10-14T00:00:00Z" if started else "0001-01-01T00:00:00Z") sys.stdout.flush() sys.exit(code) path(name, "pid").write_text(str(os.getpid())) - path(name, "state").write_text(json.dumps({"code": 0, "oom": False, "started": "2026-10-14T00:00:00Z"})) + save(status="running", code=0, oom=False, started="2026-10-14T00:00:00Z") + if program == "DETACHFAIL": + # The attached client fails while the container keeps running. + sys.stderr.write("error waiting for container: unexpected EOF") + sys.exit(1) if program == "ECHO": data = sys.stdin.read() print(" ".join(user_args)) @@ -83,16 +94,27 @@ def finish(code, oom=False, started=True): except FileNotFoundError: print("Error: No such object: " + name) sys.exit(1) - print(info["code"], "true" if info["oom"] else "false", info["started"]) + if info.get("slow_inspect"): + (state / "inspect.pid").write_text(str(os.getpid())) + time.sleep(60) + print(info["status"], info["code"], "true" if info["oom"] else "false", info["started"]) elif command == "kill": name = args[-1] pid = int(path(name, "pid").read_text()) os.kill(pid, signal.SIGKILL) info = json.loads(path(name, "state").read_text()) info["code"] = 137 + info["status"] = "exited" path(name, "state").write_text(json.dumps(info)) elif command == "rm": name = args[-1] + try: + info = json.loads(path(name, "state").read_text()) + except FileNotFoundError: + info = {} + if info.get("rm_fails"): + print("Error response from daemon: daemon unavailable") + sys.exit(1) with (state / "removed.log").open("a") as log: log.write(name + "\n") path(name, "state").unlink(missing_ok=True) @@ -266,3 +288,78 @@ def test_client_satisfies_the_sandbox_protocol(tmp_path: Path) -> None: params = client.get_identifier().params assert params["profiles"]["python-3.12"]["image"] == _IMAGE assert params["timeout_seconds"] == 5.0 + + +async def test_a_client_failure_while_the_container_runs_is_a_sandbox_error_async(fake_docker: Path) -> None: + with pytest.raises(SandboxExecutionError, match="still running"): + await _run(_client(fake_docker), "DETACHFAIL") + assert len(_removed(fake_docker)) == 1 + + +async def test_a_failed_removal_is_reported_and_stops_further_cases_async(fake_docker: Path) -> None: + run = await _run(_client(fake_docker), "RMFAIL", ProgramInvocation(), ProgramInvocation()) + + assert [event.case_index for event in run.events] == [0] + assert run.events[0].exit_code == 0 + assert "could not remove container pyrit-sandbox-" in run.warnings[0] + assert "daemon unavailable" in run.warnings[0] + assert run.reasons[0].startswith("stopped after case 0") + + +async def test_a_failed_removal_on_the_last_case_keeps_the_evidence_async(fake_docker: Path) -> None: + run = await _run(_client(fake_docker), "RMFAIL") + + assert len(run.events) == 1 + assert run.reasons == () + assert len(run.warnings) == 1 + + +async def test_cancellation_reaps_the_management_command_async(fake_docker: Path) -> None: + pid_file = fake_docker.parent / "state" / "inspect.pid" + task = asyncio.create_task(_run(_client(fake_docker), "SLOWINSPECT")) + for _ in range(200): + if pid_file.exists(): + break + await asyncio.sleep(0.05) + inspect_pid = int(pid_file.read_text()) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + with pytest.raises(ProcessLookupError): + os.kill(inspect_pid, 0) + + +async def test_the_program_directory_is_removed_after_making_it_writable_async(tmp_path: Path) -> None: + client = _client(tmp_path / "docker") + seen: dict[str, object] = {} + real_rmtree = shutil.rmtree + + async def run_invocation(**kwargs: object): # type: ignore[no-untyped-def] + program_dir = kwargs["program_dir"] + seen["dir"] = program_dir + event = ProgramExecutionEvent( + case_index=kwargs["case_index"], + runtime_profile="python-3.12", + runtime_image_digest="sha256:" + "c" * 64, + resolved_command=("python3",), + exit_code=0, + duration_ms=1, + ) + return event, None + + def recording_rmtree(path, *args, **kwargs): # type: ignore[no-untyped-def] + program = next(iter(path.iterdir())) + seen["writable"] = bool(program.stat().st_mode & stat.S_IWRITE) + return real_rmtree(path, *args, **kwargs) + + with ( + patch.object(client, "_run_invocation_async", side_effect=run_invocation), + patch("pyrit.score.observation.docker_sandbox_client.shutil.rmtree", side_effect=recording_rmtree), + ): + run = await _run(client, "print(1)") + + assert len(run.events) == 1 + assert seen["writable"] is True + assert not seen["dir"].exists() # type: ignore[union-attr] diff --git a/tests/unit/score/test_program_execution_scorer.py b/tests/unit/score/test_program_execution_scorer.py index 59cd7f81ea..cf94396a2b 100644 --- a/tests/unit/score/test_program_execution_scorer.py +++ b/tests/unit/score/test_program_execution_scorer.py @@ -3,6 +3,7 @@ from __future__ import annotations +import asyncio import hashlib import uuid from typing import TYPE_CHECKING @@ -13,6 +14,7 @@ Acquisition, ComponentIdentifier, Contains, + ContentEntryScorable, ContentScorable, ContentWritten, Equals, @@ -558,3 +560,111 @@ def test_observation_round_trip_keeps_the_program_reference() -> None: payload=ProgramExecutionObservationPayload(program_sha256="a" * 64, events=(_event(),)), ) assert Observation.model_validate(observation.model_dump(mode="json")) == observation + + +class _SlowSandbox(_FakeSandbox): + """Holds each run open until released, so concurrent callers overlap.""" + + def __init__(self) -> None: + super().__init__() + self.release = asyncio.Event() + + async def execute_async(self, *, program, runtime_profile, invocations): # type: ignore[no-untyped-def] + self.calls.append((program, runtime_profile, invocations)) + await self.release.wait() + return SandboxRun(events=tuple(_event(case_index=index) for index in range(len(invocations)))) + + +def _attempt() -> ProgramExecutionScorable: + return ProgramExecutionScorable( + execution_id=uuid.uuid4(), + program=ContentScorable(value="x"), + runtime_profile=_PROFILE, + invocations=(ProgramInvocation(),), + ) + + +async def test_concurrent_callers_share_one_run_of_an_attempt_async() -> None: + sandbox = _SlowSandbox() + source = ProgramExecutionSource(sandbox=sandbox) + execution = _attempt() + + pending = asyncio.gather(source.acquire_async(scorable=execution), source.acquire_async(scorable=execution)) + await asyncio.sleep(0) + sandbox.release.set() + first, second = await pending + + assert len(sandbox.calls) == 1 + assert first.id == second.id + + +async def test_one_cancelled_caller_does_not_cancel_the_shared_run_async() -> None: + sandbox = _SlowSandbox() + source = ProgramExecutionSource(sandbox=sandbox) + execution = _attempt() + + cancelled = asyncio.create_task(source.acquire_async(scorable=execution)) + kept = asyncio.create_task(source.acquire_async(scorable=execution)) + await asyncio.sleep(0) + cancelled.cancel() + await asyncio.sleep(0) + sandbox.release.set() + + observation = await kept + with pytest.raises(asyncio.CancelledError): + await cancelled + assert observation.acquisition is Acquisition.COMPLETE + assert len(sandbox.calls) == 1 + + +async def test_a_run_nobody_waits_for_is_cancelled_async() -> None: + sandbox = _SlowSandbox() + source = ProgramExecutionSource(sandbox=sandbox) + execution = _attempt() + + task = asyncio.create_task(source.acquire_async(scorable=execution)) + await asyncio.sleep(0) + flight = source._in_flight[execution] + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.sleep(0) + + assert flight.task.cancelled() + assert execution not in source._in_flight + + +async def test_sandbox_warnings_are_kept_on_the_observation_async(sqlite_instance: SQLiteMemory) -> None: + class _LeakySandbox(_FakeSandbox): + async def execute_async(self, *, program, runtime_profile, invocations): # type: ignore[no-untyped-def] + return SandboxRun( + events=(_event(),), warnings=("could not remove container pyrit-sandbox-1: daemon unavailable",) + ) + + score = ( + await _scorer(_LeakySandbox()).score_async( + scorable=ContentScorable(value="x"), expectation=_expectation(_case("1")) + ) + )[0] + + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert score.score_value == "true" + assert observation.acquisition is Acquisition.COMPLETE + assert "daemon unavailable" in observation.metadata["sandbox_warnings"] + + +async def test_scoring_a_saved_score_anchor_runs_the_stored_program_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "9"}, "4": {"stdout": "16"}}) + scorer = _scorer(sandbox) + program = ContentScorable(value="import sys\nprint(int(sys.argv[1]) ** 2)\n") + first = (await scorer.score_async(scorable=program, expectation=_expectation(*_SQUARES)))[0] + saved = (await sqlite_instance.get_scores_async(score_ids=[first.id]))[0] + assert isinstance(saved.scorable, ContentEntryScorable) + + again = (await scorer.score_async(scorable=saved.scorable, expectation=_expectation(*_SQUARES)))[0] + + assert again.score_value == "true" + assert [call[0] for call in sandbox.calls] == [program.value, program.value] + observation = (await sqlite_instance.get_observations_async(observation_ids=again.observation_ids))[0] + assert isinstance(observation.scorable, ProgramExecutionScorable) + assert observation.scorable.program == saved.scorable From 51be88f069b51b374e869cab6c97477533a773fd Mon Sep 17 00:00:00 2001 From: WatchTree-19 <119982314+WatchTree-19@users.noreply.github.com> Date: Fri, 16 Oct 2026 23:42:39 +0000 Subject: [PATCH 3/6] Address review: bounded client exit, cleanup on cancel, shared observations - DockerSandboxClient uses one deadline for the whole invocation, including the attached client's exit. A client still running after its streams close is ended at the deadline, and the container's own state then decides the result, so a stalled client is neither unbounded nor mistaken for a program timeout. - ProgramExecutionSource waits for the sandbox's cleanup when the last caller of an attempt is cancelled, so cancellation still means the environment is gone. - Memory links a new observation whose ID is already stored when the stored row is identical, checked under the observation row locks; different evidence under the same ID is still rejected. ProgramExecutionScorer drops its separate pre-check, so concurrent scoring calls that share one acquisition both persist their scores. --- pyrit/memory/memory_interface.py | 40 +++++++--- .../observation/docker_sandbox_client.py | 19 ++++- .../observation/program_execution_source.py | 7 +- .../true_false/program_execution_scorer.py | 23 +----- .../test_interface_observations.py | 42 ++++++++++ .../unit/score/test_docker_sandbox_client.py | 39 +++++++++ .../score/test_program_execution_scorer.py | 79 +++++++++++++++++++ 7 files changed, 217 insertions(+), 32 deletions(-) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index bdb3456e1e..3bda1faa14 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -2370,15 +2370,19 @@ def _add_scores_to_memory( SQLAlchemyError: If the score or identifier rows cannot be persisted. """ observations = [Observation.model_validate(observation.model_dump()) for observation in observations] - new_ids, referenced_ids = self._validate_score_inputs(scores=scores, observations=observations) + _, referenced_ids = self._validate_score_inputs(scores=scores, observations=observations) with closing(self._get_session()) as session: try: _begin_sqlite_write(session) - self._validate_observation_links( + already_stored = self._validate_observation_links( session=session, - new_observation_ids=new_ids, + new_observations={str(observation.id): observation for observation in observations}, referenced_observation_ids=referenced_ids, ) + # An identical observation another call stored first is linked, not inserted again. + observations = [ + observation for observation in observations if str(observation.id) not in already_stored + ] canonical_scores = self._resolve_score_message_anchors(session=session, scores=scores) content_entries, persisted_scores, persisted_observations = self._prepare_score_anchors( scores=canonical_scores, @@ -2590,24 +2594,42 @@ def _validate_observation_links( cls, *, session: Session, - new_observation_ids: set[str], + new_observations: Mapping[str, Observation], referenced_observation_ids: set[str], - ) -> None: + ) -> set[str]: """ Reject conflicting new IDs and links to observations that do not exist. + A new observation whose ID is already stored is accepted only when the stored row is + the same observation, as when concurrent scoring calls share one acquisition. The + check runs under the observation locks, so the first writer inserts and the others link. + + Returns: + set[str]: IDs of new observations already stored unchanged, to link rather than insert. + Raises: - ValueError: If a new ID conflicts or a referenced observation does not exist. + ValueError: If a new ID holds different evidence or a referenced observation does not exist. """ + new_observation_ids = set(new_observations) existing = _lock_observations( session=session, observation_ids=sorted(new_observation_ids | referenced_observation_ids) ) - conflicts = new_observation_ids & existing - if conflicts: - raise ValueError(f"New observation IDs already exist: {sorted(conflicts)}.") + already_stored = new_observation_ids & existing + if already_stored: + stored: dict[str, Observation] = {} + ids = sorted(already_stored) + for start in range(0, len(ids), cls._MAX_BIND_VARS): + entries = session.scalars( + select(ObservationEntry).where(ObservationEntry.id.in_(ids[start : start + cls._MAX_BIND_VARS])) + ) + stored.update({str(entry.id): entry.get_observation() for entry in entries}) + conflicts = sorted(i for i in already_stored if stored.get(i) != new_observations[i]) + if conflicts: + raise ValueError(f"New observation IDs already exist: {conflicts}.") missing = referenced_observation_ids - new_observation_ids - existing if missing: raise ValueError(f"Score references observations not found in memory: {sorted(missing)}.") + return already_stored @classmethod def _validate_observation_evidence( diff --git a/pyrit/score/observation/docker_sandbox_client.py b/pyrit/score/observation/docker_sandbox_client.py index f1e90d038d..31c1317088 100644 --- a/pyrit/score/observation/docker_sandbox_client.py +++ b/pyrit/score/observation/docker_sandbox_client.py @@ -89,7 +89,9 @@ class DockerSandboxClient: A container that runs past ``timeout_seconds`` is killed and reported as timed out. One that the kernel stops for exceeding its memory limit is reported as exceeding a resource limit. - Each stream keeps at most ``max_output_bytes`` bytes. The timeout includes container start-up. + Each stream keeps at most ``max_output_bytes`` bytes. The timeout includes container start-up + and the attached client's exit; a client still running once its streams have closed is + ended at the deadline, and the container's own state then decides the result. This client isolates submitted code with Docker's defaults plus the restrictions above. It is not a hardened multi-tenant sandbox; run it on a host dedicated to evaluation. @@ -311,6 +313,9 @@ async def _run_invocation_async( raise SandboxExecutionError(f"could not start {self._docker!r}: {error}") from error timed_out = False cleanup_warning: str | None = None + # One wall-clock deadline covers the whole invocation: reading the output and the + # attached client's exit. + deadline = started + self._timeout_seconds try: readers = asyncio.gather( self._read_capped_async(process.stdout), @@ -331,7 +336,17 @@ async def _run_invocation_async( with contextlib.suppress(ProcessLookupError): process.kill() (stdout, stdout_truncated), (stderr, stderr_truncated), _ = await readers - await process.wait() + try: + exit_budget = self._CONTROL_TIMEOUT_SECONDS if timed_out else max(deadline - time.monotonic(), 0.0) + await asyncio.wait_for(process.wait(), timeout=exit_budget) + except TimeoutError: + # The streams closed but the attached client did not exit in time. Ending the + # client keeps the invocation bounded; the container's own state, read below, + # decides whether the program finished, so a stalled client is not mistaken + # for a program that ran too long. + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() duration_ms = int((time.monotonic() - started) * 1000) container = await self._inspect_async(name=name, cli_stderr=stderr, timed_out=timed_out) finally: diff --git a/pyrit/score/observation/program_execution_source.py b/pyrit/score/observation/program_execution_source.py index d034d6483f..ef196b865f 100644 --- a/pyrit/score/observation/program_execution_source.py +++ b/pyrit/score/observation/program_execution_source.py @@ -134,8 +134,13 @@ async def acquire_async(self, *, scorable: ProgramExecutionScorable) -> Observat finally: flight.waiters -= 1 if flight.waiters == 0 and not flight.task.done(): - # Every caller of this attempt was cancelled, so nobody needs its result. + # Every caller of this attempt was cancelled, so nobody needs its result. Wait for + # the sandbox's cleanup before the caller's cancellation completes, so that + # cancelling still means the environment is gone. flight.task.cancel() + await asyncio.wait({flight.task}) + if not flight.task.cancelled() and flight.task.exception() is not None: + logger.debug("Cancelled execution %s ended with %r", scorable.execution_id, flight.task.exception()) async def _acquire_uncached_async(self, *, scorable: ProgramExecutionScorable) -> Observation: """ diff --git a/pyrit/score/true_false/program_execution_scorer.py b/pyrit/score/true_false/program_execution_scorer.py index 923aadba7e..5e0c09c9ac 100644 --- a/pyrit/score/true_false/program_execution_scorer.py +++ b/pyrit/score/true_false/program_execution_scorer.py @@ -183,32 +183,15 @@ async def _score_scorable_async(self, *, scorable: Scorable, expectation: Scorin raise ValueError("Program execution source changed the caller's evidence anchor.") if not isinstance(observation.payload, ProgramExecutionObservationPayload): raise ValueError("Program execution source returned incompatible evidence.") - # A reused attempt's evidence may already be stored; a new score then only links to it. - if execution is not scorable or not await self._is_stored_async(observation=observation): - _collect_observation(observation) + # A reused or shared attempt's evidence may already be stored. Memory links an identical + # stored observation instead of inserting it again, atomically under its row lock. + _collect_observation(observation) scores = self._score_observation(observation=observation, evidence=observation.payload, expectation=expectation) for score in scores: score.scorable = scorable score.message_piece_id = self._piece_id_from_scorable(scorable) return scores - async def _is_stored_async(self, *, observation: Observation) -> bool: - """ - Check whether an observation the source reused is already stored unchanged. - - Returns: - bool: True when memory holds the same observation. - - Raises: - ValueError: If memory holds a different observation under the same id. - """ - stored = await self._memory.get_observations_async(observation_ids=[observation.id]) - if not stored: - return False - if stored[0] != observation: - raise ValueError(f"Observation {observation.id} does not match its stored evidence.") - return True - def _score_observation( self, *, diff --git a/tests/unit/memory/memory_interface/test_interface_observations.py b/tests/unit/memory/memory_interface/test_interface_observations.py index 77abf077b7..e362342140 100644 --- a/tests/unit/memory/memory_interface/test_interface_observations.py +++ b/tests/unit/memory/memory_interface/test_interface_observations.py @@ -134,6 +134,48 @@ async def test_existing_observation_can_support_another_score(sqlite_instance: M assert (await sqlite_instance.get_scores_async(score_ids=[replay_score.id]))[0].observation_ids == [observation.id] +async def _stored_observation_and_score_async(memory: MemoryInterface) -> tuple[Observation, Score]: + piece = MessagePiece(role="assistant", original_value="response", conversation_id=str(uuid.uuid4())) + (await memory.add_message_pieces_to_memory_async(message_pieces=[piece])) + scorable = MessageScorable(message_piece_ids=(piece.id,)) + expectation = ScoringExpectation(objective="Judge the response") + observation = await _observation_async( + memory=memory, scorable=scorable, response_piece_id=piece.id, expectation=expectation + ) + score = Score( + score_value="true", + score_type="true_false", + scorable=scorable, + scored_expectation=expectation, + observation_ids=[observation.id], + ) + (await memory.add_scores_to_memory_async(scores=[score], observations=[observation])) + return observation, score + + +async def test_identical_new_observation_already_stored_is_linked_not_inserted(sqlite_instance: MemoryInterface): + observation, first_score = await _stored_observation_and_score_async(sqlite_instance) + second_score = first_score.model_copy(update={"id": uuid.uuid4()}) + + (await sqlite_instance.add_scores_to_memory_async(scores=[second_score], observations=[observation])) + + assert len(sqlite_instance._query_entries(ObservationEntry)) == 1 + stored = await sqlite_instance.get_scores_async(score_ids=[first_score.id, second_score.id]) + assert [score.observation_ids for score in stored] == [[observation.id], [observation.id]] + + +async def test_different_evidence_under_a_stored_observation_id_is_rejected(sqlite_instance: MemoryInterface): + observation, first_score = await _stored_observation_and_score_async(sqlite_instance) + changed = observation.model_copy(update={"metadata": {"note": "different"}}) + second_score = first_score.model_copy(update={"id": uuid.uuid4()}) + + with pytest.raises(ValueError, match="New observation IDs already exist"): + await sqlite_instance.add_scores_to_memory_async(scores=[second_score], observations=[changed]) + + assert await sqlite_instance.get_scores_async(score_ids=[second_score.id]) == [] + assert (await sqlite_instance.get_observations_async(observation_ids=[observation.id]))[0] == observation + + async def test_unreferenced_observation_is_not_persisted(sqlite_instance: MemoryInterface): piece = MessagePiece(role="assistant", original_value="response", conversation_id=str(uuid.uuid4())) (await sqlite_instance.add_message_pieces_to_memory_async(message_pieces=[piece])) diff --git a/tests/unit/score/test_docker_sandbox_client.py b/tests/unit/score/test_docker_sandbox_client.py index 4fc6e20aa1..3084aa50e6 100644 --- a/tests/unit/score/test_docker_sandbox_client.py +++ b/tests/unit/score/test_docker_sandbox_client.py @@ -57,6 +57,18 @@ def finish(code, oom=False, started=True): sys.exit(code) path(name, "pid").write_text(str(os.getpid())) save(status="running", code=0, oom=False, started="2026-10-14T00:00:00Z") + if program in ("CLOSEHANG", "CLOSEHANGRUNNING"): + # The client closes its output streams but does not exit. With CLOSEHANG the + # container has finished; with CLOSEHANGRUNNING it is still running. + print("done") + sys.stdout.flush() + if program == "CLOSEHANG": + save(status="exited", code=0, oom=False, started="2026-10-14T00:00:00Z") + for fd in (1, 2): + os.close(fd) + (state / "hang.pid").write_text(str(os.getpid())) + time.sleep(60) + sys.exit(0) if program == "DETACHFAIL": # The attached client fails while the container keeps running. sys.stderr.write("error waiting for container: unexpected EOF") @@ -216,6 +228,33 @@ async def test_timeout_kills_the_container_async(fake_docker: Path) -> None: assert len(_removed(fake_docker)) == 1 +def _alive(pid_file: Path) -> bool: + try: + os.kill(int(pid_file.read_text()), 0) + except ProcessLookupError: + return False + return True + + +async def test_a_client_that_closes_its_streams_but_never_exits_is_bounded_async(fake_docker: Path) -> None: + # The container finished, so its exit is real evidence even though the client hung. + run = await asyncio.wait_for(_run(_client(fake_docker, timeout_seconds=1.0), "CLOSEHANG"), timeout=10) + + event = run.events[0] + assert (event.exit_code, event.stdout, event.timed_out) == (0, "done\n", False) + assert not _alive(fake_docker.parent / "state" / "hang.pid") + assert len(_removed(fake_docker)) == 1 + + +async def test_a_hung_client_over_a_running_container_is_a_bounded_sandbox_error_async(fake_docker: Path) -> None: + # The container never finished, so the result is a sandbox error, not a program timeout. + with pytest.raises(SandboxExecutionError, match="still running"): + await asyncio.wait_for(_run(_client(fake_docker, timeout_seconds=1.0), "CLOSEHANGRUNNING"), timeout=10) + + assert not _alive(fake_docker.parent / "state" / "hang.pid") + assert len(_removed(fake_docker)) == 1 + + async def test_memory_kill_is_a_resource_limit_async(fake_docker: Path) -> None: event = (await _run(_client(fake_docker), "OOM")).events[0] diff --git a/tests/unit/score/test_program_execution_scorer.py b/tests/unit/score/test_program_execution_scorer.py index cf94396a2b..209bfa55b5 100644 --- a/tests/unit/score/test_program_execution_scorer.py +++ b/tests/unit/score/test_program_execution_scorer.py @@ -634,6 +634,85 @@ async def test_a_run_nobody_waits_for_is_cancelled_async() -> None: assert execution not in source._in_flight +class _SlowCleanupSandbox(_SlowSandbox): + """Holds its cleanup open after cancellation until released.""" + + def __init__(self) -> None: + super().__init__() + self.cleanup_started = asyncio.Event() + self.cleanup_release = asyncio.Event() + self.cleaned_up = False + + async def execute_async(self, *, program, runtime_profile, invocations): # type: ignore[no-untyped-def] + try: + return await super().execute_async( + program=program, runtime_profile=runtime_profile, invocations=invocations + ) + except asyncio.CancelledError: + self.cleanup_started.set() + await self.cleanup_release.wait() + self.cleaned_up = True + raise + + +async def test_cancelling_the_last_caller_waits_for_sandbox_cleanup_async() -> None: + sandbox = _SlowCleanupSandbox() + source = ProgramExecutionSource(sandbox=sandbox) + execution = _attempt() + + task = asyncio.create_task(source.acquire_async(scorable=execution)) + await asyncio.sleep(0) + task.cancel() + await asyncio.wait_for(sandbox.cleanup_started.wait(), timeout=5) + for _ in range(5): + await asyncio.sleep(0) + assert not task.done(), "cancellation finished before the sandbox cleaned up" + + sandbox.cleanup_release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert sandbox.cleaned_up + assert execution not in source._in_flight + + +async def test_concurrent_scoring_of_one_attempt_stores_its_observation_once_async( + sqlite_instance: SQLiteMemory, +) -> None: + sandbox = _SlowSandbox() + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=sandbox)) + execution = _attempt() + expectation = _expectation(ProgramCase(invocation=ProgramInvocation())) + # Hold both callers at the write until each has done every read, so neither sees the + # other's observation already stored. + arrived, both_arrived = 0, asyncio.Event() + write = sqlite_instance.add_scores_to_memory_async + + async def write_together(**kwargs): # type: ignore[no-untyped-def] + nonlocal arrived + arrived += 1 + if arrived == 2: + both_arrived.set() + await asyncio.wait_for(both_arrived.wait(), timeout=5) + return await write(**kwargs) + + sqlite_instance.add_scores_to_memory_async = write_together # type: ignore[method-assign] + + pending = asyncio.gather( + scorer.score_async(scorable=execution, expectation=expectation), + scorer.score_async(scorable=execution, expectation=expectation), + ) + await asyncio.sleep(0) + sandbox.release.set() + first, second = await pending + + assert len(sandbox.calls) == 1 + assert first[0].score_value == second[0].score_value == "true" + assert first[0].observation_ids == second[0].observation_ids + stored = await sqlite_instance.get_scores_async(score_ids=[first[0].id, second[0].id]) + assert len(stored) == 2 + assert len(await sqlite_instance.get_observations_async(observation_ids=first[0].observation_ids)) == 1 + + async def test_sandbox_warnings_are_kept_on_the_observation_async(sqlite_instance: SQLiteMemory) -> None: class _LeakySandbox(_FakeSandbox): async def execute_async(self, *, program, runtime_profile, invocations): # type: ignore[no-untyped-def] From c4d4c7695e7c381f6a4143d5b17f97eb3d5cd551 Mon Sep 17 00:00:00 2001 From: WatchTree-19 <119982314+WatchTree-19@users.noreply.github.com> Date: Sat, 17 Oct 2026 12:38:30 +0000 Subject: [PATCH 4/6] Address review: Docker logs, entrypoint, reuse after eviction, file work off the loop - DockerSandboxClient runs containers with --log-driver none, so the daemon keeps no copy of the output, and sets --entrypoint from the profile, so an image's own entrypoint never changes the command that runs. Building the run arguments and waiting on the attached client are now separate helpers. - Writing and removing the program directory runs in a worker thread; a cancelled caller waits for that work before its cancellation completes. - ProgramExecutionSource gives an attempt's reusable evidence an ID derived from the source and the attempt, and loads it from memory once the cache has let it go, so a stored attempt is never run again. Errors keep fresh IDs and are retried. --- .../observation/docker_sandbox_client.py | 230 ++++++++++++------ .../observation/program_execution_source.py | 37 ++- .../unit/score/test_docker_sandbox_client.py | 116 ++++++++- .../score/test_program_execution_scorer.py | 86 ++++++- 4 files changed, 392 insertions(+), 77 deletions(-) diff --git a/pyrit/score/observation/docker_sandbox_client.py b/pyrit/score/observation/docker_sandbox_client.py index 31c1317088..8e53942d2f 100644 --- a/pyrit/score/observation/docker_sandbox_client.py +++ b/pyrit/score/observation/docker_sandbox_client.py @@ -17,13 +17,13 @@ import uuid from dataclasses import dataclass from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ParamSpec, TypeVar from pyrit.models import ComponentIdentifier, ProgramExecutionEvent from pyrit.score.observation.sandbox_client import SandboxExecutionError, SandboxRun if TYPE_CHECKING: - from collections.abc import Mapping + from collections.abc import Callable, Mapping from pyrit.models import ProgramInvocation @@ -32,6 +32,32 @@ _DIGEST_PATTERN = re.compile(r"@(sha256:[0-9a-f]{64})$") _NEVER_STARTED = "0001-01-01T00:00:00Z" +_P = ParamSpec("_P") +_T = TypeVar("_T") + + +async def _complete_in_thread(function: Callable[_P, _T], /, *args: _P.args, **kwargs: _P.kwargs) -> _T: + """ + Run blocking file work in a worker thread, letting it finish if the caller is cancelled. + + A thread cannot be interrupted, so a cancelled caller waits for the work to end before its + cancellation completes. Cleanup that follows then sees the file system the work left. + + Returns: + _T: What the function returned. + + Raises: + CancelledError: Once the work has ended, if the caller was cancelled. + """ + work = asyncio.ensure_future(asyncio.to_thread(function, *args, **kwargs)) + try: + return await asyncio.shield(work) + except asyncio.CancelledError: + await asyncio.wait({work}) + if not work.cancelled() and work.exception() is not None: + logger.debug("File work ended with %r after its caller was cancelled", work.exception()) + raise + @dataclass(frozen=True) class DockerRuntimeProfile: @@ -74,6 +100,15 @@ class _ContainerExit: oom_killed: bool +@dataclass(frozen=True) +class _ClientOutput: + stdout: bytes + stdout_truncated: bool + stderr: bytes + stderr_truncated: bool + timed_out: bool + + class DockerSandboxClient: """ Run each invocation in its own disposable Docker container. @@ -82,7 +117,9 @@ class DockerSandboxClient: no privilege escalation, an unprivileged user, and caps on memory, CPU, and process count. The program is mounted read-only; a small ``/tmp`` is the only writable location. Each invocation gets a fresh container that is removed afterwards, so one run cannot leave state - for the next. + for the next. Docker keeps no log of the container's output, so output volume never grows + the daemon's files, and the profile's command replaces the image's entrypoint, so the + recorded command is the one that runs. Images are never pulled while scoring: a profile's image must already be present, or the attempt fails as a sandbox error. Pull the pinned images before a run. @@ -142,6 +179,7 @@ def __init__( self._cpus = cpus self._pids_limit = pids_limit self._max_output_bytes = max_output_bytes + self._temp_root = Path(tempfile.gettempdir()) def get_identifier(self) -> ComponentIdentifier: """ @@ -201,12 +239,12 @@ async def execute_async( profile = self._profiles.get(runtime_profile) if profile is None: raise ValueError(f"Runtime profile {runtime_profile!r} is not configured.") - program_dir = Path(tempfile.mkdtemp(prefix="pyrit-program-")) + program_dir = self._temp_root / f"pyrit-program-{uuid.uuid4().hex}" program_path = program_dir / profile.program_filename try: - program_path.write_text(program, encoding="utf-8") - program_path.chmod(0o444) - program_dir.chmod(0o755) + await _complete_in_thread( + self._write_program, program_dir=program_dir, program_path=program_path, program=program + ) events: list[ProgramExecutionEvent] = [] for index, invocation in enumerate(invocations): try: @@ -229,7 +267,15 @@ async def execute_async( return SandboxRun(events=tuple(events), reasons=reasons, warnings=(cleanup_warning,)) return SandboxRun(events=tuple(events)) finally: - self._remove_program_dir(program_dir=program_dir, program_path=program_path) + await _complete_in_thread(self._remove_program_dir, program_dir=program_dir, program_path=program_path) + + @staticmethod + def _write_program(*, program_dir: Path, program_path: Path, program: str) -> None: + """Create a private directory and write the program into it, then make both read-only for others.""" + program_dir.mkdir(mode=0o700) + program_path.write_text(program, encoding="utf-8") + program_path.chmod(0o444) + program_dir.chmod(0o755) @staticmethod def _remove_program_dir(*, program_dir: Path, program_path: Path) -> None: @@ -267,7 +313,67 @@ async def _run_invocation_async( """ name = f"pyrit-sandbox-{uuid.uuid4().hex}" resolved_command = (*profile.command, f"/program/{profile.program_filename}", *invocation.arguments) - argv = [ + argv = self._run_arguments( + name=name, profile=profile, program_dir=program_dir, resolved_command=resolved_command + ) + started = time.monotonic() + try: + process = await asyncio.create_subprocess_exec( + *argv, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + except OSError as error: + raise SandboxExecutionError(f"could not start {self._docker!r}: {error}") from error + cleanup_warning: str | None = None + try: + output = await self._collect_output_async( + process=process, name=name, stdin=invocation.stdin, deadline=started + self._timeout_seconds + ) + duration_ms = int((time.monotonic() - started) * 1000) + container = await self._inspect_async(name=name, cli_stderr=output.stderr, timed_out=output.timed_out) + finally: + if process.returncode is None: + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() + removal_status, removal_output = await self._control_async("rm", "--force", name) + if removal_status != 0: + detail = removal_output.strip()[:300] or "unknown Docker failure" + cleanup_warning = f"could not remove container {name}: {detail}" + logger.warning("Sandbox cleanup failed: %s", cleanup_warning) + timed_out = output.timed_out + event = ProgramExecutionEvent( + case_index=case_index, + runtime_profile=runtime_profile, + runtime_image_digest=profile.image_digest, + resolved_command=resolved_command, + exit_code=None if (timed_out or container.oom_killed) else container.exit_code, + stdout=output.stdout.decode("utf-8", errors="replace"), + stderr=output.stderr.decode("utf-8", errors="replace"), + stdout_truncated=output.stdout_truncated, + stderr_truncated=output.stderr_truncated, + timed_out=timed_out, + resource_limit_exceeded=container.oom_killed and not timed_out, + duration_ms=duration_ms, + ) + return event, cleanup_warning + + def _run_arguments( + self, *, name: str, profile: DockerRuntimeProfile, program_dir: Path, resolved_command: tuple[str, ...] + ) -> list[str]: + """ + Build the ``docker run`` command line for one invocation. + + The profile's first argument becomes the entrypoint, so an image's own entrypoint + never runs in front of ``resolved_command``. Daemon-side logging is off: the attached + client reads the output, and nothing is kept on the host beyond what it retains. + + Returns: + list[str]: The arguments, starting with the Docker client. + """ + return [ self._docker, "run", "--name", @@ -275,6 +381,8 @@ async def _run_invocation_async( "--interactive", "--pull", "never", + "--log-driver", + "none", "--network", "none", "--read-only", @@ -298,82 +406,60 @@ async def _run_invocation_async( "/tmp", "--mount", f"type=bind,source={os.fspath(program_dir)},target=/program,readonly", + "--entrypoint", + resolved_command[0], profile.image, - *resolved_command, + *resolved_command[1:], ] - started = time.monotonic() - try: - process = await asyncio.create_subprocess_exec( - *argv, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - except OSError as error: - raise SandboxExecutionError(f"could not start {self._docker!r}: {error}") from error + + async def _collect_output_async( + self, *, process: asyncio.subprocess.Process, name: str, stdin: str | None, deadline: float + ) -> _ClientOutput: + """ + Feed the input, read both streams, and wait for the attached client to exit, by one deadline. + + A container still running at the deadline is killed and the run is marked timed out. A + client whose streams closed but that has not exited by the deadline is ended; the + container's own state, inspected afterwards, then decides whether the program finished, + so a stalled client is not mistaken for a program that ran too long. + + Returns: + _ClientOutput: The retained output and whether the run timed out. + """ + readers = asyncio.gather( + self._read_capped_async(process.stdout), + self._read_capped_async(process.stderr), + self._feed_stdin_async(process.stdin, stdin), + ) timed_out = False - cleanup_warning: str | None = None - # One wall-clock deadline covers the whole invocation: reading the output and the - # attached client's exit. - deadline = started + self._timeout_seconds try: - readers = asyncio.gather( - self._read_capped_async(process.stdout), - self._read_capped_async(process.stderr), - self._feed_stdin_async(process.stdin, invocation.stdin), + (stdout, stdout_truncated), (stderr, stderr_truncated), _ = await asyncio.wait_for( + asyncio.shield(readers), timeout=max(deadline - time.monotonic(), 0.0) ) + except TimeoutError: + timed_out = True + await self._control_async("kill", name) try: - (stdout, stdout_truncated), (stderr, stderr_truncated), _ = await asyncio.wait_for( - asyncio.shield(readers), timeout=self._timeout_seconds - ) - except TimeoutError: - timed_out = True - await self._control_async("kill", name) - try: - await asyncio.wait_for(asyncio.shield(readers), timeout=self._CONTROL_TIMEOUT_SECONDS) - except TimeoutError: - # The container would not stop; ending the attached client closes its streams. - with contextlib.suppress(ProcessLookupError): - process.kill() - (stdout, stdout_truncated), (stderr, stderr_truncated), _ = await readers - try: - exit_budget = self._CONTROL_TIMEOUT_SECONDS if timed_out else max(deadline - time.monotonic(), 0.0) - await asyncio.wait_for(process.wait(), timeout=exit_budget) + await asyncio.wait_for(asyncio.shield(readers), timeout=self._CONTROL_TIMEOUT_SECONDS) except TimeoutError: - # The streams closed but the attached client did not exit in time. Ending the - # client keeps the invocation bounded; the container's own state, read below, - # decides whether the program finished, so a stalled client is not mistaken - # for a program that ran too long. + # The container would not stop; ending the attached client closes its streams. with contextlib.suppress(ProcessLookupError): process.kill() - await process.wait() - duration_ms = int((time.monotonic() - started) * 1000) - container = await self._inspect_async(name=name, cli_stderr=stderr, timed_out=timed_out) - finally: - if process.returncode is None: - with contextlib.suppress(ProcessLookupError): - process.kill() - await process.wait() - removal_status, removal_output = await self._control_async("rm", "--force", name) - if removal_status != 0: - detail = removal_output.strip()[:300] or "unknown Docker failure" - cleanup_warning = f"could not remove container {name}: {detail}" - logger.warning("Sandbox cleanup failed: %s", cleanup_warning) - event = ProgramExecutionEvent( - case_index=case_index, - runtime_profile=runtime_profile, - runtime_image_digest=profile.image_digest, - resolved_command=resolved_command, - exit_code=None if (timed_out or container.oom_killed) else container.exit_code, - stdout=stdout.decode("utf-8", errors="replace"), - stderr=stderr.decode("utf-8", errors="replace"), + (stdout, stdout_truncated), (stderr, stderr_truncated), _ = await readers + try: + exit_budget = self._CONTROL_TIMEOUT_SECONDS if timed_out else max(deadline - time.monotonic(), 0.0) + await asyncio.wait_for(process.wait(), timeout=exit_budget) + except TimeoutError: + with contextlib.suppress(ProcessLookupError): + process.kill() + await process.wait() + return _ClientOutput( + stdout=stdout, stdout_truncated=stdout_truncated, + stderr=stderr, stderr_truncated=stderr_truncated, timed_out=timed_out, - resource_limit_exceeded=container.oom_killed and not timed_out, - duration_ms=duration_ms, ) - return event, cleanup_warning async def _inspect_async(self, *, name: str, cli_stderr: bytes, timed_out: bool) -> _ContainerExit: """ diff --git a/pyrit/score/observation/program_execution_source.py b/pyrit/score/observation/program_execution_source.py index ef196b865f..b1bfafef8b 100644 --- a/pyrit/score/observation/program_execution_source.py +++ b/pyrit/score/observation/program_execution_source.py @@ -8,6 +8,7 @@ import asyncio import hashlib import logging +import uuid from collections import OrderedDict from dataclasses import dataclass from typing import TYPE_CHECKING @@ -31,6 +32,10 @@ logger = logging.getLogger(__name__) +# Namespace for the IDs of reusable execution evidence, derived from the source and the attempt. +_OBSERVATION_NAMESPACE = uuid.UUID("20acb8a7-a3ed-5df1-b292-83d20448a0fd") +_REUSABLE = (Acquisition.COMPLETE, Acquisition.PARTIAL) + class _NoProgramError(Exception): """The named evidence holds nothing that can be submitted as a program.""" @@ -66,9 +71,17 @@ class ProgramExecutionSource: Each stream keeps at most ``max_output_chars`` characters and records when it was cut. A completed or partial observation is reused when the same scorable is acquired again, and concurrent callers for one attempt share a single run, so one execution attempt is observed - once; a new ``execution_id`` runs the program again. Problems the sandbox reports after the - evidence was collected, such as a container it could not remove, are kept in the - observation's ``sandbox_warnings`` metadata. + once; a new ``execution_id`` runs the program again. The source holds recent observations, + and its evidence for an attempt always has the same ID, derived from the source + configuration and the attempt. Once memory stores that observation, as scoring does, a + later acquisition loads it instead of running the program, however many attempts the + source has seen since; every stored score therefore judges the attempt's one stored run. + Evidence that was never stored and has left the cache runs again under the same ID, and + memory refuses a second, different run of an attempt that already has stored evidence, + so contradicting results are never stored for one attempt. Acquisition reads memory, so + the central memory instance must be set. + Problems the sandbox reports after the evidence was collected, such as a container it + could not remove, are kept in the observation's ``sandbox_warnings`` metadata. """ _MAX_REUSED_OBSERVATIONS = 256 @@ -158,6 +171,12 @@ async def _acquire_uncached_async(self, *, scorable: ProgramExecutionScorable) - f"Runtime profile {scorable.runtime_profile!r} is not offered by the sandbox; " f"available profiles: {sorted(self._sandbox.get_runtime_profiles())}." ) + stored = await CentralMemory.get_memory_instance().get_observations_async( + observation_ids=[self._attempt_observation_id(scorable)] + ) + if stored: + self._remember(scorable=scorable, observation=stored[0]) + return stored[0] try: program = await self._resolve_program_async(scorable=scorable) except _NoProgramError as error: @@ -252,7 +271,10 @@ def _observation( program_sha256: str | None = None, reasons: tuple[str, ...] = (), ) -> Observation: + # Only reusable evidence takes the attempt's ID; an error is retried as a new observation. + observation_id = self._attempt_observation_id(scorable) if acquisition in _REUSABLE else uuid.uuid4() return Observation( + id=observation_id, source_identifier=self.get_identifier(), acquisition=acquisition, scorable=scorable, @@ -263,6 +285,15 @@ def _observation( ), ) + def _attempt_observation_id(self, scorable: ProgramExecutionScorable) -> uuid.UUID: + """ + Derive the ID of an attempt's reusable evidence from the source configuration and the attempt. + + Returns: + uuid.UUID: The same ID for the same source configuration and attempt. + """ + return uuid.uuid5(_OBSERVATION_NAMESPACE, f"{self.get_identifier().hash}:{scorable.model_dump_json()}") + def _remember(self, *, scorable: ProgramExecutionScorable, observation: Observation) -> None: self._observed[scorable] = observation while len(self._observed) > self._MAX_REUSED_OBSERVATIONS: diff --git a/tests/unit/score/test_docker_sandbox_client.py b/tests/unit/score/test_docker_sandbox_client.py index 3084aa50e6..d43f90aeb3 100644 --- a/tests/unit/score/test_docker_sandbox_client.py +++ b/tests/unit/score/test_docker_sandbox_client.py @@ -8,9 +8,12 @@ import asyncio import json import os +import pathlib import shutil import stat import sys +import tempfile +import threading from typing import TYPE_CHECKING from unittest.mock import patch @@ -44,6 +47,11 @@ def path(name, suffix): source = dict(part.split("=", 1) for part in mount.split(",") if "=" in part)["source"] image_index = next(i for i, a in enumerate(args) if "@sha256:" in a) inner = args[image_index + 1:] + # Docker runs the image's entrypoint before the command, unless --entrypoint replaces it. + entrypoint = os.environ.get("FAKE_IMAGE_ENTRYPOINT") + if "--entrypoint" in args: + entrypoint = args[args.index("--entrypoint") + 1] + path(name, "effective").write_text(json.dumps(([entrypoint] if entrypoint else []) + inner)) program_arg = next(a for a in inner if a.startswith("/program/")) program = (Path(source) / program_arg.rsplit("/", 1)[1]).read_text() user_args = inner[inner.index(program_arg) + 1:] @@ -201,16 +209,33 @@ async def test_containers_are_locked_down_and_removed_async(fake_docker: Path) - ("--memory-swap", "128m"), ("--cpus", "0.5"), ("--pids-limit", "8"), + ("--log-driver", "none"), ]: assert argv[argv.index(flag) + 1] == value assert "--read-only" in argv assert "--privileged" not in argv + assert "--log-opt" not in argv mount = argv[argv.index("--mount") + 1] assert mount.endswith(",target=/program,readonly") - assert argv[argv.index(_IMAGE) + 1 :] == ["python3", "-I", "/program/main.py"] + # The profile's first argument is the entrypoint, so the image's own entrypoint never runs. + assert argv[argv.index("--entrypoint") + 1] == "python3" + assert argv.index("--entrypoint") < argv.index(_IMAGE) + assert argv[argv.index(_IMAGE) + 1 :] == ["-I", "/program/main.py"] assert _removed(fake_docker) == [argv_files[0].stem] +async def test_an_image_entrypoint_does_not_change_the_command_that_runs_async( + fake_docker: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("FAKE_IMAGE_ENTRYPOINT", "python3") + + run = await _run(_client(fake_docker), "ECHO", ProgramInvocation(arguments=("x",))) + + effective = json.loads(next((fake_docker.parent / "state").glob("*.effective")).read_text()) + assert tuple(effective) == run.events[0].resolved_command == ("python3", "-I", "/program/main.py", "x") + assert run.events[0].exit_code == 0 + + async def test_program_exit_codes_are_observed_not_errors_async(fake_docker: Path) -> None: run = await _run(_client(fake_docker), "EXIT 125") @@ -402,3 +427,92 @@ def recording_rmtree(path, *args, **kwargs): # type: ignore[no-untyped-def] assert len(run.events) == 1 assert seen["writable"] is True assert not seen["dir"].exists() # type: ignore[union-attr] + + +async def _fake_run_invocation(**kwargs: object) -> tuple[ProgramExecutionEvent, None]: + event = ProgramExecutionEvent( + case_index=kwargs["case_index"], # type: ignore[arg-type] + runtime_profile="python-3.12", + runtime_image_digest="sha256:" + "c" * 64, + resolved_command=("python3",), + exit_code=0, + duration_ms=1, + ) + return event, None + + +@pytest.mark.parametrize("blocked", ["write", "remove"]) +async def test_file_work_runs_off_the_event_loop_async(blocked: str, tmp_path: Path) -> None: + client = _client(tmp_path / "docker") + release = threading.Event() + threads: list[int] = [] + real_write_text = pathlib.Path.write_text + real_rmtree = shutil.rmtree + + def wait_for_the_loop() -> None: + # Only the running event loop releases this, so file work on the loop thread deadlocks here. + threads.append(threading.get_ident()) + if not release.wait(2): + raise AssertionError("the event loop was blocked while file work ran") + + def blocking_write_text(self: pathlib.Path, *args: object, **kwargs: object) -> int: + if blocked == "write": + wait_for_the_loop() + return real_write_text(self, *args, **kwargs) # type: ignore[arg-type] + + def blocking_rmtree(path: pathlib.Path, *args: object, **kwargs: object) -> None: + if blocked == "remove": + wait_for_the_loop() + real_rmtree(path, *args, **kwargs) # type: ignore[arg-type] + + async def release_once_blocked() -> None: + while not threads: + await asyncio.sleep(0.005) + release.set() + + with ( + patch.object(client, "_run_invocation_async", side_effect=_fake_run_invocation), + patch.object(pathlib.Path, "write_text", blocking_write_text), + patch("pyrit.score.observation.docker_sandbox_client.shutil.rmtree", side_effect=blocking_rmtree), + ): + releaser = asyncio.create_task(release_once_blocked()) + run = await _run(client, "print(1)") + await asyncio.wait_for(releaser, timeout=1) + + assert len(run.events) == 1 + assert threads[0] != threading.get_ident() + + +async def test_cancelling_during_file_work_still_removes_the_program_async( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + root = tmp_path / "temp" + root.mkdir() + monkeypatch.setattr(tempfile, "tempdir", str(root)) + client = _client(tmp_path / "docker") + real_mkdir = pathlib.Path.mkdir + entered = threading.Event() + release = threading.Event() + + def gated_mkdir(self: pathlib.Path, *args: object, **kwargs: object) -> None: + entered.set() + release.wait(5) + real_mkdir(self, *args, **kwargs) # type: ignore[arg-type] + + with ( + patch.object(client, "_run_invocation_async", side_effect=_fake_run_invocation), + patch.object(pathlib.Path, "mkdir", gated_mkdir), + ): + task = asyncio.create_task(_run(client, "print(1)")) + for _ in range(200): + if entered.is_set(): + break + await asyncio.sleep(0.01) + assert entered.is_set() + task.cancel() + asyncio.get_running_loop().call_later(0.2, release.set) + with pytest.raises(asyncio.CancelledError): + await task + # The cancellation completed only after the program directory had been made and removed. + assert release.is_set() + assert list(root.iterdir()) == [] diff --git a/tests/unit/score/test_program_execution_scorer.py b/tests/unit/score/test_program_execution_scorer.py index 209bfa55b5..df1cca05df 100644 --- a/tests/unit/score/test_program_execution_scorer.py +++ b/tests/unit/score/test_program_execution_scorer.py @@ -661,7 +661,12 @@ async def test_cancelling_the_last_caller_waits_for_sandbox_cleanup_async() -> N execution = _attempt() task = asyncio.create_task(source.acquire_async(scorable=execution)) - await asyncio.sleep(0) + # The source checks memory before it runs the program, so wait for the run itself to start. + for _ in range(500): + if sandbox.calls: + break + await asyncio.sleep(0.01) + assert sandbox.calls task.cancel() await asyncio.wait_for(sandbox.cleanup_started.wait(), timeout=5) for _ in range(5): @@ -747,3 +752,82 @@ async def test_scoring_a_saved_score_anchor_runs_the_stored_program_async(sqlite observation = (await sqlite_instance.get_observations_async(observation_ids=again.observation_ids))[0] assert isinstance(observation.scorable, ProgramExecutionScorable) assert observation.scorable.program == saved.scorable + + +class _ChangingSandbox(_FakeSandbox): + """Prints something different on every run, so a second run of one attempt would show.""" + + async def execute_async(self, *, program, runtime_profile, invocations): # type: ignore[no-untyped-def] + self.calls.append((program, runtime_profile, invocations)) + stdout = "9" if len(self.calls) == 1 else f"run {len(self.calls)}" + return SandboxRun(events=tuple(_event(case_index=index, stdout=stdout) for index in range(len(invocations)))) + + +async def test_a_stored_attempt_is_reused_after_the_source_evicts_it_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _ChangingSandbox() + source = ProgramExecutionSource(sandbox=sandbox) + scorer = ProgramExecutionScorer(source=source) + expectation = _expectation(_case("3", stdout=Equals(value="9"))) + execution = ProgramExecutionScorable( + execution_id=uuid.uuid4(), + program=ContentScorable(value="x"), + runtime_profile=_PROFILE, + invocations=(ProgramInvocation(arguments=("3",)),), + ) + + first = (await scorer.score_async(scorable=execution, expectation=expectation))[0] + for _ in range(ProgramExecutionSource._MAX_REUSED_OBSERVATIONS): + await source.acquire_async(scorable=_attempt()) + assert execution not in source._observed + again = (await scorer.score_async(scorable=execution, expectation=expectation))[0] + + runs_of_execution = [call for call in sandbox.calls if call[2] == execution.invocations] + assert len(runs_of_execution) == 1 + assert (first.score_value, again.score_value) == ("true", "true") + assert first.observation_ids == again.observation_ids + stored = await sqlite_instance.get_scores_async(score_ids=[first.id, again.id]) + assert len(stored) == 2 + assert len(await sqlite_instance.get_observations_async(observation_ids=first.observation_ids)) == 1 + + +async def test_errors_are_retried_as_new_observations_not_the_attempts_evidence_async() -> None: + sandbox = _FakeSandbox() + source = ProgramExecutionSource(sandbox=sandbox) + execution = _attempt() + sandbox.error = SandboxExecutionError("no daemon") + + failed = await source.acquire_async(scorable=execution) + sandbox.error = None + ran = await source.acquire_async(scorable=execution) + + assert failed.acquisition is Acquisition.ERROR + assert ran.acquisition is Acquisition.COMPLETE + assert failed.id != ran.id + assert ran.id == ProgramExecutionSource(sandbox=sandbox)._attempt_observation_id(execution) + + +async def test_two_sources_cannot_store_two_runs_of_one_attempt_async(sqlite_instance: SQLiteMemory) -> None: + first_sandbox, second_sandbox = _SlowSandbox(), _SlowSandbox() + scorers = [_scorer(first_sandbox), _scorer(second_sandbox)] + execution = _attempt() + expectation = _expectation(ProgramCase(invocation=ProgramInvocation(), expected=ProgramResult())) + + pending = asyncio.gather( + *(scorer.score_async(scorable=execution, expectation=expectation) for scorer in scorers), + return_exceptions=True, + ) + for _ in range(500): + if first_sandbox.calls and second_sandbox.calls: + break + await asyncio.sleep(0.01) + first_sandbox.release.set() + second_sandbox.release.set() + results = await pending + + # Each source ran the attempt, so the later store is refused rather than keeping a second run. + errors = [result for result in results if isinstance(result, BaseException)] + stored = [result for result in results if not isinstance(result, BaseException)] + assert len(errors) == 1 + assert "already exist" in str(errors[0]) + assert len(stored) == 1 + assert len(await sqlite_instance.get_observations_async(observation_ids=stored[0][0].observation_ids)) == 1 From 8dc4e50ed708acfd0324a7fc1f54dd5ba5641f53 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Thu, 22 Oct 2026 22:13:34 -0700 Subject: [PATCH 5/6] Simplify and harden experimental program execution Extract reusable execution clients from scoring and make shared execution data types private. Remove implicit execution reuse, obsolete helpers, and duplicate test scaffolding. Preserve exact source bytes, bounded output, and cancellation-safe cleanup. Validate Docker resource support, reject image-declared volumes, remove owned anonymous volumes, and protect submitted source with a private host workspace. Update scoring replay, scoped regressions, real Docker integration, and paired documentation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0bdf701a-5fa6-44c0-ba7e-82a5beb6722d --- doc/code/framework.md | 24 +- .../scoring/7_program_execution_scorer.ipynb | 285 ++----- .../scoring/7_program_execution_scorer.py | 203 ++--- pyrit/_code_execution/__init__.py | 39 + pyrit/_code_execution/client.py | 65 ++ pyrit/_code_execution/docker_client.py | 749 ++++++++++++++++++ pyrit/memory/memory_interface.py | 40 +- pyrit/models/__init__.py | 4 - pyrit/models/code_execution.py | 89 +++ pyrit/models/score/__init__.py | 3 - pyrit/models/score/condition.py | 14 +- pyrit/models/score/observation.py | 28 +- pyrit/models/score/program.py | 78 -- pyrit/models/score/scorable.py | 19 +- pyrit/score/__init__.py | 7 - pyrit/score/observation/__init__.py | 7 - .../observation/docker_sandbox_client.py | 560 ------------- .../observation/program_execution_source.py | 144 +--- pyrit/score/observation/sandbox_client.py | 69 -- .../true_false/program_execution_scorer.py | 197 ++--- ...st_program_execution_scorer_integration.py | 299 +++++++ .../score/test_scorer_notebooks.py | 11 +- .../_code_execution/test_docker_client.py | 533 +++++++++++++ .../unit/common/test_lazy_package_imports.py | 61 ++ .../test_interface_observations.py | 42 - tests/unit/models/test_code_execution.py | 75 ++ tests/unit/models/test_program.py | 109 +-- tests/unit/models/test_scorable.py | 7 +- .../unit/score/test_docker_sandbox_client.py | 518 ------------ .../score/test_program_execution_scorer.py | 499 ++++-------- 30 files changed, 2417 insertions(+), 2361 deletions(-) create mode 100644 pyrit/_code_execution/__init__.py create mode 100644 pyrit/_code_execution/client.py create mode 100644 pyrit/_code_execution/docker_client.py create mode 100644 pyrit/models/code_execution.py delete mode 100644 pyrit/models/score/program.py delete mode 100644 pyrit/score/observation/docker_sandbox_client.py delete mode 100644 pyrit/score/observation/sandbox_client.py create mode 100644 tests/integration/score/test_program_execution_scorer_integration.py create mode 100644 tests/unit/_code_execution/test_docker_client.py create mode 100644 tests/unit/models/test_code_execution.py delete mode 100644 tests/unit/score/test_docker_sandbox_client.py diff --git a/doc/code/framework.md b/doc/code/framework.md index 3bbc794a57..ff9723dc28 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -290,10 +290,12 @@ If you are contributing to PyRIT, that work will most likely land in one of the root directory. `FileWriteScorer` evaluates its `ContentWritten` condition against that evidence. The caller owns workspace isolation; a snapshot does not identify its writer. - Program-execution sources run a program for a `ProgramExecutionScorable` through an injected - `SandboxClient`, such as `DockerSandboxClient`. `ProgramExecutionScorer` splits its `ProgramBehaves` - condition so the source receives only the runtime profile and invocations, never the expected - results. A program that fails is observed evidence; a sandbox that cannot run it is an - acquisition error, so the verdict is undetermined rather than false. + client from the experimental `pyrit._code_execution` module, and turn its results into scoring evidence. + `ProgramExecutionScorer` splits its `ProgramBehaves` condition so the source receives only the + runtime profile and invocations, never the expected results. A program that fails is observed + evidence; a sandbox that cannot run it is an acquisition error, so the verdict is undetermined + rather than false. + Each scoring call runs afresh; an observation identifies the run and supports explicit replay. - Observation lifecycle is optional. A scorer that implements `AttemptObserver`, whether it is a root scorer or wrapped inside one, is asked by the attack to start observing after setup, once its attempt ID is allocated and before anything is sent, and to stop when the attempt ends. A @@ -367,6 +369,20 @@ The below talks about responsibilities of most modules in the PyRIT library - Components that need credentials should go through these helpers rather than handling tokens themselves. +## Code Execution (Experimental) + +**Responsibility**: Connect to caller-configured execution backends and return typed invocation results. + +- The client API is private and experimental in `pyrit._code_execution`; it may change without a deprecation cycle. +- Targets, converters, dataset generators, and scorers can inject the same `SandboxClient` contract + from `pyrit._code_execution`. Execution inputs and results are private types in `pyrit.models.code_execution`. +- `runtime_profile_name` selects a client-owned configuration by name, not by supplying an image or command. + Clients advertise the available names through `get_runtime_profile_names()`. +- The caller supplies the backend and trusted runtimes. The optional `DockerSandboxClient` uses an + existing daemon and pre-pulled images; it creates and removes only its disposable containers. +- **Does not own**: interpreting results, scoring criteria, message or seed creation, memory + persistence, or infrastructure setup. Those remain with the consumer or caller. + ## [Exceptions](../contributing/9_exception) **Responsibility**: Define PyRIT's exception hierarchy and the retry behavior built around it. diff --git a/doc/code/scoring/7_program_execution_scorer.ipynb b/doc/code/scoring/7_program_execution_scorer.ipynb index 9089018cb3..2c79fe18c4 100644 --- a/doc/code/scoring/7_program_execution_scorer.ipynb +++ b/doc/code/scoring/7_program_execution_scorer.ipynb @@ -5,31 +5,21 @@ "id": "0", "metadata": {}, "source": [ - "# Program-execution scoring\n", + "# Program execution and scoring\n", "\n", - "`ProgramExecutionScorer` answers \"Does this program behave as intended for these inputs?\" It runs\n", - "the program in an isolated sandbox and judges what the program did, not what a response claims\n", - "it would do.\n", + "`pyrit._code_execution` is experimental; its API may change without a deprecation cycle.\n", + "Its shared data types in `pyrit.models.code_execution` are also private and experimental.\n", + "The caller provides a Docker daemon and this exact pre-pulled image:\n", + "`docker pull python@sha256:a6e34c598f2467ed0e9a8d349809fcd8b5c603269512df273a0bb1784edc11b1`.\n", + "PyRIT does not start Docker or pull images. This example needs no model credentials.\n", + "Docker must support cgroup memory, swap, CPU, and process limits. Runtime images must not\n", + "declare volumes. Unsupported backends and images are rejected before execution.\n", "\n", - "The `ProgramBehaves` condition lists cases. Each `ProgramCase` pairs a `ProgramInvocation`\n", - "(arguments and standard input) with a `ProgramResult` (exit code, and optional `stdout` and\n", - "`stderr` matchers). The scorer splits them before anything runs: the runtime profile and the\n", - "invocations go to the source, and the expected results stay with the scorer. The code that runs\n", - "the program never sees which output passes.\n", + "## Execute without scoring or memory\n", "\n", - "A `ProgramExecutionSource` asks an injected `SandboxClient` to run the program.\n", - "`DockerSandboxClient` runs each invocation in a fresh container with no network, a read-only\n", - "file system, no capabilities, an unprivileged user, and memory, CPU and process limits. A runtime\n", - "profile is a trusted name for a digest-pinned image and a command; the program cannot pick either.\n", - "Images are never pulled while scoring, so pull them first.\n", - "\n", - "A program that fails to compile, exits with the wrong code, prints the wrong output, times out, or\n", - "exceeds a resource limit is a `False` verdict: the sandbox observed it. A sandbox that cannot run\n", - "the program, or stops before every case ran, leaves the verdict undetermined, so broken\n", - "infrastructure never looks like a failing program.\n", - "\n", - "This walkthrough needs a Docker daemon and the pinned image below\n", - "(`docker pull python:3.12-slim`), and PyRIT's in-memory storage. It needs no model or credentials." + "A runtime profile name selects a trusted client configuration, not an arbitrary image.\n", + "The client runs code verbatim and returns bounded output and exit information.\n", + "Targets, converters, and generators can use this same interface." ] }, { @@ -37,254 +27,129 @@ "execution_count": null, "id": "1", "metadata": {}, - "outputs": [], - "source": [ - "import httpx\n", - "\n", - "from pyrit.executor.attack import AttackScoringConfig, PromptSendingAttack\n", - "from pyrit.memory import CentralMemory\n", - "from pyrit.models import (\n", - " ContentScorable,\n", - " Equals,\n", - " ProgramBehaves,\n", - " ProgramCase,\n", - " ProgramInvocation,\n", - " ProgramResult,\n", - " ScoringExpectation,\n", - ")\n", - "from pyrit.prompt_target import HTTPTarget\n", - "from pyrit.score import DockerRuntimeProfile, DockerSandboxClient, ProgramExecutionScorer, ProgramExecutionSource\n", - "from pyrit.setup import IN_MEMORY, initialize_pyrit_async\n", - "\n", - "await initialize_pyrit_async( # type: ignore\n", - " memory_db_type=IN_MEMORY,\n", - " load_defaults=False,\n", - " env_files=[],\n", - " silent=True,\n", - ")\n", - "memory = CentralMemory.get_memory_instance()\n", - "\n", - "python_312 = DockerRuntimeProfile(\n", - " image=\"python@sha256:a6e34c598f2467ed0e9a8d349809fcd8b5c603269512df273a0bb1784edc11b1\",\n", - " command=(\"python3\", \"-I\"),\n", - ")\n", - "sandbox = DockerSandboxClient(profiles={\"python-3.12\": python_312}, timeout_seconds=5)\n", - "scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=sandbox))\n", - "\n", - "squares = ScoringExpectation(\n", - " conditions=(\n", - " ProgramBehaves(\n", - " runtime_profile=\"python-3.12\",\n", - " cases=(\n", - " ProgramCase(\n", - " invocation=ProgramInvocation(arguments=(\"3\",)), expected=ProgramResult(stdout=Equals(value=\"9\"))\n", - " ),\n", - " ProgramCase(\n", - " invocation=ProgramInvocation(stdin=\"12\\n\"), expected=ProgramResult(stdout=Equals(value=\"144\"))\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - ")" - ] - }, - { - "cell_type": "markdown", - "id": "2", - "metadata": {}, - "source": [ - "## Score a program directly\n", - "\n", - "A `ContentScorable` holds the program text. The program is run verbatim, once per case." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "3", - "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "square: True - All 2 case(s) produced the expected result under python-3.12.\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "doubles: False - The program did not produce the expected result under python-3.12: case 0 output did not match; case 1 output did not match.\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "loops: False - The program did not produce the expected result under python-3.12: case 0 timed out; case 1 timed out.\n" + "Direct execution: 9\n" ] } ], "source": [ - "square = \"import sys\\nn = int(sys.argv[1]) if len(sys.argv) > 1 else int(sys.stdin.read())\\nprint(n * n)\\n\"\n", - "doubles = \"import sys\\nn = int(sys.argv[1]) if len(sys.argv) > 1 else int(sys.stdin.read())\\nprint(n * 2)\\n\"\n", - "loops = \"while True:\\n pass\\n\"\n", - "\n", - "for name, program in ((\"square\", square), (\"doubles\", doubles), (\"loops\", loops)):\n", - " score = (await scorer.score_async(scorable=ContentScorable(value=program), expectation=squares))[0] # type: ignore\n", - " print(f\"{name}: {score.get_value()} - {score.score_rationale}\")" + "from pyrit._code_execution import DockerRuntimeProfile, DockerSandboxClient\n", + "from pyrit.models.code_execution import _ProgramInvocation\n", + "\n", + "client = DockerSandboxClient(\n", + " profiles={\n", + " \"python\": DockerRuntimeProfile(\n", + " image=\"python@sha256:a6e34c598f2467ed0e9a8d349809fcd8b5c603269512df273a0bb1784edc11b1\",\n", + " command=(\"python3\", \"-I\"),\n", + " )\n", + " }\n", + ")\n", + "program = \"import sys\\nprint(int(sys.argv[1]) ** 2)\\n\"\n", + "invocation = _ProgramInvocation(arguments=(\"3\",))\n", + "assert \"python\" in client.get_runtime_profile_names()\n", + "run = await client.execute_async( # type: ignore\n", + " program=program, runtime_profile_name=\"python\", invocations=(invocation,)\n", + ")\n", + "assert run.events[0].exit_code == 0\n", + "assert run.events[0].stdout == \"9\\n\"\n", + "print(f\"Direct execution: {run.events[0].stdout.strip()}\")" ] }, { "cell_type": "markdown", - "id": "4", + "id": "2", "metadata": {}, "source": [ - "## A sandbox failure is not a program failure\n", + "## Score and save the evidence\n", "\n", - "Here the profile names an image that is not on the machine, so no container starts. The\n", - "observation is an acquisition error and the score is undetermined, not `False`." + "The source receives only code and invocation inputs; expected results stay with the scorer.\n", + "Every `score_async` call starts a fresh execution, identified by its saved observation ID.\n", + "Wrong output, unexpected exits, and enforced resource/time limits fail a case. Infrastructure\n", + "failures or insufficient evidence leave the verdict undetermined." ] }, { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "3", "metadata": {}, "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "Sandbox could not run program for execution d96c13b0-eb8f-4b1d-ad01-a7c2b160230d: the container did not start: docker: Error response from daemon: No such image: python@sha256:0000000000000000000000000000000000000000000000000000000000000000\n" - ] - }, { "name": "stdout", "output_type": "stream", "text": [ - "undetermined: The program's behavior under python-3.12 could not be observed (sandbox error: the container did not start: docker: Error response from daemon: No such image: python@sha256:0000000000000000000000000000000000000000000000000000000000000000).\n" + "Score: True - All 1 case(s) produced the expected result under python.\n" ] } ], "source": [ - "missing_image = DockerSandboxClient(\n", - " profiles={\"python-3.12\": DockerRuntimeProfile(image=\"python@sha256:\" + \"0\" * 64, command=(\"python3\", \"-I\"))}\n", + "from pyrit.memory import CentralMemory\n", + "from pyrit.models import ContentScorable, Equals, ProgramBehaves, ProgramCase, ProgramResult, ScoringExpectation\n", + "from pyrit.score import ProgramExecutionScorer, ProgramExecutionSource\n", + "from pyrit.setup import IN_MEMORY, initialize_pyrit_async\n", + "\n", + "await initialize_pyrit_async( # type: ignore\n", + " memory_db_type=IN_MEMORY, load_defaults=False, env_files=[], silent=True\n", ")\n", - "unavailable = (\n", - " await ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=missing_image)).score_async( # type: ignore\n", - " scorable=ContentScorable(value=square), expectation=squares\n", - " )\n", - ")[0]\n", - "print(f\"{unavailable.status.value}: {unavailable.score_rationale}\")" + "case = ProgramCase(invocation=invocation, expected=ProgramResult(stdout=Equals(value=\"9\")))\n", + "expectation = ScoringExpectation(conditions=(ProgramBehaves(runtime_profile_name=\"python\", cases=(case,)),))\n", + "scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=client))\n", + "score = (await scorer.score_async(scorable=ContentScorable(value=program), expectation=expectation))[0] # type: ignore\n", + "assert score.get_value() is True\n", + "print(f\"Score: {score.get_value()} - {score.score_rationale}\")" ] }, { "cell_type": "markdown", - "id": "6", + "id": "4", "metadata": {}, "source": [ - "## Re-judge saved evidence\n", + "## Replay without another execution\n", "\n", - "The observation keeps each invocation's exit state, a bounded copy of both streams, the resolved\n", - "command, and the image digest. Replay judges that evidence against new expected results without\n", - "running the program again. The invocations must match the ones that ran." + "Replay uses the saved outputs, not the Docker client. Expected results may change, but the\n", + "runtime profile name and invocation inputs must still match the saved run." ] }, { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "5", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "case 0: ('python3', '-I', '/program/main.py', '3') exit 0, stdout '9\\n'\n", - "case 1: ('python3', '-I', '/program/main.py') exit 0, stdout '144\\n'\n", - "Replayed against a stricter expectation: False - The program did not produce the expected result under python-3.12: case 1 output did not match.\n" + "Replay: False - The program did not produce the expected result under python: case 0 output did not match.\n" ] } ], "source": [ - "passing = (await scorer.score_async(scorable=ContentScorable(value=square), expectation=squares))[0] # type: ignore\n", - "saved = (await memory.get_observations_async(observation_ids=passing.observation_ids))[0]\n", - "for event in saved.payload.events:\n", - " print(f\"case {event.case_index}: {event.resolved_command} exit {event.exit_code}, stdout {event.stdout!r}\")\n", - "\n", - "stricter = ScoringExpectation(\n", - " conditions=(\n", - " ProgramBehaves(\n", - " runtime_profile=\"python-3.12\",\n", - " cases=(\n", - " ProgramCase(\n", - " invocation=ProgramInvocation(arguments=(\"3\",)), expected=ProgramResult(stdout=Equals(value=\"9\"))\n", - " ),\n", - " ProgramCase(\n", - " invocation=ProgramInvocation(stdin=\"12\\n\"), expected=ProgramResult(stdout=Equals(value=\"145\"))\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", + "memory = CentralMemory.get_memory_instance()\n", + "observation = (await memory.get_observations_async(observation_ids=score.observation_ids))[0]\n", + "different_case = ProgramCase(invocation=invocation, expected=ProgramResult(stdout=Equals(value=\"8\")))\n", + "different_expectation = ScoringExpectation(\n", + " conditions=(ProgramBehaves(runtime_profile_name=\"python\", cases=(different_case,)),)\n", ")\n", - "replayed = (await scorer.score_observation_async(observation=saved, expectation=stricter))[0] # type: ignore\n", - "print(f\"Replayed against a stricter expectation: {replayed.get_value()} - {replayed.score_rationale}\")" + "replayed = (\n", + " await scorer.score_observation_async(observation=observation, expectation=different_expectation) # type: ignore\n", + ")[0]\n", + "assert replayed.get_value() is False\n", + "assert replayed.observation_ids == score.observation_ids\n", + "print(f\"Replay: {replayed.get_value()} - {replayed.score_rationale}\")" ] }, { "cell_type": "markdown", - "id": "8", - "metadata": { - "lines_to_next_cell": 2 - }, - "source": [ - "## Score a generated program through an attack\n", - "\n", - "The attack scores its response as a program. This local target returns code as its whole response.\n", - "A response that wraps code in explanation is run as it is and fails its cases; the scorer does not\n", - "guess where the code starts." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "9", + "id": "6", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Target writes square -> outcome success\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Target writes doubles -> outcome failure\n" - ] - } - ], "source": [ - "def coder(*, program: str) -> HTTPTarget:\n", - " \"\"\"Return a local target that answers every prompt with the same program.\"\"\"\n", - " return HTTPTarget(\n", - " http_request=\"POST / HTTP/1.1\\nHost: coder.test\\n\\n{PROMPT}\",\n", - " transport=httpx.MockTransport(lambda request: httpx.Response(200, text=program)),\n", - " )\n", - "\n", - "\n", - "for name, program in ((\"square\", square), (\"doubles\", doubles)):\n", - " attack = PromptSendingAttack(\n", - " objective_target=coder(program=program),\n", - " attack_scoring_config=AttackScoringConfig(objective_scorer=scorer),\n", - " max_attempts_on_failure=0,\n", - " )\n", - " result = await attack.execute_async(objective=\"write a program that squares its input\", expectation=squares) # type: ignore\n", - " print(f\"Target writes {name} -> outcome {result.outcome.value}\")" + "Real execution and fault coverage is in `tests/integration/score/test_program_execution_scorer_integration.py`.\n", + "Opt in with `RUN_ALL_TESTS=true`; the tests require the daemon and pinned image above." ] } ], @@ -302,7 +167,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.16" + "version": "3.14.2" } }, "nbformat": 4, diff --git a/doc/code/scoring/7_program_execution_scorer.py b/doc/code/scoring/7_program_execution_scorer.py index 536aaf751b..12775a87ec 100644 --- a/doc/code/scoring/7_program_execution_scorer.py +++ b/doc/code/scoring/7_program_execution_scorer.py @@ -9,165 +9,88 @@ # --- # %% [markdown] -# # Program-execution scoring +# # Program execution and scoring # -# `ProgramExecutionScorer` answers "Does this program behave as intended for these inputs?" It runs -# the program in an isolated sandbox and judges what the program did, not what a response claims -# it would do. +# `pyrit._code_execution` is experimental; its API may change without a deprecation cycle. +# Its shared data types in `pyrit.models.code_execution` are also private and experimental. +# The caller provides a Docker daemon and this exact pre-pulled image: +# `docker pull python@sha256:a6e34c598f2467ed0e9a8d349809fcd8b5c603269512df273a0bb1784edc11b1`. +# PyRIT does not start Docker or pull images. This example needs no model credentials. +# Docker must support cgroup memory, swap, CPU, and process limits. Runtime images must not +# declare volumes. Unsupported backends and images are rejected before execution. # -# The `ProgramBehaves` condition lists cases. Each `ProgramCase` pairs a `ProgramInvocation` -# (arguments and standard input) with a `ProgramResult` (exit code, and optional `stdout` and -# `stderr` matchers). The scorer splits them before anything runs: the runtime profile and the -# invocations go to the source, and the expected results stay with the scorer. The code that runs -# the program never sees which output passes. +# ## Execute without scoring or memory # -# A `ProgramExecutionSource` asks an injected `SandboxClient` to run the program. -# `DockerSandboxClient` runs each invocation in a fresh container with no network, a read-only -# file system, no capabilities, an unprivileged user, and memory, CPU and process limits. A runtime -# profile is a trusted name for a digest-pinned image and a command; the program cannot pick either. -# Images are never pulled while scoring, so pull them first. -# -# A program that fails to compile, exits with the wrong code, prints the wrong output, times out, or -# exceeds a resource limit is a `False` verdict: the sandbox observed it. A sandbox that cannot run -# the program, or stops before every case ran, leaves the verdict undetermined, so broken -# infrastructure never looks like a failing program. -# -# This walkthrough needs a Docker daemon and the pinned image below -# (`docker pull python:3.12-slim`), and PyRIT's in-memory storage. It needs no model or credentials. +# A runtime profile name selects a trusted client configuration, not an arbitrary image. +# The client runs code verbatim and returns bounded output and exit information. +# Targets, converters, and generators can use this same interface. # %% -import httpx - -from pyrit.executor.attack import AttackScoringConfig, PromptSendingAttack -from pyrit.memory import CentralMemory -from pyrit.models import ( - ContentScorable, - Equals, - ProgramBehaves, - ProgramCase, - ProgramInvocation, - ProgramResult, - ScoringExpectation, +from pyrit._code_execution import DockerRuntimeProfile, DockerSandboxClient +from pyrit.models.code_execution import _ProgramInvocation + +client = DockerSandboxClient( + profiles={ + "python": DockerRuntimeProfile( + image="python@sha256:a6e34c598f2467ed0e9a8d349809fcd8b5c603269512df273a0bb1784edc11b1", + command=("python3", "-I"), + ) + } ) -from pyrit.prompt_target import HTTPTarget -from pyrit.score import DockerRuntimeProfile, DockerSandboxClient, ProgramExecutionScorer, ProgramExecutionSource -from pyrit.setup import IN_MEMORY, initialize_pyrit_async - -await initialize_pyrit_async( # type: ignore - memory_db_type=IN_MEMORY, - load_defaults=False, - env_files=[], - silent=True, -) -memory = CentralMemory.get_memory_instance() - -python_312 = DockerRuntimeProfile( - image="python@sha256:a6e34c598f2467ed0e9a8d349809fcd8b5c603269512df273a0bb1784edc11b1", - command=("python3", "-I"), -) -sandbox = DockerSandboxClient(profiles={"python-3.12": python_312}, timeout_seconds=5) -scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=sandbox)) - -squares = ScoringExpectation( - conditions=( - ProgramBehaves( - runtime_profile="python-3.12", - cases=( - ProgramCase( - invocation=ProgramInvocation(arguments=("3",)), expected=ProgramResult(stdout=Equals(value="9")) - ), - ProgramCase( - invocation=ProgramInvocation(stdin="12\n"), expected=ProgramResult(stdout=Equals(value="144")) - ), - ), - ), - ) +program = "import sys\nprint(int(sys.argv[1]) ** 2)\n" +invocation = _ProgramInvocation(arguments=("3",)) +assert "python" in client.get_runtime_profile_names() +run = await client.execute_async( # type: ignore + program=program, runtime_profile_name="python", invocations=(invocation,) ) +assert run.events[0].exit_code == 0 +assert run.events[0].stdout == "9\n" +print(f"Direct execution: {run.events[0].stdout.strip()}") # %% [markdown] -# ## Score a program directly +# ## Score and save the evidence # -# A `ContentScorable` holds the program text. The program is run verbatim, once per case. +# The source receives only code and invocation inputs; expected results stay with the scorer. +# Every `score_async` call starts a fresh execution, identified by its saved observation ID. +# Wrong output, unexpected exits, and enforced resource/time limits fail a case. Infrastructure +# failures or insufficient evidence leave the verdict undetermined. # %% -square = "import sys\nn = int(sys.argv[1]) if len(sys.argv) > 1 else int(sys.stdin.read())\nprint(n * n)\n" -doubles = "import sys\nn = int(sys.argv[1]) if len(sys.argv) > 1 else int(sys.stdin.read())\nprint(n * 2)\n" -loops = "while True:\n pass\n" - -for name, program in (("square", square), ("doubles", doubles), ("loops", loops)): - score = (await scorer.score_async(scorable=ContentScorable(value=program), expectation=squares))[0] # type: ignore - print(f"{name}: {score.get_value()} - {score.score_rationale}") - -# %% [markdown] -# ## A sandbox failure is not a program failure -# -# Here the profile names an image that is not on the machine, so no container starts. The -# observation is an acquisition error and the score is undetermined, not `False`. +from pyrit.memory import CentralMemory +from pyrit.models import ContentScorable, Equals, ProgramBehaves, ProgramCase, ProgramResult, ScoringExpectation +from pyrit.score import ProgramExecutionScorer, ProgramExecutionSource +from pyrit.setup import IN_MEMORY, initialize_pyrit_async -# %% -missing_image = DockerSandboxClient( - profiles={"python-3.12": DockerRuntimeProfile(image="python@sha256:" + "0" * 64, command=("python3", "-I"))} +await initialize_pyrit_async( # type: ignore + memory_db_type=IN_MEMORY, load_defaults=False, env_files=[], silent=True ) -unavailable = ( - await ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=missing_image)).score_async( # type: ignore - scorable=ContentScorable(value=square), expectation=squares - ) -)[0] -print(f"{unavailable.status.value}: {unavailable.score_rationale}") +case = ProgramCase(invocation=invocation, expected=ProgramResult(stdout=Equals(value="9"))) +expectation = ScoringExpectation(conditions=(ProgramBehaves(runtime_profile_name="python", cases=(case,)),)) +scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=client)) +score = (await scorer.score_async(scorable=ContentScorable(value=program), expectation=expectation))[0] # type: ignore +assert score.get_value() is True +print(f"Score: {score.get_value()} - {score.score_rationale}") # %% [markdown] -# ## Re-judge saved evidence +# ## Replay without another execution # -# The observation keeps each invocation's exit state, a bounded copy of both streams, the resolved -# command, and the image digest. Replay judges that evidence against new expected results without -# running the program again. The invocations must match the ones that ran. +# Replay uses the saved outputs, not the Docker client. Expected results may change, but the +# runtime profile name and invocation inputs must still match the saved run. # %% -passing = (await scorer.score_async(scorable=ContentScorable(value=square), expectation=squares))[0] # type: ignore -saved = (await memory.get_observations_async(observation_ids=passing.observation_ids))[0] -for event in saved.payload.events: - print(f"case {event.case_index}: {event.resolved_command} exit {event.exit_code}, stdout {event.stdout!r}") - -stricter = ScoringExpectation( - conditions=( - ProgramBehaves( - runtime_profile="python-3.12", - cases=( - ProgramCase( - invocation=ProgramInvocation(arguments=("3",)), expected=ProgramResult(stdout=Equals(value="9")) - ), - ProgramCase( - invocation=ProgramInvocation(stdin="12\n"), expected=ProgramResult(stdout=Equals(value="145")) - ), - ), - ), - ) +memory = CentralMemory.get_memory_instance() +observation = (await memory.get_observations_async(observation_ids=score.observation_ids))[0] +different_case = ProgramCase(invocation=invocation, expected=ProgramResult(stdout=Equals(value="8"))) +different_expectation = ScoringExpectation( + conditions=(ProgramBehaves(runtime_profile_name="python", cases=(different_case,)),) ) -replayed = (await scorer.score_observation_async(observation=saved, expectation=stricter))[0] # type: ignore -print(f"Replayed against a stricter expectation: {replayed.get_value()} - {replayed.score_rationale}") +replayed = ( + await scorer.score_observation_async(observation=observation, expectation=different_expectation) # type: ignore +)[0] +assert replayed.get_value() is False +assert replayed.observation_ids == score.observation_ids +print(f"Replay: {replayed.get_value()} - {replayed.score_rationale}") # %% [markdown] -# ## Score a generated program through an attack -# -# The attack scores its response as a program. This local target returns code as its whole response. -# A response that wraps code in explanation is run as it is and fails its cases; the scorer does not -# guess where the code starts. - - -# %% -def coder(*, program: str) -> HTTPTarget: - """Return a local target that answers every prompt with the same program.""" - return HTTPTarget( - http_request="POST / HTTP/1.1\nHost: coder.test\n\n{PROMPT}", - transport=httpx.MockTransport(lambda request: httpx.Response(200, text=program)), - ) - - -for name, program in (("square", square), ("doubles", doubles)): - attack = PromptSendingAttack( - objective_target=coder(program=program), - attack_scoring_config=AttackScoringConfig(objective_scorer=scorer), - max_attempts_on_failure=0, - ) - result = await attack.execute_async(objective="write a program that squares its input", expectation=squares) # type: ignore - print(f"Target writes {name} -> outcome {result.outcome.value}") +# Real execution and fault coverage is in `tests/integration/score/test_program_execution_scorer_integration.py`. +# Opt in with `RUN_ALL_TESTS=true`; the tests require the daemon and pinned image above. diff --git a/pyrit/_code_execution/__init__.py b/pyrit/_code_execution/__init__.py new file mode 100644 index 0000000000..84ed3c54a7 --- /dev/null +++ b/pyrit/_code_execution/__init__.py @@ -0,0 +1,39 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +# ruff: noqa: F401 + +""" +Experimental clients for caller-configured code-execution backends. + +This private package is an explicit opt-in. Its API may change without a deprecation cycle. +""" + +from typing import TYPE_CHECKING + +from pyrit.common.lazy_imports import get_lazy_dir, resolve_lazy_export + +if TYPE_CHECKING: + from pyrit._code_execution.client import SandboxClient, SandboxExecutionError + from pyrit._code_execution.docker_client import DockerRuntimeProfile, DockerSandboxClient + +_LAZY_EXPORTS: dict[str, str] = { + "SandboxClient": "pyrit._code_execution.client", + "SandboxExecutionError": "pyrit._code_execution.client", + "DockerRuntimeProfile": "pyrit._code_execution.docker_client", + "DockerSandboxClient": "pyrit._code_execution.docker_client", +} + +__all__ = list(_LAZY_EXPORTS) + + +def __getattr__(name: str) -> object: + return resolve_lazy_export( + name=name, + module_name=__name__, + module_globals=globals(), + exports=_LAZY_EXPORTS, + ) + + +def __dir__() -> list[str]: + return get_lazy_dir(module_globals=globals(), exports=_LAZY_EXPORTS) diff --git a/pyrit/_code_execution/client.py b/pyrit/_code_execution/client.py new file mode 100644 index 0000000000..ff1fdf51f8 --- /dev/null +++ b/pyrit/_code_execution/client.py @@ -0,0 +1,65 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""The experimental contract for connecting to an isolated code-execution backend.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +if TYPE_CHECKING: + from pyrit.models import ComponentIdentifier + from pyrit.models.code_execution import _ProgramInvocation, _RuntimeProfileName, _SandboxRun + + +class SandboxExecutionError(Exception): + """ + The sandbox could not run the program. + + This is an infrastructure failure, such as an unreachable service, a failed + authentication, or an environment that could not be provisioned. It is never raised for a + program that fails, times out, or exceeds a limit: those are observed behavior. + """ + + +@runtime_checkable +class SandboxClient(Protocol): + """ + Connect to a caller-configured backend that runs code in isolation. + + Each invocation runs in a fresh environment. The client requests execution and cleanup, + and waits for cleanup before a cancelled call returns. It does not set up the backend + service, build images, or retain environments between calls. + A runtime profile is a trusted name the client maps to its own configuration; the program + cannot choose an image or a command. The caller supplies the running backend and trusted + runtimes, and decides how to use the returned results. + """ + + def get_identifier(self) -> ComponentIdentifier: + """Return the nonsecret client configuration.""" + ... + + def get_runtime_profile_names(self) -> frozenset[_RuntimeProfileName]: + """Return the names that select this client's configured runtimes.""" + ... + + async def execute_async( + self, *, program: str, runtime_profile_name: _RuntimeProfileName, invocations: tuple[_ProgramInvocation, ...] + ) -> _SandboxRun: + """ + Run each invocation of the program once, in order, in a fresh isolated environment. + + Args: + program (str): The code to run, without extracting it from surrounding prose. + runtime_profile_name (_RuntimeProfileName): One of the names returned by + ``get_runtime_profile_names``. The client owns the configuration selected by this name. + invocations (tuple[_ProgramInvocation, ...]): A nonempty sequence of arguments and input. + + Returns: + _SandboxRun: The invocation results and any execution gaps or cleanup warnings. + + Raises: + ValueError: If the runtime profile is unsupported or there are no invocations. + SandboxExecutionError: If the sandbox could not run the program at all. + """ + ... diff --git a/pyrit/_code_execution/docker_client.py b/pyrit/_code_execution/docker_client.py new file mode 100644 index 0000000000..67ec96763b --- /dev/null +++ b/pyrit/_code_execution/docker_client.py @@ -0,0 +1,749 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Run programs in disposable, locked-down Docker containers.""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +import os +import re +import shutil +import stat +import tempfile +import time +import uuid +from dataclasses import dataclass, replace +from pathlib import Path +from typing import TYPE_CHECKING, TypeVar + +from pyrit._code_execution.client import SandboxExecutionError +from pyrit.common.task_utils import gather_with_cleanup_async +from pyrit.models import ComponentIdentifier +from pyrit.models.code_execution import _ProgramExecutionEvent, _SandboxRun + +if TYPE_CHECKING: + from collections.abc import Awaitable, Mapping + + from pyrit.models.code_execution import _ProgramInvocation, _RuntimeProfileName + +logger = logging.getLogger(__name__) + +_DIGEST_PATTERN = re.compile(r"@(sha256:[0-9a-f]{64})$") +_NEVER_STARTED = "0001-01-01T00:00:00Z" + +_T = TypeVar("_T") + + +async def _finish_work_async(work: Awaitable[_T]) -> _T: + """ + Drain owned file work or cleanup even if the caller is cancelled repeatedly. + + Returns: + _T: The completed work's result. + + Raises: + CancelledError: Once the work has ended, if the caller was cancelled. + """ + task = asyncio.ensure_future(work) + cancellation: asyncio.CancelledError | None = None + while not task.done(): + try: + await asyncio.shield(task) + except asyncio.CancelledError as error: + cancellation = cancellation or error + except Exception: + break + try: + result = task.result() + except Exception as error: + if cancellation is not None: + raise cancellation from error + raise + if cancellation is not None: + raise cancellation + return result + + +async def _close_stdin_async(stream: asyncio.StreamWriter) -> None: + """ + Close input without letting cancellation poison the writer's close notification. + + Raises: + CancelledError: After the buffered writes have been aborted and closure has finished. + """ + stream.close() + closing = asyncio.create_task(stream.wait_closed()) + try: + await asyncio.shield(closing) + except asyncio.CancelledError: + stream.transport.abort() + with contextlib.suppress(BrokenPipeError, ConnectionResetError): + await _finish_work_async(closing) + raise + except (BrokenPipeError, ConnectionResetError): + pass + + +@dataclass(frozen=True) +class DockerRuntimeProfile: + """ + Configure a trusted, digest-pinned image and its program interpreter. + + The command is followed by the program's path and the invocation's arguments. + """ + + image: str + command: tuple[str, ...] + program_filename: str = "main.py" + + def __post_init__(self) -> None: + """ + Validate the image pin and the program file name. + + Raises: + ValueError: If the image is not pinned by digest, the command is empty, or the + file name is not a plain name. + """ + if _DIGEST_PATTERN.search(self.image) is None: + raise ValueError(f"Runtime image {self.image!r} must be pinned by digest (name@sha256:<64 hex>).") + if not self.command or any(not part for part in self.command): + raise ValueError("A runtime profile command must be a nonempty argument list.") + if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]*", self.program_filename): + raise ValueError(f"Program file name {self.program_filename!r} must be a plain file name.") + + @property + def image_digest(self) -> str: + """The ``sha256:`` digest the image is pinned to; ``__post_init__`` guarantees its presence.""" + return self.image.rsplit("@", 1)[1] + + +@dataclass(frozen=True) +class _ContainerState: + status: str + exit_code: int | None + oom_killed: bool + + +@dataclass(frozen=True) +class _ClientOutput: + stdout: bytes + stdout_truncated: bool + stderr: bytes + stderr_truncated: bool + + +@dataclass(frozen=True) +class _InvocationResult: + output: _ClientOutput + container: _ContainerState + timed_out: bool = False + + +class DockerSandboxClient: + """ + Run each invocation in a disposable container on a caller-provided Docker daemon. + + Profiles require pre-pulled, pinned images with no declared volumes. The daemon must support + cgroup memory, swap, CPU, and process limits; unsupported backends are rejected before execution. + Containers have no network or privileges, run as an unprivileged user, and have memory, CPU, + and process limits. The root and program are read-only; ``/tmp`` is a bounded tmpfs mount. + The profile replaces the image's entrypoint. + Daemon logging is disabled, and each stream retains at most ``max_output_bytes`` bytes. + The program mount sits inside a private host directory that is never mounted. + + One deadline covers startup, output capture, and the attached client's exit. The container's + state distinguishes program termination from client failure. Cleanup removes each container + before another invocation starts. + + This is not a hardened multi-tenant sandbox. Use a host dedicated to evaluation. + """ + + _READ_CHUNK_BYTES = 1 << 16 + _CONTROL_TIMEOUT_SECONDS = 30.0 + + def __init__( + self, + *, + profiles: Mapping[_RuntimeProfileName, DockerRuntimeProfile], + docker_executable: str = "docker", + timeout_seconds: float = 10.0, + memory: str = "256m", + cpus: float = 1.0, + pids_limit: int = 64, + max_output_bytes: int = 64_000, + ) -> None: + """ + Initialize with the allowlisted runtime profiles and per-container limits. + + Args: + profiles (Mapping[_RuntimeProfileName, DockerRuntimeProfile]): Names and their trusted + runtime configurations. + docker_executable (str): The Docker command line client to call. + timeout_seconds (float): Wall-clock limit for each invocation. + memory (str): Memory limit for each container, in Docker's notation, such as ``256m``. + cpus (float): CPU limit for each container. + pids_limit (int): Most processes a container may hold. + max_output_bytes (int): Most bytes of each stream kept per invocation. + + Raises: + ValueError: If no profile is given, a profile name is invalid, or a limit is not positive. + """ + if not profiles: + raise ValueError("DockerSandboxClient requires at least one runtime profile.") + invalid = [name for name in profiles if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]{0,63}", name)] + if invalid: + raise ValueError(f"Invalid runtime profile names: {invalid}.") + if timeout_seconds <= 0 or cpus <= 0 or pids_limit < 1 or max_output_bytes < 1: + raise ValueError("timeout_seconds, cpus, pids_limit and max_output_bytes must be positive.") + if not re.fullmatch(r"[1-9][0-9]*[bkmg]?", memory): + raise ValueError(f"Memory limit {memory!r} must look like 256m or 1g.") + self._profiles = dict(profiles) + self._docker = docker_executable + self._timeout_seconds = timeout_seconds + self._memory = memory + self._cpus = cpus + self._pids_limit = pids_limit + self._max_output_bytes = max_output_bytes + self._temp_root = Path(tempfile.gettempdir()) + + def get_identifier(self) -> ComponentIdentifier: + """ + Return the profiles and limits. + + Returns: + ComponentIdentifier: The nonsecret client configuration. + """ + return ComponentIdentifier.of( + self, + params={ + "client_version": 3, + "profiles": { + name: { + "image": profile.image, + "command": list(profile.command), + "program_filename": profile.program_filename, + } + for name, profile in sorted(self._profiles.items()) + }, + "timeout_seconds": self._timeout_seconds, + "memory": self._memory, + "cpus": self._cpus, + "pids_limit": self._pids_limit, + "max_output_bytes": self._max_output_bytes, + }, + ) + + def get_runtime_profile_names(self) -> frozenset[_RuntimeProfileName]: + """ + Return the configured runtime profile names. + + Returns: + frozenset[_RuntimeProfileName]: Names that select the configured runtimes. + """ + return frozenset(self._profiles) + + async def execute_async( + self, *, program: str, runtime_profile_name: _RuntimeProfileName, invocations: tuple[_ProgramInvocation, ...] + ) -> _SandboxRun: + """ + Run each invocation in a fresh container, in order. + + Args: + program (str): The program text, written to a read-only file in the container. + runtime_profile_name (_RuntimeProfileName): A key from ``profiles``, also returned by + ``get_runtime_profile_names``. The key selects a trusted image and command. + invocations (tuple[_ProgramInvocation, ...]): The arguments and input of each run. + + Returns: + _SandboxRun: One event per invocation that ran. If the sandbox fails after some + invocations ran, the run holds those events and the reason it stopped. + + Raises: + ValueError: If the runtime profile is not configured or there are no invocations. + SandboxExecutionError: If the backend cannot enforce the limits, the image declares + volumes, or no invocation could be run. + """ + profile = self._profiles.get(runtime_profile_name) + if profile is None: + raise ValueError(f"Runtime profile name {runtime_profile_name!r} is not configured.") + if not invocations: + raise ValueError("At least one program invocation is required.") + await self._validate_backend_async(profile=profile) + workspace_dir = self._temp_root / f"pyrit-program-{uuid.uuid4().hex}" + program_dir = workspace_dir / "program" + program_path = program_dir / profile.program_filename + try: + await _finish_work_async( + asyncio.to_thread( + self._write_program, workspace_dir=workspace_dir, program_path=program_path, program=program + ) + ) + events: list[_ProgramExecutionEvent] = [] + for index, invocation in enumerate(invocations): + try: + event, cleanup_warning = await self._run_invocation_async( + invocation_index=index, + runtime_profile_name=runtime_profile_name, + profile=profile, + program_dir=program_dir, + invocation=invocation, + ) + except SandboxExecutionError as error: + if not events: + raise + return _SandboxRun(events=tuple(events), reasons=(f"sandbox error at invocation {index}: {error}",)) + events.append(event) + if cleanup_warning is not None: + # The environment may still be running, so no further invocation starts beside it. + remaining = len(invocations) - len(events) + reasons = (f"stopped after invocation {index}: {cleanup_warning}",) if remaining else () + return _SandboxRun(events=tuple(events), reasons=reasons, warnings=(cleanup_warning,)) + return _SandboxRun(events=tuple(events)) + finally: + await _finish_work_async( + asyncio.to_thread(self._remove_workspace, workspace_dir=workspace_dir, program_path=program_path) + ) + + async def _validate_backend_async(self, *, profile: DockerRuntimeProfile) -> None: + """ + Reject unsupported resource controls and image-declared writable volumes. + + Raises: + SandboxExecutionError: If the backend or image cannot meet the isolation requirements. + """ + status, output = await self._control_async( + "info", + "--format", + "{{.CgroupDriver}} {{.MemoryLimit}} {{.SwapLimit}} {{.CPUCfsPeriod}} {{.CPUCfsQuota}} {{.PidsLimit}}", + ) + if status != 0: + raise SandboxExecutionError(f"could not inspect Docker resource support: {output.strip()[:300]}") + fields = output.split() + if len(fields) != 6 or fields[0] not in ("cgroupfs", "systemd") or fields[1:] != ["true"] * 5: + raise SandboxExecutionError( + "Docker must support cgroup memory, swap, CPU, and process limits; " + f"check controller delegation on rootless hosts (reported: {output.strip()[:300]})." + ) + status, output = await self._control_async( + "image", "inspect", "--format", "{{json .Config.Volumes}}", profile.image + ) + if status != 0: + raise SandboxExecutionError(f"could not inspect the runtime image: {output.strip()[:300]}") + if output.strip() not in ("null", "{}"): + raise SandboxExecutionError( + "Runtime images must not declare volumes, which can bypass the bounded tmpfs storage policy." + ) + + @staticmethod + def _write_program(*, workspace_dir: Path, program_path: Path, program: str) -> None: + """Keep the host workspace private while making its inner program mount readable.""" + workspace_dir.mkdir(mode=0o700) + program_path.parent.mkdir(mode=0o700) + program_path.write_bytes(program.encode("utf-8")) + program_path.chmod(0o444) + program_path.parent.chmod(0o755) + + @staticmethod + def _remove_workspace(*, workspace_dir: Path, program_path: Path) -> None: + """Restore file write permission for Windows and remove the entire private workspace.""" + with contextlib.suppress(OSError): + program_path.chmod(stat.S_IREAD | stat.S_IWRITE) + shutil.rmtree(workspace_dir, ignore_errors=True) + if workspace_dir.exists(): + logger.warning("Could not remove the program workspace %s", workspace_dir) + + async def _run_invocation_async( + self, + *, + invocation_index: int, + runtime_profile_name: _RuntimeProfileName, + profile: DockerRuntimeProfile, + program_dir: Path, + invocation: _ProgramInvocation, + ) -> tuple[_ProgramExecutionEvent, str | None]: + """ + Run one invocation in a new container and remove the container afterwards. + + Returns: + tuple[_ProgramExecutionEvent, str | None]: What the invocation did, and a warning when + the container could not be removed. + + Raises: + SandboxExecutionError: If Docker cannot start or inspect the container, or the + container had not finished when the attached client exited. + """ + name = f"pyrit-sandbox-{uuid.uuid4().hex}" + resolved_command = (*profile.command, f"/program/{profile.program_filename}", *invocation.arguments) + argv = self._run_arguments( + name=name, profile=profile, program_dir=program_dir, resolved_command=resolved_command + ) + started = time.monotonic() + deadline = started + self._timeout_seconds + startup = asyncio.create_task(self._start_client_async(argv=argv)) + cleanup_warning: str | None = None + try: + try: + process = await asyncio.wait_for(asyncio.shield(startup), timeout=self._timeout_seconds) + except TimeoutError as error: + raise SandboxExecutionError("the Docker client did not finish starting before its deadline") from error + result = await self._wait_for_invocation_async( + process=process, name=name, stdin=invocation.stdin, deadline=deadline + ) + duration_ms = int((time.monotonic() - started) * 1000) + finally: + cleanup_warning = await _finish_work_async(self._cleanup_invocation_async(startup=startup, name=name)) + output, container = result.output, result.container + event = _ProgramExecutionEvent( + invocation_index=invocation_index, + runtime_profile_name=runtime_profile_name, + runtime_image_digest=profile.image_digest, + resolved_command=resolved_command, + exit_code=None if (result.timed_out or container.oom_killed) else container.exit_code, + stdout=output.stdout.decode("utf-8", errors="replace"), + stderr=output.stderr.decode("utf-8", errors="replace"), + stdout_truncated=output.stdout_truncated, + stderr_truncated=output.stderr_truncated, + timed_out=result.timed_out, + resource_limit_exceeded=container.oom_killed and not result.timed_out, + duration_ms=duration_ms, + ) + return event, cleanup_warning + + async def _start_client_async(self, *, argv: list[str]) -> asyncio.subprocess.Process: + """ + Start the attached Docker client for an invocation that already owns cleanup. + + Returns: + asyncio.subprocess.Process: The started client. + + Raises: + SandboxExecutionError: If the client cannot be started. + """ + try: + return await asyncio.create_subprocess_exec( + *argv, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + except OSError as error: + raise SandboxExecutionError(f"could not start {self._docker!r}: {error}") from error + + async def _cleanup_invocation_async( + self, *, startup: asyncio.Task[asyncio.subprocess.Process], name: str + ) -> str | None: + """ + Wait for startup before reaping the client and removing its container. + + Returns: + str | None: A warning if cleanup could not be confirmed. + """ + process: asyncio.subprocess.Process | None = None + try: + process = await startup + except SandboxExecutionError: + status, output = await self._control_async( + "container", "ls", "--all", "--filter", f"name=^/{name}$", "--format", "{{.Names}}" + ) + if status != 0: + warning = f"could not verify cleanup of container {name}: {output.strip()[:300]}" + logger.warning("Sandbox cleanup failed: %s", warning) + return warning + if not output.strip(): + return None + return await self._remove_container_async(process=process, name=name) + + async def _remove_container_async(self, *, process: asyncio.subprocess.Process | None, name: str) -> str | None: + """ + Reap the attached client and remove its container. + + Returns: + str | None: A warning if removal failed, otherwise None. + """ + warnings = [] + if process is not None: + try: + self._kill_client(process) + await process.communicate() + except OSError as error: + warnings.append(f"could not reap the Docker client for container {name}: {error}") + status, output = await self._control_async("rm", "--force", "--volumes", name) + if status != 0: + detail = output.strip()[:300] or "unknown Docker failure" + warnings.append(f"could not remove container {name}: {detail}") + if not warnings: + return None + warning = "; ".join(warnings) + logger.warning("Sandbox cleanup failed: %s", warning) + return warning + + def _run_arguments( + self, *, name: str, profile: DockerRuntimeProfile, program_dir: Path, resolved_command: tuple[str, ...] + ) -> list[str]: + """ + Build the ``docker run`` command line for one invocation. + + Returns: + list[str]: The arguments, starting with the Docker client. + """ + return [ + self._docker, + "run", + "--name", + name, + "--interactive", + "--pull", + "never", + "--log-driver", + "none", + "--network", + "none", + "--read-only", + "--cap-drop", + "ALL", + "--security-opt", + "no-new-privileges", + "--user", + "65534:65534", + "--memory", + self._memory, + "--memory-swap", + self._memory, + "--cpus", + str(self._cpus), + "--pids-limit", + str(self._pids_limit), + "--tmpfs", + "/tmp:rw,nosuid,nodev,size=16m", + "--workdir", + "/tmp", + "--mount", + f"type=bind,source={os.fspath(program_dir)},target=/program,readonly", + "--entrypoint", + resolved_command[0], + profile.image, + *resolved_command[1:], + ] + + async def _wait_for_invocation_async( + self, *, process: asyncio.subprocess.Process, name: str, stdin: str | None, deadline: float + ) -> _InvocationResult: + """ + Apply one deadline and distinguish a client failure from the program's termination. + + Returns: + _InvocationResult: The captured output and the container's termination facts. + + Raises: + SandboxExecutionError: If capture failed or the client and container ended inconsistently. + """ + capture = asyncio.create_task(self._collect_output_async(process=process, stdin=stdin)) + try: + try: + output = await asyncio.wait_for(asyncio.shield(capture), timeout=max(deadline - time.monotonic(), 0.0)) + except TimeoutError: + return await self._finish_timed_out_capture_async(process=process, name=name, capture=capture) + client_stopped = False + if process.returncode is None: + try: + await asyncio.wait_for(process.wait(), timeout=max(deadline - time.monotonic(), 0.0)) + except TimeoutError: + self._kill_client(process) + await process.wait() + client_stopped = True + output = replace(output, stdout_truncated=True, stderr_truncated=True) + logger.warning("Docker client for %s stalled after EOF; output completeness is unknown", name) + container = await self._inspect_async(name=name, cli_stderr=output.stderr) + if container.exit_code is None: + raise SandboxExecutionError( + f"the Docker client exited while the container was still {container.status}, " + "so the program did not finish" + ) + if not client_stopped and process.returncode != container.exit_code: + raise SandboxExecutionError( + f"the Docker client exited with {process.returncode}, but the container exited with " + f"{container.exit_code}; output capture may be incomplete" + ) + return _InvocationResult(output=output, container=container) + finally: + if not capture.done(): + capture.cancel() + await _finish_work_async(asyncio.gather(capture, return_exceptions=True)) + + async def _finish_timed_out_capture_async( + self, + *, + process: asyncio.subprocess.Process, + name: str, + capture: asyncio.Task[_ClientOutput], + ) -> _InvocationResult: + """ + Enforce the program deadline only when the container is still running. + + Returns: + _InvocationResult: A timed-out invocation with its retained output. + + Raises: + SandboxExecutionError: If the deadline expired because capture stalled after execution. + """ + container = await self._inspect_async(name=name, cli_stderr=b"") + if container.status != "running": + raise SandboxExecutionError( + f"output capture did not finish before the deadline, but the container was {container.status}; " + "a program timeout could not be established" + ) + status, detail = await self._control_async("kill", name) + if status != 0: + logger.warning("Could not kill timed-out container %s: %s", name, detail.strip()[:300]) + interrupted = not capture.done() + self._kill_client(process) + try: + output = await asyncio.wait_for(asyncio.shield(capture), timeout=self._CONTROL_TIMEOUT_SECONDS) + except TimeoutError as error: + raise SandboxExecutionError("output capture did not finish after stopping the Docker client") from error + if interrupted: + output = replace(output, stdout_truncated=True, stderr_truncated=True) + await process.wait() + container = await self._inspect_async(name=name, cli_stderr=output.stderr) + return _InvocationResult(output=output, container=container, timed_out=True) + + async def _collect_output_async(self, *, process: asyncio.subprocess.Process, stdin: str | None) -> _ClientOutput: + """ + Feed standard input and capture bounded output without managing the container. + + Returns: + _ClientOutput: The retained streams and their truncation markers. + + Raises: + SandboxExecutionError: If the client streams cannot be read. + """ + feeder = asyncio.create_task(self._feed_stdin_async(stream=process.stdin, data=stdin)) + try: + (stdout, stdout_truncated), (stderr, stderr_truncated) = await gather_with_cleanup_async( + (self._read_capped_async(process.stdout), self._read_capped_async(process.stderr)) + ) + await feeder + return _ClientOutput( + stdout=stdout, + stdout_truncated=stdout_truncated, + stderr=stderr, + stderr_truncated=stderr_truncated, + ) + except OSError as error: + raise SandboxExecutionError(f"could not capture Docker client output: {error}") from error + finally: + if not feeder.done(): + feeder.cancel() + await _finish_work_async(asyncio.gather(feeder, return_exceptions=True)) + + async def _inspect_async(self, *, name: str, cli_stderr: bytes) -> _ContainerState: + """ + Read the container's state without inferring anything from the attached client. + + Returns: + _ContainerState: The state, with an exit code only for a finished container. + + Raises: + SandboxExecutionError: If the container does not exist, never started, or has invalid state. + """ + returncode, output = await self._control_async( + "inspect", + "--format", + "{{.State.Status}} {{.State.ExitCode}} {{.State.OOMKilled}} {{.State.StartedAt}}", + name, + ) + fields = output.split() + if returncode != 0 or len(fields) != 4 or fields[3] == _NEVER_STARTED: + lines = [line.strip() for line in cli_stderr.decode("utf-8", errors="replace").splitlines() if line.strip()] + detail = lines[0][:500] if lines else output.strip()[:500] + raise SandboxExecutionError(f"the container did not start: {detail or 'unknown Docker failure'}") + try: + exit_code = int(fields[1]) + except ValueError as error: + raise SandboxExecutionError(f"the container reported an invalid exit code: {fields[1]!r}") from error + if fields[2] not in ("true", "false"): + raise SandboxExecutionError(f"the container reported an invalid memory-kill state: {fields[2]!r}") + return _ContainerState( + status=fields[0], exit_code=exit_code if fields[0] == "exited" else None, oom_killed=fields[2] == "true" + ) + + async def _control_async(self, *arguments: str) -> tuple[int, str]: + """ + Run a short Docker management command, bounded in time. + + Returns: + tuple[int, str]: The exit status and combined output; status -1 when it failed to run. + """ + try: + process = await asyncio.create_subprocess_exec( + self._docker, + *arguments, + stdin=asyncio.subprocess.DEVNULL, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + ) + except OSError as error: + logger.warning("Docker %s failed to start: %s", arguments[0], error) + return -1, str(error) + try: + output, _ = await asyncio.wait_for(process.communicate(), timeout=self._CONTROL_TIMEOUT_SECONDS) + except TimeoutError: + logger.warning("Docker %s timed out", arguments[0]) + return -1, "timed out" + finally: + # Runs on timeout and on cancellation alike, so no management command outlives its caller. + if process.returncode is None: + with contextlib.suppress(ProcessLookupError): + process.kill() + await _finish_work_async(process.communicate()) + return process.returncode if process.returncode is not None else -1, output.decode("utf-8", errors="replace") + + async def _read_capped_async(self, stream: asyncio.StreamReader | None) -> tuple[bytes, bool]: + """ + Read a stream to its end, keeping at most the configured number of bytes. + + Returns: + tuple[bytes, bool]: The retained bytes and whether more were discarded. + """ + if stream is None: + return b"", False + kept = bytearray() + truncated = False + while chunk := await stream.read(self._READ_CHUNK_BYTES): + room = self._max_output_bytes - len(kept) + if len(chunk) > room: + truncated = True + kept.extend(chunk[: max(room, 0)]) + return bytes(kept), truncated + + @staticmethod + async def _feed_stdin_async(*, stream: asyncio.StreamWriter | None, data: str | None) -> None: + """ + Write the invocation's input, then close standard input so the program sees its end. + + Raises: + CancelledError: After input has been aborted and its writer closed. + """ + if stream is None: + return + try: + if data: + stream.write(data.encode("utf-8")) + await stream.drain() + except (BrokenPipeError, ConnectionResetError): + pass + except asyncio.CancelledError: + stream.transport.abort() + raise + finally: + await _close_stdin_async(stream) + + @staticmethod + def _kill_client(process: asyncio.subprocess.Process) -> None: + """Stop only this invocation's client, leaving container removal to its owner.""" + if process.returncode is None: + with contextlib.suppress(ProcessLookupError): + process.kill() diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 3bda1faa14..bdb3456e1e 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -2370,19 +2370,15 @@ def _add_scores_to_memory( SQLAlchemyError: If the score or identifier rows cannot be persisted. """ observations = [Observation.model_validate(observation.model_dump()) for observation in observations] - _, referenced_ids = self._validate_score_inputs(scores=scores, observations=observations) + new_ids, referenced_ids = self._validate_score_inputs(scores=scores, observations=observations) with closing(self._get_session()) as session: try: _begin_sqlite_write(session) - already_stored = self._validate_observation_links( + self._validate_observation_links( session=session, - new_observations={str(observation.id): observation for observation in observations}, + new_observation_ids=new_ids, referenced_observation_ids=referenced_ids, ) - # An identical observation another call stored first is linked, not inserted again. - observations = [ - observation for observation in observations if str(observation.id) not in already_stored - ] canonical_scores = self._resolve_score_message_anchors(session=session, scores=scores) content_entries, persisted_scores, persisted_observations = self._prepare_score_anchors( scores=canonical_scores, @@ -2594,42 +2590,24 @@ def _validate_observation_links( cls, *, session: Session, - new_observations: Mapping[str, Observation], + new_observation_ids: set[str], referenced_observation_ids: set[str], - ) -> set[str]: + ) -> None: """ Reject conflicting new IDs and links to observations that do not exist. - A new observation whose ID is already stored is accepted only when the stored row is - the same observation, as when concurrent scoring calls share one acquisition. The - check runs under the observation locks, so the first writer inserts and the others link. - - Returns: - set[str]: IDs of new observations already stored unchanged, to link rather than insert. - Raises: - ValueError: If a new ID holds different evidence or a referenced observation does not exist. + ValueError: If a new ID conflicts or a referenced observation does not exist. """ - new_observation_ids = set(new_observations) existing = _lock_observations( session=session, observation_ids=sorted(new_observation_ids | referenced_observation_ids) ) - already_stored = new_observation_ids & existing - if already_stored: - stored: dict[str, Observation] = {} - ids = sorted(already_stored) - for start in range(0, len(ids), cls._MAX_BIND_VARS): - entries = session.scalars( - select(ObservationEntry).where(ObservationEntry.id.in_(ids[start : start + cls._MAX_BIND_VARS])) - ) - stored.update({str(entry.id): entry.get_observation() for entry in entries}) - conflicts = sorted(i for i in already_stored if stored.get(i) != new_observations[i]) - if conflicts: - raise ValueError(f"New observation IDs already exist: {conflicts}.") + conflicts = new_observation_ids & existing + if conflicts: + raise ValueError(f"New observation IDs already exist: {sorted(conflicts)}.") missing = referenced_observation_ids - new_observation_ids - existing if missing: raise ValueError(f"Score references observations not found in memory: {sorted(missing)}.") - return already_stored @classmethod def _validate_observation_evidence( diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index 9852c8f1b8..dbac7de10f 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -205,10 +205,8 @@ OutputMatches, ProgramBehaves, ProgramCase, - ProgramExecutionEvent, ProgramExecutionObservationPayload, ProgramExecutionScorable, - ProgramInvocation, ProgramResult, Regex, Scorable, @@ -288,10 +286,8 @@ "OutputMatches": "pyrit.models.score", "ProgramBehaves": "pyrit.models.score", "ProgramCase": "pyrit.models.score", - "ProgramExecutionEvent": "pyrit.models.score", "ProgramExecutionObservationPayload": "pyrit.models.score", "ProgramExecutionScorable": "pyrit.models.score", - "ProgramInvocation": "pyrit.models.score", "ProgramResult": "pyrit.models.score", "ConversationScorable": "pyrit.models.score", "ConversationObservationPayload": "pyrit.models.score", diff --git a/pyrit/models/code_execution.py b/pyrit/models/code_execution.py new file mode 100644 index 0000000000..10a86ca0e2 --- /dev/null +++ b/pyrit/models/code_execution.py @@ -0,0 +1,89 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Private shared data for experimental code execution; these contracts may change.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Annotated + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +#: A trusted runtime profile name, such as ``python-3.12-restricted``. It names a sandbox +#: configuration the caller allowlisted; it is never an image name or a command. +_RuntimeProfileName = Annotated[str, Field(min_length=1, max_length=64, pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*$")] + + +class _ProgramInvocation(BaseModel): + """The arguments and standard input for one invocation of a program.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + arguments: tuple[Annotated[str, Field(pattern=r"^[^\x00]*$")], ...] = () + stdin: str | None = None + + +class _ProgramExecutionEvent(BaseModel): + """ + What one invocation did when the sandbox ran it. + + An event records an outcome, not a guarantee that the next execution will produce it. + A program can read the clock, use randomness, or depend on its runtime. + + Each output stream records whether its retained text is incomplete, including interrupted capture. + ``runtime_image_digest`` and ``resolved_command`` are optional provenance: a backend that cannot + expose them leaves them unset. + ``exit_code`` is ``None`` when the program did not exit on its own, for example because the + sandbox stopped it at its time or resource limit. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + invocation_index: int = Field(ge=0) + runtime_profile_name: _RuntimeProfileName + runtime_image_digest: str | None = Field(default=None, min_length=1) + resolved_command: tuple[str, ...] | None = Field(default=None, min_length=1) + exit_code: int | None + stdout: str = "" + stderr: str = "" + stdout_truncated: bool = False + stderr_truncated: bool = False + timed_out: bool = False + resource_limit_exceeded: bool = False + duration_ms: int = Field(ge=0) + + @model_validator(mode="after") + def _validate_exit_state(self) -> _ProgramExecutionEvent: + """ + Keep the exit state consistent with how the run ended. + + Returns: + _ProgramExecutionEvent: The validated event. + + Raises: + ValueError: If a run the sandbox stopped also reports an exit code of its own. + """ + if self.timed_out and self.exit_code is not None: + raise ValueError("A timed-out invocation did not exit on its own and has no exit code.") + return self + + @property + def output_truncated(self) -> bool: + """Whether either output stream could not be retained in full.""" + return self.stdout_truncated or self.stderr_truncated + + +@dataclass(frozen=True) +class _SandboxRun: + """ + Results from one request to execute a program's invocations. + + ``events`` holds one event per invocation that ran, in invocation order. ``reasons`` explains + why any invocations did not run. ``warnings`` reports cleanup problems that do not invalidate + the results already collected. + """ + + events: tuple[_ProgramExecutionEvent, ...] = () + reasons: tuple[str, ...] = () + warnings: tuple[str, ...] = () diff --git a/pyrit/models/score/__init__.py b/pyrit/models/score/__init__.py index 6d82a06631..44baf4d0e0 100644 --- a/pyrit/models/score/__init__.py +++ b/pyrit/models/score/__init__.py @@ -42,7 +42,6 @@ SurfaceObservationPayload, ToolEventsObservationPayload, ) - from pyrit.models.score.program import ProgramExecutionEvent, ProgramInvocation from pyrit.models.score.scorable import ( ContentEntryScorable, ContentScorable, @@ -77,10 +76,8 @@ _LAZY_EXPORTS: dict[str, str] = { "ProgramBehaves": "pyrit.models.score.condition", "ProgramCase": "pyrit.models.score.condition", - "ProgramExecutionEvent": "pyrit.models.score.program", "ProgramExecutionObservationPayload": "pyrit.models.score.observation", "ProgramExecutionScorable": "pyrit.models.score.scorable", - "ProgramInvocation": "pyrit.models.score.program", "ProgramResult": "pyrit.models.score.condition", "ConversationScorable": "pyrit.models.score.scorable", "ConversationObservationPayload": "pyrit.models.score.observation", diff --git a/pyrit/models/score/condition.py b/pyrit/models/score/condition.py index 196aa01e74..8ca274ed95 100644 --- a/pyrit/models/score/condition.py +++ b/pyrit/models/score/condition.py @@ -8,11 +8,11 @@ from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, SerializeAsAny, TypeAdapter, model_validator -from pyrit.models.score._trace_validation import ToolName # noqa: TC001 (runtime-required by Pydantic) -from pyrit.models.score.program import ( # noqa: TC001 (runtime-required by Pydantic) - ProgramInvocation, - RuntimeProfileName, +from pyrit.models.code_execution import ( # noqa: TC001 (runtime-required by Pydantic) + _ProgramInvocation, + _RuntimeProfileName, ) +from pyrit.models.score._trace_validation import ToolName # noqa: TC001 (runtime-required by Pydantic) from pyrit.models.score.surface import SurfaceMatch # noqa: TC001 (runtime-required by Pydantic) from pyrit.models.score.text_matcher import TextMatcher # noqa: TC001 (runtime-required by Pydantic) @@ -264,7 +264,7 @@ class ProgramCase(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") - invocation: ProgramInvocation = Field(default_factory=ProgramInvocation) + invocation: _ProgramInvocation = Field(default_factory=_ProgramInvocation) expected: ProgramResult = Field(default_factory=ProgramResult) @@ -279,7 +279,7 @@ class ProgramBehaves(Condition): """ condition_type: Literal["program_behaves"] = "program_behaves" - runtime_profile: RuntimeProfileName + runtime_profile_name: _RuntimeProfileName cases: tuple[ProgramCase, ...] @model_validator(mode="after") @@ -298,7 +298,7 @@ def _validate_cases(self) -> ProgramBehaves: return self @property - def invocations(self) -> tuple[ProgramInvocation, ...]: + def invocations(self) -> tuple[_ProgramInvocation, ...]: """The stimuli of every case, in case order, without their expected results.""" return tuple(case.invocation for case in self.cases) diff --git a/pyrit/models/score/observation.py b/pyrit/models/score/observation.py index 0cde1a803e..7a4034a58d 100644 --- a/pyrit/models/score/observation.py +++ b/pyrit/models/score/observation.py @@ -12,11 +12,11 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, field_validator, model_validator +from pyrit.models.code_execution import _ProgramExecutionEvent # noqa: TC001 (runtime-required by Pydantic) from pyrit.models.identifiers.component_identifier import ( ComponentIdentifier, # noqa: TC001 (runtime-required by Pydantic field annotations) ) from pyrit.models.literals import MEDIA_PATH_DATA_TYPES -from pyrit.models.score.program import ProgramExecutionEvent # noqa: TC001 (runtime-required by Pydantic) from pyrit.models.score.scorable import ( ContentEntryScorable, ContentScorable, @@ -436,7 +436,7 @@ class ProgramExecutionObservationPayload(BaseModel): kind: Literal["program_execution"] = "program_execution" schema_version: Literal[1] = 1 program_sha256: str | None = Field(default=None, min_length=64, max_length=64, pattern=r"^[0-9a-f]{64}$") - events: tuple[ProgramExecutionEvent, ...] = () + events: tuple[_ProgramExecutionEvent, ...] = () reasons: tuple[str, ...] = () @field_validator("schema_version", mode="before") @@ -458,33 +458,21 @@ def _validate_schema_version(cls, value: Any) -> Any: @model_validator(mode="after") def _validate_events(self) -> ProgramExecutionObservationPayload: """ - Keep one event per case, in case order. + Keep one event per invocation, in invocation order. Returns: ProgramExecutionObservationPayload: The validated payload. Raises: - ValueError: If a case repeats, events are out of order, or events lack a program digest. + ValueError: If an invocation repeats, events are out of order, or events lack a program digest. """ - indices = [event.case_index for event in self.events] + indices = [event.invocation_index for event in self.events] if indices != sorted(set(indices)): - raise ValueError("Program execution events must name each case once, in case order.") + raise ValueError("Program execution events must name each invocation once, in invocation order.") if self.events and self.program_sha256 is None: raise ValueError("Program execution events require the digest of the program that ran.") return self - def event_for(self, case_index: int) -> ProgramExecutionEvent | None: - """ - Return the event for one case. - - Args: - case_index (int): The case's position in the condition. - - Returns: - ProgramExecutionEvent | None: The event, or None when that case did not run. - """ - return next((event for event in self.events if event.case_index == case_index), None) - ObservationPayload = Annotated[ ScorerTargetResponsePayload @@ -612,9 +600,9 @@ def _validate_program_acquisition(self) -> Observation: if not isinstance(self.scorable, ProgramExecutionScorable): raise ValueError("Program execution observations require a ProgramExecutionScorable.") events = self.payload.events - if any(event.case_index >= len(self.scorable.invocations) for event in events): + if any(event.invocation_index >= len(self.scorable.invocations) for event in events): raise ValueError("Program execution events must belong to the scorable's invocations.") - if any(event.runtime_profile != self.scorable.runtime_profile for event in events): + if any(event.runtime_profile_name != self.scorable.runtime_profile_name for event in events): raise ValueError("Program execution events must use the scorable's runtime profile.") if self.acquisition is Acquisition.COMPLETE and ( len(events) != len(self.scorable.invocations) or self.payload.reasons diff --git a/pyrit/models/score/program.py b/pyrit/models/score/program.py deleted file mode 100644 index 84248dd6f2..0000000000 --- a/pyrit/models/score/program.py +++ /dev/null @@ -1,78 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -"""Backend-neutral evidence about how a program behaved when an isolated sandbox ran it.""" - -from __future__ import annotations - -from typing import Annotated - -from pydantic import BaseModel, ConfigDict, Field, model_validator - -#: A trusted runtime profile name, such as ``python-3.12-restricted``. It names a sandbox -#: configuration the caller allowlisted; it is never an image name or a command. -RuntimeProfileName = Annotated[str, Field(min_length=1, max_length=64, pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*$")] - - -class ProgramInvocation(BaseModel): - """ - One bounded invocation of a program: its arguments and standard input. - - An invocation is a test stimulus. It carries no expected result, so a source that runs it - cannot see which output passes. - """ - - model_config = ConfigDict(frozen=True, extra="forbid") - - arguments: tuple[Annotated[str, Field(pattern=r"^[^\x00]*$")], ...] = () - stdin: str | None = None - - -class ProgramExecutionEvent(BaseModel): - """ - What one invocation did when the sandbox ran it. - - The event records enough to audit the result and to judge it again. It does not make - execution reproducible: a program can read the clock, use randomness, or depend on its - runtime. Replay judges this stored evidence again without running the program. - - ``stdout`` and ``stderr`` hold at most the source's configured limit. Each stream has its - own truncation marker, so a matcher never treats a missing suffix as proof of absence. - ``exit_code`` is ``None`` when the program did not exit on its own, for example because the - sandbox stopped it at its time or resource limit. - """ - - model_config = ConfigDict(frozen=True, extra="forbid") - - case_index: int = Field(ge=0) - runtime_profile: RuntimeProfileName - runtime_image_digest: str = Field(min_length=1) - resolved_command: tuple[str, ...] = Field(min_length=1) - exit_code: int | None - stdout: str = "" - stderr: str = "" - stdout_truncated: bool = False - stderr_truncated: bool = False - timed_out: bool = False - resource_limit_exceeded: bool = False - duration_ms: int = Field(ge=0) - - @model_validator(mode="after") - def _validate_exit_state(self) -> ProgramExecutionEvent: - """ - Keep the exit state consistent with how the run ended. - - Returns: - ProgramExecutionEvent: The validated event. - - Raises: - ValueError: If a run the sandbox stopped also reports an exit code of its own. - """ - if self.timed_out and self.exit_code is not None: - raise ValueError("A timed-out invocation did not exit on its own and has no exit code.") - return self - - @property - def output_truncated(self) -> bool: - """Whether either output stream was cut at the retention limit.""" - return self.stdout_truncated or self.stderr_truncated diff --git a/pyrit/models/score/scorable.py b/pyrit/models/score/scorable.py index b711cd6112..d9b649e2b3 100644 --- a/pyrit/models/score/scorable.py +++ b/pyrit/models/score/scorable.py @@ -9,12 +9,12 @@ from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, model_validator +from pyrit.models.code_execution import ( # noqa: TC001 (runtime-required by Pydantic) + _ProgramInvocation, + _RuntimeProfileName, +) from pyrit.models.literals import PromptDataType # noqa: TC001 (runtime-required by Pydantic field annotations) from pyrit.models.score._trace_validation import TraceId # noqa: TC001 (runtime-required by Pydantic) -from pyrit.models.score.program import ( # noqa: TC001 (runtime-required by Pydantic) - ProgramInvocation, - RuntimeProfileName, -) from pyrit.models.score.surface import SurfaceMatch # noqa: TC001 (runtime-required by Pydantic) if TYPE_CHECKING: @@ -177,22 +177,21 @@ class SurfaceScorable(Scorable): class ProgramExecutionScorable(Scorable): """ - One isolated execution attempt of a program: the program, its runtime, and its invocations. + A program, its runtime, and the invocations to execute. This is an acquisition request, not a result: it holds no expected output, so the source that runs it cannot see which output passes. ``program`` names the code to run: a stored message with one text piece, loose text, or stored loose text. It is used verbatim; code is never extracted from surrounding prose. - ``execution_id`` identifies the attempt. Reusing a scorable reuses one observed execution, - and a caller that wants a fresh run mints a new id. + Each acquisition runs afresh. The resulting observation's ID identifies that run; + replay uses the stored observation, not this input description. """ scorable_type: Literal["program_execution"] = "program_execution" - execution_id: uuid.UUID program: Annotated[MessageScorable | ContentScorable | ContentEntryScorable, Field(discriminator="scorable_type")] - runtime_profile: RuntimeProfileName - invocations: tuple[ProgramInvocation, ...] + runtime_profile_name: _RuntimeProfileName + invocations: tuple[_ProgramInvocation, ...] @model_validator(mode="after") def _validate_execution(self) -> ProgramExecutionScorable: diff --git a/pyrit/score/__init__.py b/pyrit/score/__init__.py index 121e615100..48cea04304 100644 --- a/pyrit/score/__init__.py +++ b/pyrit/score/__init__.py @@ -42,7 +42,6 @@ from pyrit.score.float_scale.video_float_scale_scorer import VideoFloatScaleScorer from pyrit.score.message_scorable_resolver import MessageScorableResolver from pyrit.score.message_scorer import MessageScorer - from pyrit.score.observation.docker_sandbox_client import DockerRuntimeProfile, DockerSandboxClient from pyrit.score.observation.execution import NonReplayableObservationError from pyrit.score.observation.lifecycle import AttemptObserver, SurfaceLifecycleSource from pyrit.score.observation.local_file_surface_source import LocalFileSurfaceSource @@ -50,7 +49,6 @@ from pyrit.score.observation.otel_span_exporter import InMemoryTraceExporter from pyrit.score.observation.otel_trace_source import OtelTraceSource from pyrit.score.observation.program_execution_source import ProgramExecutionSource - from pyrit.score.observation.sandbox_client import SandboxClient, SandboxExecutionError, SandboxRun from pyrit.score.observation.trace_client import InMemoryTraceClient, TraceAcquisitionError, TraceClient from pyrit.score.response_handler import ( CallableResponseHandler, @@ -165,13 +163,8 @@ from pyrit.score.true_false.wildguard_scorer import WildGuardScorer, render_wildguard_prompt _LAZY_EXPORTS: dict[str, str | tuple[str, str | None]] = { - "DockerRuntimeProfile": "pyrit.score.observation.docker_sandbox_client", - "DockerSandboxClient": "pyrit.score.observation.docker_sandbox_client", "ProgramExecutionScorer": "pyrit.score.true_false.program_execution_scorer", "ProgramExecutionSource": "pyrit.score.observation.program_execution_source", - "SandboxClient": "pyrit.score.observation.sandbox_client", - "SandboxExecutionError": "pyrit.score.observation.sandbox_client", - "SandboxRun": "pyrit.score.observation.sandbox_client", "OutputMatchesScorer": "pyrit.score.true_false.output_matches_scorer", "AnsiEscapeOutputScorer": "pyrit.score.true_false.regex.ansi_escape_output_scorer", "AnthraxKeywordScorer": "pyrit.score.true_false.regex.anthrax_keyword_scorer", diff --git a/pyrit/score/observation/__init__.py b/pyrit/score/observation/__init__.py index 9d9ac31b08..ddcf100056 100644 --- a/pyrit/score/observation/__init__.py +++ b/pyrit/score/observation/__init__.py @@ -10,7 +10,6 @@ if TYPE_CHECKING: from pyrit.score.observation.conversation_source import ConversationSource - from pyrit.score.observation.docker_sandbox_client import DockerRuntimeProfile, DockerSandboxClient from pyrit.score.observation.execution import NonReplayableObservationError from pyrit.score.observation.lifecycle import AttemptObserver, SurfaceLifecycleSource from pyrit.score.observation.local_file_surface_source import LocalFileSurfaceSource @@ -18,16 +17,10 @@ from pyrit.score.observation.otel_span_exporter import InMemoryTraceExporter from pyrit.score.observation.otel_trace_source import OtelTraceSource from pyrit.score.observation.program_execution_source import ProgramExecutionSource - from pyrit.score.observation.sandbox_client import SandboxClient, SandboxExecutionError, SandboxRun from pyrit.score.observation.trace_client import InMemoryTraceClient, TraceAcquisitionError, TraceClient _LAZY_EXPORTS: dict[str, str] = { - "DockerRuntimeProfile": "pyrit.score.observation.docker_sandbox_client", - "DockerSandboxClient": "pyrit.score.observation.docker_sandbox_client", "ProgramExecutionSource": "pyrit.score.observation.program_execution_source", - "SandboxClient": "pyrit.score.observation.sandbox_client", - "SandboxExecutionError": "pyrit.score.observation.sandbox_client", - "SandboxRun": "pyrit.score.observation.sandbox_client", "AttemptObserver": "pyrit.score.observation.lifecycle", "ConversationSource": "pyrit.score.observation.conversation_source", "InMemoryTraceClient": "pyrit.score.observation.trace_client", diff --git a/pyrit/score/observation/docker_sandbox_client.py b/pyrit/score/observation/docker_sandbox_client.py deleted file mode 100644 index 8e53942d2f..0000000000 --- a/pyrit/score/observation/docker_sandbox_client.py +++ /dev/null @@ -1,560 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -"""Run programs in disposable, locked-down Docker containers.""" - -from __future__ import annotations - -import asyncio -import contextlib -import logging -import os -import re -import shutil -import stat -import tempfile -import time -import uuid -from dataclasses import dataclass -from pathlib import Path -from typing import TYPE_CHECKING, ParamSpec, TypeVar - -from pyrit.models import ComponentIdentifier, ProgramExecutionEvent -from pyrit.score.observation.sandbox_client import SandboxExecutionError, SandboxRun - -if TYPE_CHECKING: - from collections.abc import Callable, Mapping - - from pyrit.models import ProgramInvocation - -logger = logging.getLogger(__name__) - -_DIGEST_PATTERN = re.compile(r"@(sha256:[0-9a-f]{64})$") -_NEVER_STARTED = "0001-01-01T00:00:00Z" - -_P = ParamSpec("_P") -_T = TypeVar("_T") - - -async def _complete_in_thread(function: Callable[_P, _T], /, *args: _P.args, **kwargs: _P.kwargs) -> _T: - """ - Run blocking file work in a worker thread, letting it finish if the caller is cancelled. - - A thread cannot be interrupted, so a cancelled caller waits for the work to end before its - cancellation completes. Cleanup that follows then sees the file system the work left. - - Returns: - _T: What the function returned. - - Raises: - CancelledError: Once the work has ended, if the caller was cancelled. - """ - work = asyncio.ensure_future(asyncio.to_thread(function, *args, **kwargs)) - try: - return await asyncio.shield(work) - except asyncio.CancelledError: - await asyncio.wait({work}) - if not work.cancelled() and work.exception() is not None: - logger.debug("File work ended with %r after its caller was cancelled", work.exception()) - raise - - -@dataclass(frozen=True) -class DockerRuntimeProfile: - """ - One allowlisted runtime: a digest-pinned image and the command that runs a program in it. - - The container runs ``command``, then the program's path, then the invocation's arguments. - For example ``command=("python3", "-I")`` runs ``python3 -I /program/main.py ``. - ``image`` must be pinned by digest, so the profile always names the same runtime. - """ - - image: str - command: tuple[str, ...] - program_filename: str = "main.py" - - def __post_init__(self) -> None: - """ - Validate the image pin and the program file name. - - Raises: - ValueError: If the image is not pinned by digest, the command is empty, or the - file name is not a plain name. - """ - if _DIGEST_PATTERN.search(self.image) is None: - raise ValueError(f"Runtime image {self.image!r} must be pinned by digest (name@sha256:<64 hex>).") - if not self.command or any(not part for part in self.command): - raise ValueError("A runtime profile command must be a nonempty argument list.") - if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]*", self.program_filename): - raise ValueError(f"Program file name {self.program_filename!r} must be a plain file name.") - - @property - def image_digest(self) -> str: - """The ``sha256:`` digest the image is pinned to; ``__post_init__`` guarantees its presence.""" - return self.image.rsplit("@", 1)[1] - - -@dataclass(frozen=True) -class _ContainerExit: - exit_code: int | None - oom_killed: bool - - -@dataclass(frozen=True) -class _ClientOutput: - stdout: bytes - stdout_truncated: bool - stderr: bytes - stderr_truncated: bool - timed_out: bool - - -class DockerSandboxClient: - """ - Run each invocation in its own disposable Docker container. - - Every container has no network, a read-only root file system, all capabilities dropped, - no privilege escalation, an unprivileged user, and caps on memory, CPU, and process count. - The program is mounted read-only; a small ``/tmp`` is the only writable location. Each - invocation gets a fresh container that is removed afterwards, so one run cannot leave state - for the next. Docker keeps no log of the container's output, so output volume never grows - the daemon's files, and the profile's command replaces the image's entrypoint, so the - recorded command is the one that runs. - - Images are never pulled while scoring: a profile's image must already be present, or the - attempt fails as a sandbox error. Pull the pinned images before a run. - - A container that runs past ``timeout_seconds`` is killed and reported as timed out. One that - the kernel stops for exceeding its memory limit is reported as exceeding a resource limit. - Each stream keeps at most ``max_output_bytes`` bytes. The timeout includes container start-up - and the attached client's exit; a client still running once its streams have closed is - ended at the deadline, and the container's own state then decides the result. - - This client isolates submitted code with Docker's defaults plus the restrictions above. It - is not a hardened multi-tenant sandbox; run it on a host dedicated to evaluation. - """ - - _READ_CHUNK_BYTES = 1 << 16 - _CONTROL_TIMEOUT_SECONDS = 30.0 - - def __init__( - self, - *, - profiles: Mapping[str, DockerRuntimeProfile], - docker_executable: str = "docker", - timeout_seconds: float = 10.0, - memory: str = "256m", - cpus: float = 1.0, - pids_limit: int = 64, - max_output_bytes: int = 64_000, - ) -> None: - """ - Initialize with the allowlisted runtime profiles and per-container limits. - - Args: - profiles (Mapping[str, DockerRuntimeProfile]): Runtime profile names and their runtimes. - docker_executable (str): The Docker command line client to call. - timeout_seconds (float): Wall-clock limit for each invocation. - memory (str): Memory limit for each container, in Docker's notation, such as ``256m``. - cpus (float): CPU limit for each container. - pids_limit (int): Most processes a container may hold. - max_output_bytes (int): Most bytes of each stream kept per invocation. - - Raises: - ValueError: If no profile is given, a profile name is invalid, or a limit is not positive. - """ - if not profiles: - raise ValueError("DockerSandboxClient requires at least one runtime profile.") - invalid = [name for name in profiles if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]{0,63}", name)] - if invalid: - raise ValueError(f"Invalid runtime profile names: {invalid}.") - if timeout_seconds <= 0 or cpus <= 0 or pids_limit < 1 or max_output_bytes < 1: - raise ValueError("timeout_seconds, cpus, pids_limit and max_output_bytes must be positive.") - if not re.fullmatch(r"[1-9][0-9]*[bkmg]?", memory): - raise ValueError(f"Memory limit {memory!r} must look like 256m or 1g.") - self._profiles = dict(profiles) - self._docker = docker_executable - self._timeout_seconds = timeout_seconds - self._memory = memory - self._cpus = cpus - self._pids_limit = pids_limit - self._max_output_bytes = max_output_bytes - self._temp_root = Path(tempfile.gettempdir()) - - def get_identifier(self) -> ComponentIdentifier: - """ - Return the profiles and limits. - - Returns: - ComponentIdentifier: The nonsecret client configuration. - """ - return ComponentIdentifier.of( - self, - params={ - "client_version": 1, - "profiles": { - name: { - "image": profile.image, - "command": list(profile.command), - "program_filename": profile.program_filename, - } - for name, profile in sorted(self._profiles.items()) - }, - "timeout_seconds": self._timeout_seconds, - "memory": self._memory, - "cpus": self._cpus, - "pids_limit": self._pids_limit, - "max_output_bytes": self._max_output_bytes, - }, - ) - - def get_runtime_profiles(self) -> frozenset[str]: - """ - Return the configured runtime profile names. - - Returns: - frozenset[str]: The profile names this client can run. - """ - return frozenset(self._profiles) - - async def execute_async( - self, *, program: str, runtime_profile: str, invocations: tuple[ProgramInvocation, ...] - ) -> SandboxRun: - """ - Run each invocation in a fresh container, in order. - - Args: - program (str): The program text, written to a read-only file in the container. - runtime_profile (str): The allowlisted profile to run under. - invocations (tuple[ProgramInvocation, ...]): The arguments and input of each run. - - Returns: - SandboxRun: One event per invocation that ran. If the sandbox fails after some - invocations ran, the run holds those events and the reason it stopped. - - Raises: - ValueError: If the runtime profile is not configured. - SandboxExecutionError: If no invocation could be run. - """ - profile = self._profiles.get(runtime_profile) - if profile is None: - raise ValueError(f"Runtime profile {runtime_profile!r} is not configured.") - program_dir = self._temp_root / f"pyrit-program-{uuid.uuid4().hex}" - program_path = program_dir / profile.program_filename - try: - await _complete_in_thread( - self._write_program, program_dir=program_dir, program_path=program_path, program=program - ) - events: list[ProgramExecutionEvent] = [] - for index, invocation in enumerate(invocations): - try: - event, cleanup_warning = await self._run_invocation_async( - case_index=index, - runtime_profile=runtime_profile, - profile=profile, - program_dir=program_dir, - invocation=invocation, - ) - except SandboxExecutionError as error: - if not events: - raise - return SandboxRun(events=tuple(events), reasons=(f"sandbox error at case {index}: {error}",)) - events.append(event) - if cleanup_warning is not None: - # The environment may still be running, so no further case starts beside it. - remaining = len(invocations) - len(events) - reasons = (f"stopped after case {index}: {cleanup_warning}",) if remaining else () - return SandboxRun(events=tuple(events), reasons=reasons, warnings=(cleanup_warning,)) - return SandboxRun(events=tuple(events)) - finally: - await _complete_in_thread(self._remove_program_dir, program_dir=program_dir, program_path=program_path) - - @staticmethod - def _write_program(*, program_dir: Path, program_path: Path, program: str) -> None: - """Create a private directory and write the program into it, then make both read-only for others.""" - program_dir.mkdir(mode=0o700) - program_path.write_text(program, encoding="utf-8") - program_path.chmod(0o444) - program_dir.chmod(0o755) - - @staticmethod - def _remove_program_dir(*, program_dir: Path, program_path: Path) -> None: - """ - Delete the submitted program, restoring write permission first. - - On Windows the read-only mode is a file attribute that blocks deletion, so the file - is made writable again before the directory is removed. - """ - with contextlib.suppress(OSError): - program_path.chmod(stat.S_IREAD | stat.S_IWRITE) - shutil.rmtree(program_dir, ignore_errors=True) - if program_dir.exists(): - logger.warning("Could not remove the program directory %s", program_dir) - - async def _run_invocation_async( - self, - *, - case_index: int, - runtime_profile: str, - profile: DockerRuntimeProfile, - program_dir: Path, - invocation: ProgramInvocation, - ) -> tuple[ProgramExecutionEvent, str | None]: - """ - Run one invocation in a new container and remove the container afterwards. - - Returns: - tuple[ProgramExecutionEvent, str | None]: What the invocation did, and a warning when - the container could not be removed. - - Raises: - SandboxExecutionError: If Docker cannot start or inspect the container, or the - container had not finished when the attached client exited. - """ - name = f"pyrit-sandbox-{uuid.uuid4().hex}" - resolved_command = (*profile.command, f"/program/{profile.program_filename}", *invocation.arguments) - argv = self._run_arguments( - name=name, profile=profile, program_dir=program_dir, resolved_command=resolved_command - ) - started = time.monotonic() - try: - process = await asyncio.create_subprocess_exec( - *argv, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - except OSError as error: - raise SandboxExecutionError(f"could not start {self._docker!r}: {error}") from error - cleanup_warning: str | None = None - try: - output = await self._collect_output_async( - process=process, name=name, stdin=invocation.stdin, deadline=started + self._timeout_seconds - ) - duration_ms = int((time.monotonic() - started) * 1000) - container = await self._inspect_async(name=name, cli_stderr=output.stderr, timed_out=output.timed_out) - finally: - if process.returncode is None: - with contextlib.suppress(ProcessLookupError): - process.kill() - await process.wait() - removal_status, removal_output = await self._control_async("rm", "--force", name) - if removal_status != 0: - detail = removal_output.strip()[:300] or "unknown Docker failure" - cleanup_warning = f"could not remove container {name}: {detail}" - logger.warning("Sandbox cleanup failed: %s", cleanup_warning) - timed_out = output.timed_out - event = ProgramExecutionEvent( - case_index=case_index, - runtime_profile=runtime_profile, - runtime_image_digest=profile.image_digest, - resolved_command=resolved_command, - exit_code=None if (timed_out or container.oom_killed) else container.exit_code, - stdout=output.stdout.decode("utf-8", errors="replace"), - stderr=output.stderr.decode("utf-8", errors="replace"), - stdout_truncated=output.stdout_truncated, - stderr_truncated=output.stderr_truncated, - timed_out=timed_out, - resource_limit_exceeded=container.oom_killed and not timed_out, - duration_ms=duration_ms, - ) - return event, cleanup_warning - - def _run_arguments( - self, *, name: str, profile: DockerRuntimeProfile, program_dir: Path, resolved_command: tuple[str, ...] - ) -> list[str]: - """ - Build the ``docker run`` command line for one invocation. - - The profile's first argument becomes the entrypoint, so an image's own entrypoint - never runs in front of ``resolved_command``. Daemon-side logging is off: the attached - client reads the output, and nothing is kept on the host beyond what it retains. - - Returns: - list[str]: The arguments, starting with the Docker client. - """ - return [ - self._docker, - "run", - "--name", - name, - "--interactive", - "--pull", - "never", - "--log-driver", - "none", - "--network", - "none", - "--read-only", - "--cap-drop", - "ALL", - "--security-opt", - "no-new-privileges", - "--user", - "65534:65534", - "--memory", - self._memory, - "--memory-swap", - self._memory, - "--cpus", - str(self._cpus), - "--pids-limit", - str(self._pids_limit), - "--tmpfs", - "/tmp:rw,nosuid,nodev,size=16m", - "--workdir", - "/tmp", - "--mount", - f"type=bind,source={os.fspath(program_dir)},target=/program,readonly", - "--entrypoint", - resolved_command[0], - profile.image, - *resolved_command[1:], - ] - - async def _collect_output_async( - self, *, process: asyncio.subprocess.Process, name: str, stdin: str | None, deadline: float - ) -> _ClientOutput: - """ - Feed the input, read both streams, and wait for the attached client to exit, by one deadline. - - A container still running at the deadline is killed and the run is marked timed out. A - client whose streams closed but that has not exited by the deadline is ended; the - container's own state, inspected afterwards, then decides whether the program finished, - so a stalled client is not mistaken for a program that ran too long. - - Returns: - _ClientOutput: The retained output and whether the run timed out. - """ - readers = asyncio.gather( - self._read_capped_async(process.stdout), - self._read_capped_async(process.stderr), - self._feed_stdin_async(process.stdin, stdin), - ) - timed_out = False - try: - (stdout, stdout_truncated), (stderr, stderr_truncated), _ = await asyncio.wait_for( - asyncio.shield(readers), timeout=max(deadline - time.monotonic(), 0.0) - ) - except TimeoutError: - timed_out = True - await self._control_async("kill", name) - try: - await asyncio.wait_for(asyncio.shield(readers), timeout=self._CONTROL_TIMEOUT_SECONDS) - except TimeoutError: - # The container would not stop; ending the attached client closes its streams. - with contextlib.suppress(ProcessLookupError): - process.kill() - (stdout, stdout_truncated), (stderr, stderr_truncated), _ = await readers - try: - exit_budget = self._CONTROL_TIMEOUT_SECONDS if timed_out else max(deadline - time.monotonic(), 0.0) - await asyncio.wait_for(process.wait(), timeout=exit_budget) - except TimeoutError: - with contextlib.suppress(ProcessLookupError): - process.kill() - await process.wait() - return _ClientOutput( - stdout=stdout, - stdout_truncated=stdout_truncated, - stderr=stderr, - stderr_truncated=stderr_truncated, - timed_out=timed_out, - ) - - async def _inspect_async(self, *, name: str, cli_stderr: bytes, timed_out: bool) -> _ContainerExit: - """ - Read how the container ended, distinguishing a program exit from a Docker failure. - - An exit code is only reported for a container that has finished. A container still - running when the attached client exited means the run was cut short by Docker, not by - the program, unless the sandbox stopped it at its time limit. - - Returns: - _ContainerExit: The container's exit code and whether the kernel killed it for memory. - - Raises: - SandboxExecutionError: If the container does not exist, never started, or had not - finished when the client exited. - """ - returncode, output = await self._control_async( - "inspect", - "--format", - "{{.State.Status}} {{.State.ExitCode}} {{.State.OOMKilled}} {{.State.StartedAt}}", - name, - ) - fields = output.split() - if returncode != 0 or len(fields) != 4 or fields[3] == _NEVER_STARTED: - lines = [line.strip() for line in cli_stderr.decode("utf-8", errors="replace").splitlines() if line.strip()] - detail = lines[0][:500] if lines else output.strip()[:500] - raise SandboxExecutionError(f"the container did not start: {detail or 'unknown Docker failure'}") - status = fields[0] - if status != "exited" and not timed_out: - raise SandboxExecutionError( - f"the Docker client exited while the container was still {status}, so the program did not finish" - ) - return _ContainerExit(exit_code=int(fields[1]), oom_killed=fields[2] == "true") - - async def _control_async(self, *arguments: str) -> tuple[int, str]: - """ - Run a short Docker management command, bounded in time. - - Returns: - tuple[int, str]: The exit status and combined output; status -1 when it failed to run. - """ - try: - process = await asyncio.create_subprocess_exec( - self._docker, - *arguments, - stdin=asyncio.subprocess.DEVNULL, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.STDOUT, - ) - except OSError as error: - logger.warning("Docker %s failed to start: %s", arguments[0], error) - return -1, str(error) - try: - output, _ = await asyncio.wait_for(process.communicate(), timeout=self._CONTROL_TIMEOUT_SECONDS) - except TimeoutError: - logger.warning("Docker %s timed out", arguments[0]) - return -1, "timed out" - finally: - # Runs on timeout and on cancellation alike, so no management command outlives its caller. - if process.returncode is None: - with contextlib.suppress(ProcessLookupError): - process.kill() - await process.wait() - return process.returncode if process.returncode is not None else -1, output.decode("utf-8", errors="replace") - - async def _read_capped_async(self, stream: asyncio.StreamReader | None) -> tuple[bytes, bool]: - """ - Read a stream to its end, keeping at most the configured number of bytes. - - Returns: - tuple[bytes, bool]: The retained bytes and whether more were discarded. - """ - if stream is None: - return b"", False - kept = bytearray() - truncated = False - while chunk := await stream.read(self._READ_CHUNK_BYTES): - room = self._max_output_bytes - len(kept) - if len(chunk) > room: - truncated = True - kept.extend(chunk[: max(room, 0)]) - return bytes(kept), truncated - - @staticmethod - async def _feed_stdin_async(stream: asyncio.StreamWriter | None, data: str | None) -> None: - """Write the invocation's input, then close standard input so the program sees its end.""" - if stream is None: - return - try: - if data: - stream.write(data.encode("utf-8")) - await stream.drain() - except (BrokenPipeError, ConnectionResetError): - pass - finally: - stream.close() - with contextlib.suppress(BrokenPipeError, ConnectionResetError): - await stream.wait_closed() diff --git a/pyrit/score/observation/program_execution_source.py b/pyrit/score/observation/program_execution_source.py index b1bfafef8b..878027a3af 100644 --- a/pyrit/score/observation/program_execution_source.py +++ b/pyrit/score/observation/program_execution_source.py @@ -5,14 +5,11 @@ from __future__ import annotations -import asyncio import hashlib import logging -import uuid -from collections import OrderedDict -from dataclasses import dataclass from typing import TYPE_CHECKING +from pyrit._code_execution import SandboxExecutionError from pyrit.memory import CentralMemory from pyrit.models import ( Acquisition, @@ -20,72 +17,36 @@ ContentScorable, MessageScorable, Observation, - ProgramExecutionEvent, ProgramExecutionObservationPayload, ) from pyrit.score.message_scorable_resolver import MessageScorableResolver -from pyrit.score.observation.sandbox_client import SandboxExecutionError if TYPE_CHECKING: + from pyrit._code_execution import SandboxClient from pyrit.models import ProgramExecutionScorable - from pyrit.score.observation.sandbox_client import SandboxClient + from pyrit.models.code_execution import _ProgramExecutionEvent logger = logging.getLogger(__name__) -# Namespace for the IDs of reusable execution evidence, derived from the source and the attempt. -_OBSERVATION_NAMESPACE = uuid.UUID("20acb8a7-a3ed-5df1-b292-83d20448a0fd") -_REUSABLE = (Acquisition.COMPLETE, Acquisition.PARTIAL) - class _NoProgramError(Exception): """The named evidence holds nothing that can be submitted as a program.""" -@dataclass -class _InFlight: - """One running acquisition and how many callers are waiting for it.""" - - task: asyncio.Future[Observation] - waiters: int = 0 - - class ProgramExecutionSource: """ - Run the program a ``ProgramExecutionScorable`` names, once per invocation, in a sandbox. - - The source receives only the derived scorable: the program, the runtime profile, and the - invocations. It never sees expected results, so it can run the program but cannot tell - which output passes. An injected ``SandboxClient`` owns the disposable environment; the - source only submits the attempt and waits for it. - - The program is the named text, used verbatim. A message must name exactly one text piece; - code is never extracted from surrounding prose. A message that holds a blocked or failed - response has no program to run, so the evidence is unavailable. + Submit verbatim program text and invocation inputs to an injected ``SandboxClient``. - Program failure is not acquisition failure. A compile error, a nonzero exit, wrong output, - a timeout, or an enforced resource limit is complete evidence that the program misbehaved. - Only a sandbox that cannot run the program makes the acquisition an error, and only a - sandbox that stops early leaves it partial, so broken infrastructure never looks like a - program that failed its tests. + Expected results never reach the source. Memory is needed only to resolve stored program + references. A message must contain exactly one text piece; blocked or failed responses + leave evidence unavailable. Code is not extracted from prose. - Each stream keeps at most ``max_output_chars`` characters and records when it was cut. - A completed or partial observation is reused when the same scorable is acquired again, and - concurrent callers for one attempt share a single run, so one execution attempt is observed - once; a new ``execution_id`` runs the program again. The source holds recent observations, - and its evidence for an attempt always has the same ID, derived from the source - configuration and the attempt. Once memory stores that observation, as scoring does, a - later acquisition loads it instead of running the program, however many attempts the - source has seen since; every stored score therefore judges the attempt's one stored run. - Evidence that was never stored and has left the cache runs again under the same ID, and - memory refuses a second, different run of an attempt that already has stored evidence, - so contradicting results are never stored for one attempt. Acquisition reads memory, so - the central memory instance must be set. - Problems the sandbox reports after the evidence was collected, such as a container it - could not remove, are kept in the observation's ``sandbox_warnings`` metadata. + Each acquisition runs afresh and returns a new observation. Program failures are execution + evidence; infrastructure failures are acquisition errors. A run that stops after some + invocations leaves partial evidence. Streams retain at most ``max_output_chars`` characters + with truncation flags, and cleanup warnings remain in ``sandbox_warnings`` metadata. """ - _MAX_REUSED_OBSERVATIONS = 256 - def __init__(self, *, sandbox: SandboxClient, max_output_chars: int = 64_000) -> None: """ Initialize with the sandbox that runs submitted programs. @@ -101,8 +62,6 @@ def __init__(self, *, sandbox: SandboxClient, max_output_chars: int = 64_000) -> raise ValueError("max_output_chars must be positive.") self._sandbox = sandbox self._max_output_chars = max_output_chars - self._observed: OrderedDict[ProgramExecutionScorable, Observation] = OrderedDict() - self._in_flight: dict[ProgramExecutionScorable, _InFlight] = {} def get_identifier(self) -> ComponentIdentifier: """ @@ -113,13 +72,13 @@ def get_identifier(self) -> ComponentIdentifier: """ return ComponentIdentifier.of( self, - params={"acquisition_version": 1, "max_output_chars": self._max_output_chars}, + params={"acquisition_version": 3, "max_output_chars": self._max_output_chars}, children={"sandbox": self._sandbox.get_identifier()}, ) async def acquire_async(self, *, scorable: ProgramExecutionScorable) -> Observation: """ - Run every invocation once and return what the program did. + Run every invocation afresh and return what the program did. Args: scorable (ProgramExecutionScorable): The program, runtime profile, and invocations. @@ -132,51 +91,11 @@ async def acquire_async(self, *, scorable: ProgramExecutionScorable) -> Observat returns events for invocations it was not given, or the program reference names missing or several pieces. """ - reused = self._observed.get(scorable) - if reused is not None: - self._observed.move_to_end(scorable) - return reused - flight = self._in_flight.get(scorable) - if flight is None: - flight = _InFlight(task=asyncio.ensure_future(self._acquire_uncached_async(scorable=scorable))) - self._in_flight[scorable] = flight - flight.task.add_done_callback(lambda _task, key=scorable: self._in_flight.pop(key, None)) - flight.waiters += 1 - try: - return await asyncio.shield(flight.task) - finally: - flight.waiters -= 1 - if flight.waiters == 0 and not flight.task.done(): - # Every caller of this attempt was cancelled, so nobody needs its result. Wait for - # the sandbox's cleanup before the caller's cancellation completes, so that - # cancelling still means the environment is gone. - flight.task.cancel() - await asyncio.wait({flight.task}) - if not flight.task.cancelled() and flight.task.exception() is not None: - logger.debug("Cancelled execution %s ended with %r", scorable.execution_id, flight.task.exception()) - - async def _acquire_uncached_async(self, *, scorable: ProgramExecutionScorable) -> Observation: - """ - Run the attempt once; ``acquire_async`` shares this work between concurrent callers. - - Returns: - Observation: The execution evidence, with its acquisition state. - - Raises: - ValueError: If the runtime profile is not offered, the program reference is - unusable, or the sandbox returns events for invocations it was not given. - """ - if scorable.runtime_profile not in self._sandbox.get_runtime_profiles(): + if scorable.runtime_profile_name not in self._sandbox.get_runtime_profile_names(): raise ValueError( - f"Runtime profile {scorable.runtime_profile!r} is not offered by the sandbox; " - f"available profiles: {sorted(self._sandbox.get_runtime_profiles())}." + f"Runtime profile name {scorable.runtime_profile_name!r} is not offered by the sandbox; " + f"available names: {sorted(self._sandbox.get_runtime_profile_names())}." ) - stored = await CentralMemory.get_memory_instance().get_observations_async( - observation_ids=[self._attempt_observation_id(scorable)] - ) - if stored: - self._remember(scorable=scorable, observation=stored[0]) - return stored[0] try: program = await self._resolve_program_async(scorable=scorable) except _NoProgramError as error: @@ -184,16 +103,16 @@ async def _acquire_uncached_async(self, *, scorable: ProgramExecutionScorable) - try: run = await self._sandbox.execute_async( program=program, - runtime_profile=scorable.runtime_profile, + runtime_profile_name=scorable.runtime_profile_name, invocations=scorable.invocations, ) except SandboxExecutionError as error: - logger.warning("Sandbox could not run program for execution %s: %s", scorable.execution_id, error) + logger.warning("Sandbox could not run program under %s: %s", scorable.runtime_profile_name, error) return self._observation( scorable=scorable, acquisition=Acquisition.ERROR, reasons=(f"sandbox error: {error}",) ) events = tuple(self._bounded(event) for event in run.events) - if any(event.case_index >= len(scorable.invocations) for event in events): + if any(event.invocation_index >= len(scorable.invocations) for event in events): raise ValueError("The sandbox returned an event for an invocation it was not given.") program_sha256 = hashlib.sha256(program.encode("utf-8")).hexdigest() if len(events) == len(scorable.invocations): @@ -212,8 +131,6 @@ async def _acquire_uncached_async(self, *, scorable: ProgramExecutionScorable) - ) if run.warnings: observation = observation.model_copy(update={"metadata": {"sandbox_warnings": "; ".join(run.warnings)}}) - if observation.acquisition in (Acquisition.COMPLETE, Acquisition.PARTIAL): - self._remember(scorable=scorable, observation=observation) return observation async def _resolve_program_async(self, *, scorable: ProgramExecutionScorable) -> str: @@ -243,12 +160,12 @@ async def _resolve_program_async(self, *, scorable: ProgramExecutionScorable) -> raise ValueError(f"A program must be text, not {piece.converted_value_data_type}.") return piece.converted_value - def _bounded(self, event: ProgramExecutionEvent) -> ProgramExecutionEvent: + def _bounded(self, event: _ProgramExecutionEvent) -> _ProgramExecutionEvent: """ Cut each stream to the retention limit, keeping any truncation the sandbox reported. Returns: - ProgramExecutionEvent: The event as it is retained. + _ProgramExecutionEvent: The event as it is retained. """ limit = self._max_output_chars if len(event.stdout) <= limit and len(event.stderr) <= limit: @@ -267,14 +184,11 @@ def _observation( *, scorable: ProgramExecutionScorable, acquisition: Acquisition, - events: tuple[ProgramExecutionEvent, ...] = (), + events: tuple[_ProgramExecutionEvent, ...] = (), program_sha256: str | None = None, reasons: tuple[str, ...] = (), ) -> Observation: - # Only reusable evidence takes the attempt's ID; an error is retried as a new observation. - observation_id = self._attempt_observation_id(scorable) if acquisition in _REUSABLE else uuid.uuid4() return Observation( - id=observation_id, source_identifier=self.get_identifier(), acquisition=acquisition, scorable=scorable, @@ -284,17 +198,3 @@ def _observation( reasons=reasons, ), ) - - def _attempt_observation_id(self, scorable: ProgramExecutionScorable) -> uuid.UUID: - """ - Derive the ID of an attempt's reusable evidence from the source configuration and the attempt. - - Returns: - uuid.UUID: The same ID for the same source configuration and attempt. - """ - return uuid.uuid5(_OBSERVATION_NAMESPACE, f"{self.get_identifier().hash}:{scorable.model_dump_json()}") - - def _remember(self, *, scorable: ProgramExecutionScorable, observation: Observation) -> None: - self._observed[scorable] = observation - while len(self._observed) > self._MAX_REUSED_OBSERVATIONS: - self._observed.popitem(last=False) diff --git a/pyrit/score/observation/sandbox_client.py b/pyrit/score/observation/sandbox_client.py deleted file mode 100644 index bd66a62c81..0000000000 --- a/pyrit/score/observation/sandbox_client.py +++ /dev/null @@ -1,69 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -"""The injected service that runs a program in a disposable, isolated environment.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import TYPE_CHECKING, Protocol, runtime_checkable - -if TYPE_CHECKING: - from pyrit.models import ComponentIdentifier, ProgramExecutionEvent, ProgramInvocation - - -class SandboxExecutionError(Exception): - """ - The sandbox could not run the program. - - This is an infrastructure failure, such as an unreachable service, a failed - authentication, or an environment that could not be provisioned. It is never raised for a - program that fails, times out, or exceeds a limit: those are observed behavior. - """ - - -@dataclass(frozen=True) -class SandboxRun: - """ - The invocations a sandbox ran for one execution attempt. - - ``events`` holds one event for each invocation that ran, in case order. When the sandbox - stopped before running every invocation, ``reasons`` says why, so the missing cases stay - undetermined instead of failing. ``warnings`` reports problems that do not invalidate the - evidence collected, such as an environment the sandbox could not destroy afterwards. - """ - - events: tuple[ProgramExecutionEvent, ...] = () - reasons: tuple[str, ...] = () - warnings: tuple[str, ...] = () - - -@runtime_checkable -class SandboxClient(Protocol): - """ - Run a program under an allowlisted runtime profile in a disposable environment. - - The client owns the environment: it creates it for the attempt and destroys it afterwards. - It never retains an environment between attempts, builds runtime images, or starts a daemon. - A runtime profile is a trusted name the client maps to its own configuration; the program - cannot choose an image or a command. - """ - - def get_identifier(self) -> ComponentIdentifier: - """Return the nonsecret client configuration.""" - ... - - def get_runtime_profiles(self) -> frozenset[str]: - """Return the runtime profile names this client is configured to run.""" - ... - - async def execute_async( - self, *, program: str, runtime_profile: str, invocations: tuple[ProgramInvocation, ...] - ) -> SandboxRun: - """ - Run each invocation of the program once, in order, in a fresh isolated environment. - - Raises: - SandboxExecutionError: If the sandbox could not run the program at all. - """ - ... diff --git a/pyrit/score/true_false/program_execution_scorer.py b/pyrit/score/true_false/program_execution_scorer.py index 5e0c09c9ac..90e3309e21 100644 --- a/pyrit/score/true_false/program_execution_scorer.py +++ b/pyrit/score/true_false/program_execution_scorer.py @@ -5,7 +5,7 @@ from __future__ import annotations -import uuid +from dataclasses import dataclass from typing import TYPE_CHECKING from pyrit.models import ( @@ -28,16 +28,22 @@ from pyrit.models import ( ComponentIdentifier, Observation, - ProgramExecutionEvent, ProgramResult, Scorable, ScoringExpectation, TextMatcher, ) + from pyrit.models.code_execution import _ProgramExecutionEvent from pyrit.score.observation.execution import _ObservationEvidence from pyrit.score.observation.observation_source import ObservationSource +@dataclass(frozen=True) +class _Evaluation: + value: bool | None + rationale: str + + def _stream_matches(*, matcher: TextMatcher, text: str, truncated: bool) -> bool | None: """ Match one output stream, treating a cut stream as unknown unless a prefix proves a match. @@ -52,90 +58,86 @@ def _stream_matches(*, matcher: TextMatcher, text: str, truncated: bool) -> bool return None if truncated else False -def match_case(*, expected: ProgramResult, event: ProgramExecutionEvent) -> bool | None: +def _evaluate_case(*, expected: ProgramResult, event: _ProgramExecutionEvent | None) -> _Evaluation: """ - Decide whether one invocation produced its expected result. - - A run the sandbox stopped at its time or resource limit fails, because the program did not - finish. Otherwise the exit code and each stream that has a matcher must all agree. - - Args: - expected (ProgramResult): The expected result. - event (ProgramExecutionEvent): What the invocation did. + Evaluate one case once, keeping its explanation with the decision. Returns: - bool | None: True when every check passes, False when one fails, None when a cut - stream leaves a check open. + _Evaluation: The case's verdict and explanation. """ - if event.timed_out or event.resource_limit_exceeded: - return False + if event is None: + return _Evaluation(value=None, rationale="did not run") + if event.timed_out: + return _Evaluation(value=False, rationale="timed out") + if event.resource_limit_exceeded: + return _Evaluation(value=False, rationale="exceeded a resource limit") if event.exit_code is None: - return None + return _Evaluation(value=None, rationale="did not report an exit code") if expected.exit_code is not None and event.exit_code != expected.exit_code: - return False + return _Evaluation(value=False, rationale=f"exited with {event.exit_code}, expected {expected.exit_code}") checks: list[bool | None] = [] if expected.stdout is not None: checks.append(_stream_matches(matcher=expected.stdout, text=event.stdout, truncated=event.stdout_truncated)) if expected.stderr is not None: checks.append(_stream_matches(matcher=expected.stderr, text=event.stderr, truncated=event.stderr_truncated)) if False in checks: - return False - return None if None in checks else True + return _Evaluation(value=False, rationale="output did not match") + if None in checks: + return _Evaluation(value=None, rationale="output was cut before the check could be settled") + return _Evaluation(value=True, rationale="") -def match_program_behaves( +def _evaluate_program( *, condition: ProgramBehaves, payload: ProgramExecutionObservationPayload, acquisition: Acquisition -) -> bool | None: +) -> _Evaluation: """ - Match every case against the execution evidence, without I/O. - - One failing case is enough for a false verdict, even when other cases did not run. A true - verdict needs every case to have run and passed. - - Args: - condition (ProgramBehaves): The cases and their expected results. - payload (ProgramExecutionObservationPayload): The acquired execution evidence. - acquisition (Acquisition): Whether every invocation ran. + Index evidence once and reuse case evaluations for the verdict and rationale. Returns: - bool | None: True when every case passes, False when one fails, otherwise None. + _Evaluation: The program's verdict and explanation. """ + profile_name = condition.runtime_profile_name if acquisition in (Acquisition.ERROR, Acquisition.UNAVAILABLE): - return None - unknown = False - for index, case in enumerate(condition.cases): - event = payload.event_for(index) - if event is None: - unknown = True - continue - result = match_case(expected=case.expected, event=event) - if result is False: - return False - unknown = unknown or result is None - return None if unknown else True + gaps = ", ".join(payload.reasons) or "no execution evidence" + return _Evaluation( + value=None, rationale=f"The program's behavior under {profile_name} could not be observed ({gaps})." + ) + events = {event.invocation_index: event for event in payload.events} + evaluations = [ + _evaluate_case(expected=case.expected, event=events.get(index)) for index, case in enumerate(condition.cases) + ] + if all(evaluation.value is True for evaluation in evaluations): + return _Evaluation( + value=True, rationale=f"All {len(evaluations)} case(s) produced the expected result under {profile_name}." + ) + summary = "; ".join( + f"case {index} {evaluation.rationale}" + for index, evaluation in enumerate(evaluations) + if evaluation.value is not True + ) + if any(evaluation.value is False for evaluation in evaluations): + return _Evaluation( + value=False, rationale=f"The program did not produce the expected result under {profile_name}: {summary}." + ) + gaps = ", ".join(payload.reasons) + suffix = f" ({gaps})" if gaps else "" + return _Evaluation( + value=None, + rationale=f"The program's behavior under {profile_name} could not be fully established: {summary}{suffix}.", + ) class ProgramExecutionScorer(TrueFalseScorer): """ - Score whether a program behaves as expected, judged from what an isolated sandbox observed. - - The ``ProgramBehaves`` condition carries each case's stimulus and expected result. The scorer - splits them: it builds a ``ProgramExecutionScorable`` from the program it is given, the - runtime profile and the invocations, and sends only that to its source. Each expected result - stays with the scorer, so the code that runs the program never learns which output passes. + Compare execution evidence with the expected results in a ``ProgramBehaves`` condition. - The program is the scored content, run verbatim: loose text, stored loose text, or a stored - message with one text piece. Code is not extracted from surrounding prose, so a response that wraps code in - explanation fails its cases rather than being guessed at. + The source receives verbatim program text, a runtime profile name, and invocation inputs, + but never expected results. Each case checks its expected exit and output; a timeout or + resource limit fails the case. All cases passing yields true, any proven failure yields + false, and insufficient evidence leaves the verdict undetermined. - The verdict is true when every case passes and false when any case fails. A compile error, - nonzero exit, wrong output, timeout, or resource limit is a failure the sandbox observed. A - sandbox that could not run the program leaves the verdict undetermined, as does a case that - did not run or a check that cut output cannot settle. - - Every scoring call is a new execution attempt. Scoring an explicit ``ProgramExecutionScorable`` - reuses that attempt's evidence when the source has already observed it, and the stored - observation can be judged again without running the program. + Every scoring call runs afresh. ``score_observation_async`` replays saved evidence without + executing the program. """ CONDITION_TYPE = ProgramBehaves @@ -159,7 +161,10 @@ def _build_identifier(self) -> ComponentIdentifier: async def _score_scorable_async(self, *, scorable: Scorable, expectation: ScoringExpectation | None) -> list[Score]: condition = self._get_required_condition(expectation=expectation, condition_type=ProgramBehaves) if isinstance(scorable, ProgramExecutionScorable): - if (scorable.runtime_profile, scorable.invocations) != (condition.runtime_profile, condition.invocations): + if (scorable.runtime_profile_name, scorable.invocations) != ( + condition.runtime_profile_name, + condition.invocations, + ): raise ValueError( "A ProgramExecutionScorable must run the runtime profile and invocations of the " "ProgramBehaves condition." @@ -167,9 +172,8 @@ async def _score_scorable_async(self, *, scorable: Scorable, expectation: Scorin execution = scorable elif isinstance(scorable, (MessageScorable, ContentScorable, ContentEntryScorable)): execution = ProgramExecutionScorable( - execution_id=uuid.uuid4(), program=scorable, - runtime_profile=condition.runtime_profile, + runtime_profile_name=condition.runtime_profile_name, invocations=condition.invocations, ) else: @@ -183,8 +187,6 @@ async def _score_scorable_async(self, *, scorable: Scorable, expectation: Scorin raise ValueError("Program execution source changed the caller's evidence anchor.") if not isinstance(observation.payload, ProgramExecutionObservationPayload): raise ValueError("Program execution source returned incompatible evidence.") - # A reused or shared attempt's evidence may already be stored. Memory links an identical - # stored observation instead of inserting it again, atomically under its row lock. _collect_observation(observation) scores = self._score_observation(observation=observation, evidence=observation.payload, expectation=expectation) for score in scores: @@ -205,20 +207,21 @@ def _score_observation( raise NonReplayableObservationError("Program behavior scoring requires a stored program execution.") condition = self._get_required_condition(expectation=expectation, condition_type=ProgramBehaves) execution = observation.scorable - if (execution.runtime_profile, execution.invocations) != (condition.runtime_profile, condition.invocations): + if (execution.runtime_profile_name, execution.invocations) != ( + condition.runtime_profile_name, + condition.invocations, + ): raise NonReplayableObservationError( "The stored execution ran different invocations or a different runtime profile than the " "ProgramBehaves condition." ) - value = match_program_behaves(condition=condition, payload=evidence, acquisition=observation.acquisition) + evaluation = _evaluate_program(condition=condition, payload=evidence, acquisition=observation.acquisition) return [ Score( - score_value=None if value is None else str(value).lower(), - status=ScoreStatus.UNDETERMINED if value is None else ScoreStatus.COMPLETE, + score_value=None if evaluation.value is None else str(evaluation.value).lower(), + status=ScoreStatus.UNDETERMINED if evaluation.value is None else ScoreStatus.COMPLETE, score_type="true_false", - score_rationale=self._rationale( - condition=condition, payload=evidence, acquisition=observation.acquisition, value=value - ), + score_rationale=evaluation.rationale, score_value_description="The program produced the expected result for every case.", scorer_class_identifier=self.get_identifier(), scorable=observation.scorable, @@ -226,53 +229,3 @@ def _score_observation( observation_ids=[observation.id], ) ] - - @staticmethod - def _rationale( - *, - condition: ProgramBehaves, - payload: ProgramExecutionObservationPayload, - acquisition: Acquisition, - value: bool | None, - ) -> str: - """ - Explain the verdict in terms of the cases. - - Returns: - str: The rationale. - """ - total = len(condition.cases) - profile = condition.runtime_profile - if value is True: - return f"All {total} case(s) produced the expected result under {profile}." - if acquisition in (Acquisition.ERROR, Acquisition.UNAVAILABLE): - gaps = ", ".join(payload.reasons) or "no execution evidence" - return f"The program's behavior under {profile} could not be observed ({gaps})." - outcomes = [] - for index, case in enumerate(condition.cases): - event = payload.event_for(index) - if event is None: - outcomes.append(f"case {index} did not run") - continue - result = match_case(expected=case.expected, event=event) - if result is True: - continue - if event.timed_out: - detail = "timed out" - elif event.resource_limit_exceeded: - detail = "exceeded a resource limit" - elif event.exit_code is None: - detail = "did not report an exit code" - elif case.expected.exit_code is not None and event.exit_code != case.expected.exit_code: - detail = f"exited with {event.exit_code}, expected {case.expected.exit_code}" - elif result is False: - detail = "output did not match" - else: - detail = "output was cut before the check could be settled" - outcomes.append(f"case {index} {detail}") - summary = "; ".join(outcomes) - if value is False: - return f"The program did not produce the expected result under {profile}: {summary}." - gaps = ", ".join(payload.reasons) - suffix = f" ({gaps})" if gaps else "" - return f"The program's behavior under {profile} could not be fully established: {summary}{suffix}." diff --git a/tests/integration/score/test_program_execution_scorer_integration.py b/tests/integration/score/test_program_execution_scorer_integration.py new file mode 100644 index 0000000000..0261fd21ea --- /dev/null +++ b/tests/integration/score/test_program_execution_scorer_integration.py @@ -0,0 +1,299 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Opt-in shared execution and scoring against a caller-provided Docker daemon and image.""" + +from __future__ import annotations + +import asyncio +import hashlib +from typing import TYPE_CHECKING +from unittest.mock import patch + +import pytest + +from pyrit._code_execution import DockerRuntimeProfile, DockerSandboxClient +from pyrit.models import ( + Acquisition, + Contains, + ContentScorable, + Equals, + ProgramBehaves, + ProgramCase, + ProgramExecutionObservationPayload, + ProgramExecutionScorable, + ProgramResult, + ScoreStatus, + ScoringExpectation, +) +from pyrit.models.code_execution import _ProgramInvocation +from pyrit.score import ProgramExecutionScorer, ProgramExecutionSource + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator + from pathlib import Path + + from pyrit.memory import SQLiteMemory + +pytestmark = [pytest.mark.run_only_if_all_tests, pytest.mark.usefixtures("patch_central_database")] + +_IMAGE = "python@sha256:a6e34c598f2467ed0e9a8d349809fcd8b5c603269512df273a0bb1784edc11b1" +_PROFILE = "python-3.12" +_PROGRAM = "import sys\nn = int(sys.argv[1]) if len(sys.argv) > 1 else int(sys.stdin.read())\nprint(n * n)\n" + + +@pytest.fixture +async def docker_sandbox(tmp_path: Path) -> AsyncGenerator[DockerSandboxClient]: + client = DockerSandboxClient( + profiles={_PROFILE: DockerRuntimeProfile(image=_IMAGE, command=("python3", "-I"))}, + timeout_seconds=5, + max_output_bytes=128, + ) + client._temp_root = tmp_path + for command in (("version", "--format", "{{.Server.Version}}"), ("image", "inspect", _IMAGE)): + status, output = await client._control_async(*command) + if status != 0: + pytest.fail( + f"Live execution requires a running Docker daemon and a pre-pulled {_IMAGE}. " + f"Docker {' '.join(command)} failed: {output}" + ) + with patch.object(client, "_run_arguments", wraps=client._run_arguments) as arguments: + try: + yield client + finally: + problems = [] + for call in arguments.call_args_list: + name = call.kwargs["name"] + status, output = await client._control_async( + "container", "ls", "--all", "--filter", f"name=^/{name}$", "--format", "{{.Names}}" + ) + if status != 0: + problems.append(f"Could not verify container cleanup for {name}: {output}") + elif output.strip(): + problems.append(f"Container {name} remained after execution") + removed, detail = await client._control_async("rm", "--force", "--volumes", name) + if removed != 0: + problems.append(f"Could not remove test container {name}: {detail}") + remaining = await asyncio.to_thread(lambda: list(tmp_path.iterdir())) + for directory in remaining: + problems.append(f"Program directory remained after execution: {directory}") + if directory.name.startswith("pyrit-program-"): + await asyncio.to_thread( + client._remove_workspace, + workspace_dir=directory, + program_path=directory / "program" / "main.py", + ) + assert not problems, "\n".join(problems) + + +def _expectation(*, stdout: str = "ok") -> ScoringExpectation: + return ScoringExpectation( + conditions=( + ProgramBehaves( + runtime_profile_name=_PROFILE, + cases=(ProgramCase(expected=ProgramResult(stdout=Equals(value=stdout))),), + ), + ) + ) + + +async def test_real_client_returns_results_without_a_scorer_async(docker_sandbox: DockerSandboxClient) -> None: + run = await docker_sandbox.execute_async( + program="import sys\nprint(sys.argv[1] + ':' + sys.stdin.read())\n", + runtime_profile_name=_PROFILE, + invocations=( + _ProgramInvocation(arguments=("first",), stdin="one"), + _ProgramInvocation(arguments=("second",), stdin="two"), + ), + ) + assert [(event.invocation_index, event.exit_code, event.stdout) for event in run.events] == [ + (0, 0, "first:one\n"), + (1, 0, "second:two\n"), + ] + assert run.reasons == run.warnings == () + + +async def test_recorded_digest_matches_the_exact_executed_bytes_async( + docker_sandbox: DockerSandboxClient, sqlite_instance: SQLiteMemory +) -> None: + program = ( + "import hashlib\nfrom pathlib import Path\nprint(hashlib.sha256(Path(__file__).read_bytes()).hexdigest())\n" + ) + digest = hashlib.sha256(program.encode("utf-8")).hexdigest() + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=docker_sandbox)) + score = ( + await scorer.score_async(scorable=ContentScorable(value=program), expectation=_expectation(stdout=digest)) + )[0] + assert score.get_value() is True + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + assert observation.payload.program_sha256 == digest + assert observation.payload.events[0].stdout.strip() == digest + + +async def test_real_execution_persistence_and_replay_async( + docker_sandbox: DockerSandboxClient, sqlite_instance: SQLiteMemory +) -> None: + condition = ProgramBehaves( + runtime_profile_name=_PROFILE, + cases=( + ProgramCase( + invocation=_ProgramInvocation(arguments=("3",)), expected=ProgramResult(stdout=Equals(value="9")) + ), + ProgramCase( + invocation=_ProgramInvocation(stdin="12\n"), expected=ProgramResult(stdout=Equals(value="144")) + ), + ), + ) + expectation = ScoringExpectation(objective="Square each input", conditions=(condition,)) + execution = ProgramExecutionScorable( + program=ContentScorable(value=_PROGRAM), runtime_profile_name=_PROFILE, invocations=condition.invocations + ) + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=docker_sandbox)) + with patch.object(docker_sandbox, "execute_async", wraps=docker_sandbox.execute_async) as execute: + first = (await scorer.score_async(scorable=execution, expectation=expectation))[0] + second = (await scorer.score_async(scorable=execution, expectation=expectation))[0] + assert execute.await_count == 2 + assert first.observation_ids != second.observation_ids + stored = await sqlite_instance.get_scores_async(score_ids=[str(first.id), str(second.id)]) + assert len(stored) == 2 + for score in stored: + assert score.get_value() is True + assert score.status is ScoreStatus.COMPLETE + assert score.scored_expectation == expectation + assert score.scorable == execution + observation = (await sqlite_instance.get_observations_async(observation_ids=first.observation_ids))[0] + assert observation.acquisition is Acquisition.COMPLETE + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + assert observation.payload.program_sha256 == hashlib.sha256(_PROGRAM.encode("utf-8")).hexdigest() + assert [event.stdout for event in observation.payload.events] == ["9\n", "144\n"] + for event, invocation in zip(observation.payload.events, condition.invocations, strict=True): + assert event.runtime_profile_name == _PROFILE + assert event.runtime_image_digest == _IMAGE.rsplit("@", 1)[1] + assert event.resolved_command == ("python3", "-I", "/program/main.py", *invocation.arguments) + assert event.exit_code == 0 + assert not event.timed_out and not event.output_truncated + stricter = condition.model_copy( + update={ + "cases": ( + condition.cases[0], + ProgramCase( + invocation=condition.cases[1].invocation, expected=ProgramResult(stdout=Equals(value="145")) + ), + ) + } + ) + replayed = ( + await scorer.score_observation_async( + observation=observation, expectation=ScoringExpectation(conditions=(stricter,)) + ) + )[0] + assert replayed.get_value() is False + assert replayed.observation_ids == first.observation_ids + assert execute.await_count == 2 + saved_replay = (await sqlite_instance.get_scores_async(score_ids=[str(replayed.id)]))[0] + assert saved_replay.get_value() is False + assert saved_replay.observation_ids == first.observation_ids + + +@pytest.mark.parametrize(("program", "exit_code"), [("print('wrong')", 0), ("raise SystemExit(7)", 7), ("def (", 1)]) +async def test_real_program_failure_is_false_async( + docker_sandbox: DockerSandboxClient, sqlite_instance: SQLiteMemory, program: str, exit_code: int +) -> None: + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=docker_sandbox)) + score = (await scorer.score_async(scorable=ContentScorable(value=program), expectation=_expectation()))[0] + assert score.status is ScoreStatus.COMPLETE + assert score.get_value() is False + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert observation.acquisition is Acquisition.COMPLETE + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + assert observation.payload.events[0].exit_code == exit_code + + +async def test_real_timeout_is_false_async(docker_sandbox: DockerSandboxClient, sqlite_instance: SQLiteMemory) -> None: + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=docker_sandbox)) + score = ( + await scorer.score_async(scorable=ContentScorable(value="while True:\n pass\n"), expectation=_expectation()) + )[0] + assert score.status is ScoreStatus.COMPLETE + assert score.get_value() is False + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + assert observation.payload.events[0].timed_out + assert observation.payload.events[0].exit_code is None + + +async def test_real_output_is_bounded_async(docker_sandbox: DockerSandboxClient, sqlite_instance: SQLiteMemory) -> None: + expectation = ScoringExpectation( + conditions=( + ProgramBehaves( + runtime_profile_name=_PROFILE, + cases=(ProgramCase(expected=ProgramResult(stdout=Contains(value="xxx"))),), + ), + ) + ) + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=docker_sandbox)) + score = ( + await scorer.score_async( + scorable=ContentScorable(value="import sys\nsys.stdout.write('x' * 8192)\nsys.stderr.write('e' * 8192)\n"), + expectation=expectation, + ) + )[0] + assert score.get_value() is True + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + event = observation.payload.events[0] + assert event.stdout == "x" * 128 and event.stdout_truncated + assert event.stderr == "e" * 128 and event.stderr_truncated + + +async def test_missing_image_is_undetermined_async( + docker_sandbox: DockerSandboxClient, sqlite_instance: SQLiteMemory +) -> None: + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=docker_sandbox)) + with patch.object( + docker_sandbox, + "_profiles", + {_PROFILE: DockerRuntimeProfile(image="python@sha256:" + "0" * 64, command=("python3", "-I"))}, + ): + score = (await scorer.score_async(scorable=ContentScorable(value="print('ok')"), expectation=_expectation()))[0] + assert score.status is ScoreStatus.UNDETERMINED + assert score.score_value is None + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert observation.acquisition is Acquisition.ERROR + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + assert observation.payload.events == () + assert observation.payload.reasons + + +async def test_cancelling_real_execution_cleans_up_async(docker_sandbox: DockerSandboxClient) -> None: + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=docker_sandbox)) + with ( + patch.object(docker_sandbox, "_timeout_seconds", 30), + patch.object(docker_sandbox, "_run_arguments", wraps=docker_sandbox._run_arguments) as arguments, + ): + task = asyncio.create_task( + scorer.score_async(scorable=ContentScorable(value="while True:\n pass\n"), expectation=_expectation()) + ) + + async def wait_until_running_async() -> None: + while not task.done(): + if arguments.call_args is not None: + name = arguments.call_args.kwargs["name"] + status, output = await docker_sandbox._control_async( + "inspect", "--format", "{{.State.Running}}", name + ) + if status == 0 and output.strip() == "true": + return + await asyncio.sleep(0.01) + pytest.fail(f"Execution ended before cancellation: {task.result()}") + + try: + await asyncio.wait_for(wait_until_running_async(), timeout=15) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=15) + finally: + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) diff --git a/tests/integration/score/test_scorer_notebooks.py b/tests/integration/score/test_scorer_notebooks.py index b8444482ac..1b0ff9db2c 100644 --- a/tests/integration/score/test_scorer_notebooks.py +++ b/tests/integration/score/test_scorer_notebooks.py @@ -17,9 +17,16 @@ @pytest.mark.parametrize( "file_name", - [file for file in os.listdir(nb_directory_path) if file.endswith(".ipynb") and file not in skipped_files], + [ + pytest.param( + file, + marks=pytest.mark.run_only_if_all_tests if file == "7_program_execution_scorer.ipynb" else (), + ) + for file in os.listdir(nb_directory_path) + if file.endswith(".ipynb") and file not in skipped_files + ], ) -def test_execute_notebooks(file_name): +def test_execute_notebooks(file_name: str) -> None: nb_path = pathlib.Path(nb_directory_path, file_name).resolve() with open(nb_path, encoding="utf-8") as f: nb = nbformat.read(f, as_version=4) diff --git a/tests/unit/_code_execution/test_docker_client.py b/tests/unit/_code_execution/test_docker_client.py new file mode 100644 index 0000000000..9f65d8219b --- /dev/null +++ b/tests/unit/_code_execution/test_docker_client.py @@ -0,0 +1,533 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Deterministic Docker fault injection; real execution is covered by integration tests.""" + +from __future__ import annotations + +import asyncio +import os +import stat +import threading +from contextlib import nullcontext +from typing import TYPE_CHECKING +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from pyrit._code_execution import DockerRuntimeProfile, DockerSandboxClient, SandboxExecutionError +from pyrit.models.code_execution import _ProgramInvocation, _SandboxRun + +if TYPE_CHECKING: + from pathlib import Path + +_IMAGE = "python@sha256:" + "c" * 64 +_STARTED = "2026-10-21T00:00:00Z" + + +@pytest.fixture +def client(tmp_path: Path) -> DockerSandboxClient: + instance = DockerSandboxClient( + profiles={"python": DockerRuntimeProfile(image=_IMAGE, command=("python3", "-I"))}, + docker_executable=str(tmp_path / "no-docker"), + ) + instance._temp_root = tmp_path + return instance + + +def _process(*, returncode: int | None = 0, output: bytes = b"ok", eof: bool = True) -> MagicMock: + process = MagicMock(spec=asyncio.subprocess.Process) + process.returncode, process.stdin = returncode, None + readers = [asyncio.StreamReader(), asyncio.StreamReader()] + for reader, data in zip(readers, (output, b""), strict=True): + reader.feed_data(data) + if eof: + reader.feed_eof() + process.stdout, process.stderr = readers + exited = asyncio.Event() + if returncode is not None: + exited.set() + + async def wait_async() -> int: + await exited.wait() + code = process.returncode + assert isinstance(code, int) + return code + + def kill() -> None: + process.returncode = -9 + exited.set() + for reader in readers: + reader.feed_eof() + + process.wait = AsyncMock(side_effect=wait_async) + process.kill.side_effect = kill + process.communicate = AsyncMock(return_value=(b"", b"")) + return process + + +def _control( + *, + state: str = "exited", + code: int = 0, + oom: bool = False, + removal_fails: bool = False, + capabilities: str = "cgroupfs true true true true true", + volumes: str = "null", +) -> AsyncMock: + async def command_async(*arguments: str) -> tuple[int, str]: + nonlocal state, code + if arguments[0] == "info": + return 0, capabilities + if arguments[:2] == ("image", "inspect"): + return 0, volumes + if arguments[0] == "inspect": + started = "0001-01-01T00:00:00Z" if state == "created" else _STARTED + return 0, f"{state} {code} {str(oom).lower()} {started}" + if arguments[0] == "kill": + state, code = "exited", 137 + return 0, "" + if arguments[:2] == ("container", "ls"): + return 0, "" + assert arguments[:3] == ("rm", "--force", "--volumes"), arguments + return (1, "daemon unavailable") if removal_fails else (0, "") + + return AsyncMock(side_effect=command_async) + + +async def _run_async(client: DockerSandboxClient) -> _SandboxRun: + return await client.execute_async( + program="print('ok')", runtime_profile_name="python", invocations=(_ProgramInvocation(),) + ) + + +@pytest.mark.parametrize( + "capabilities", + [ + "none true true true true true", + "cgroupfs false true true true true", + "cgroupfs true false true true true", + "systemd true true false true true", + "systemd true true true false true", + "systemd true true true true false", + "", + "unexpected response", + ], +) +async def test_unsupported_resource_controls_prevent_execution_async( + client: DockerSandboxClient, tmp_path: Path, capabilities: str +) -> None: + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as launch, + patch.object(client, "_control_async", _control(capabilities=capabilities)), + pytest.raises(SandboxExecutionError, match="must support cgroup"), + ): + await _run_async(client) + launch.assert_not_awaited() + assert await asyncio.to_thread(lambda: list(tmp_path.iterdir())) == [] + + +@pytest.mark.parametrize("volumes", ['{"/data":{}}', '{"/tmp":{}}', "invalid metadata"]) +async def test_image_volumes_prevent_execution_async(client: DockerSandboxClient, tmp_path: Path, volumes: str) -> None: + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as launch, + patch.object(client, "_control_async", _control(volumes=volumes)), + pytest.raises(SandboxExecutionError, match="must not declare volumes"), + ): + await _run_async(client) + launch.assert_not_awaited() + assert await asyncio.to_thread(lambda: list(tmp_path.iterdir())) == [] + + +async def test_systemd_controls_and_empty_volume_map_allow_execution_async(client: DockerSandboxClient) -> None: + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=_process()), + patch.object(client, "_control_async", _control(capabilities="systemd true true true true true", volumes="{}")), + ): + assert (await _run_async(client)).events[0].exit_code == 0 + + +@pytest.mark.parametrize(("command", "message"), [("info", "resource support"), ("image", "runtime image")]) +async def test_preflight_failures_prevent_execution_async( + client: DockerSandboxClient, command: str, message: str +) -> None: + healthy = _control() + + async def control_async(*arguments: str) -> tuple[int, str]: + if arguments[0] == command: + return 1, "unavailable" + response: tuple[int, str] = await healthy(*arguments) + return response + + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as launch, + patch.object(client, "_control_async", side_effect=control_async), + pytest.raises(SandboxExecutionError, match=f"could not inspect.*{message}"), + ): + await _run_async(client) + launch.assert_not_awaited() + + +def test_run_arguments_enforce_runtime_and_isolation(client: DockerSandboxClient, tmp_path: Path) -> None: + command = ("python3", "-I", "/program/main.py", "argument") + argv = client._run_arguments( + name="test-container", profile=client._profiles["python"], program_dir=tmp_path, resolved_command=command + ) + for flag, value in { + "--pull": "never", + "--log-driver": "none", + "--network": "none", + "--cap-drop": "ALL", + "--security-opt": "no-new-privileges", + "--user": "65534:65534", + "--memory": "256m", + "--memory-swap": "256m", + "--cpus": "1.0", + "--pids-limit": "64", + }.items(): + assert argv[argv.index(flag) + 1] == value + assert "--read-only" in argv and "--privileged" not in argv + assert argv[argv.index("--entrypoint") + 1] == command[0] + assert argv.index("--entrypoint") < argv.index(_IMAGE) + assert argv[argv.index(_IMAGE) + 1 :] == list(command[1:]) + assert argv[argv.index("--mount") + 1].endswith(",target=/program,readonly") + + +@pytest.mark.parametrize( + ("state", "code", "oom", "client_code", "error"), + [ + ("exited", 0, False, 0, None), + ("exited", 125, False, 125, None), + ("exited", 137, True, 137, None), + ("running", 0, False, 1, "still running"), + ("created", 125, False, 125, "did not start"), + ("exited", 0, False, 1, "output capture may be incomplete"), + ], +) +async def test_container_and_client_outcomes_async( + client: DockerSandboxClient, + tmp_path: Path, + state: str, + code: int, + oom: bool, + client_code: int, + error: str | None, +) -> None: + process = _process(returncode=client_code) + control = _control(state=state, code=code, oom=oom) + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=process), + patch.object(client, "_control_async", control), + patch.object(client, "_run_arguments", wraps=client._run_arguments) as arguments, + ): + if error is not None: + with pytest.raises(SandboxExecutionError, match=error): + await _run_async(client) + else: + event = (await _run_async(client)).events[0] + assert event.exit_code == (None if oom else code) + assert event.resource_limit_exceeded is oom + assert event.runtime_image_digest == "sha256:" + "c" * 64 + assert control.call_args_list[-1].args[:3] == ("rm", "--force", "--volumes") + mounted_dir = arguments.call_args.kwargs["program_dir"] + assert mounted_dir.name == "program" + assert mounted_dir.parent.parent == tmp_path + process.communicate.assert_awaited_once() + assert await asyncio.to_thread(lambda: list(tmp_path.iterdir())) == [] + + +@pytest.mark.parametrize( + ("eof", "state", "error"), + [ + (False, "exited", "a program timeout could not be established"), + (True, "running", "still running"), + (True, "exited", None), + (False, "running", None), + ], +) +async def test_deadline_distinguishes_stalled_capture_client_and_program_async( + client: DockerSandboxClient, eof: bool, state: str, error: str | None +) -> None: + process, control = _process(returncode=None, eof=eof), _control(state=state) + with ( + patch.object(client, "_timeout_seconds", 0.05), + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=process), + patch.object(client, "_control_async", control), + ): + if error is not None: + with pytest.raises(SandboxExecutionError, match=error): + await _run_async(client) + else: + event = (await _run_async(client)).events[0] + assert event.timed_out is (not eof) + assert event.exit_code == (0 if eof else None) + assert event.stdout_truncated and event.stderr_truncated + commands = [call.args[0] for call in control.call_args_list] + assert ("kill" in commands) is (state == "running" and not eof) + assert commands[-1] == "rm" + process.kill.assert_called_once() + + +@pytest.mark.parametrize("count", [1, 2]) +async def test_cleanup_failure_keeps_evidence_and_stops_further_invocations_async( + client: DockerSandboxClient, count: int +) -> None: + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=_process()) as launch, + patch.object(client, "_control_async", _control(removal_fails=True)), + ): + run = await client.execute_async( + program="print('ok')", runtime_profile_name="python", invocations=(_ProgramInvocation(),) * count + ) + assert len(run.events) == 1 and run.events[0].exit_code == 0 + assert "could not remove container" in run.warnings[0] + assert bool(run.reasons) is (count > 1) + launch.assert_awaited_once() + + +async def test_failure_after_a_completed_invocation_returns_partial_results_async(client: DockerSandboxClient) -> None: + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock, side_effect=[_process(), OSError("offline")]), + patch.object(client, "_control_async", _control()), + ): + run = await client.execute_async( + program="print('ok')", runtime_profile_name="python", invocations=(_ProgramInvocation(),) * 2 + ) + assert [event.invocation_index for event in run.events] == [0] + assert "invocation 1" in run.reasons[0] and "offline" in run.reasons[0] + + +@pytest.mark.parametrize("failed", [False, True]) +async def test_startup_is_owned_until_cleanup_async(client: DockerSandboxClient, tmp_path: Path, failed: bool) -> None: + process, entered, release = _process(returncode=None), asyncio.Event(), asyncio.Event() + created: list[str] = [] + healthy = _control() + + async def start_async(*args: str, **kwargs: object) -> MagicMock: + created.append(args[args.index("--name") + 1]) + entered.set() + await release.wait() + if failed: + raise OSError("stdio connection failed after container creation") + return process + + async def control_async(*arguments: str) -> tuple[int, str]: + if arguments[0] in ("info", "image"): + response: tuple[int, str] = await healthy(*arguments) + return response + if arguments[:2] == ("container", "ls"): + assert f"name=^/{created[0]}$" in arguments + return 0, created[0] + assert arguments == ("rm", "--force", "--volumes", created[0]) + return 0, "" + + with ( + patch("asyncio.create_subprocess_exec", side_effect=start_async), + patch.object(client, "_control_async", side_effect=control_async) as control, + ): + task = asyncio.create_task(_run_async(client)) + try: + await asyncio.wait_for(entered.wait(), timeout=5) + if not failed: + for _ in range(2): + task.cancel() + await asyncio.sleep(0.01) + assert not task.done() + assert all(call.args[0] in ("info", "image") for call in control.call_args_list) + release.set() + with pytest.raises(SandboxExecutionError if failed else asyncio.CancelledError): + await asyncio.wait_for(task, timeout=5) + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) + assert control.call_args.args == ("rm", "--force", "--volumes", created[0]) + assert await asyncio.to_thread(lambda: list(tmp_path.iterdir())) == [] + if not failed: + process.kill.assert_called_once() + process.communicate.assert_awaited_once() + + +@pytest.mark.parametrize("newline", ["\n", "\r\n"]) +def test_program_bytes_and_readonly_file_cleanup(client: DockerSandboxClient, tmp_path: Path, newline: str) -> None: + workspace = tmp_path / "workspace" + path = workspace / "program" / "main.py" + program = "print('\u00e9')\nprint(42)\n".replace("\n", newline) + try: + client._write_program(workspace_dir=workspace, program_path=path, program=program) + assert path.read_bytes() == program.encode("utf-8") + if os.name == "posix": + assert stat.S_IMODE(workspace.stat().st_mode) == 0o700 + assert stat.S_IMODE(path.parent.stat().st_mode) == 0o755 + assert stat.S_IMODE(path.stat().st_mode) == 0o444 + finally: + client._remove_workspace(workspace_dir=workspace, program_path=path) + assert not workspace.exists() + + +async def test_partial_workspace_is_removed_after_write_failure_async( + client: DockerSandboxClient, tmp_path: Path +) -> None: + with ( + patch.object(client, "_control_async", _control()), + patch("pathlib.Path.write_bytes", side_effect=OSError("disk full")), + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as launch, + pytest.raises(OSError, match="disk full"), + ): + await _run_async(client) + launch.assert_not_awaited() + assert await asyncio.to_thread(lambda: list(tmp_path.iterdir())) == [] + + +@pytest.mark.parametrize("stage", ["drain", "close"]) +async def test_input_cancellation_aborts_buffered_writes_async(stage: str) -> None: + stream = MagicMock(spec=asyncio.StreamWriter) + stream.transport = MagicMock(spec=asyncio.WriteTransport) + blocked, released = asyncio.Event(), asyncio.Event() + + async def drain_async() -> None: + if stage == "drain": + blocked.set() + await asyncio.Event().wait() + + async def close_async() -> None: + if stage == "close": + blocked.set() + await released.wait() + + stream.transport.abort.side_effect = released.set + stream.drain, stream.wait_closed = AsyncMock(side_effect=drain_async), AsyncMock(side_effect=close_async) + task = asyncio.create_task(DockerSandboxClient._feed_stdin_async(stream=stream, data="input")) + await asyncio.wait_for(blocked.wait(), timeout=5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + stream.transport.abort.assert_called_once() + stream.wait_closed.assert_awaited_once() + + +@pytest.mark.parametrize("phase", ["write", "remove"]) +@pytest.mark.parametrize("cancel", [False, True]) +async def test_file_work_does_not_block_or_outlive_the_call_async( + client: DockerSandboxClient, tmp_path: Path, phase: str, cancel: bool +) -> None: + entered, release = threading.Event(), threading.Event() + operation = "_write_program" if phase == "write" else "_remove_workspace" + original = getattr(client, operation) + loop_thread = threading.get_ident() + + def gated(**kwargs: object) -> None: + assert threading.get_ident() != loop_thread + entered.set() + assert release.wait(5), "file work blocked the event loop" + original(**kwargs) + + with ( + patch.object(client, operation, side_effect=gated), + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=_process()), + patch.object(client, "_control_async", _control()), + ): + task = asyncio.create_task(_run_async(client)) + try: + for _ in range(200): + if entered.is_set(): + break + await asyncio.sleep(0.01) + assert entered.is_set() + if cancel: + for _ in range(2): + task.cancel() + await asyncio.sleep(0.01) + assert not task.done() + release.set() + if cancel: + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=5) + else: + assert (await asyncio.wait_for(task, timeout=5)).events[0].exit_code == 0 + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) + assert await asyncio.to_thread(lambda: list(tmp_path.iterdir())) == [] + + +@pytest.mark.parametrize("phase", ["container", "management"]) +async def test_repeated_cancellation_drains_process_cleanup_async(client: DockerSandboxClient, phase: str) -> None: + process = _process(returncode=None, eof=False) + running, cleaning, release = asyncio.Event(), asyncio.Event(), asyncio.Event() + calls = 0 + + async def communicate_async() -> tuple[bytes, bytes]: + nonlocal calls + calls += 1 + if phase == "management" and calls == 1: + running.set() + await asyncio.Event().wait() + cleaning.set() + await release.wait() + return b"", b"" + + async def collect_async(**kwargs: object) -> None: + running.set() + await asyncio.Event().wait() + + process.communicate = AsyncMock(side_effect=communicate_async) + control = _control() + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=process), + patch.object(client, "_collect_output_async", side_effect=collect_async), + patch.object(client, "_control_async", control) if phase == "container" else nullcontext(), + ): + task = asyncio.create_task( + _run_async(client) if phase == "container" else client._control_async("inspect", "test") + ) + try: + await asyncio.wait_for(running.wait(), timeout=5) + task.cancel() + await asyncio.wait_for(cleaning.wait(), timeout=5) + task.cancel() + await asyncio.sleep(0.01) + assert not task.done() + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=5) + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) + process.kill.assert_called_once() + assert process.communicate.await_count == (2 if phase == "management" else 1) + if phase == "container": + removals = [call for call in control.call_args_list if call.args[0] == "rm"] + assert len(removals) == 1 + assert removals[0].args[:3] == ("rm", "--force", "--volumes") + + +@pytest.mark.parametrize("case", ["empty", "unknown", "missing-client"]) +async def test_invalid_requests_fail_before_execution_async(client: DockerSandboxClient, case: str) -> None: + with pytest.raises(SandboxExecutionError if case == "missing-client" else ValueError): + await client.execute_async( + program="print(1)", + runtime_profile_name="unknown" if case == "unknown" else "python", + invocations=() if case == "empty" else (_ProgramInvocation(),), + ) + + +@pytest.mark.parametrize( + "settings", [{"timeout_seconds": 0}, {"pids_limit": 0}, {"max_output_bytes": 0}, {"memory": "bad"}] +) +def test_invalid_limits(settings: dict[str, object]) -> None: + with pytest.raises(ValueError): + DockerSandboxClient(profiles={"python": DockerRuntimeProfile(image=_IMAGE, command=("python3",))}, **settings) # type: ignore[arg-type] + + +def test_profiles_are_explicit_trusted_configuration() -> None: + for settings in ( + {"image": "python:latest", "command": ("python3",)}, + {"image": _IMAGE, "command": ()}, + {"image": _IMAGE, "command": ("python3",), "program_filename": "../main.py"}, + ): + with pytest.raises(ValueError): + DockerRuntimeProfile(**settings) # type: ignore[arg-type] + with pytest.raises(ValueError): + DockerSandboxClient(profiles={}) + with pytest.raises(ValueError): + DockerSandboxClient(profiles={"bad name": DockerRuntimeProfile(image=_IMAGE, command=("python3",))}) diff --git a/tests/unit/common/test_lazy_package_imports.py b/tests/unit/common/test_lazy_package_imports.py index f8cb7a709a..655f73bed9 100644 --- a/tests/unit/common/test_lazy_package_imports.py +++ b/tests/unit/common/test_lazy_package_imports.py @@ -22,6 +22,18 @@ ) _LAZY_IMPORT_SPOT_CHECKS = [ + ( + "pyrit._code_execution", + "SandboxClient", + "pyrit._code_execution.client", + "pyrit._code_execution.docker_client", + ), + ( + "pyrit._code_execution", + "DockerSandboxClient", + "pyrit._code_execution.docker_client", + "pyrit.score", + ), ( "pyrit.analytics", "compute_technique_stats_async", @@ -456,6 +468,55 @@ def test_analytics_foundations_do_not_load_higher_layers() -> None: ) +def test_code_execution_is_independent_of_component_layers() -> None: + _assert_subprocess_succeeds( + """ + import importlib.util + import sys + import pyrit.models + import pyrit.models.code_execution as execution_models + from pyrit._code_execution import DockerRuntimeProfile, DockerSandboxClient, SandboxClient + from pyrit.models.code_execution import ( + _ProgramExecutionEvent, _ProgramInvocation, _RuntimeProfileName, _SandboxRun, + ) + + for name in ("_ProgramExecutionEvent", "_ProgramInvocation", "_RuntimeProfileName", "_SandboxRun"): + assert name not in pyrit.models.__all__ + assert not hasattr(pyrit.models, name) + public_name = name.removeprefix("_") + assert public_name not in pyrit.models.__all__ + assert not hasattr(pyrit.models, public_name) + assert not hasattr(execution_models, public_name) + + client = DockerSandboxClient( + profiles={"python": DockerRuntimeProfile(image="python@sha256:" + "a" * 64, command=("python3",))} + ) + assert isinstance(client, SandboxClient) + assert client.get_runtime_profile_names() == frozenset({"python"}) + assert client.get_identifier().class_module == "pyrit._code_execution.docker_client" + assert importlib.util.find_spec("pyrit.code_execution") is None + event = _ProgramExecutionEvent( + invocation_index=0, runtime_profile_name="python", + exit_code=0, stdout="result", duration_ms=1, + ) + assert event.resolved_command is None + assert _SandboxRun(events=(event,)).events[0].stdout == "result" + assert _ProgramInvocation(stdin="input").stdin == "input" + + forbidden = ( + "pyrit.score", "pyrit.models.score", "pyrit.memory", "pyrit.executor", + "pyrit.prompt_target", "pyrit.converter", "pyrit.datasets", "pyrit.scenario", + "pyrit.backend", "sqlalchemy", + ) + loaded = [ + name for name in sys.modules + if any(name == prefix or name.startswith(prefix + ".") for prefix in forbidden) + ] + assert not loaded, loaded + """ + ) + + def test_observation_acquisition_and_scorer_do_not_load_otel_sdk() -> None: _assert_subprocess_succeeds( """ diff --git a/tests/unit/memory/memory_interface/test_interface_observations.py b/tests/unit/memory/memory_interface/test_interface_observations.py index e362342140..77abf077b7 100644 --- a/tests/unit/memory/memory_interface/test_interface_observations.py +++ b/tests/unit/memory/memory_interface/test_interface_observations.py @@ -134,48 +134,6 @@ async def test_existing_observation_can_support_another_score(sqlite_instance: M assert (await sqlite_instance.get_scores_async(score_ids=[replay_score.id]))[0].observation_ids == [observation.id] -async def _stored_observation_and_score_async(memory: MemoryInterface) -> tuple[Observation, Score]: - piece = MessagePiece(role="assistant", original_value="response", conversation_id=str(uuid.uuid4())) - (await memory.add_message_pieces_to_memory_async(message_pieces=[piece])) - scorable = MessageScorable(message_piece_ids=(piece.id,)) - expectation = ScoringExpectation(objective="Judge the response") - observation = await _observation_async( - memory=memory, scorable=scorable, response_piece_id=piece.id, expectation=expectation - ) - score = Score( - score_value="true", - score_type="true_false", - scorable=scorable, - scored_expectation=expectation, - observation_ids=[observation.id], - ) - (await memory.add_scores_to_memory_async(scores=[score], observations=[observation])) - return observation, score - - -async def test_identical_new_observation_already_stored_is_linked_not_inserted(sqlite_instance: MemoryInterface): - observation, first_score = await _stored_observation_and_score_async(sqlite_instance) - second_score = first_score.model_copy(update={"id": uuid.uuid4()}) - - (await sqlite_instance.add_scores_to_memory_async(scores=[second_score], observations=[observation])) - - assert len(sqlite_instance._query_entries(ObservationEntry)) == 1 - stored = await sqlite_instance.get_scores_async(score_ids=[first_score.id, second_score.id]) - assert [score.observation_ids for score in stored] == [[observation.id], [observation.id]] - - -async def test_different_evidence_under_a_stored_observation_id_is_rejected(sqlite_instance: MemoryInterface): - observation, first_score = await _stored_observation_and_score_async(sqlite_instance) - changed = observation.model_copy(update={"metadata": {"note": "different"}}) - second_score = first_score.model_copy(update={"id": uuid.uuid4()}) - - with pytest.raises(ValueError, match="New observation IDs already exist"): - await sqlite_instance.add_scores_to_memory_async(scores=[second_score], observations=[changed]) - - assert await sqlite_instance.get_scores_async(score_ids=[second_score.id]) == [] - assert (await sqlite_instance.get_observations_async(observation_ids=[observation.id]))[0] == observation - - async def test_unreferenced_observation_is_not_persisted(sqlite_instance: MemoryInterface): piece = MessagePiece(role="assistant", original_value="response", conversation_id=str(uuid.uuid4())) (await sqlite_instance.add_message_pieces_to_memory_async(message_pieces=[piece])) diff --git a/tests/unit/models/test_code_execution.py b/tests/unit/models/test_code_execution.py new file mode 100644 index 0000000000..6f3b55c64e --- /dev/null +++ b/tests/unit/models/test_code_execution.py @@ -0,0 +1,75 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Execution data is usable without scoring, messages, or memory.""" + +import pytest +from pydantic import TypeAdapter, ValidationError + +from pyrit.models.code_execution import _ProgramExecutionEvent, _ProgramInvocation, _RuntimeProfileName, _SandboxRun + + +def _event(**updates: object) -> _ProgramExecutionEvent: + return _ProgramExecutionEvent.model_validate( + { + "invocation_index": 0, + "runtime_profile_name": "python-3.12", + "resolved_command": ("python3", "-I", "/program/main.py"), + "exit_code": 0, + "duration_ms": 1, + **updates, + } + ) + + +def test_invocation_arguments_reject_nul() -> None: + with pytest.raises(ValidationError): + _ProgramInvocation(arguments=("a\x00b",)) + + +@pytest.mark.parametrize("profile", ["", "-leading-dash", "has space", "image:tag", "a/b", "x" * 65]) +def test_runtime_profile_name_rejects_invalid_names(profile: str) -> None: + with pytest.raises(ValidationError): + TypeAdapter(_RuntimeProfileName).validate_python(profile) + + +def test_timed_out_event_cannot_report_an_exit_code() -> None: + with pytest.raises(ValidationError, match="timed-out"): + _event(timed_out=True, exit_code=0) + assert _event(timed_out=True, exit_code=None).exit_code is None + + +def test_runtime_image_provenance_is_optional_but_not_blank() -> None: + assert _event().runtime_image_digest is None + assert _event(runtime_image_digest="sha256:" + "a" * 64).runtime_image_digest == "sha256:" + "a" * 64 + with pytest.raises(ValidationError): + _event(runtime_image_digest="") + + +def test_command_provenance_is_optional_but_not_empty() -> None: + assert _event(resolved_command=None).resolved_command is None + with pytest.raises(ValidationError): + _event(resolved_command=()) + + +@pytest.mark.parametrize(("stdout", "stderr"), [(False, False), (True, False), (False, True), (True, True)]) +def test_output_truncation_reports_either_stream(stdout: bool, stderr: bool) -> None: + assert _event(stdout_truncated=stdout, stderr_truncated=stderr).output_truncated is (stdout or stderr) + + +def test_invocation_index_is_nonnegative() -> None: + with pytest.raises(ValidationError): + _event(invocation_index=-1) + + +def test_sandbox_run_round_trips_shared_execution_data() -> None: + run = _SandboxRun( + events=(_event(stdout="result"),), + reasons=("the backend stopped before invocation 1",), + warnings=("cleanup was incomplete",), + ) + adapter = TypeAdapter(_SandboxRun) + serialized = adapter.dump_json(run) + assert b'"invocation_index":0' in serialized + assert b"case_index" not in serialized + assert adapter.validate_json(serialized) == run diff --git a/tests/unit/models/test_program.py b/tests/unit/models/test_program.py index 1ea1e730c1..ba06c355b1 100644 --- a/tests/unit/models/test_program.py +++ b/tests/unit/models/test_program.py @@ -19,26 +19,25 @@ Observation, ProgramBehaves, ProgramCase, - ProgramExecutionEvent, ProgramExecutionObservationPayload, ProgramExecutionScorable, - ProgramInvocation, ProgramResult, Score, ScoringExpectation, SurfaceScorable, scorable_from_dict, ) +from pyrit.models.code_execution import _ProgramExecutionEvent, _ProgramInvocation _SOURCE = ComponentIdentifier(class_name="FakeSource", class_module="tests.fake") _DIGEST = "a" * 64 _IMAGE_DIGEST = "sha256:" + "b" * 64 -def _event(*, case_index: int = 0, profile: str = "python-3.12", **updates: object) -> ProgramExecutionEvent: +def _event(*, invocation_index: int = 0, profile: str = "python-3.12", **updates: object) -> _ProgramExecutionEvent: values: dict[str, object] = { - "case_index": case_index, - "runtime_profile": profile, + "invocation_index": invocation_index, + "runtime_profile_name": profile, "runtime_image_digest": _IMAGE_DIGEST, "resolved_command": ("python3", "/program/main.py"), "exit_code": 0, @@ -46,15 +45,14 @@ def _event(*, case_index: int = 0, profile: str = "python-3.12", **updates: obje "duration_ms": 5, } values.update(updates) - return ProgramExecutionEvent.model_validate(values) + return _ProgramExecutionEvent.model_validate(values) def _execution(*, invocations: int = 2, profile: str = "python-3.12") -> ProgramExecutionScorable: return ProgramExecutionScorable( - execution_id=uuid.uuid4(), program=ContentScorable(value="print('ok')"), - runtime_profile=profile, - invocations=tuple(ProgramInvocation(arguments=(str(index),)) for index in range(invocations)), + runtime_profile_name=profile, + invocations=tuple(_ProgramInvocation(arguments=(str(index),)) for index in range(invocations)), ) @@ -62,7 +60,7 @@ def _observation( *, scorable: ProgramExecutionScorable, acquisition: Acquisition, - events: tuple[ProgramExecutionEvent, ...] = (), + events: tuple[_ProgramExecutionEvent, ...] = (), reasons: tuple[str, ...] = (), ) -> Observation: return Observation( @@ -77,10 +75,10 @@ def _observation( def test_program_behaves_round_trips_through_condition_registry() -> None: condition = ProgramBehaves( - runtime_profile="python-3.12", + runtime_profile_name="python-3.12", cases=( ProgramCase( - invocation=ProgramInvocation(arguments=("3",), stdin="x"), + invocation=_ProgramInvocation(arguments=("3",), stdin="x"), expected=ProgramResult(exit_code=None, stdout=Equals(value="9")), ), ), @@ -92,50 +90,36 @@ def test_program_behaves_round_trips_through_condition_registry() -> None: def test_program_behaves_invocations_drop_expected_results() -> None: condition = ProgramBehaves( - runtime_profile="python-3.12", + runtime_profile_name="python-3.12", cases=( - ProgramCase(invocation=ProgramInvocation(arguments=("1",)), expected=ProgramResult(exit_code=3)), - ProgramCase(invocation=ProgramInvocation(stdin="in")), + ProgramCase(invocation=_ProgramInvocation(arguments=("1",)), expected=ProgramResult(exit_code=3)), + ProgramCase(invocation=_ProgramInvocation(stdin="in")), ), ) - assert condition.invocations == (ProgramInvocation(arguments=("1",)), ProgramInvocation(stdin="in")) + assert condition.invocations == (_ProgramInvocation(arguments=("1",)), _ProgramInvocation(stdin="in")) def test_program_case_defaults_expect_a_clean_exit_with_no_arguments() -> None: case = ProgramCase() - assert case.invocation == ProgramInvocation() + assert case.invocation == _ProgramInvocation() assert case.expected == ProgramResult(exit_code=0) def test_program_behaves_requires_a_case() -> None: with pytest.raises(ValidationError, match="at least one case"): - ProgramBehaves(runtime_profile="python-3.12", cases=()) + ProgramBehaves(runtime_profile_name="python-3.12", cases=()) -@pytest.mark.parametrize("profile", ["", "-leading-dash", "has space", "image:tag", "a/b", "x" * 65]) -def test_runtime_profile_must_be_a_plain_name(profile: str) -> None: +def test_condition_validates_its_runtime_profile_name() -> None: with pytest.raises(ValidationError): - ProgramBehaves(runtime_profile=profile, cases=(ProgramCase(),)) - - -def test_invocation_arguments_reject_nul() -> None: - with pytest.raises(ValidationError): - ProgramInvocation(arguments=("a\x00b",)) - - -def test_timed_out_event_cannot_report_an_exit_code() -> None: - with pytest.raises(ValidationError, match="timed-out"): - _event(timed_out=True, exit_code=0) - assert _event(timed_out=True, exit_code=None).output_truncated is False - assert _event(stderr_truncated=True).output_truncated is True + ProgramBehaves(runtime_profile_name="bad name", cases=(ProgramCase(),)) def test_execution_scorable_round_trips_as_a_stored_scorable() -> None: message_program = ProgramExecutionScorable( - execution_id=uuid.uuid4(), program=MessageScorable(message_piece_ids=(uuid.uuid4(),)), - runtime_profile="python-3.12", - invocations=(ProgramInvocation(),), + runtime_profile_name="python-3.12", + invocations=(_ProgramInvocation(),), ) for scorable in (message_program, _execution()): assert scorable_from_dict(scorable.model_dump(mode="json")) == scorable @@ -143,34 +127,30 @@ def test_execution_scorable_round_trips_as_a_stored_scorable() -> None: def test_execution_scorable_accepts_stored_text_content() -> None: stored = ProgramExecutionScorable( - execution_id=uuid.uuid4(), program=ContentEntryScorable(content_id=uuid.uuid4()), - runtime_profile="python-3.12", - invocations=(ProgramInvocation(),), + runtime_profile_name="python-3.12", + invocations=(_ProgramInvocation(),), ) assert scorable_from_dict(stored.model_dump(mode="json")) == stored with pytest.raises(ValidationError, match="text"): ProgramExecutionScorable( - execution_id=uuid.uuid4(), program=ContentEntryScorable(content_id=uuid.uuid4(), data_type="image_path"), - runtime_profile="python-3.12", - invocations=(ProgramInvocation(),), + runtime_profile_name="python-3.12", + invocations=(_ProgramInvocation(),), ) def test_execution_scorable_requires_text_and_an_invocation() -> None: with pytest.raises(ValidationError, match="text"): ProgramExecutionScorable( - execution_id=uuid.uuid4(), program=ContentScorable(value="image.png", data_type="image_path"), - runtime_profile="python-3.12", - invocations=(ProgramInvocation(),), + runtime_profile_name="python-3.12", + invocations=(_ProgramInvocation(),), ) with pytest.raises(ValidationError, match="at least one invocation"): ProgramExecutionScorable( - execution_id=uuid.uuid4(), program=ContentScorable(value="print(1)"), - runtime_profile="python-3.12", + runtime_profile_name="python-3.12", invocations=(), ) @@ -179,9 +159,8 @@ def test_execution_scorable_rejects_a_program_that_is_not_text_or_a_message() -> with pytest.raises(ValidationError): ProgramExecutionScorable.model_validate( { - "execution_id": str(uuid.uuid4()), "program": SurfaceScorable(uri="/x").model_dump(), - "runtime_profile": "python-3.12", + "runtime_profile_name": "python-3.12", "invocations": [{}], } ) @@ -204,10 +183,10 @@ def test_score_round_trips_an_execution_anchor() -> None: "indices", [(1, 0), (0, 0)], ) -def test_payload_requires_events_once_in_case_order(indices: tuple[int, ...]) -> None: - with pytest.raises(ValidationError, match="case order"): +def test_payload_requires_events_once_in_invocation_order(indices: tuple[int, ...]) -> None: + with pytest.raises(ValidationError, match="invocation order"): ProgramExecutionObservationPayload( - program_sha256=_DIGEST, events=tuple(_event(case_index=index) for index in indices) + program_sha256=_DIGEST, events=tuple(_event(invocation_index=index) for index in indices) ) @@ -221,35 +200,35 @@ def test_payload_rejects_a_coerced_schema_version() -> None: ProgramExecutionObservationPayload.model_validate({"schema_version": "1"}) -def test_payload_finds_events_by_case() -> None: - payload = ProgramExecutionObservationPayload(program_sha256=_DIGEST, events=(_event(case_index=1),)) - assert payload.event_for(1) is payload.events[0] - assert payload.event_for(0) is None - - def test_complete_observation_holds_one_event_per_invocation() -> None: scorable = _execution() observation = _observation( - scorable=scorable, acquisition=Acquisition.COMPLETE, events=(_event(case_index=0), _event(case_index=1)) + scorable=scorable, + acquisition=Acquisition.COMPLETE, + events=(_event(invocation_index=0), _event(invocation_index=1)), ) assert Observation.model_validate(observation.model_dump(mode="json")) == observation with pytest.raises(ValidationError, match="Complete program acquisition"): - _observation(scorable=scorable, acquisition=Acquisition.COMPLETE, events=(_event(case_index=0),)) + _observation(scorable=scorable, acquisition=Acquisition.COMPLETE, events=(_event(invocation_index=0),)) with pytest.raises(ValidationError, match="Complete program acquisition"): _observation( scorable=scorable, acquisition=Acquisition.COMPLETE, - events=(_event(case_index=0), _event(case_index=1)), + events=(_event(invocation_index=0), _event(invocation_index=1)), reasons=("gap",), ) def test_partial_observation_is_missing_an_invocation() -> None: scorable = _execution() - _observation(scorable=scorable, acquisition=Acquisition.PARTIAL, events=(_event(case_index=0),), reasons=("x",)) + _observation( + scorable=scorable, acquisition=Acquisition.PARTIAL, events=(_event(invocation_index=0),), reasons=("x",) + ) with pytest.raises(ValidationError, match="Partial program acquisition"): _observation( - scorable=scorable, acquisition=Acquisition.PARTIAL, events=(_event(case_index=0), _event(case_index=1)) + scorable=scorable, + acquisition=Acquisition.PARTIAL, + events=(_event(invocation_index=0), _event(invocation_index=1)), ) @@ -258,13 +237,13 @@ def test_failed_observation_holds_no_events(acquisition: Acquisition) -> None: scorable = _execution() _observation(scorable=scorable, acquisition=acquisition, reasons=("sandbox down",)) with pytest.raises(ValidationError, match="cannot contain events"): - _observation(scorable=scorable, acquisition=acquisition, events=(_event(case_index=0),)) + _observation(scorable=scorable, acquisition=acquisition, events=(_event(invocation_index=0),)) def test_observation_events_must_belong_to_the_scorable() -> None: scorable = _execution(invocations=1) with pytest.raises(ValidationError, match="belong to the scorable"): - _observation(scorable=scorable, acquisition=Acquisition.PARTIAL, events=(_event(case_index=1),)) + _observation(scorable=scorable, acquisition=Acquisition.PARTIAL, events=(_event(invocation_index=1),)) with pytest.raises(ValidationError, match="runtime profile"): _observation(scorable=scorable, acquisition=Acquisition.COMPLETE, events=(_event(profile="node-20"),)) diff --git a/tests/unit/models/test_scorable.py b/tests/unit/models/test_scorable.py index 2fb1b2980b..532b85d49b 100644 --- a/tests/unit/models/test_scorable.py +++ b/tests/unit/models/test_scorable.py @@ -14,13 +14,13 @@ MessagePiece, MessageScorable, ProgramExecutionScorable, - ProgramInvocation, Scorable, Score, SurfaceScorable, TraceScorable, scorable_from_dict, ) +from pyrit.models.code_execution import _ProgramInvocation from pyrit.models.score.scorable import SCORABLE_TYPES @@ -158,10 +158,9 @@ def test_every_union_member_round_trips_to_its_own_type(): SurfaceScorable(uri="/data/out.txt"), ConversationScorable(conversation_id="whole-conversation"), ProgramExecutionScorable( - execution_id=uuid.uuid4(), program=ContentScorable(value="print(1)"), - runtime_profile="python-3.12", - invocations=(ProgramInvocation(),), + runtime_profile_name="python-3.12", + invocations=(_ProgramInvocation(),), ), ] diff --git a/tests/unit/score/test_docker_sandbox_client.py b/tests/unit/score/test_docker_sandbox_client.py deleted file mode 100644 index d43f90aeb3..0000000000 --- a/tests/unit/score/test_docker_sandbox_client.py +++ /dev/null @@ -1,518 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -"""Exercise DockerSandboxClient against a fake Docker command line client.""" - -from __future__ import annotations - -import asyncio -import json -import os -import pathlib -import shutil -import stat -import sys -import tempfile -import threading -from typing import TYPE_CHECKING -from unittest.mock import patch - -import pytest - -from pyrit.models import ProgramExecutionEvent, ProgramInvocation -from pyrit.score import DockerRuntimeProfile, DockerSandboxClient, SandboxClient, SandboxExecutionError - -if TYPE_CHECKING: - from pathlib import Path - -_IMAGE = "python@sha256:" + "c" * 64 - -# The fake runs the "program" as a scenario keyword, and keeps container state in files so -# that inspect, kill and rm see what run did. -_FAKE_DOCKER = r""" -import json, os, signal, sys, time -from pathlib import Path - -state = Path(os.environ["FAKE_DOCKER_STATE"]) -args = sys.argv[1:] -command = args[0] - -def path(name, suffix): - return state / f"{name}.{suffix}" - -if command == "run": - name = args[args.index("--name") + 1] - path(name, "argv").write_text(json.dumps(args)) - mount = args[args.index("--mount") + 1] - source = dict(part.split("=", 1) for part in mount.split(",") if "=" in part)["source"] - image_index = next(i for i, a in enumerate(args) if "@sha256:" in a) - inner = args[image_index + 1:] - # Docker runs the image's entrypoint before the command, unless --entrypoint replaces it. - entrypoint = os.environ.get("FAKE_IMAGE_ENTRYPOINT") - if "--entrypoint" in args: - entrypoint = args[args.index("--entrypoint") + 1] - path(name, "effective").write_text(json.dumps(([entrypoint] if entrypoint else []) + inner)) - program_arg = next(a for a in inner if a.startswith("/program/")) - program = (Path(source) / program_arg.rsplit("/", 1)[1]).read_text() - user_args = inner[inner.index(program_arg) + 1:] - flags = {"rm_fails": program == "RMFAIL", "slow_inspect": program == "SLOWINSPECT"} - def save(**info): - path(name, "state").write_text(json.dumps({**flags, **info})) - def finish(code, oom=False, started=True): - save(status="exited" if started else "created", code=code, oom=oom, - started="2026-10-14T00:00:00Z" if started else "0001-01-01T00:00:00Z") - sys.stdout.flush() - sys.exit(code) - path(name, "pid").write_text(str(os.getpid())) - save(status="running", code=0, oom=False, started="2026-10-14T00:00:00Z") - if program in ("CLOSEHANG", "CLOSEHANGRUNNING"): - # The client closes its output streams but does not exit. With CLOSEHANG the - # container has finished; with CLOSEHANGRUNNING it is still running. - print("done") - sys.stdout.flush() - if program == "CLOSEHANG": - save(status="exited", code=0, oom=False, started="2026-10-14T00:00:00Z") - for fd in (1, 2): - os.close(fd) - (state / "hang.pid").write_text(str(os.getpid())) - time.sleep(60) - sys.exit(0) - if program == "DETACHFAIL": - # The attached client fails while the container keeps running. - sys.stderr.write("error waiting for container: unexpected EOF") - sys.exit(1) - if program == "ECHO": - data = sys.stdin.read() - print(" ".join(user_args)) - sys.stderr.write(data) - finish(0) - elif program.startswith("EXIT "): - finish(int(program.split()[1])) - elif program == "SLEEP": - time.sleep(60) - finish(0) - elif program == "OOM": - finish(137, oom=True) - elif program == "BIG": - sys.stdout.write("a" * 200_000) - finish(0) - elif program == "NOSTART": - sys.stderr.write("docker: Error response from daemon: No such image") - finish(125, started=False) - elif program == "FAILSECOND": - marker = state / "ran-once" - if marker.exists(): - sys.stderr.write("docker: daemon went away") - finish(125, started=False) - marker.write_text("x") - finish(0) - finish(0) -elif command == "inspect": - name = args[-1] - try: - info = json.loads(path(name, "state").read_text()) - except FileNotFoundError: - print("Error: No such object: " + name) - sys.exit(1) - if info.get("slow_inspect"): - (state / "inspect.pid").write_text(str(os.getpid())) - time.sleep(60) - print(info["status"], info["code"], "true" if info["oom"] else "false", info["started"]) -elif command == "kill": - name = args[-1] - pid = int(path(name, "pid").read_text()) - os.kill(pid, signal.SIGKILL) - info = json.loads(path(name, "state").read_text()) - info["code"] = 137 - info["status"] = "exited" - path(name, "state").write_text(json.dumps(info)) -elif command == "rm": - name = args[-1] - try: - info = json.loads(path(name, "state").read_text()) - except FileNotFoundError: - info = {} - if info.get("rm_fails"): - print("Error response from daemon: daemon unavailable") - sys.exit(1) - with (state / "removed.log").open("a") as log: - log.write(name + "\n") - path(name, "state").unlink(missing_ok=True) -""" - - -@pytest.fixture -def fake_docker(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: - if sys.platform == "win32": - pytest.skip("The fake Docker client is a POSIX script.") - state = tmp_path / "state" - state.mkdir() - script = tmp_path / "docker" - script.write_text(f"#!{sys.executable}\n{_FAKE_DOCKER}") - script.chmod(script.stat().st_mode | stat.S_IXUSR) - monkeypatch.setenv("FAKE_DOCKER_STATE", str(state)) - return script - - -def _client(docker: Path, **overrides: object) -> DockerSandboxClient: - settings: dict[str, object] = { - "profiles": {"python-3.12": DockerRuntimeProfile(image=_IMAGE, command=("python3", "-I"))}, - "docker_executable": str(docker), - "timeout_seconds": 5.0, - } - settings.update(overrides) - return DockerSandboxClient(**settings) # type: ignore[arg-type] - - -async def _run(client: DockerSandboxClient, program: str, *invocations: ProgramInvocation): # type: ignore[no-untyped-def] - return await client.execute_async( - program=program, runtime_profile="python-3.12", invocations=invocations or (ProgramInvocation(),) - ) - - -def _removed(docker: Path) -> list[str]: - log = docker.parent / "state" / "removed.log" - return log.read_text().split() if log.exists() else [] - - -async def test_runs_each_invocation_with_its_arguments_and_input_async(fake_docker: Path) -> None: - run = await _run( - _client(fake_docker), - "ECHO", - ProgramInvocation(arguments=("a", "b c"), stdin="input one"), - ProgramInvocation(arguments=("z",)), - ) - - assert [event.case_index for event in run.events] == [0, 1] - first, second = run.events - assert (first.stdout, first.stderr, first.exit_code) == ("a b c\n", "input one", 0) - assert (second.stdout, second.stderr) == ("z\n", "") - assert first.resolved_command == ("python3", "-I", "/program/main.py", "a", "b c") - assert first.runtime_image_digest == "sha256:" + "c" * 64 - assert first.runtime_profile == "python-3.12" - assert run.reasons == () - - -async def test_containers_are_locked_down_and_removed_async(fake_docker: Path) -> None: - await _run(_client(fake_docker, memory="128m", cpus=0.5, pids_limit=8), "ECHO") - - argv_files = list((fake_docker.parent / "state").glob("*.argv")) - assert len(argv_files) == 1 - argv = json.loads(argv_files[0].read_text()) - for flag, value in [ - ("--network", "none"), - ("--pull", "never"), - ("--cap-drop", "ALL"), - ("--security-opt", "no-new-privileges"), - ("--user", "65534:65534"), - ("--memory", "128m"), - ("--memory-swap", "128m"), - ("--cpus", "0.5"), - ("--pids-limit", "8"), - ("--log-driver", "none"), - ]: - assert argv[argv.index(flag) + 1] == value - assert "--read-only" in argv - assert "--privileged" not in argv - assert "--log-opt" not in argv - mount = argv[argv.index("--mount") + 1] - assert mount.endswith(",target=/program,readonly") - # The profile's first argument is the entrypoint, so the image's own entrypoint never runs. - assert argv[argv.index("--entrypoint") + 1] == "python3" - assert argv.index("--entrypoint") < argv.index(_IMAGE) - assert argv[argv.index(_IMAGE) + 1 :] == ["-I", "/program/main.py"] - assert _removed(fake_docker) == [argv_files[0].stem] - - -async def test_an_image_entrypoint_does_not_change_the_command_that_runs_async( - fake_docker: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setenv("FAKE_IMAGE_ENTRYPOINT", "python3") - - run = await _run(_client(fake_docker), "ECHO", ProgramInvocation(arguments=("x",))) - - effective = json.loads(next((fake_docker.parent / "state").glob("*.effective")).read_text()) - assert tuple(effective) == run.events[0].resolved_command == ("python3", "-I", "/program/main.py", "x") - assert run.events[0].exit_code == 0 - - -async def test_program_exit_codes_are_observed_not_errors_async(fake_docker: Path) -> None: - run = await _run(_client(fake_docker), "EXIT 125") - - assert run.events[0].exit_code == 125 - assert not run.events[0].timed_out - - -async def test_timeout_kills_the_container_async(fake_docker: Path) -> None: - run = await _run(_client(fake_docker, timeout_seconds=0.5), "SLEEP") - - event = run.events[0] - assert event.timed_out - assert event.exit_code is None - assert not event.resource_limit_exceeded - assert len(_removed(fake_docker)) == 1 - - -def _alive(pid_file: Path) -> bool: - try: - os.kill(int(pid_file.read_text()), 0) - except ProcessLookupError: - return False - return True - - -async def test_a_client_that_closes_its_streams_but_never_exits_is_bounded_async(fake_docker: Path) -> None: - # The container finished, so its exit is real evidence even though the client hung. - run = await asyncio.wait_for(_run(_client(fake_docker, timeout_seconds=1.0), "CLOSEHANG"), timeout=10) - - event = run.events[0] - assert (event.exit_code, event.stdout, event.timed_out) == (0, "done\n", False) - assert not _alive(fake_docker.parent / "state" / "hang.pid") - assert len(_removed(fake_docker)) == 1 - - -async def test_a_hung_client_over_a_running_container_is_a_bounded_sandbox_error_async(fake_docker: Path) -> None: - # The container never finished, so the result is a sandbox error, not a program timeout. - with pytest.raises(SandboxExecutionError, match="still running"): - await asyncio.wait_for(_run(_client(fake_docker, timeout_seconds=1.0), "CLOSEHANGRUNNING"), timeout=10) - - assert not _alive(fake_docker.parent / "state" / "hang.pid") - assert len(_removed(fake_docker)) == 1 - - -async def test_memory_kill_is_a_resource_limit_async(fake_docker: Path) -> None: - event = (await _run(_client(fake_docker), "OOM")).events[0] - - assert event.resource_limit_exceeded - assert event.exit_code is None - - -async def test_output_is_capped_per_stream_async(fake_docker: Path) -> None: - event = (await _run(_client(fake_docker, max_output_bytes=1000), "BIG")).events[0] - - assert event.stdout == "a" * 1000 - assert event.stdout_truncated - assert not event.stderr_truncated - - -async def test_a_container_that_never_starts_is_a_sandbox_error_async(fake_docker: Path) -> None: - with pytest.raises(SandboxExecutionError, match="No such image"): - await _run(_client(fake_docker), "NOSTART") - assert len(_removed(fake_docker)) == 1 - - -async def test_a_failure_after_some_runs_returns_what_ran_async(fake_docker: Path) -> None: - run = await _run(_client(fake_docker), "FAILSECOND", ProgramInvocation(), ProgramInvocation()) - - assert [event.case_index for event in run.events] == [0] - assert "case 1" in run.reasons[0] - - -async def test_a_missing_docker_client_is_a_sandbox_error_async(tmp_path: Path) -> None: - with pytest.raises(SandboxExecutionError, match="could not start"): - await _run(_client(tmp_path / "no-such-docker"), "ECHO") - - -async def test_unknown_profile_is_rejected_async(fake_docker: Path) -> None: - with pytest.raises(ValueError, match="not configured"): - await _client(fake_docker).execute_async( - program="ECHO", runtime_profile="node-20", invocations=(ProgramInvocation(),) - ) - - -def test_profile_image_must_be_pinned_by_digest() -> None: - with pytest.raises(ValueError, match="pinned by digest"): - DockerRuntimeProfile(image="python:3.12-slim", command=("python3",)) - with pytest.raises(ValueError, match="nonempty"): - DockerRuntimeProfile(image=_IMAGE, command=()) - with pytest.raises(ValueError, match="plain file name"): - DockerRuntimeProfile(image=_IMAGE, command=("python3",), program_filename="../main.py") - - -@pytest.mark.parametrize( - "overrides", - [ - {"profiles": {}}, - {"profiles": {"bad name": DockerRuntimeProfile(image=_IMAGE, command=("python3",))}}, - {"timeout_seconds": 0}, - {"pids_limit": 0}, - {"max_output_bytes": 0}, - {"memory": "lots"}, - ], -) -def test_client_rejects_invalid_settings(overrides: dict[str, object], tmp_path: Path) -> None: - with pytest.raises(ValueError): - _client(tmp_path / "docker", **overrides) - - -def test_client_satisfies_the_sandbox_protocol(tmp_path: Path) -> None: - client = _client(tmp_path / "docker") - assert isinstance(client, SandboxClient) - assert client.get_runtime_profiles() == frozenset({"python-3.12"}) - params = client.get_identifier().params - assert params["profiles"]["python-3.12"]["image"] == _IMAGE - assert params["timeout_seconds"] == 5.0 - - -async def test_a_client_failure_while_the_container_runs_is_a_sandbox_error_async(fake_docker: Path) -> None: - with pytest.raises(SandboxExecutionError, match="still running"): - await _run(_client(fake_docker), "DETACHFAIL") - assert len(_removed(fake_docker)) == 1 - - -async def test_a_failed_removal_is_reported_and_stops_further_cases_async(fake_docker: Path) -> None: - run = await _run(_client(fake_docker), "RMFAIL", ProgramInvocation(), ProgramInvocation()) - - assert [event.case_index for event in run.events] == [0] - assert run.events[0].exit_code == 0 - assert "could not remove container pyrit-sandbox-" in run.warnings[0] - assert "daemon unavailable" in run.warnings[0] - assert run.reasons[0].startswith("stopped after case 0") - - -async def test_a_failed_removal_on_the_last_case_keeps_the_evidence_async(fake_docker: Path) -> None: - run = await _run(_client(fake_docker), "RMFAIL") - - assert len(run.events) == 1 - assert run.reasons == () - assert len(run.warnings) == 1 - - -async def test_cancellation_reaps_the_management_command_async(fake_docker: Path) -> None: - pid_file = fake_docker.parent / "state" / "inspect.pid" - task = asyncio.create_task(_run(_client(fake_docker), "SLOWINSPECT")) - for _ in range(200): - if pid_file.exists(): - break - await asyncio.sleep(0.05) - inspect_pid = int(pid_file.read_text()) - - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - - with pytest.raises(ProcessLookupError): - os.kill(inspect_pid, 0) - - -async def test_the_program_directory_is_removed_after_making_it_writable_async(tmp_path: Path) -> None: - client = _client(tmp_path / "docker") - seen: dict[str, object] = {} - real_rmtree = shutil.rmtree - - async def run_invocation(**kwargs: object): # type: ignore[no-untyped-def] - program_dir = kwargs["program_dir"] - seen["dir"] = program_dir - event = ProgramExecutionEvent( - case_index=kwargs["case_index"], - runtime_profile="python-3.12", - runtime_image_digest="sha256:" + "c" * 64, - resolved_command=("python3",), - exit_code=0, - duration_ms=1, - ) - return event, None - - def recording_rmtree(path, *args, **kwargs): # type: ignore[no-untyped-def] - program = next(iter(path.iterdir())) - seen["writable"] = bool(program.stat().st_mode & stat.S_IWRITE) - return real_rmtree(path, *args, **kwargs) - - with ( - patch.object(client, "_run_invocation_async", side_effect=run_invocation), - patch("pyrit.score.observation.docker_sandbox_client.shutil.rmtree", side_effect=recording_rmtree), - ): - run = await _run(client, "print(1)") - - assert len(run.events) == 1 - assert seen["writable"] is True - assert not seen["dir"].exists() # type: ignore[union-attr] - - -async def _fake_run_invocation(**kwargs: object) -> tuple[ProgramExecutionEvent, None]: - event = ProgramExecutionEvent( - case_index=kwargs["case_index"], # type: ignore[arg-type] - runtime_profile="python-3.12", - runtime_image_digest="sha256:" + "c" * 64, - resolved_command=("python3",), - exit_code=0, - duration_ms=1, - ) - return event, None - - -@pytest.mark.parametrize("blocked", ["write", "remove"]) -async def test_file_work_runs_off_the_event_loop_async(blocked: str, tmp_path: Path) -> None: - client = _client(tmp_path / "docker") - release = threading.Event() - threads: list[int] = [] - real_write_text = pathlib.Path.write_text - real_rmtree = shutil.rmtree - - def wait_for_the_loop() -> None: - # Only the running event loop releases this, so file work on the loop thread deadlocks here. - threads.append(threading.get_ident()) - if not release.wait(2): - raise AssertionError("the event loop was blocked while file work ran") - - def blocking_write_text(self: pathlib.Path, *args: object, **kwargs: object) -> int: - if blocked == "write": - wait_for_the_loop() - return real_write_text(self, *args, **kwargs) # type: ignore[arg-type] - - def blocking_rmtree(path: pathlib.Path, *args: object, **kwargs: object) -> None: - if blocked == "remove": - wait_for_the_loop() - real_rmtree(path, *args, **kwargs) # type: ignore[arg-type] - - async def release_once_blocked() -> None: - while not threads: - await asyncio.sleep(0.005) - release.set() - - with ( - patch.object(client, "_run_invocation_async", side_effect=_fake_run_invocation), - patch.object(pathlib.Path, "write_text", blocking_write_text), - patch("pyrit.score.observation.docker_sandbox_client.shutil.rmtree", side_effect=blocking_rmtree), - ): - releaser = asyncio.create_task(release_once_blocked()) - run = await _run(client, "print(1)") - await asyncio.wait_for(releaser, timeout=1) - - assert len(run.events) == 1 - assert threads[0] != threading.get_ident() - - -async def test_cancelling_during_file_work_still_removes_the_program_async( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - root = tmp_path / "temp" - root.mkdir() - monkeypatch.setattr(tempfile, "tempdir", str(root)) - client = _client(tmp_path / "docker") - real_mkdir = pathlib.Path.mkdir - entered = threading.Event() - release = threading.Event() - - def gated_mkdir(self: pathlib.Path, *args: object, **kwargs: object) -> None: - entered.set() - release.wait(5) - real_mkdir(self, *args, **kwargs) # type: ignore[arg-type] - - with ( - patch.object(client, "_run_invocation_async", side_effect=_fake_run_invocation), - patch.object(pathlib.Path, "mkdir", gated_mkdir), - ): - task = asyncio.create_task(_run(client, "print(1)")) - for _ in range(200): - if entered.is_set(): - break - await asyncio.sleep(0.01) - assert entered.is_set() - task.cancel() - asyncio.get_running_loop().call_later(0.2, release.set) - with pytest.raises(asyncio.CancelledError): - await task - # The cancellation completed only after the program directory had been made and removed. - assert release.is_set() - assert list(root.iterdir()) == [] diff --git a/tests/unit/score/test_program_execution_scorer.py b/tests/unit/score/test_program_execution_scorer.py index df1cca05df..e72dcd2a42 100644 --- a/tests/unit/score/test_program_execution_scorer.py +++ b/tests/unit/score/test_program_execution_scorer.py @@ -7,9 +7,11 @@ import hashlib import uuid from typing import TYPE_CHECKING +from unittest.mock import AsyncMock, MagicMock, patch import pytest +from pyrit._code_execution import DockerRuntimeProfile, DockerSandboxClient, SandboxExecutionError from pyrit.models import ( Acquisition, ComponentIdentifier, @@ -20,29 +22,25 @@ Equals, MessagePiece, MessageScorable, - Observation, ProgramBehaves, ProgramCase, - ProgramExecutionEvent, ProgramExecutionObservationPayload, ProgramExecutionScorable, - ProgramInvocation, ProgramResult, Regex, ScoreStatus, ScoringExpectation, SurfaceScorable, ) +from pyrit.models.code_execution import _ProgramExecutionEvent, _ProgramInvocation, _SandboxRun from pyrit.score import ( NonReplayableObservationError, ProgramExecutionScorer, ProgramExecutionSource, - SandboxClient, - SandboxExecutionError, - SandboxRun, TrueFalseInverterScorer, ) -from pyrit.score.true_false.program_execution_scorer import match_case, match_program_behaves +from pyrit.score.text_matching import match_text +from pyrit.score.true_false.program_execution_scorer import _evaluate_case, _evaluate_program if TYPE_CHECKING: from pyrit.memory import SQLiteMemory @@ -54,10 +52,10 @@ _IMAGE_DIGEST = "sha256:" + "b" * 64 -def _event(*, case_index: int = 0, **updates: object) -> ProgramExecutionEvent: +def _event(*, invocation_index: int = 0, **updates: object) -> _ProgramExecutionEvent: values: dict[str, object] = { - "case_index": case_index, - "runtime_profile": _PROFILE, + "invocation_index": invocation_index, + "runtime_profile_name": _PROFILE, "runtime_image_digest": _IMAGE_DIGEST, "resolved_command": ("python3", "/program/main.py"), "exit_code": 0, @@ -65,7 +63,7 @@ def _event(*, case_index: int = 0, **updates: object) -> ProgramExecutionEvent: "duration_ms": 1, } values.update(updates) - return ProgramExecutionEvent.model_validate(values) + return _ProgramExecutionEvent.model_validate(values) class _FakeSandbox: @@ -74,40 +72,42 @@ class _FakeSandbox: def __init__(self, *, outputs: dict[str, dict[str, object]] | None = None, profiles: frozenset[str] | None = None): self.outputs = outputs or {} self.profiles = profiles or frozenset({_PROFILE}) - self.calls: list[tuple[str, str, tuple[ProgramInvocation, ...]]] = [] + self.calls: list[tuple[str, str, tuple[_ProgramInvocation, ...]]] = [] self.error: SandboxExecutionError | None = None self.stop_after: int | None = None def get_identifier(self) -> ComponentIdentifier: return ComponentIdentifier.of(self, params={"fake": True}) - def get_runtime_profiles(self) -> frozenset[str]: + def get_runtime_profile_names(self) -> frozenset[str]: return self.profiles async def execute_async( - self, *, program: str, runtime_profile: str, invocations: tuple[ProgramInvocation, ...] - ) -> SandboxRun: - self.calls.append((program, runtime_profile, invocations)) + self, *, program: str, runtime_profile_name: str, invocations: tuple[_ProgramInvocation, ...] + ) -> _SandboxRun: + self.calls.append((program, runtime_profile_name, invocations)) if self.error is not None: raise self.error ran = invocations if self.stop_after is None else invocations[: self.stop_after] events = tuple( - _event(case_index=index, **self.outputs.get(invocation.arguments[0] if invocation.arguments else "", {})) + _event( + invocation_index=index, **self.outputs.get(invocation.arguments[0] if invocation.arguments else "", {}) + ) for index, invocation in enumerate(ran) ) reasons = ("the sandbox lost its worker",) if len(ran) < len(invocations) else () - return SandboxRun(events=events, reasons=reasons) + return _SandboxRun(events=events, reasons=reasons) def _case(argument: str, *, exit_code: int | None = 0, stdout: TextMatcher | None = None) -> ProgramCase: return ProgramCase( - invocation=ProgramInvocation(arguments=(argument,)), + invocation=_ProgramInvocation(arguments=(argument,)), expected=ProgramResult(exit_code=exit_code, stdout=stdout), ) def _expectation(*cases: ProgramCase, profile: str = _PROFILE) -> ScoringExpectation: - return ScoringExpectation(conditions=(ProgramBehaves(runtime_profile=profile, cases=cases),)) + return ScoringExpectation(conditions=(ProgramBehaves(runtime_profile_name=profile, cases=cases),)) _SQUARES = (_case("3", stdout=Equals(value="9")), _case("4", stdout=Equals(value="16"))) @@ -144,27 +144,27 @@ def _scorer(sandbox: _FakeSandbox, *, max_output_chars: int = 64_000) -> Program ), ], ) -def test_match_case(*, expected: ProgramResult, event: dict[str, object], result: bool | None) -> None: - assert match_case(expected=expected, event=_event(**event)) is result +def test_evaluate_case(*, expected: ProgramResult, event: dict[str, object], result: bool | None) -> None: + assert _evaluate_case(expected=expected, event=_event(**event)).value is result -def test_match_program_behaves_needs_every_case() -> None: - condition = ProgramBehaves(runtime_profile=_PROFILE, cases=_SQUARES) - passing = (_event(case_index=0, stdout="9"), _event(case_index=1, stdout="16")) +def test_evaluate_program_needs_every_case() -> None: + condition = ProgramBehaves(runtime_profile_name=_PROFILE, cases=_SQUARES) + passing = (_event(invocation_index=0, stdout="9"), _event(invocation_index=1, stdout="16")) payload = ProgramExecutionObservationPayload(program_sha256="a" * 64, events=passing) - assert match_program_behaves(condition=condition, payload=payload, acquisition=Acquisition.COMPLETE) is True + assert _evaluate_program(condition=condition, payload=payload, acquisition=Acquisition.COMPLETE).value is True one_ran = ProgramExecutionObservationPayload(program_sha256="a" * 64, events=passing[:1]) - assert match_program_behaves(condition=condition, payload=one_ran, acquisition=Acquisition.PARTIAL) is None + assert _evaluate_program(condition=condition, payload=one_ran, acquisition=Acquisition.PARTIAL).value is None failing_first = ProgramExecutionObservationPayload( - program_sha256="a" * 64, events=(_event(case_index=0, stdout="8"),) + program_sha256="a" * 64, events=(_event(invocation_index=0, stdout="8"),) ) - assert match_program_behaves(condition=condition, payload=failing_first, acquisition=Acquisition.PARTIAL) is False + assert _evaluate_program(condition=condition, payload=failing_first, acquisition=Acquisition.PARTIAL).value is False empty = ProgramExecutionObservationPayload(reasons=("down",)) for acquisition in (Acquisition.ERROR, Acquisition.UNAVAILABLE): - assert match_program_behaves(condition=condition, payload=empty, acquisition=acquisition) is None + assert _evaluate_program(condition=condition, payload=empty, acquisition=acquisition).value is None async def test_scores_true_when_every_case_passes_async(sqlite_instance: SQLiteMemory) -> None: @@ -188,7 +188,7 @@ async def test_sandbox_receives_stimuli_but_not_expected_results_async() -> None secret = "the-expected-answer" expectation = _expectation( ProgramCase( - invocation=ProgramInvocation(arguments=("3",), stdin="in"), + invocation=_ProgramInvocation(arguments=("3",), stdin="in"), expected=ProgramResult(stdout=Equals(value=secret)), ) ) @@ -197,10 +197,97 @@ async def test_sandbox_receives_stimuli_but_not_expected_results_async() -> None program, profile, invocations = sandbox.calls[0] assert (program, profile) == ("print(1)", _PROFILE) - assert invocations == (ProgramInvocation(arguments=("3",), stdin="in"),) + assert invocations == (_ProgramInvocation(arguments=("3",), stdin="in"),) assert secret not in repr(sandbox.calls) +async def test_scoring_and_replay_do_not_require_runtime_provenance_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "9", "runtime_image_digest": None, "resolved_command": None}}) + scorer = _scorer(sandbox) + expectation = _expectation(_case("3", stdout=Equals(value="9"))) + score = (await scorer.score_async(scorable=ContentScorable(value="print(9)"), expectation=expectation))[0] + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + + assert score.get_value() is True + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + assert observation.payload.events[0].runtime_image_digest is None + assert observation.payload.events[0].resolved_command is None + replayed = (await scorer.score_observation_async(observation=observation, expectation=expectation))[0] + assert replayed.get_value() is True + assert replayed.observation_ids == score.observation_ids + assert len(sandbox.calls) == 1 + saved = (await sqlite_instance.get_scores_async(score_ids=[str(replayed.id)]))[0] + assert saved.get_value() is True + assert saved.scored_expectation == expectation + + +async def test_cases_are_matched_once_for_verdict_and_rationale_async() -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "9"}, "4": {"stdout": "wrong"}}) + with patch("pyrit.score.true_false.program_execution_scorer.match_text", wraps=match_text) as matching: + score = ( + await _scorer(sandbox).score_async( + scorable=ContentScorable(value="print(1)"), expectation=_expectation(*_SQUARES) + ) + )[0] + assert score.get_value() is False + assert "case 1 output did not match" in score.score_rationale + assert matching.call_count == 2 + + +@pytest.mark.parametrize( + ("failure", "rationale"), + [ + ("lost-output", "output capture may be incomplete"), + ("resource-limits", "must support cgroup"), + ("image-volumes", "must not declare volumes"), + ], +) +async def test_unsafe_docker_execution_cannot_produce_a_passing_score_async( + sqlite_instance: SQLiteMemory, failure: str, rationale: str +) -> None: + client = DockerSandboxClient( + profiles={_PROFILE: DockerRuntimeProfile(image="python@" + _IMAGE_DIGEST, command=("python3", "-I"))} + ) + process = MagicMock(spec=asyncio.subprocess.Process) + process.returncode = 1 + process.stdin = None + for name, data in (("stdout", b"ok"), ("stderr", b"attachment failed")): + reader = asyncio.StreamReader() + reader.feed_data(data) + reader.feed_eof() + setattr(process, name, reader) + process.communicate = AsyncMock(return_value=(b"", b"")) + + async def control_async(*arguments: str) -> tuple[int, str]: + if arguments[0] == "info": + driver = "none" if failure == "resource-limits" else "cgroupfs" + return 0, f"{driver} true true true true true" + if arguments[:2] == ("image", "inspect"): + return 0, '{"/data":{}}' if failure == "image-volumes" else "null" + if arguments[0] == "inspect": + return 0, "exited 0 false 2026-10-21T00:00:00Z" + assert arguments[:3] == ("rm", "--force", "--volumes") + return 0, "" + + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=client)) + expectation = _expectation(_case("3", stdout=Equals(value="ok"))) + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=process) as launch, + patch.object(client, "_control_async", side_effect=control_async), + ): + score = (await scorer.score_async(scorable=ContentScorable(value="print('okextra')"), expectation=expectation))[ + 0 + ] + assert score.status is ScoreStatus.UNDETERMINED + assert score.score_value is None + assert rationale in score.score_rationale + assert launch.await_count == (1 if failure == "lost-output" else 0) + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert observation.acquisition is Acquisition.ERROR + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + assert observation.payload.events == () + + async def test_scores_false_when_one_case_fails_async() -> None: sandbox = _FakeSandbox(outputs={"3": {"stdout": "9"}, "4": {"stdout": "15"}}) @@ -339,80 +426,61 @@ async def test_replay_rejects_different_stimuli_async(sqlite_instance: SQLiteMem ) -async def test_each_scoring_call_is_a_new_execution_async() -> None: +@pytest.mark.parametrize("explicit", [False, True]) +@pytest.mark.parametrize("concurrent", [False, True]) +async def test_scoring_calls_store_independent_executions_async( + sqlite_instance: SQLiteMemory, explicit: bool, concurrent: bool +) -> None: sandbox = _FakeSandbox() scorer = _scorer(sandbox) program = ContentScorable(value="x") - first = (await scorer.score_async(scorable=program, expectation=_expectation(_case("1"))))[0] - second = (await scorer.score_async(scorable=program, expectation=_expectation(_case("1"))))[0] - - assert len(sandbox.calls) == 2 - assert first.observation_ids != second.observation_ids - - -async def test_explicit_execution_reuses_its_observed_run_async() -> None: - sandbox = _FakeSandbox() - scorer = _scorer(sandbox) execution = ProgramExecutionScorable( - execution_id=uuid.uuid4(), - program=ContentScorable(value="x"), - runtime_profile=_PROFILE, - invocations=(ProgramInvocation(arguments=("1",)),), + program=program, + runtime_profile_name=_PROFILE, + invocations=(_ProgramInvocation(arguments=("1",)),), ) - first = (await scorer.score_async(scorable=execution, expectation=_expectation(_case("1"))))[0] - second = (await scorer.score_async(scorable=execution, expectation=_expectation(_case("1", exit_code=None))))[0] - - assert len(sandbox.calls) == 1 - assert first.observation_ids == second.observation_ids - assert first.scorable == execution - - -async def test_sandbox_errors_are_not_reused_async() -> None: - sandbox = _FakeSandbox() - sandbox.error = SandboxExecutionError("down") - source = ProgramExecutionSource(sandbox=sandbox) - execution = ProgramExecutionScorable( - execution_id=uuid.uuid4(), - program=ContentScorable(value="x"), - runtime_profile=_PROFILE, - invocations=(ProgramInvocation(),), - ) + scorable = execution if explicit else program + expectation = _expectation(_case("1")) + if concurrent: + first_scores, second_scores = await asyncio.gather( + scorer.score_async(scorable=scorable, expectation=expectation), + scorer.score_async(scorable=scorable, expectation=expectation), + ) + else: + first_scores = await scorer.score_async(scorable=scorable, expectation=expectation) + second_scores = await scorer.score_async(scorable=scorable, expectation=expectation) + first, second = first_scores[0], second_scores[0] - assert (await source.acquire_async(scorable=execution)).acquisition is Acquisition.ERROR - sandbox.error = None - assert (await source.acquire_async(scorable=execution)).acquisition is Acquisition.COMPLETE assert len(sandbox.calls) == 2 + assert first.observation_ids != second.observation_ids + stored = await sqlite_instance.get_scores_async(score_ids=[str(first.id), str(second.id)]) + assert len(stored) == 2 + assert all(score.get_value() is True for score in stored) -async def test_a_run_with_no_events_is_an_error_and_not_reused_async() -> None: +async def test_a_run_with_no_events_is_an_acquisition_error_async() -> None: sandbox = _FakeSandbox() sandbox.stop_after = 0 source = ProgramExecutionSource(sandbox=sandbox) execution = ProgramExecutionScorable( - execution_id=uuid.uuid4(), program=ContentScorable(value="x"), - runtime_profile=_PROFILE, - invocations=(ProgramInvocation(),), + runtime_profile_name=_PROFILE, + invocations=(_ProgramInvocation(),), ) failed = await source.acquire_async(scorable=execution) - sandbox.stop_after = None - retried = await source.acquire_async(scorable=execution) - assert failed.acquisition is Acquisition.ERROR assert isinstance(failed.payload, ProgramExecutionObservationPayload) assert failed.payload.reasons == ("the sandbox lost its worker",) - assert retried.acquisition is Acquisition.COMPLETE async def test_explicit_execution_must_match_the_condition_async() -> None: execution = ProgramExecutionScorable( - execution_id=uuid.uuid4(), program=ContentScorable(value="x"), - runtime_profile=_PROFILE, - invocations=(ProgramInvocation(arguments=("2",)),), + runtime_profile_name=_PROFILE, + invocations=(_ProgramInvocation(arguments=("2",)),), ) with pytest.raises(RuntimeError, match="runtime profile and invocations"): @@ -514,10 +582,6 @@ async def test_inverter_keeps_undetermined_async() -> None: assert inverted.get_value() is False -def test_fake_sandbox_satisfies_the_protocol() -> None: - assert isinstance(_FakeSandbox(), SandboxClient) - - def test_source_rejects_a_nonpositive_output_limit() -> None: with pytest.raises(ValueError, match="positive"): ProgramExecutionSource(sandbox=_FakeSandbox(), max_output_chars=0) @@ -533,195 +597,22 @@ def test_identifier_names_the_source_and_sandbox() -> None: async def test_source_rejects_events_for_unknown_invocations_async() -> None: class _Rogue(_FakeSandbox): - async def execute_async(self, *, program, runtime_profile, invocations): # type: ignore[no-untyped-def] - return SandboxRun(events=(_event(case_index=5),)) + async def execute_async(self, *, program, runtime_profile_name, invocations): # type: ignore[no-untyped-def] + return _SandboxRun(events=(_event(invocation_index=5),)) execution = ProgramExecutionScorable( - execution_id=uuid.uuid4(), program=ContentScorable(value="x"), - runtime_profile=_PROFILE, - invocations=(ProgramInvocation(),), + runtime_profile_name=_PROFILE, + invocations=(_ProgramInvocation(),), ) with pytest.raises(ValueError, match="not given"): await ProgramExecutionSource(sandbox=_Rogue()).acquire_async(scorable=execution) -def test_observation_round_trip_keeps_the_program_reference() -> None: - execution = ProgramExecutionScorable( - execution_id=uuid.uuid4(), - program=ContentScorable(value="x"), - runtime_profile=_PROFILE, - invocations=(ProgramInvocation(),), - ) - observation = Observation( - source_identifier=ComponentIdentifier(class_name="S", class_module="m"), - acquisition=Acquisition.COMPLETE, - scorable=execution, - payload=ProgramExecutionObservationPayload(program_sha256="a" * 64, events=(_event(),)), - ) - assert Observation.model_validate(observation.model_dump(mode="json")) == observation - - -class _SlowSandbox(_FakeSandbox): - """Holds each run open until released, so concurrent callers overlap.""" - - def __init__(self) -> None: - super().__init__() - self.release = asyncio.Event() - - async def execute_async(self, *, program, runtime_profile, invocations): # type: ignore[no-untyped-def] - self.calls.append((program, runtime_profile, invocations)) - await self.release.wait() - return SandboxRun(events=tuple(_event(case_index=index) for index in range(len(invocations)))) - - -def _attempt() -> ProgramExecutionScorable: - return ProgramExecutionScorable( - execution_id=uuid.uuid4(), - program=ContentScorable(value="x"), - runtime_profile=_PROFILE, - invocations=(ProgramInvocation(),), - ) - - -async def test_concurrent_callers_share_one_run_of_an_attempt_async() -> None: - sandbox = _SlowSandbox() - source = ProgramExecutionSource(sandbox=sandbox) - execution = _attempt() - - pending = asyncio.gather(source.acquire_async(scorable=execution), source.acquire_async(scorable=execution)) - await asyncio.sleep(0) - sandbox.release.set() - first, second = await pending - - assert len(sandbox.calls) == 1 - assert first.id == second.id - - -async def test_one_cancelled_caller_does_not_cancel_the_shared_run_async() -> None: - sandbox = _SlowSandbox() - source = ProgramExecutionSource(sandbox=sandbox) - execution = _attempt() - - cancelled = asyncio.create_task(source.acquire_async(scorable=execution)) - kept = asyncio.create_task(source.acquire_async(scorable=execution)) - await asyncio.sleep(0) - cancelled.cancel() - await asyncio.sleep(0) - sandbox.release.set() - - observation = await kept - with pytest.raises(asyncio.CancelledError): - await cancelled - assert observation.acquisition is Acquisition.COMPLETE - assert len(sandbox.calls) == 1 - - -async def test_a_run_nobody_waits_for_is_cancelled_async() -> None: - sandbox = _SlowSandbox() - source = ProgramExecutionSource(sandbox=sandbox) - execution = _attempt() - - task = asyncio.create_task(source.acquire_async(scorable=execution)) - await asyncio.sleep(0) - flight = source._in_flight[execution] - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - await asyncio.sleep(0) - - assert flight.task.cancelled() - assert execution not in source._in_flight - - -class _SlowCleanupSandbox(_SlowSandbox): - """Holds its cleanup open after cancellation until released.""" - - def __init__(self) -> None: - super().__init__() - self.cleanup_started = asyncio.Event() - self.cleanup_release = asyncio.Event() - self.cleaned_up = False - - async def execute_async(self, *, program, runtime_profile, invocations): # type: ignore[no-untyped-def] - try: - return await super().execute_async( - program=program, runtime_profile=runtime_profile, invocations=invocations - ) - except asyncio.CancelledError: - self.cleanup_started.set() - await self.cleanup_release.wait() - self.cleaned_up = True - raise - - -async def test_cancelling_the_last_caller_waits_for_sandbox_cleanup_async() -> None: - sandbox = _SlowCleanupSandbox() - source = ProgramExecutionSource(sandbox=sandbox) - execution = _attempt() - - task = asyncio.create_task(source.acquire_async(scorable=execution)) - # The source checks memory before it runs the program, so wait for the run itself to start. - for _ in range(500): - if sandbox.calls: - break - await asyncio.sleep(0.01) - assert sandbox.calls - task.cancel() - await asyncio.wait_for(sandbox.cleanup_started.wait(), timeout=5) - for _ in range(5): - await asyncio.sleep(0) - assert not task.done(), "cancellation finished before the sandbox cleaned up" - - sandbox.cleanup_release.set() - with pytest.raises(asyncio.CancelledError): - await task - assert sandbox.cleaned_up - assert execution not in source._in_flight - - -async def test_concurrent_scoring_of_one_attempt_stores_its_observation_once_async( - sqlite_instance: SQLiteMemory, -) -> None: - sandbox = _SlowSandbox() - scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=sandbox)) - execution = _attempt() - expectation = _expectation(ProgramCase(invocation=ProgramInvocation())) - # Hold both callers at the write until each has done every read, so neither sees the - # other's observation already stored. - arrived, both_arrived = 0, asyncio.Event() - write = sqlite_instance.add_scores_to_memory_async - - async def write_together(**kwargs): # type: ignore[no-untyped-def] - nonlocal arrived - arrived += 1 - if arrived == 2: - both_arrived.set() - await asyncio.wait_for(both_arrived.wait(), timeout=5) - return await write(**kwargs) - - sqlite_instance.add_scores_to_memory_async = write_together # type: ignore[method-assign] - - pending = asyncio.gather( - scorer.score_async(scorable=execution, expectation=expectation), - scorer.score_async(scorable=execution, expectation=expectation), - ) - await asyncio.sleep(0) - sandbox.release.set() - first, second = await pending - - assert len(sandbox.calls) == 1 - assert first[0].score_value == second[0].score_value == "true" - assert first[0].observation_ids == second[0].observation_ids - stored = await sqlite_instance.get_scores_async(score_ids=[first[0].id, second[0].id]) - assert len(stored) == 2 - assert len(await sqlite_instance.get_observations_async(observation_ids=first[0].observation_ids)) == 1 - - async def test_sandbox_warnings_are_kept_on_the_observation_async(sqlite_instance: SQLiteMemory) -> None: class _LeakySandbox(_FakeSandbox): - async def execute_async(self, *, program, runtime_profile, invocations): # type: ignore[no-untyped-def] - return SandboxRun( + async def execute_async(self, *, program, runtime_profile_name, invocations): # type: ignore[no-untyped-def] + return _SandboxRun( events=(_event(),), warnings=("could not remove container pyrit-sandbox-1: daemon unavailable",) ) @@ -754,80 +645,16 @@ async def test_scoring_a_saved_score_anchor_runs_the_stored_program_async(sqlite assert observation.scorable.program == saved.scorable -class _ChangingSandbox(_FakeSandbox): - """Prints something different on every run, so a second run of one attempt would show.""" - - async def execute_async(self, *, program, runtime_profile, invocations): # type: ignore[no-untyped-def] - self.calls.append((program, runtime_profile, invocations)) - stdout = "9" if len(self.calls) == 1 else f"run {len(self.calls)}" - return SandboxRun(events=tuple(_event(case_index=index, stdout=stdout) for index in range(len(invocations)))) - - -async def test_a_stored_attempt_is_reused_after_the_source_evicts_it_async(sqlite_instance: SQLiteMemory) -> None: - sandbox = _ChangingSandbox() - source = ProgramExecutionSource(sandbox=sandbox) - scorer = ProgramExecutionScorer(source=source) - expectation = _expectation(_case("3", stdout=Equals(value="9"))) - execution = ProgramExecutionScorable( - execution_id=uuid.uuid4(), - program=ContentScorable(value="x"), - runtime_profile=_PROFILE, - invocations=(ProgramInvocation(arguments=("3",)),), - ) - - first = (await scorer.score_async(scorable=execution, expectation=expectation))[0] - for _ in range(ProgramExecutionSource._MAX_REUSED_OBSERVATIONS): - await source.acquire_async(scorable=_attempt()) - assert execution not in source._observed - again = (await scorer.score_async(scorable=execution, expectation=expectation))[0] - - runs_of_execution = [call for call in sandbox.calls if call[2] == execution.invocations] - assert len(runs_of_execution) == 1 - assert (first.score_value, again.score_value) == ("true", "true") - assert first.observation_ids == again.observation_ids - stored = await sqlite_instance.get_scores_async(score_ids=[first.id, again.id]) - assert len(stored) == 2 - assert len(await sqlite_instance.get_observations_async(observation_ids=first.observation_ids)) == 1 - - -async def test_errors_are_retried_as_new_observations_not_the_attempts_evidence_async() -> None: +async def test_each_acquisition_returns_new_evidence_without_memory_async() -> None: sandbox = _FakeSandbox() source = ProgramExecutionSource(sandbox=sandbox) - execution = _attempt() - sandbox.error = SandboxExecutionError("no daemon") - - failed = await source.acquire_async(scorable=execution) - sandbox.error = None - ran = await source.acquire_async(scorable=execution) - - assert failed.acquisition is Acquisition.ERROR - assert ran.acquisition is Acquisition.COMPLETE - assert failed.id != ran.id - assert ran.id == ProgramExecutionSource(sandbox=sandbox)._attempt_observation_id(execution) - - -async def test_two_sources_cannot_store_two_runs_of_one_attempt_async(sqlite_instance: SQLiteMemory) -> None: - first_sandbox, second_sandbox = _SlowSandbox(), _SlowSandbox() - scorers = [_scorer(first_sandbox), _scorer(second_sandbox)] - execution = _attempt() - expectation = _expectation(ProgramCase(invocation=ProgramInvocation(), expected=ProgramResult())) - - pending = asyncio.gather( - *(scorer.score_async(scorable=execution, expectation=expectation) for scorer in scorers), - return_exceptions=True, + execution = ProgramExecutionScorable( + program=ContentScorable(value="x"), runtime_profile_name=_PROFILE, invocations=(_ProgramInvocation(),) ) - for _ in range(500): - if first_sandbox.calls and second_sandbox.calls: - break - await asyncio.sleep(0.01) - first_sandbox.release.set() - second_sandbox.release.set() - results = await pending - - # Each source ran the attempt, so the later store is refused rather than keeping a second run. - errors = [result for result in results if isinstance(result, BaseException)] - stored = [result for result in results if not isinstance(result, BaseException)] - assert len(errors) == 1 - assert "already exist" in str(errors[0]) - assert len(stored) == 1 - assert len(await sqlite_instance.get_observations_async(observation_ids=stored[0][0].observation_ids)) == 1 + with patch("pyrit.score.observation.program_execution_source.CentralMemory.get_memory_instance") as memory: + memory.side_effect = AssertionError("Loose program acquisition must not read memory") + first = await source.acquire_async(scorable=execution) + second = await source.acquire_async(scorable=execution) + assert first.id != second.id + assert first.scorable == second.scorable == execution + assert len(sandbox.calls) == 2 From 8df91faef6e50683a9c2feb6dfa749fff9cc35cb Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Fri, 23 Oct 2026 10:08:23 -0700 Subject: [PATCH 6/6] DOC: clarify execution environment ownership Keep the single-image Docker example as the only exception to caller-owned execution environments without expanding the framework section. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0bdf701a-5fa6-44c0-ba7e-82a5beb6722d --- doc/code/framework.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/doc/code/framework.md b/doc/code/framework.md index ff9723dc28..209ae5a77e 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -378,8 +378,8 @@ The below talks about responsibilities of most modules in the PyRIT library from `pyrit._code_execution`. Execution inputs and results are private types in `pyrit.models.code_execution`. - `runtime_profile_name` selects a client-owned configuration by name, not by supplying an image or command. Clients advertise the available names through `get_runtime_profile_names()`. -- The caller supplies the backend and trusted runtimes. The optional `DockerSandboxClient` uses an - existing daemon and pre-pulled images; it creates and removes only its disposable containers. +- PyRIT does not own environments. The one-image Docker example is the only exception: + `DockerSandboxClient` manages disposable containers; callers supply the daemon and pinned image. - **Does not own**: interpreting results, scoring criteria, message or seed creation, memory persistence, or infrastructure setup. Those remain with the consumer or caller.