This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 65a83820fef Let SandboxToolset be deep-copied again (#74525)
65a83820fef is described below
commit 65a83820fefdf693153c2f1016af4daa6c1fbe61
Author: PoAn Yang <[email protected]>
AuthorDate: Sat Oct 10 16:04:54 2026 +0800
Let SandboxToolset be deep-copied again (#74525)
* Let SandboxToolset be deep-copied again
Signed-off-by: PoAn Yang <[email protected]>
* Describe what the shared sandbox create lock guards
The lock is process-wide and also protects close and the open re-check,
which the previous comment did not say.
Generated-by: Claude Opus 5
Claude-Session: https://claude.ai/code/session_014np1fVAxHmGaBFt3qebJ8x
---------
Signed-off-by: PoAn Yang <[email protected]>
Co-authored-by: Jarek Potiuk <[email protected]>
---
.../providers/common/ai/toolsets/sandbox.py | 16 ++++++++++-----
.../tests/unit/common/ai/toolsets/test_sandbox.py | 24 +++++++++++++++++++++-
2 files changed, 34 insertions(+), 6 deletions(-)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py
index 2b4561c4098..a03c968b390 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py
@@ -74,6 +74,13 @@ _RELEASE_RETRY_DELAY = 1.0
# Runs backend.create off the event loop. Its threads start on first use, not
at import.
_provisioning =
concurrent.futures.ThreadPoolExecutor(thread_name_prefix="sandbox-create")
+# One process-wide lock for every SandboxToolset. It guards each toolset's
pending creation,
+# the ``_open`` flag that ``_close`` clears, and the ``_open`` re-check before
a newly created
+# sandbox is published. Airflow deep-copies tasks and default_args together
with the toolsets
+# they hold, and a lock cannot be deep-copied, so toolsets share this lock
instead of holding
+# their own. It is never held during a backend call, so one toolset never
waits for another
+# toolset's sandbox.
+_create_lock = threading.Lock()
RUN_COMMAND = "run_command"
@@ -350,7 +357,6 @@ class SandboxToolset(AirflowToolset):
self._max_export_bytes = int(max_export_bytes)
self._tool_prefix = tool_prefix
self._sandbox: str | None = None
- self._create_lock = threading.Lock()
self._create_future: concurrent.futures.Future[str] | None = None
# Set while attached: who this run claimed the sandbox as, when the
sandbox ends
# on this process's clock (None when the creator recorded no
lifetime), and what
@@ -581,7 +587,7 @@ class SandboxToolset(AirflowToolset):
return None
def _close(self, *, run_failed: bool) -> None:
- with self._create_lock:
+ with _create_lock:
self._open = False
pending = self._create_future
sandbox = self._sandbox
@@ -747,7 +753,7 @@ class SandboxToolset(AirflowToolset):
# A native framework can make the first calls from several threads,
each with an
# event loop of its own, so the one creation they share is a
thread-safe future
# that any loop can await, not a task bound to the loop that started
it.
- with self._create_lock:
+ with _create_lock:
# Checked again under the lock: another thread may have finished
creating it.
if self._sandbox is not None:
return self._sandbox
@@ -775,14 +781,14 @@ class SandboxToolset(AirflowToolset):
f"Could not provision a sandbox on backend
{self._backend.name!r}: {e}"
) from e
else:
- with self._create_lock:
+ with _create_lock:
if not self._open:
# The block ended while this call waited; closing destroys
the sandbox.
raise SandboxTerminalError("The sandbox was closed while
it was being provisioned.")
self._sandbox = sandbox
return sandbox
finally:
- with self._create_lock:
+ with _create_lock:
if self._create_future is create_future:
self._create_future = None
diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
index 4023628b0dd..4c291cd2f76 100644
--- a/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
@@ -17,6 +17,7 @@
from __future__ import annotations
import asyncio
+import copy
import re
import threading
import time
@@ -50,7 +51,7 @@ from airflow.providers.common.ai.sandbox.base import (
from airflow.providers.common.ai.tools import ToolCallError
from airflow.providers.common.ai.toolsets import sandbox as sandbox_module
from airflow.providers.common.ai.toolsets.sandbox import SandboxToolset
-from airflow.providers.common.compat.sdk import ObjectStoragePath
+from airflow.providers.common.compat.sdk import DAG, ObjectStoragePath,
TaskGroup
from unit.common.ai.sandbox.fake_tags import InMemoryTagStore
@@ -798,6 +799,27 @@ class TestForRun:
assert isinstance(forked, CustomToolset)
+class TestDeepCopy:
+ """Airflow deep-copies tasks and default_args together with the toolsets
they hold."""
+
+ @pytest.mark.asyncio
+ async def test_a_deep_copied_toolset_runs_its_tools(self):
+ copied = copy.deepcopy(SandboxToolset(_RecordingBackend()))
+
+ async with copied:
+ result = await _call(copied, "run_command", {"command": "ls"})
+
+ assert result == "[stdout]\nout"
+
+ def test_task_group_default_args_can_hold_a_sandbox_toolset(self):
+ with DAG(dag_id="sandbox_in_task_group", schedule=None):
+ group = TaskGroup(
+ group_id="agents", default_args={"toolsets":
[SandboxToolset(_RecordingBackend())]}
+ )
+
+ assert [type(toolset) for toolset in group.default_args["toolsets"]]
== [SandboxToolset]
+
+
class TestExports:
@pytest.mark.parametrize(
("kwargs", "match"),
