kaxil commented on code in PR #73990:
URL: https://github.com/apache/airflow/pull/73990#discussion_r4153929727
##########
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:
Renamed to `_print_export_size` (builds the shell that prints it) and
`_parse_export_size` (reads it back).
##########
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:
Done.
--
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]
