kaxil commented on code in PR #71672: URL: https://github.com/apache/airflow/pull/71672#discussion_r3963844130
########## providers/common/ai/src/airflow/providers/common/ai/sandbox/islo.py: ########## @@ -0,0 +1,397 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""islo.dev microVM backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import math +import shlex +import time +from contextlib import contextmanager, suppress +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _new_sandbox_name, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from islo import Islo + from islo.errors import NotFoundError + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + +_TERMINAL_EXEC_STATUSES = frozenset({"completed", "failed", "timeout"}) +_POLL_INITIAL = 0.2 +_POLL_MAX = 2.0 +_POLL_BACKOFF = 1.5 +_FILE_OP_TIMEOUT = 120.0 +_HELPER_OUTPUT_CAP = 1024 * 1024 +_COMMAND_WRAPPER = """\ +umask 077 +dir="${TMPDIR:-/tmp}/airflow-sandbox-$$" +mkdir "$dir" || exit 70 +trap 'rm -rf "$dir"' EXIT HUP INT TERM +mkfifo "$dir/out" "$dir/err" "$dir/out-tail" "$dir/err-tail" || exit 70 +tail -c "$2" <"$dir/out-tail" >"$dir/out-result" & +out_tail_pid=$! +tail -c "$2" <"$dir/err-tail" >"$dir/err-result" & +err_tail_pid=$! +tee "$dir/out-tail" <"$dir/out" | wc -c >"$dir/out-count" & +out_drain_pid=$! +tee "$dir/err-tail" <"$dir/err" | wc -c >"$dir/err-count" & +err_drain_pid=$! +sh -lc "$1" >"$dir/out" 2>"$dir/err" +status=$? +wait "$out_drain_pid" "$err_drain_pid" "$out_tail_pid" "$err_tail_pid" || exit 70 Review Comment: This `wait` only returns once the tee/wc drains hit EOF, and any process the command backgrounds inherits the fifo write end, so EOF never comes. Running the wrapper under `sh` in a Debian slim image, `sleep 25 & echo started` took 25s to return and `nohup sleep 300 & echo server-started` never returned at all. Past the deadline `_destroy_after_timeout` deletes the microVM, so the model is told its sandbox was replaced and its files are gone, for a command whose foreground part finished instantly. `sbx.py` hits the same wall and bounds it with daemon drain threads plus a deadline-limited join; the wrapper needs a similar bound, such as a watchdog that kills the drains once `$status` is captured. ########## providers/common/ai/src/airflow/providers/common/ai/sandbox/islo.py: ########## @@ -0,0 +1,397 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""islo.dev microVM backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import math +import shlex +import time +from contextlib import contextmanager, suppress +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _new_sandbox_name, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from islo import Islo + from islo.errors import NotFoundError + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + +_TERMINAL_EXEC_STATUSES = frozenset({"completed", "failed", "timeout"}) +_POLL_INITIAL = 0.2 +_POLL_MAX = 2.0 +_POLL_BACKOFF = 1.5 +_FILE_OP_TIMEOUT = 120.0 +_HELPER_OUTPUT_CAP = 1024 * 1024 +_COMMAND_WRAPPER = """\ +umask 077 Review Comment: `umask 077` guards the scratch dir but it is still in effect for `sh -lc "$1"`, so it changes the mode of everything the agent creates. Same image and command: through the wrapper `touch f; mkdir d` gives `-rw-------` and `drwx------`; without it, `-rw-r--r--` and `drwxr-xr-x`. Resetting the umask before running the command, or dropping it and doing `chmod 700 "$dir"` after `mkdir`, keeps the scratch dir private without that side effect. ########## providers/common/ai/src/airflow/providers/common/ai/sandbox/islo.py: ########## @@ -0,0 +1,397 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""islo.dev microVM backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import math +import shlex +import time +from contextlib import contextmanager, suppress +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _new_sandbox_name, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from islo import Islo + from islo.errors import NotFoundError + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + +_TERMINAL_EXEC_STATUSES = frozenset({"completed", "failed", "timeout"}) +_POLL_INITIAL = 0.2 +_POLL_MAX = 2.0 +_POLL_BACKOFF = 1.5 +_FILE_OP_TIMEOUT = 120.0 +_HELPER_OUTPUT_CAP = 1024 * 1024 +_COMMAND_WRAPPER = """\ +umask 077 +dir="${TMPDIR:-/tmp}/airflow-sandbox-$$" +mkdir "$dir" || exit 70 +trap 'rm -rf "$dir"' EXIT HUP INT TERM Review Comment: A POSIX trap handler with no `exit` hands control back to the script, so on `TERM` this removes `$dir` and then keeps running. Sending `TERM` mid-command gave me stdout of just the flag line, stderr filled with `cat: .../out-count: No such file or directory` and `[: Illegal number:`, and an exit status of 0. Whether that reaches the model depends on the status Islo reports for a `timeout_secs` kill, but `trap 'rm -rf "$dir"; exit 143' HUP INT TERM`, keeping the bare `EXIT` trap for cleanup, would take the question off the table. ########## providers/common/ai/src/airflow/providers/common/ai/sandbox/islo.py: ########## @@ -0,0 +1,397 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""islo.dev microVM backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import math +import shlex +import time +from contextlib import contextmanager, suppress +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _new_sandbox_name, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from islo import Islo + from islo.errors import NotFoundError + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + +_TERMINAL_EXEC_STATUSES = frozenset({"completed", "failed", "timeout"}) +_POLL_INITIAL = 0.2 +_POLL_MAX = 2.0 +_POLL_BACKOFF = 1.5 +_FILE_OP_TIMEOUT = 120.0 +_HELPER_OUTPUT_CAP = 1024 * 1024 +_COMMAND_WRAPPER = """\ +umask 077 +dir="${TMPDIR:-/tmp}/airflow-sandbox-$$" +mkdir "$dir" || exit 70 +trap 'rm -rf "$dir"' EXIT HUP INT TERM +mkfifo "$dir/out" "$dir/err" "$dir/out-tail" "$dir/err-tail" || exit 70 +tail -c "$2" <"$dir/out-tail" >"$dir/out-result" & +out_tail_pid=$! +tail -c "$2" <"$dir/err-tail" >"$dir/err-result" & +err_tail_pid=$! +tee "$dir/out-tail" <"$dir/out" | wc -c >"$dir/out-count" & +out_drain_pid=$! +tee "$dir/err-tail" <"$dir/err" | wc -c >"$dir/err-count" & +err_drain_pid=$! +sh -lc "$1" >"$dir/out" 2>"$dir/err" +status=$? +wait "$out_drain_pid" "$err_drain_pid" "$out_tail_pid" "$err_tail_pid" || exit 70 +if [ "$(cat "$dir/out-count")" -gt "$2" ]; then printf '1\\n'; else printf '0\\n'; fi +cat "$dir/out-result" +if [ "$(cat "$dir/err-count")" -gt "$2" ]; then printf '1\\n' >&2; else printf '0\\n' >&2; fi +cat "$dir/err-result" >&2 +exit "$status" +""" + + +@contextmanager +def _translate_islo_errors(operation: str) -> Iterator[None]: + try: + yield + except SandboxError: + raise + except Exception as e: + try: + from islo.core.api_error import ApiError + except ImportError: + raise SandboxTerminalError( + 'The Islo SDK is not installed. Install "apache-airflow-providers-common-ai[sandbox-islo]".' + ) from e + if isinstance(e, ApiError): + status = f" (HTTP {e.status_code})" if e.status_code is not None else "" + raise SandboxTerminalError(f"Islo could not {operation}{status}.") from e + raise SandboxTerminalError(f"Islo could not {operation}: {type(e).__name__}.") from e + + +def _bound_result_stream(text: str, max_bytes: int, *, server_truncated: bool) -> tuple[str, bool]: + flag, separator, payload = text.partition("\n") + if separator and flag in {"0", "1"}: + truncated = flag == "1" or server_truncated + else: + payload = text + truncated = True + + encoded = payload.encode("utf-8") + if len(encoded) > max_bytes: + payload = encoded[-max_bytes:].decode("utf-8", errors="ignore") + truncated = True + return payload, truncated + + +class IsloSandboxBackend(SandboxBackend): + """ + Sandbox backend that runs agent commands in an `islo.dev <https://islo.dev>`__ microVM. + + Islo is a hosted API with no local daemon or host-virtualization requirement, + so this backend works from an Airflow worker running in a container. + Credentials resolve lazily from an Airflow connection on first use. + + Connection fields: ``password`` is the Islo API key (required), ``host`` the + compute URL (optional), and the extra may set ``base_url`` and ``timeout`` + (request timeout in seconds). + + File reads and writes use Islo's native streaming APIs. Directory listings + and command-output bounding require common Unix command-line tools in the + sandbox image, including ``sh``, ``mkfifo``, ``tail``, ``tee`` and a ``find`` + implementation with ``-printf`` support. + + :param islo_conn_id: Airflow connection ID for Islo. ``None`` lets the SDK + resolve credentials from its own environment variables (``ISLO_API_KEY``). + :param image: Sandbox image. ``None`` (default) uses the server default. + :param vcpus: Number of virtual CPUs. ``None`` uses the server default. + :param memory_mb: Memory in MB. ``None`` uses the server default. + :param delete_after: Server-side TTL in seconds after which the sandbox is + deleted even if the worker never got to destroy it. Default ``3600``. + """ + + name = "islo" + + def __init__( + self, + islo_conn_id: str | None = "islo_default", + *, + image: str | None = None, + vcpus: int | None = None, + memory_mb: int | None = None, + delete_after: int = 3600, + ) -> None: + _validate_positive_finite(delete_after, "delete_after") + if vcpus is not None: + _validate_positive_finite(vcpus, "vcpus") + if memory_mb is not None: + _validate_positive_finite(memory_mb, "memory_mb") + if image == "": + raise ValueError("image must not be empty.") + self._islo_conn_id = islo_conn_id + self._image = image + self._vcpus = vcpus + self._memory_mb = memory_mb + self._delete_after = delete_after + self._client: Islo | None = None + + def _get_client(self) -> Islo: + if self._client is not None: + return self._client + with _translate_islo_errors("initialize its client"): + from islo import Islo + + if self._islo_conn_id is None: + self._client = Islo() + return self._client + conn = BaseHook.get_connection(self._islo_conn_id) + api_key = (conn.password or "").strip() + if not api_key: + raise SandboxTerminalError( + f"Connection {self._islo_conn_id!r} has no password; set it to the Islo API key." + ) + kwargs: dict[str, Any] = {"api_key": api_key} + if conn.host: + kwargs["compute_url"] = conn.host + extra = conn.extra_dejson + if extra.get("base_url"): + kwargs["base_url"] = extra["base_url"] + if extra.get("timeout") is not None: + try: + request_timeout = float(extra["timeout"]) + _validate_positive_finite(request_timeout, "connection extra timeout") + except (TypeError, ValueError) as e: + raise SandboxTerminalError( + "The Islo connection extra timeout must be a positive finite number." + ) from e + kwargs["timeout"] = request_timeout + self._client = Islo(**kwargs) + return self._client + + @staticmethod + def _request_options(*, timeout: float, chunk_size: int | None = None) -> dict[str, int]: + options = {"timeout_in_seconds": max(1, math.ceil(timeout)), "max_retries": 0} + if chunk_size is not None: + options["chunk_size"] = chunk_size + return options + + def create(self, *, spec: SandboxSpec | None = None) -> str: + if spec is not None and spec.allow_egress_to: + raise SandboxTerminalError( + "The Islo backend cannot apply a per-domain egress allowlist; it can only turn " + "outbound access on or off. Drop allow_egress_to, or use a backend with " + "per-domain network rules." + ) + with _translate_islo_errors("create a sandbox"): + from islo.types import LifecyclePolicy + + kwargs: dict[str, Any] = { + "internet_enabled": False if spec is None else not spec.block_network, + "lifecycle": LifecyclePolicy(delete_after=self._delete_after), + } + if self._image is not None: + kwargs["image"] = self._image + if self._vcpus is not None: + kwargs["vcpus"] = self._vcpus + if self._memory_mb is not None: + kwargs["memory_mb"] = self._memory_mb + if spec is not None and spec.env: + kwargs["env"] = dict(spec.env) Review Comment: Is `env` set at creation visible to later `exec_in_sandbox` calls, or only to the sandbox's own init? `exec_in_sandbox` takes its own `env`, so carrying `spec.env` on the backend and passing it per exec would remove the doubt. Worth pinning down, because `base.py` treats silently dropping a `SandboxSpec` field as a contract violation, and neither the unit tests nor the system test ever set one. ########## providers/common/ai/src/airflow/providers/common/ai/sandbox/islo.py: ########## @@ -0,0 +1,397 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""islo.dev microVM backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import math +import shlex +import time +from contextlib import contextmanager, suppress +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _new_sandbox_name, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from islo import Islo + from islo.errors import NotFoundError + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + +_TERMINAL_EXEC_STATUSES = frozenset({"completed", "failed", "timeout"}) +_POLL_INITIAL = 0.2 +_POLL_MAX = 2.0 +_POLL_BACKOFF = 1.5 +_FILE_OP_TIMEOUT = 120.0 +_HELPER_OUTPUT_CAP = 1024 * 1024 +_COMMAND_WRAPPER = """\ +umask 077 +dir="${TMPDIR:-/tmp}/airflow-sandbox-$$" +mkdir "$dir" || exit 70 +trap 'rm -rf "$dir"' EXIT HUP INT TERM +mkfifo "$dir/out" "$dir/err" "$dir/out-tail" "$dir/err-tail" || exit 70 +tail -c "$2" <"$dir/out-tail" >"$dir/out-result" & Review Comment: `tail -c` cuts at a byte boundary, so the first record it keeps is normally a fragment, and neither `_bound_result_stream` nor `truncate_output` drops it. Pushing 1000 numbered lines through the wrapper with a 100 byte cap, the model ends up with `88` on a line of its own, which is the tail of `line988`. `output.py`'s own header says truncation never emits a partial line, and `sbx.py`'s drain drops the leading partial line for exactly this reason. The second cut at line 106 has the same gap, and `list_directory` shares the path, where a mid-record cut can yield a mangled entry name. ########## providers/common/ai/tests/unit/common/ai/sandbox/test_islo.py: ########## @@ -0,0 +1,459 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from types import SimpleNamespace +from unittest import mock + +import pytest + +pytest.importorskip("islo") + +from islo.core.api_error import ApiError +from islo.errors import NotFoundError + +from airflow.providers.common.ai.sandbox.base import ( + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxSpec, + SandboxTerminalError, +) +from airflow.providers.common.ai.sandbox.islo import IsloSandboxBackend + +_BASE_HOOK_PATH = "airflow.providers.common.ai.sandbox.islo.BaseHook" +_ISLO_PATH = "islo.Islo" + + +def _connection(password="secret-key", host=None, extra=None): + return SimpleNamespace(password=password, host=host, extra_dejson=extra or {}) + + +def _exec_result(status="completed", exit_code=0, stdout="0\n", stderr="0\n", truncated=False): + return SimpleNamespace( + status=status, exit_code=exit_code, stdout=stdout, stderr=stderr, truncated=truncated + ) + + +def _backend_with_client(**kwargs) -> tuple[IsloSandboxBackend, mock.MagicMock]: + backend = IsloSandboxBackend(**kwargs) + client = mock.MagicMock(spec=["sandboxes"]) + client.sandboxes = mock.MagicMock( + spec=[ + "create_sandbox", + "delete_sandbox", + "download_file", + "exec_in_sandbox", + "get_exec_result", + "get_sandbox", + "upload_file", + ] + ) + client.sandboxes.exec_in_sandbox.return_value = SimpleNamespace(exec_id="exec-1") + client.sandboxes.create_sandbox.return_value = SimpleNamespace(name="box-1") + client.sandboxes.get_exec_result.return_value = _exec_result() + backend._client = client + return backend, client + + +class TestCredentials: + @mock.patch(_ISLO_PATH, autospec=True) + @mock.patch(_BASE_HOOK_PATH, autospec=True) + def test_api_key_and_allowlisted_connection_options_are_forwarded(self, hook, islo): + backend = IsloSandboxBackend(islo_conn_id="my_islo") + hook.get_connection.return_value = _connection( + password=" key ", + host="https://compute";, + extra={"base_url": "https://api";, "timeout": 12}, + ) + + backend._get_client() + + hook.get_connection.assert_called_once_with("my_islo") + islo.assert_called_once_with( + api_key="key", compute_url="https://compute";, base_url="https://api";, timeout=12.0 + ) + + @mock.patch(_ISLO_PATH, autospec=True) + @mock.patch(_BASE_HOOK_PATH, autospec=True) + def test_client_is_resolved_once_and_cached(self, hook, _islo): + backend = IsloSandboxBackend() + hook.get_connection.return_value = _connection() + + backend._get_client() + backend._get_client() + + hook.get_connection.assert_called_once_with("islo_default") + + @mock.patch(_BASE_HOOK_PATH, autospec=True) + def test_missing_api_key_is_terminal(self, hook): + backend = IsloSandboxBackend() + hook.get_connection.return_value = _connection(password="") + + with pytest.raises(SandboxTerminalError, match="has no password"): + backend._get_client() + + @mock.patch(_ISLO_PATH, autospec=True) + def test_none_conn_id_defers_to_the_sdk_environment(self, islo): + backend = IsloSandboxBackend(islo_conn_id=None) + + backend._get_client() + + islo.assert_called_once_with() + + @mock.patch(_BASE_HOOK_PATH, autospec=True) + def test_connection_resolution_failure_is_terminal(self, hook): + backend = IsloSandboxBackend() + hook.get_connection.side_effect = RuntimeError("secret backend down") + + with pytest.raises(SandboxTerminalError, match="initialize its client"): + backend._get_client() + + @mock.patch(_BASE_HOOK_PATH, autospec=True) + def test_invalid_connection_timeout_is_terminal_and_actionable(self, hook): + backend = IsloSandboxBackend() + hook.get_connection.return_value = _connection(extra={"timeout": "never"}) + + with pytest.raises(SandboxTerminalError, match="timeout must be a positive finite number"): + backend._get_client() + + [email protected]( + ("kwargs", "message"), + [ + ({"image": ""}, "image"), + ({"vcpus": 0}, "vcpus"), + ({"memory_mb": 0}, "memory_mb"), + ({"delete_after": 0}, "delete_after"), + ], +) +def test_constructor_rejects_invalid_values(kwargs, message): + with pytest.raises(ValueError, match=message): + IsloSandboxBackend(**kwargs) + + +class TestCreate: + def test_refuses_a_per_domain_egress_allowlist(self): + backend, _ = _backend_with_client() + + with pytest.raises(SandboxTerminalError, match="per-domain egress allowlist"): + backend.create(spec=SandboxSpec(allow_egress_to=["example.com"])) + + @pytest.mark.parametrize( + ("spec", "expected"), + [ + (None, False), + (SandboxSpec(), False), + (SandboxSpec(block_network=True), False), + (SandboxSpec(block_network=False), True), + ], + ) + def test_block_network_maps_to_internet_enabled(self, spec, expected): + backend, client = _backend_with_client() + + backend.create(spec=spec) + + assert client.sandboxes.create_sandbox.call_args.kwargs["internet_enabled"] is expected + + def test_spec_and_sizing_are_passed_at_creation(self): + backend, client = _backend_with_client(image="python", vcpus=2, memory_mb=1024, delete_after=120) + + name = backend.create(spec=SandboxSpec(env={"TOKEN": "value"})) + + assert name == "box-1" + kwargs = client.sandboxes.create_sandbox.call_args.kwargs + assert kwargs["image"] == "python" + assert kwargs["vcpus"] == 2 + assert kwargs["memory_mb"] == 1024 + assert kwargs["env"] == {"TOKEN": "value"} + assert kwargs["lifecycle"].delete_after == 120 + assert kwargs["request_options"] == {"timeout_in_seconds": 120, "max_retries": 0} + + def test_omitted_sizing_is_left_to_the_server(self): + backend, client = _backend_with_client() + + backend.create() + + assert not {"image", "vcpus", "memory_mb"} & client.sandboxes.create_sandbox.call_args.kwargs.keys() + + def test_api_failure_is_terminal(self): + backend, client = _backend_with_client() + client.sandboxes.create_sandbox.side_effect = ApiError(status_code=503) + + with pytest.raises(SandboxTerminalError, match="HTTP 503"): + backend.create() + + +class TestRunCommand: Review Comment: `_backend_with_client` mocks `exec_in_sandbox` outright, so `_COMMAND_WRAPPER` is never executed by anything except the system test, which needs a live API key and so will not run in ordinary CI. It is the riskiest code in the PR (four background jobs, four fifos, a trap, and a `wait`), and running it through a local `sh` subprocess needs no Islo access at all. A test like that would have caught the drain that blocks on a backgrounded child, the umask reaching the agent's command, and the partial leading line. ########## providers/common/ai/src/airflow/providers/common/ai/sandbox/islo.py: ########## @@ -0,0 +1,397 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""islo.dev microVM backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import math +import shlex +import time +from contextlib import contextmanager, suppress +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _new_sandbox_name, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from islo import Islo + from islo.errors import NotFoundError + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + +_TERMINAL_EXEC_STATUSES = frozenset({"completed", "failed", "timeout"}) +_POLL_INITIAL = 0.2 +_POLL_MAX = 2.0 +_POLL_BACKOFF = 1.5 +_FILE_OP_TIMEOUT = 120.0 +_HELPER_OUTPUT_CAP = 1024 * 1024 +_COMMAND_WRAPPER = """\ +umask 077 +dir="${TMPDIR:-/tmp}/airflow-sandbox-$$" +mkdir "$dir" || exit 70 +trap 'rm -rf "$dir"' EXIT HUP INT TERM +mkfifo "$dir/out" "$dir/err" "$dir/out-tail" "$dir/err-tail" || exit 70 +tail -c "$2" <"$dir/out-tail" >"$dir/out-result" & +out_tail_pid=$! +tail -c "$2" <"$dir/err-tail" >"$dir/err-result" & +err_tail_pid=$! +tee "$dir/out-tail" <"$dir/out" | wc -c >"$dir/out-count" & +out_drain_pid=$! +tee "$dir/err-tail" <"$dir/err" | wc -c >"$dir/err-count" & +err_drain_pid=$! +sh -lc "$1" >"$dir/out" 2>"$dir/err" +status=$? +wait "$out_drain_pid" "$err_drain_pid" "$out_tail_pid" "$err_tail_pid" || exit 70 +if [ "$(cat "$dir/out-count")" -gt "$2" ]; then printf '1\\n'; else printf '0\\n'; fi +cat "$dir/out-result" +if [ "$(cat "$dir/err-count")" -gt "$2" ]; then printf '1\\n' >&2; else printf '0\\n' >&2; fi +cat "$dir/err-result" >&2 +exit "$status" +""" + + +@contextmanager +def _translate_islo_errors(operation: str) -> Iterator[None]: + try: + yield + except SandboxError: + raise + except Exception as e: + try: + from islo.core.api_error import ApiError + except ImportError: + raise SandboxTerminalError( + 'The Islo SDK is not installed. Install "apache-airflow-providers-common-ai[sandbox-islo]".' + ) from e + if isinstance(e, ApiError): + status = f" (HTTP {e.status_code})" if e.status_code is not None else "" + raise SandboxTerminalError(f"Islo could not {operation}{status}.") from e + raise SandboxTerminalError(f"Islo could not {operation}: {type(e).__name__}.") from e + + +def _bound_result_stream(text: str, max_bytes: int, *, server_truncated: bool) -> tuple[str, bool]: + flag, separator, payload = text.partition("\n") + if separator and flag in {"0", "1"}: + truncated = flag == "1" or server_truncated + else: + payload = text + truncated = True + + encoded = payload.encode("utf-8") + if len(encoded) > max_bytes: + payload = encoded[-max_bytes:].decode("utf-8", errors="ignore") + truncated = True + return payload, truncated + + +class IsloSandboxBackend(SandboxBackend): + """ + Sandbox backend that runs agent commands in an `islo.dev <https://islo.dev>`__ microVM. + + Islo is a hosted API with no local daemon or host-virtualization requirement, + so this backend works from an Airflow worker running in a container. + Credentials resolve lazily from an Airflow connection on first use. + + Connection fields: ``password`` is the Islo API key (required), ``host`` the + compute URL (optional), and the extra may set ``base_url`` and ``timeout`` + (request timeout in seconds). + + File reads and writes use Islo's native streaming APIs. Directory listings + and command-output bounding require common Unix command-line tools in the + sandbox image, including ``sh``, ``mkfifo``, ``tail``, ``tee`` and a ``find`` + implementation with ``-printf`` support. + + :param islo_conn_id: Airflow connection ID for Islo. ``None`` lets the SDK + resolve credentials from its own environment variables (``ISLO_API_KEY``). + :param image: Sandbox image. ``None`` (default) uses the server default. + :param vcpus: Number of virtual CPUs. ``None`` uses the server default. + :param memory_mb: Memory in MB. ``None`` uses the server default. + :param delete_after: Server-side TTL in seconds after which the sandbox is + deleted even if the worker never got to destroy it. Default ``3600``. + """ + + name = "islo" + + def __init__( + self, + islo_conn_id: str | None = "islo_default", + *, + image: str | None = None, + vcpus: int | None = None, + memory_mb: int | None = None, + delete_after: int = 3600, + ) -> None: + _validate_positive_finite(delete_after, "delete_after") + if vcpus is not None: + _validate_positive_finite(vcpus, "vcpus") + if memory_mb is not None: + _validate_positive_finite(memory_mb, "memory_mb") + if image == "": + raise ValueError("image must not be empty.") + self._islo_conn_id = islo_conn_id + self._image = image + self._vcpus = vcpus + self._memory_mb = memory_mb + self._delete_after = delete_after + self._client: Islo | None = None + + def _get_client(self) -> Islo: + if self._client is not None: + return self._client + with _translate_islo_errors("initialize its client"): + from islo import Islo + + if self._islo_conn_id is None: + self._client = Islo() + return self._client + conn = BaseHook.get_connection(self._islo_conn_id) + api_key = (conn.password or "").strip() + if not api_key: + raise SandboxTerminalError( + f"Connection {self._islo_conn_id!r} has no password; set it to the Islo API key." + ) + kwargs: dict[str, Any] = {"api_key": api_key} + if conn.host: + kwargs["compute_url"] = conn.host + extra = conn.extra_dejson + if extra.get("base_url"): + kwargs["base_url"] = extra["base_url"] + if extra.get("timeout") is not None: + try: + request_timeout = float(extra["timeout"]) + _validate_positive_finite(request_timeout, "connection extra timeout") + except (TypeError, ValueError) as e: + raise SandboxTerminalError( + "The Islo connection extra timeout must be a positive finite number." + ) from e + kwargs["timeout"] = request_timeout + self._client = Islo(**kwargs) + return self._client + + @staticmethod + def _request_options(*, timeout: float, chunk_size: int | None = None) -> dict[str, int]: + options = {"timeout_in_seconds": max(1, math.ceil(timeout)), "max_retries": 0} Review Comment: `max_retries: 0` also lands on `delete_sandbox`, which is idempotent, and that is the one call whose failure `_destroy_after_timeout` converts into a task failure. So a single transient error during cleanup fails a run whose command merely timed out, even though `delete_after` is going to reclaim the microVM regardless. `sbx.py` only logs a warning on the same path, and it has no TTL to fall back on. Letting the delete keep the SDK's retries, at least on the cleanup path, would fit the backstop this backend already has. ########## providers/common/ai/src/airflow/providers/common/ai/sandbox/islo.py: ########## @@ -0,0 +1,397 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""islo.dev microVM backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import math +import shlex +import time +from contextlib import contextmanager, suppress +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _new_sandbox_name, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from islo import Islo + from islo.errors import NotFoundError + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + +_TERMINAL_EXEC_STATUSES = frozenset({"completed", "failed", "timeout"}) +_POLL_INITIAL = 0.2 +_POLL_MAX = 2.0 +_POLL_BACKOFF = 1.5 +_FILE_OP_TIMEOUT = 120.0 +_HELPER_OUTPUT_CAP = 1024 * 1024 +_COMMAND_WRAPPER = """\ +umask 077 +dir="${TMPDIR:-/tmp}/airflow-sandbox-$$" +mkdir "$dir" || exit 70 +trap 'rm -rf "$dir"' EXIT HUP INT TERM +mkfifo "$dir/out" "$dir/err" "$dir/out-tail" "$dir/err-tail" || exit 70 +tail -c "$2" <"$dir/out-tail" >"$dir/out-result" & +out_tail_pid=$! +tail -c "$2" <"$dir/err-tail" >"$dir/err-result" & +err_tail_pid=$! +tee "$dir/out-tail" <"$dir/out" | wc -c >"$dir/out-count" & +out_drain_pid=$! +tee "$dir/err-tail" <"$dir/err" | wc -c >"$dir/err-count" & +err_drain_pid=$! +sh -lc "$1" >"$dir/out" 2>"$dir/err" +status=$? +wait "$out_drain_pid" "$err_drain_pid" "$out_tail_pid" "$err_tail_pid" || exit 70 +if [ "$(cat "$dir/out-count")" -gt "$2" ]; then printf '1\\n'; else printf '0\\n'; fi +cat "$dir/out-result" +if [ "$(cat "$dir/err-count")" -gt "$2" ]; then printf '1\\n' >&2; else printf '0\\n' >&2; fi +cat "$dir/err-result" >&2 +exit "$status" +""" + + +@contextmanager +def _translate_islo_errors(operation: str) -> Iterator[None]: + try: + yield + except SandboxError: + raise + except Exception as e: + try: + from islo.core.api_error import ApiError + except ImportError: + raise SandboxTerminalError( + 'The Islo SDK is not installed. Install "apache-airflow-providers-common-ai[sandbox-islo]".' + ) from e + if isinstance(e, ApiError): + status = f" (HTTP {e.status_code})" if e.status_code is not None else "" + raise SandboxTerminalError(f"Islo could not {operation}{status}.") from e + raise SandboxTerminalError(f"Islo could not {operation}: {type(e).__name__}.") from e + + +def _bound_result_stream(text: str, max_bytes: int, *, server_truncated: bool) -> tuple[str, bool]: + flag, separator, payload = text.partition("\n") + if separator and flag in {"0", "1"}: + truncated = flag == "1" or server_truncated + else: + payload = text + truncated = True + + encoded = payload.encode("utf-8") + if len(encoded) > max_bytes: + payload = encoded[-max_bytes:].decode("utf-8", errors="ignore") + truncated = True + return payload, truncated + + +class IsloSandboxBackend(SandboxBackend): + """ + Sandbox backend that runs agent commands in an `islo.dev <https://islo.dev>`__ microVM. + + Islo is a hosted API with no local daemon or host-virtualization requirement, + so this backend works from an Airflow worker running in a container. + Credentials resolve lazily from an Airflow connection on first use. + + Connection fields: ``password`` is the Islo API key (required), ``host`` the + compute URL (optional), and the extra may set ``base_url`` and ``timeout`` + (request timeout in seconds). + + File reads and writes use Islo's native streaming APIs. Directory listings + and command-output bounding require common Unix command-line tools in the + sandbox image, including ``sh``, ``mkfifo``, ``tail``, ``tee`` and a ``find`` + implementation with ``-printf`` support. + + :param islo_conn_id: Airflow connection ID for Islo. ``None`` lets the SDK + resolve credentials from its own environment variables (``ISLO_API_KEY``). + :param image: Sandbox image. ``None`` (default) uses the server default. + :param vcpus: Number of virtual CPUs. ``None`` uses the server default. + :param memory_mb: Memory in MB. ``None`` uses the server default. + :param delete_after: Server-side TTL in seconds after which the sandbox is + deleted even if the worker never got to destroy it. Default ``3600``. + """ + + name = "islo" + + def __init__( + self, + islo_conn_id: str | None = "islo_default", + *, + image: str | None = None, + vcpus: int | None = None, + memory_mb: int | None = None, + delete_after: int = 3600, + ) -> None: + _validate_positive_finite(delete_after, "delete_after") + if vcpus is not None: + _validate_positive_finite(vcpus, "vcpus") + if memory_mb is not None: + _validate_positive_finite(memory_mb, "memory_mb") + if image == "": + raise ValueError("image must not be empty.") + self._islo_conn_id = islo_conn_id + self._image = image + self._vcpus = vcpus + self._memory_mb = memory_mb + self._delete_after = delete_after + self._client: Islo | None = None + + def _get_client(self) -> Islo: + if self._client is not None: + return self._client + with _translate_islo_errors("initialize its client"): + from islo import Islo + + if self._islo_conn_id is None: + self._client = Islo() + return self._client + conn = BaseHook.get_connection(self._islo_conn_id) + api_key = (conn.password or "").strip() + if not api_key: + raise SandboxTerminalError( + f"Connection {self._islo_conn_id!r} has no password; set it to the Islo API key." + ) + kwargs: dict[str, Any] = {"api_key": api_key} + if conn.host: + kwargs["compute_url"] = conn.host + extra = conn.extra_dejson + if extra.get("base_url"): + kwargs["base_url"] = extra["base_url"] + if extra.get("timeout") is not None: + try: + request_timeout = float(extra["timeout"]) + _validate_positive_finite(request_timeout, "connection extra timeout") + except (TypeError, ValueError) as e: + raise SandboxTerminalError( + "The Islo connection extra timeout must be a positive finite number." + ) from e + kwargs["timeout"] = request_timeout + self._client = Islo(**kwargs) + return self._client + + @staticmethod + def _request_options(*, timeout: float, chunk_size: int | None = None) -> dict[str, int]: + options = {"timeout_in_seconds": max(1, math.ceil(timeout)), "max_retries": 0} + if chunk_size is not None: + options["chunk_size"] = chunk_size + return options + + def create(self, *, spec: SandboxSpec | None = None) -> str: + if spec is not None and spec.allow_egress_to: + raise SandboxTerminalError( + "The Islo backend cannot apply a per-domain egress allowlist; it can only turn " + "outbound access on or off. Drop allow_egress_to, or use a backend with " + "per-domain network rules." + ) + with _translate_islo_errors("create a sandbox"): + from islo.types import LifecyclePolicy + + kwargs: dict[str, Any] = { + "internet_enabled": False if spec is None else not spec.block_network, + "lifecycle": LifecyclePolicy(delete_after=self._delete_after), + } + if self._image is not None: + kwargs["image"] = self._image + if self._vcpus is not None: + kwargs["vcpus"] = self._vcpus + if self._memory_mb is not None: + kwargs["memory_mb"] = self._memory_mb + if spec is not None and spec.env: + kwargs["env"] = dict(spec.env) + sandbox = self._get_client().sandboxes.create_sandbox( + name=_new_sandbox_name(), Review Comment: `_new_sandbox_name()` is generated inline, so if `create_sandbox` raises after the server has already provisioned the microVM (a response timeout, a reset, a 5xx after creation), the handle is gone and nothing can delete it or even log which sandbox leaked. `sbx.py` binds the name first and best-effort removes it under `except BaseException`. `delete_after` does bound the damage, but it is documented as the backstop for a worker killed mid-run, and here the worker is alive and able to clean up immediately. ########## providers/common/ai/src/airflow/providers/common/ai/sandbox/islo.py: ########## @@ -0,0 +1,397 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""islo.dev microVM backend for :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.""" + +from __future__ import annotations + +import math +import shlex +import time +from contextlib import contextmanager, suppress +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.sandbox.base import ( + SandboxBackend, + SandboxError, + SandboxExecResult, + SandboxFileTooLargeError, + SandboxTerminalError, + _new_sandbox_name, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterator + + from islo import Islo + from islo.errors import NotFoundError + + from airflow.providers.common.ai.sandbox.base import SandboxSpec + +_TERMINAL_EXEC_STATUSES = frozenset({"completed", "failed", "timeout"}) Review Comment: The SDK types the exec `status` as a plain `str` with `extra="allow"`, and the vendor's other status enums carry values like `cancelled` and `dead`, so these three may not be the full vocabulary. Anything outside the set reads as still-running here, so `_await_exec` polls to the deadline and then destroys the sandbox and reports a timeout. Treating an unrecognised status as terminal-failed instead would keep a mislabelled status from costing the agent its files. ########## providers/common/ai/docs/toolsets.rst: ########## @@ -793,6 +793,60 @@ Constructor parameters: guarantee this backend cannot make. Set ``"deny-all"`` after running ``sbx policy init deny-all``, or ``"allow-all"`` to state that egress is open. +Islo backend Review Comment: The `sbx` warning box just above (line 753) still says a hosted backend plugs in through `SandboxBackend` but none ships with the provider yet, and `sbx.py`'s class docstring says the same; both are now false. The credentials paragraph below also reads as unconditional, but `islo_conn_id=None` sends the SDK to its own environment variable, which is the worker environment the sentence says the key stays out of, and it is the path this PR's own system test and README use. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
