Lee-W commented on code in PR #73990:
URL: https://github.com/apache/airflow/pull/73990#discussion_r4151471288
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py:
##########
@@ -506,8 +594,77 @@ def _close(self) -> None:
if holder is not None:
self._release(self._attached_handle, holder)
return
+ export = bool(self._exports) and not run_failed
if sandbox is None:
+ if export and not self._forked:
+ # The run never called a tool, so it never provisioned a
sandbox, or the
+ # sandbox ended under its last command and nothing replaced
it. Either
+ # way the files it was to leave behind do not exist.
+ raise SandboxTerminalError(
+ "The run ended with no sandbox standing, so none of the
files it was to export "
+ f"exist: {', '.join(repr(path) for path in
self._exports)}."
+ )
return
+ try:
+ if export:
+ self._export(sandbox)
+ finally:
+ self._destroy(sandbox)
+
+ def _export(self, sandbox: str) -> None:
+ """
+ Copy every file in ``exports`` out of the sandbox, failing the task on
the first that cannot be.
+
+ Before teardown, and never best effort: a task that promised a file
and did not
+ deliver it must fail, or its downstream task finds nothing and cannot
tell why.
+ """
+ written: list[ObjectStoragePath] = []
+ for path, destination in self._exports.items():
+ target = self._export_target(destination)
Review Comment:
It is opened for writing and added to `written` before the source is
validated, so a failed export deletes whatever was already at that path.
e.g., having a local destination containing "ORIGINAL": the export failed
partway and the file was gone.
With 2 exports where the second fails, the first destination (already
overwritten) is deleted too.
We could try to check the source before opening the dest and write a temp
key and rename it on success or doc that a failed run removes every destination
it processed
##########
providers/common/ai/src/airflow/providers/common/ai/sandbox/sbx.py:
##########
@@ -55,6 +58,10 @@
# Helpers return a status or a directory listing, never bulk file content, so a
# small cap is enough to bound what a hostile guest can push into worker
memory.
_HELPER_OUTPUT_CAP = 1024 * 1024
+# Seconds an export may go without a byte arriving before it is ended. A
stall, not a
+# budget for the whole file: measured, 200 MB streams out of a local microVM
in under
+# three seconds, and a larger file only takes longer.
+_EXPORT_STALL_TIMEOUT = 120.0
Review Comment:
IIUC, this 120 is per chunk timeout, so a sending a byte every ~119 s can
keep the export going indefinitely. Is `execution_timeout` the intended outer
bound here, or should export have its own total deadline?
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py:
##########
@@ -469,19 +541,35 @@ def _identity(self) -> _Identity:
return _Identity(owner=self._owner if self._owner is not None else
run, holder=holder)
async def __aexit__(self, *args: Any) -> bool | None:
- await asyncio.to_thread(self._close)
+ # Read here, not in the thread: the exception being handled belongs to
this frame.
+ run_failed = args[0] is not None or self._run_failed()
+ await asyncio.to_thread(self._close, run_failed=run_failed)
return None
+ def _run_failed(self) -> bool:
Review Comment:
maybe making it a property?
##########
providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py:
##########
@@ -440,6 +452,125 @@ def list_directory(self, sandbox: str, path: str) ->
list[tuple[str, bool]]:
entries.append((name, kind == "d"))
return entries
+ def export_file(self, sandbox: str, path: str, dest: BinaryIO, *,
max_bytes: int) -> int:
+ """
+ Copy a regular file out of the sandbox into ``dest`` and return the
bytes written.
+
+ ``dest`` is a writable binary stream, typically an object-storage
file, and the
+ copy goes through it without the whole file ever being held in worker
memory,
+ so a file far over :meth:`read_file`'s budget can leave the sandbox.
Only a
+ regular file is exported: a directory, a device, or a FIFO is refused,
since
+ none of them has a size to promise a caller. Raise
+ :class:`SandboxFileTooLargeError` instead of copying a file larger than
+ ``max_bytes``, and :class:`SandboxError` when the file changed size
while it
+ was being copied, which means a process in the sandbox is still
writing it.
+
+ This default reads the file in slices through :meth:`run_command`, one
+ command per slice, and needs ``stat``, ``tail``, ``head`` and
``base64`` in the
+ guest. It relies on ``run_command`` returning each slice's output
intact, or
+ setting ``stdout_truncated`` when it could not, and on nothing but the
command's
+ own output reaching stdout. Override it when the vendor can stream a
file out.
+ """
+ quoted = shlex.quote(path)
+ check = self.run_command(
+ sandbox,
+ f"{self._export_checks(quoted, max_bytes)}
{self._report_export_size()}",
+ timeout=_FILE_OP_TIMEOUT,
+ max_output_bytes=_FILE_OP_OUTPUT_CAP,
+ )
+ self._raise_for_export_status(path, check, max_bytes)
+ size = self._reported_export_size(check.stdout)
+ written = 0
+ while True:
+ # ``tail -c +N`` seeks on a regular file, so each slice costs its
own
+ # length rather than a read from the start.
+ result = self.run_command(
+ sandbox,
+ f"tail -c +{written + 1} -- {quoted} | head -c
{_EXPORT_CHUNK_BYTES} | base64",
+ timeout=_FILE_OP_TIMEOUT,
+ max_output_bytes=_EXPORT_CHUNK_BYTES * 2 + 4096,
+ )
+ if result.sandbox_terminated:
+ raise SandboxTerminalError(f"The sandbox ended while {path!r}
was being exported.")
+ if result.exit_code or result.stdout_truncated:
+ # A truncated slice decodes cleanly into the wrong bytes, so
it is an
+ # error, not a short read.
+ raise SandboxError(result.stderr.strip() or f"Could not export
{path!r}.")
+ try:
+ data = base64.b64decode(result.stdout, validate=False)
+ except (binascii.Error, ValueError) as e:
+ raise SandboxError(f"Could not decode {path!r} from the
sandbox.") from e
+ if written + len(data) > max_bytes:
+ raise SandboxFileTooLargeError(path, written + len(data),
max_bytes)
+ dest.write(data)
+ written += len(data)
+ if len(data) < _EXPORT_CHUNK_BYTES:
+ break
+ self._check_export_size(path, expected=size, written=written)
+ return written
+
+ @classmethod
+ def _export_checks(cls, quoted: str, max_bytes: int) -> str:
+ """
+ Shell that refuses what :meth:`export_file` will not copy, leaving the
size in ``$sz``.
+
+ For a backend overriding :meth:`export_file` that still drives the
guest
+ through a shell, so its refusals mean what the default's do.
+ """
+ return (
+ f"sz=$(stat -Lc %s -- {quoted} 2>/dev/null) || exit
{cls._MISSING_PATH_STATUS}; "
+ f"[ -d {quoted} ] && exit {cls._IS_DIRECTORY_STATUS}; "
+ f"[ -f {quoted} ] || exit {cls._NOT_REGULAR_FILE_STATUS}; "
+ f'[ "$sz" -gt {max_bytes} ] && {{
{cls._report_export_size(to_stderr=True)} exit {cls._TOO_LARGE_STATUS}; }};'
+ )
+
+ @staticmethod
+ def _report_export_size(*, to_stderr: bool = False) -> str:
+ """
+ Shell that prints ``$sz`` on a line of its own, tagged.
+
+ Tagged because the stream it lands on may also carry the vendor CLI's
own notices.
+ """
+ redirect = " >&2" if to_stderr else ""
+ return f'printf "\\n{_EXPORT_SIZE_TAG}%s\\n" "$sz"{redirect};'
+
+ @staticmethod
+ def _reported_export_size(text: str) -> int:
+ sizes = re.findall(rf"^{_EXPORT_SIZE_TAG}(\d+)$", text,
flags=re.MULTILINE)
+ if not sizes:
+ raise SandboxError("The sandbox did not report the size of the
file being exported.")
+ return int(sizes[-1])
Review Comment:
we should rename one of them. it's a bit hard to distinguish
--
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]
