Lee-W commented on code in PR #73559:
URL: https://github.com/apache/airflow/pull/73559#discussion_r4091927936
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py:
##########
@@ -230,26 +311,146 @@ async def for_run(self, ctx: RunContext[Any]) ->
AbstractToolset[Any]:
# not silently degrade to this class on every run.
return type(self)(
self._backend,
- spec=self._spec,
+ # Attach mode refuses a spec, and the default one filled in above
is
+ # not the author's, so it is not handed back.
+ spec=None if self._attach_mode else self._spec,
default_command_timeout=self._default_command_timeout,
max_command_timeout=self._max_command_timeout,
max_output_lines=self._max_output_lines,
max_output_bytes=self._max_output_bytes,
max_read_bytes=self._max_read_bytes,
tool_prefix=self._tool_prefix,
+ attach_to=self._attached_handle() if self._attach_mode else None,
+ owner=self._owner,
)
+ def _attached_handle(self) -> str:
+ """
+ Return the rendered handle, checked, because the templater bypasses
the constructor.
+
+ ``attach_to`` is templated so the handle can come from an upstream
XCom, and
+ an XCom that was never pushed renders to ``None`` under native
rendering or to
+ the string ``"None"`` otherwise. Neither is a sandbox, and silently
falling back
+ to provisioning one would run the agent in an empty workspace with no
error.
+ """
+ handle = self.attach_to
+ if not is_sandbox_handle(handle):
+ raise SandboxTerminalError(
+ f"attach_to rendered to {handle!r}, which is not a sandbox
handle. The task that "
+ "provisions the sandbox pushed nothing, or this task does not
depend on it and ran "
+ "first; check the upstream task and the XCom it returns."
+ )
+ return handle
+
async def __aenter__(self) -> Self:
- # The sandbox is provisioned lazily on first use, not here: a durable
+ # An owned sandbox is provisioned lazily on first use, not here: a
durable
# replay that only serves cached tool results must not provision one,
and
- # nothing leaks if the run fails before any tool executes.
+ # nothing leaks if the run fails before any tool executes. An attached
one is
+ # claimed now, so a wrong handle or a held sandbox fails the run
before the
+ # model has spent anything, and the tool descriptions can state the
lifetime.
+ if self._attach_mode:
+ await self._attach(self._attached_handle())
return self
+ async def _attach(self, handle: str) -> None:
+ owner, holder = self._identity()
+ backend = self._attachable_backend()
+ try:
+ attached = await asyncio.to_thread(backend.attach, handle,
owner=owner, holder=holder)
+ except SandboxTerminalError:
+ raise
+ except SandboxError as e:
+ # Nothing the model does can change whether this sandbox can be
attached
+ # to, so a recoverable label here is one it could not act on.
+ raise SandboxTerminalError(
+ f"Could not attach to sandbox {handle!r} on backend
{backend.name!r}: {e}"
+ ) from e
+ self._sandbox = handle
+ self._holder = holder
+ remaining = attached.remaining_lifetime
+ self._expires_at = None if remaining is None else time.monotonic() +
remaining
+ self._attach_note = self._describe_attached(attached.network,
remaining)
+ log.info(
+ "Attached to sandbox %s on backend %s as %s; %s of its lifetime
remain",
+ handle,
+ backend.name,
+ holder,
+ "an unknown number of seconds" if remaining is None else
f"{remaining:.0f}s",
+ )
+
+ @classmethod
+ def _describe_attached(cls, network: SandboxSpec | None,
remaining_lifetime: float | None) -> str:
+ whose = (
+ "This sandbox was set up by an earlier task, and your files stay
in it after this run "
+ "for a later task to collect."
+ )
+ policy = (
+ cls._describe_network(network)
+ if network is not None
+ else "Its network access is whatever the task that set it up
allowed; test before relying on it."
+ )
+ if remaining_lifetime is None:
+ clock = "How long it has left is not known."
+ elif remaining_lifetime < 60:
+ clock = "Under a minute of its lifetime remained when this run
began, so finish up."
+ else:
+ minutes = round(remaining_lifetime / 60)
+ unit = "minute" if minutes == 1 else "minutes"
+ clock = f"About {minutes} {unit} of its lifetime remained when
this run began."
+ return f"{whose} {policy} {clock}"
+
+ def _attachable_backend(self) -> AttachableSandboxBackend:
+ if self._attachable is None:
+ # The constructor refuses attach_to on any other backend, so this
is a
+ # programming error, not a run-time condition.
+ raise RuntimeError("attach mode on a backend that cannot attach")
+ return self._attachable
+
+ def _identity(self) -> tuple[str, str]:
Review Comment:
should we make the return type a named tuple?
--
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]
