Lee-W commented on code in PR #73559:
URL: https://github.com/apache/airflow/pull/73559#discussion_r4091902153
##########
providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py:
##########
@@ -109,12 +137,89 @@ class SandboxSpec:
than a hostname list gives, so it needs no opt-in. Both lists may be
set
together; how a backend combines them, and what that costs, is the
backend's to document.
+ :param owner: Who the sandbox is for, when a task provisions it for an
agent
+ task to attach to later. A
:class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`
+ attaching to the sandbox has to present the same value, and by default
it
+ presents the Dag run it is part of, so the provisioning task in the
same
+ run writes ``owner=dag_run_owner(context)``. Unset for a sandbox nobody
+ will attach to. A backend that cannot record it must refuse it.
"""
env: Mapping[str, str] | None = None
block_network: bool = True
allow_egress_to: Sequence[str] | None = None
allow_egress_to_cidrs: Sequence[str] | None = None
+ owner: str | None = None
+
+
+def dag_run_owner(context: Mapping[str, Any]) -> str:
+ """
+ Return the owner token naming the Dag run a task is part of:
``"<dag_id>/<run_id>"``.
+
+ This is what a ``SandboxToolset`` presents when it attaches to a sandbox
+ without an explicit ``owner``, so a task provisioning a sandbox for an
agent
+ task in the same Dag run stamps it with
``SandboxSpec(owner=dag_run_owner(context))``.
+ The pair is unique across the deployment where a bare ``run_id`` is not:
two
+ Dags on the same schedule share their run ids. ``context`` is the task
context,
+ as a ``@task`` receives it in ``**context`` or ``get_current_context``
returns it.
+ """
+ ti = context["ti"]
+ return f"{ti.dag_id}/{ti.run_id}"
+
+
+def encode_network_policy(spec: SandboxSpec) -> str:
+ """
+ Serialize a spec's network policy for a sandbox tag, so an attaching
toolset can read it back.
+
+ The toolset tells the model what the sandbox can reach, because a model
that has to
+ discover a denied network by failing wastes a turn, or a whole command
budget. An
+ attached sandbox was provisioned under a spec the toolset never sees, so
the backend
+ records the policy on the sandbox at create and
:func:`decode_network_policy` turns
+ it back into a spec. Compact JSON with sorted keys, so the same policy
always encodes
+ the same way.
+ """
+ return json.dumps(
+ {
+ "block_network": spec.block_network,
+ "allow_egress_to": list(spec.allow_egress_to or ()),
+ "allow_egress_to_cidrs": list(spec.allow_egress_to_cidrs or ()),
+ },
+ separators=(",", ":"),
+ sort_keys=True,
+ )
+
+
+def decode_network_policy(value: str | None) -> SandboxSpec | None:
+ """Read a :data:`NETWORK_TAG` value back into a spec carrying only the
network fields, or ``None``."""
+ if not value:
+ return None
Review Comment:
Should we raise instead of returning none? handling the exception like
```python
try:
decode...
except ...:
...
```
make more sense to me
##########
providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py:
##########
@@ -109,12 +137,89 @@ class SandboxSpec:
than a hostname list gives, so it needs no opt-in. Both lists may be
set
together; how a backend combines them, and what that costs, is the
backend's to document.
+ :param owner: Who the sandbox is for, when a task provisions it for an
agent
+ task to attach to later. A
:class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`
+ attaching to the sandbox has to present the same value, and by default
it
+ presents the Dag run it is part of, so the provisioning task in the
same
+ run writes ``owner=dag_run_owner(context)``. Unset for a sandbox nobody
+ will attach to. A backend that cannot record it must refuse it.
"""
env: Mapping[str, str] | None = None
block_network: bool = True
allow_egress_to: Sequence[str] | None = None
allow_egress_to_cidrs: Sequence[str] | None = None
+ owner: str | None = None
+
+
+def dag_run_owner(context: Mapping[str, Any]) -> str:
Review Comment:
```suggestion
def extract_dag_run_owner(context: Mapping[str, Any]) -> str:
```
I kinda feel i saw it somewhere. if this is the convention, then let's keep
it
--
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]
