From 84408ebe27d5d2308bdcb5b0ffbbd3aa0c0d1b8a Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Thu, 18 Jun 2026 14:54:49 +0530 Subject: [PATCH 1/4] fix: add depth guard to sync_node_execution to prevent RecursionError - Add _depth and _max_depth parameters (default 50) to _sync_execution and sync_node_execution - Raise FlyteAssertion when nesting exceeds limit instead of crashing with RecursionError - Add unit tests verifying depth guard and regression safety Fixes #7338 Signed-off-by: Abhishek Shinde --- flytekit/remote/remote.py | 48 +++++++- .../unit/remote/test_recursion_guard.py | 115 ++++++++++++++++++ 2 files changed, 157 insertions(+), 6 deletions(-) create mode 100644 tests/flytekit/unit/remote/test_recursion_guard.py diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 041c701ec4..0120b88f73 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -2587,10 +2587,17 @@ def _sync_execution( self, execution: FlyteWorkflowExecution, sync_nodes: bool = False, + _depth: int = 0, + _max_depth: int = 50, ) -> FlyteWorkflowExecution: """ Sync a FlyteWorkflowExecution object with its corresponding remote state. """ + if _depth > _max_depth: + raise FlyteAssertion( + f"Nesting depth {_depth} exceeds _max_depth={_max_depth} for execution " + f"{execution.id}. Refusing to recurse further to avoid RecursionError." + ) # Update closure, and then data, because we don't want the execution to finish between when we get the data, # and then for the closure to have is_done to be true. execution._closure = self.client.get_execution(execution.id).closure @@ -2642,7 +2649,9 @@ def _sync_execution( if sync_nodes: node_execs = {} for n in underlying_node_executions: - node_execs[n.id.node_id] = self.sync_node_execution(n, node_mapping) # noqa + node_execs[n.id.node_id] = self.sync_node_execution( # noqa + n, node_mapping, _depth=_depth + 1, _max_depth=_max_depth + ) execution._node_executions = node_execs return self._assign_inputs_and_outputs(execution, execution_data, node_interface) @@ -2650,6 +2659,8 @@ def sync_node_execution( self, execution: FlyteNodeExecution, node_mapping: typing.Dict[str, FlyteNode], + _depth: int = 0, + _max_depth: int = 50, ) -> FlyteNodeExecution: """ Get data backing a node execution. These FlyteNodeExecution objects should've come from Admin with the model @@ -2669,6 +2680,11 @@ def sync_node_execution( The data model is complicated, so ascertaining which of these happened is a bit tricky. That logic is encapsulated in this function. """ + if _depth > _max_depth: + raise FlyteAssertion( + f"Nesting depth {_depth} exceeds _max_depth={_max_depth} for node " + f"{execution.id}. Refusing to recurse further to avoid RecursionError." + ) # For single task execution - the metadata spec node id is missing. In these cases, revert to regular node id node_id = execution.metadata.spec_node_id # This case supports single-task execution compiled workflows. @@ -2712,7 +2728,7 @@ def sync_node_execution( launched_exec = self.fetch_execution( project=launched_exec_id.project, domain=launched_exec_id.domain, name=launched_exec_id.name ) - self.sync_execution(launched_exec, sync_nodes=True) + self._sync_execution(launched_exec, sync_nodes=True, _depth=_depth + 1, _max_depth=_max_depth) if launched_exec.is_done: # The synced underlying execution should've had these populated. execution._inputs = launched_exec.inputs @@ -2743,7 +2759,12 @@ def sync_node_execution( dynamic_flyte_wf = FlyteWorkflow.promote_from_closure(compiled_wf, node_launch_plans) execution._underlying_node_executions = [ - self.sync_node_execution(FlyteNodeExecution.promote_from_model(cne), dynamic_flyte_wf._node_map) + self.sync_node_execution( + FlyteNodeExecution.promote_from_model(cne), + dynamic_flyte_wf._node_map, + _depth=_depth + 1, + _max_depth=_max_depth, + ) for cne in child_node_executions ] execution._task_executions = [ @@ -2757,7 +2778,12 @@ def sync_node_execution( sub_flyte_workflow = execution._node.flyte_entity sub_node_mapping = {n.id: n for n in sub_flyte_workflow.flyte_nodes} execution._underlying_node_executions = [ - self.sync_node_execution(FlyteNodeExecution.promote_from_model(cne), sub_node_mapping) + self.sync_node_execution( + FlyteNodeExecution.promote_from_model(cne), + sub_node_mapping, + _depth=_depth + 1, + _max_depth=_max_depth, + ) for cne in child_node_executions ] execution._interface = sub_flyte_workflow.interface @@ -2778,7 +2804,12 @@ def sync_node_execution( sub_node_mapping[else_node.id] = else_node execution._underlying_node_executions = [ - self.sync_node_execution(FlyteNodeExecution.promote_from_model(cne), sub_node_mapping) + self.sync_node_execution( + FlyteNodeExecution.promote_from_model(cne), + sub_node_mapping, + _depth=_depth + 1, + _max_depth=_max_depth, + ) for cne in child_node_executions ] else: @@ -2851,7 +2882,12 @@ def sync_node_execution( sub_node_mapping[else_node.id] = else_node execution._underlying_node_executions = [ - self.sync_node_execution(FlyteNodeExecution.promote_from_model(cne), sub_node_mapping) + self.sync_node_execution( + FlyteNodeExecution.promote_from_model(cne), + sub_node_mapping, + _depth=_depth + 1, + _max_depth=_max_depth, + ) for cne in child_node_executions ] diff --git a/tests/flytekit/unit/remote/test_recursion_guard.py b/tests/flytekit/unit/remote/test_recursion_guard.py new file mode 100644 index 0000000000..2413d38e5b --- /dev/null +++ b/tests/flytekit/unit/remote/test_recursion_guard.py @@ -0,0 +1,115 @@ +from datetime import datetime, timedelta, timezone +from unittest.mock import MagicMock, patch + +import pytest + +from flytekit.configuration import Config +from flytekit.exceptions.user import FlyteAssertion +from flytekit.models.core.identifier import NodeExecutionIdentifier, WorkflowExecutionIdentifier +from flytekit.models.node_execution import ( + NodeExecutionClosure, + WorkflowNodeMetadata, +) +from flytekit.remote.executions import FlyteNodeExecution, FlyteWorkflowExecution +from flytekit.remote.interface import TypedInterface +from flytekit.remote.remote import FlyteRemote + + +def _mock_launched_exec(): + exec = MagicMock(spec=FlyteWorkflowExecution) + wf = MagicMock() + wf.interface = MagicMock(spec=TypedInterface) + exec._flyte_workflow = wf + exec.is_done = False + exec.inputs = {"a": 1} + exec.outputs = {"b": 2} + return exec + + +def _make_node_execution(node_id: str, execution_name: str, with_workflow_node_metadata: bool = False): + wf_exec_id = WorkflowExecutionIdentifier("p1", "d1", execution_name) + ne_id = NodeExecutionIdentifier(node_id, wf_exec_id) + meta = MagicMock() + meta.is_parent_node = False + meta.is_array = False + meta.spec_node_id = node_id + + wf_node_meta = None + if with_workflow_node_metadata: + wf_node_meta = WorkflowNodeMetadata( + execution_id=WorkflowExecutionIdentifier("p1", "d1", f"launched_{execution_name}") + ) + + closure = NodeExecutionClosure( + phase=0, + started_at=datetime.now(timezone.utc), + duration=timedelta(seconds=1), + workflow_node_metadata=wf_node_meta, + ) + return FlyteNodeExecution(id=ne_id, input_uri="s3://bucket/input", closure=closure, metadata=meta) + + +@pytest.fixture +def remote(): + with patch("flytekit.clients.friendly.SynchronousFlyteClient"): + flyte_remote = FlyteRemote( + config=Config.auto(), + default_project="p1", + default_domain="d1", + ) + flyte_remote._client_initialized = True + flyte_remote._client = MagicMock() + return flyte_remote + + +def test_max_depth_raises_flyte_assertion(remote): + ne = _make_node_execution("n1", "exec1", with_workflow_node_metadata=True) + + launched_exec = _mock_launched_exec() + remote.fetch_execution = MagicMock(return_value=launched_exec) + remote.client.get_node_execution_data = MagicMock( + return_value=MagicMock(dynamic_workflow=None) + ) + + # Don't mock _sync_execution — the guard fires inside it when _depth exceeds _max_depth. + # With _max_depth=0, the initial sync_node_execution enters at _depth=0 (passes guard). + # It reaches the launched LP path, which calls _sync_execution with _depth=1. + # _sync_execution then sees 1 > 0 and raises FlyteAssertion. + with pytest.raises(FlyteAssertion, match="Nesting depth"): + remote.sync_node_execution(ne, {"n1": MagicMock()}, _max_depth=0) + + +def test_reasonable_depth_does_not_raise(remote): + ne = _make_node_execution("n1", "exec1", with_workflow_node_metadata=True) + + launched_exec = _mock_launched_exec() + remote.fetch_execution = MagicMock(return_value=launched_exec) + remote.client.get_node_execution_data = MagicMock( + return_value=MagicMock(dynamic_workflow=None) + ) + remote._sync_execution = MagicMock() + + result = remote.sync_node_execution(ne, {"n1": MagicMock()}) + assert result is ne + + +def test_nested_under_default_limit(remote): + ne = _make_node_execution("n1", "exec1", with_workflow_node_metadata=True) + + launched_exec = _mock_launched_exec() + remote.fetch_execution = MagicMock(return_value=launched_exec) + remote.client.get_node_execution_data = MagicMock( + return_value=MagicMock(dynamic_workflow=None) + ) + remote._sync_execution = MagicMock() + + result = remote.sync_node_execution(ne, {"n1": MagicMock()}, _depth=1, _max_depth=50) + assert result is ne + + +def test_sync_execution_depth_guard(remote): + wf_exec = MagicMock(spec=FlyteWorkflowExecution) + wf_exec.id = WorkflowExecutionIdentifier("p1", "d1", "deep_exec") + + with pytest.raises(FlyteAssertion, match="Nesting depth"): + remote._sync_execution(wf_exec, sync_nodes=True, _depth=51, _max_depth=50) From f4e49096055cec9230b7264c1801a8f1595e5ba6 Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Wed, 24 Jun 2026 18:36:20 +0530 Subject: [PATCH 2/4] feat: add flytekit-asqav plugin for cryptographic audit trails Implements an decorator (ClassDecorator pattern) that wraps Flyte tasks with cryptographically signed receipts at started/finished/failed lifecycle points using the Asqav SDK. Receipts are rendered as a Flyte Deck card with verification links. Closes #7085 Signed-off-by: Abhishek Shinde --- plugins/flytekit-asqav/README.md | 65 +++++ .../flytekitplugins/asqav/__init__.py | 19 ++ .../flytekitplugins/asqav/tracking.py | 238 ++++++++++++++++++ plugins/flytekit-asqav/setup.py | 37 +++ plugins/flytekit-asqav/tests/__init__.py | 0 .../tests/test_asqav_tracking.py | 201 +++++++++++++++ uv.lock | 14 +- 7 files changed, 567 insertions(+), 7 deletions(-) create mode 100644 plugins/flytekit-asqav/README.md create mode 100644 plugins/flytekit-asqav/flytekitplugins/asqav/__init__.py create mode 100644 plugins/flytekit-asqav/flytekitplugins/asqav/tracking.py create mode 100644 plugins/flytekit-asqav/setup.py create mode 100644 plugins/flytekit-asqav/tests/__init__.py create mode 100644 plugins/flytekit-asqav/tests/test_asqav_tracking.py diff --git a/plugins/flytekit-asqav/README.md b/plugins/flytekit-asqav/README.md new file mode 100644 index 0000000000..f11c9fb21a --- /dev/null +++ b/plugins/flytekit-asqav/README.md @@ -0,0 +1,65 @@ +# Flytekit Asqav Plugin + +Cryptographic audit trails for Flyte tasks using [Asqav](https://asqav.com). + +Each task execution is wrapped with cryptographically signed receipts at lifecycle +points (started, finished, failed), providing provable, tamper-proof records for +AI workflows in regulated environments (finance, healthcare, etc.). + +## Installation + +```bash +pip install flytekitplugins-asqav +``` + +Requires `flytekit>=1.12.0` and `asqav>=0.1.0`. + +## Quick Start + +```python +from flytekit import task, Secret, workflow +from flytekitplugins.asqav import asqav_audit + +asqav_secret = Secret(key="asqav-api-key", group="asqav") + +@task(secret_requests=[asqav_secret]) +@asqav_audit(agent_name="model-trainer", secret=asqav_secret) +def train_model() -> float: + # ... training code ... + return 0.95 + +@workflow +def main() -> float: + return train_model() +``` + +## Configuration + +### `@asqav_audit(...)` + +| Parameter | Type | Required | Default | Description | +|---|---|---|---|---| +| `agent_name` | `str` | ✅ | — | Asqav agent name for this task | +| `secret` | `Secret` or `Callable` | ❌ | `ASQAV_API_KEY` env var | How to resolve the API key | +| `receipt_type` | `str` | ❌ | `protectmcp:lifecycle` | Asqav receipt type | +| `risk_class` | `str` | ❌ | `None` | Risk classification (`low`, `medium`, `high`, `unknown`) | +| `compliance_mode` | `bool` | ❌ | `True` | Generate IETF-compliant compliance receipts | + +### Secret Resolution (priority order) + +1. **`Secret` object** — resolved via Flyte's `SecretsManager` (for remote deployment) +2. **Callable** — zero-argument function returning the API key string +3. **`ASQAV_API_KEY`** environment variable — fallback + +## Audit Trail Output + +Each receipt is logged and rendered as a Flyte Deck card with: +- Receipt type (started / finished / failed) +- Execution ID and timestamp +- Duration (for finished/failed receipts) +- Error info (for failed receipts) +- Clickable verification link to Asqav + +## License + +Apache 2.0 diff --git a/plugins/flytekit-asqav/flytekitplugins/asqav/__init__.py b/plugins/flytekit-asqav/flytekitplugins/asqav/__init__.py new file mode 100644 index 0000000000..1e7bfc5600 --- /dev/null +++ b/plugins/flytekit-asqav/flytekitplugins/asqav/__init__.py @@ -0,0 +1,19 @@ +""" +.. currentmodule:: flytekitplugins.asqav + +This plugin provides cryptographic audit trails for Flyte tasks using +`Asqav `_. Each task execution is wrapped with +cryptographically signed receipts at lifecycle points (started, finished, +failed), enabling provable, tamper-proof audit records for AI workflows in +regulated environments. + +.. autosummary:: + :template: custom.rst + :toctree: generated/ + + asqav_audit +""" + +from .tracking import asqav_audit + +__all__ = ["asqav_audit"] diff --git a/plugins/flytekit-asqav/flytekitplugins/asqav/tracking.py b/plugins/flytekit-asqav/flytekitplugins/asqav/tracking.py new file mode 100644 index 0000000000..e95de635ac --- /dev/null +++ b/plugins/flytekit-asqav/flytekitplugins/asqav/tracking.py @@ -0,0 +1,238 @@ +import datetime +import os +from typing import Any, Callable, Dict, Optional, Union + +import asqav +from flytekit import Deck, Secret +from flytekit.core.context_manager import FlyteContextManager +from flytekit.core.utils import ClassDecorator +from flytekit.loggers import logger + +ASQAV_LINK_TYPE_KEY = "link_type" +ASQAV_DEFAULT_RECEIPT_TYPE = "protectmcp:lifecycle" +ASQAV_ENV_API_KEY = "ASQAV_API_KEY" + +ASQAV_SCOPE_KEY = "scope" +ASQAV_AGENT_NAME_KEY = "agent_name" +ASQAV_RECEIPT_TYPE_KEY = "receipt_type" +ASQAV_RISK_CLASS_KEY = "risk_class" +ASQAV_COMPLIANCE_MODE_KEY = "compliance_mode" + + +def _receipt_html(receipt_type: str, sig_id: str, verification_url: str, metadata: Dict[str, Any]) -> str: + """Render a single audit receipt as an HTML row.""" + status_icon = "✅" if receipt_type == "finished" else "❌" if receipt_type == "failed" else "⏳" + rows = "".join( + f"{k}" + f"{v}" + for k, v in metadata.items() + ) + return f""" +
+
+ {status_icon} Receipt: {receipt_type} + ID: {sig_id} +
+
+ + {rows} +
+

+ + 🔗 Verify on Asqav + +

+
+
""" + + +class asqav_audit(ClassDecorator): + """Cryptographic audit decorator for Flyte tasks using Asqav. + + Wraps a ``@task``-decorated function to sign cryptographically verifiable + receipts at each lifecycle point (started, finished, failed). Receipts are + rendered as a Flyte Deck card with verification links. + + Usage:: + + from flytekit import task, Secret, workflow + from flytekitplugins.asqav import asqav_audit + + asqav_secret = Secret(key="asqav-api-key", group="asqav") + + @task(secret_requests=[asqav_secret]) + @asqav_audit(agent_name="model-trainer", secret=asqav_secret) + def train_model() -> float: + ... + """ + + def __init__( + self, + task_function=None, + *, + agent_name: str, + secret: Optional[Union[Secret, Callable]] = None, + receipt_type: str = ASQAV_DEFAULT_RECEIPT_TYPE, + risk_class: Optional[str] = None, + compliance_mode: bool = True, + **init_kwargs, + ): + if not agent_name: + raise ValueError("agent_name must be set") + super().__init__( + task_function, + agent_name=agent_name, + secret=secret, + receipt_type=receipt_type, + risk_class=risk_class, + compliance_mode=compliance_mode, + **init_kwargs, + ) + self.agent_name = agent_name + self.secret = secret + self.receipt_type = receipt_type + self.risk_class = risk_class + self.compliance_mode = compliance_mode + + def _resolve_api_key(self) -> str: + """Resolve the Asqav API key from the most specific source available.""" + # 1. Callable — highest priority for programmatic injection + if callable(self.secret): + return self.secret() + # 2. Environment variable — convenient for local development + api_key = os.environ.get(ASQAV_ENV_API_KEY) + if api_key: + return api_key + # 3. Flyte Secret — for remote execution with secret_requests + if isinstance(self.secret, Secret): + ctx = FlyteContextManager.current_context() + return ctx.user_space_params.secrets.get(key=self.secret.key, group=self.secret.group) + raise ValueError( + f"No Asqav API key found. Provide a `Secret` or callable to " + f"`asqav_audit(secret=...)`, or set the {ASQAV_ENV_API_KEY} env var." + ) + + def get_extra_config(self): + extra_config = { + ASQAV_LINK_TYPE_KEY: "asqav-audit", + ASQAV_AGENT_NAME_KEY: self.agent_name, + ASQAV_SCOPE_KEY: "flyte-task-lifecycle", + ASQAV_RECEIPT_TYPE_KEY: self.receipt_type, + ASQAV_COMPLIANCE_MODE_KEY: str(self.compliance_mode), + } + if self.risk_class: + extra_config[ASQAV_RISK_CLASS_KEY] = self.risk_class + return extra_config + + def execute(self, *args, **kwargs): + ctx = FlyteContextManager.current_context() + is_local = ctx.execution_state.is_local_execution() + execution_id = str(ctx.user_space_params.execution_id) if not is_local else "local" + + api_key = self._resolve_api_key() + + # Initialise the Asqav SDK and create an agent + asqav.init(api_key=api_key) + agent = asqav.Agent.create(name=self.agent_name) + + receipts = [] + start_time = datetime.datetime.utcnow() + + # --- started receipt --- + started_context = { + "event_type": "flyte.task.started", + "execution_id": execution_id, + "agent_name": self.agent_name, + "timestamp": datetime.datetime.utcnow().isoformat() + "Z", + } + started_resp = agent.sign( + action_type="flyte.task.started", + context=started_context, + compliance_mode=self.compliance_mode, + receipt_type=self.receipt_type, + risk_class=self.risk_class, + ) + receipts.append(("started", started_resp)) + logger.info(f"Signed asqav started receipt: {started_resp.signature_id}") + + # --- task execution --- + output = None + exception: Optional[Exception] = None + try: + output = self.task_function(*args, **kwargs) + except Exception as e: + exception = e + raise + finally: + elapsed = (datetime.datetime.utcnow() - start_time).total_seconds() + + if exception is not None: + error_info = f"{type(exception).__name__}: {exception}" + failed_context = { + "event_type": "flyte.task.failed", + "execution_id": execution_id, + "agent_name": self.agent_name, + "duration_seconds": round(elapsed, 3), + "error": error_info, + "timestamp": datetime.datetime.utcnow().isoformat() + "Z", + } + try: + failed_resp = agent.sign( + action_type="flyte.task.failed", + context=failed_context, + compliance_mode=self.compliance_mode, + receipt_type=self.receipt_type, + risk_class=self.risk_class, + ) + receipts.append(("failed", failed_resp)) + logger.info(f"Signed asqav failed receipt: {failed_resp.signature_id}") + except Exception as sign_err: + logger.error(f"Failed to sign asqav receipt for error: {sign_err}") + else: + finished_context = { + "event_type": "flyte.task.finished", + "execution_id": execution_id, + "agent_name": self.agent_name, + "duration_seconds": round(elapsed, 3), + "output_type": type(output).__name__ if output is not None else "None", + "timestamp": datetime.datetime.utcnow().isoformat() + "Z", + } + try: + finished_resp = agent.sign( + action_type="flyte.task.finished", + context=finished_context, + compliance_mode=self.compliance_mode, + receipt_type=self.receipt_type, + risk_class=self.risk_class, + ) + receipts.append(("finished", finished_resp)) + logger.info(f"Signed asqav finished receipt: {finished_resp.signature_id}") + except Exception as sign_err: + logger.error(f"Failed to sign asqav receipt for success: {sign_err}") + + # Render Flyte Deck with all receipts + if receipts: + deck_html_parts = [ + "

Asqav Cryptographic Audit Trail

", + f"

Agent: {self.agent_name} | " + f"Execution ID: {execution_id}

", + ] + for rtype, resp in receipts: + metadata = { + "Action Type": resp.signature_id, + "Algorithm": resp.algorithm or "unknown", + "Compliance Mode": str(resp.compliance_mode), + } + if resp.receipt_type: + metadata["Receipt Type"] = resp.receipt_type + if resp.risk_class: + metadata["Risk Class"] = resp.risk_class + deck_html_parts.append( + _receipt_html(rtype, resp.signature_id, resp.verification_url, metadata) + ) + deck = Deck("Asqav Audit Trail", "\n".join(deck_html_parts)) + + return output diff --git a/plugins/flytekit-asqav/setup.py b/plugins/flytekit-asqav/setup.py new file mode 100644 index 0000000000..18dfde8540 --- /dev/null +++ b/plugins/flytekit-asqav/setup.py @@ -0,0 +1,37 @@ +from setuptools import setup + +PLUGIN_NAME = "asqav" + +microlib_name = f"flytekitplugins-{PLUGIN_NAME}" + +plugin_requires = ["flytekit>=1.12.0", "asqav>=0.1.0"] + +__version__ = "0.0.0+develop" + +setup( + title="Asqav", + title_expanded="Flytekit Asqav Audit Trail Plugin", + name=microlib_name, + version=__version__, + author="flyteorg", + author_email="admin@flyte.org", + description="Cryptographic audit trails for Flyte tasks using Asqav", + namespace_packages=["flytekitplugins"], + packages=[f"flytekitplugins.{PLUGIN_NAME}"], + install_requires=plugin_requires, + license="apache2", + python_requires=">=3.10", + classifiers=[ + "Intended Audience :: Science/Research", + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: Scientific/Engineering", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Software Development", + "Topic :: Software Development :: Libraries", + "Topic :: Software Development :: Libraries :: Python Modules", + ], +) diff --git a/plugins/flytekit-asqav/tests/__init__.py b/plugins/flytekit-asqav/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/plugins/flytekit-asqav/tests/test_asqav_tracking.py b/plugins/flytekit-asqav/tests/test_asqav_tracking.py new file mode 100644 index 0000000000..fb7f57d082 --- /dev/null +++ b/plugins/flytekit-asqav/tests/test_asqav_tracking.py @@ -0,0 +1,201 @@ +from unittest.mock import MagicMock, Mock, PropertyMock, patch + +import pytest +from flytekitplugins.asqav import asqav_audit +from flytekitplugins.asqav.tracking import ( + ASQAV_AGENT_NAME_KEY, + ASQAV_COMPLIANCE_MODE_KEY, + ASQAV_LINK_TYPE_KEY, + ASQAV_RECEIPT_TYPE_KEY, + ASQAV_RISK_CLASS_KEY, + ASQAV_SCOPE_KEY, + ASQAV_DEFAULT_RECEIPT_TYPE, +) + +from flytekit import Secret, task + +secret = Secret(key="asqav-api-key", group="asqav") + + +def _make_signature_response(signature_id: str, **overrides) -> Mock: + """Create a mock SignatureResponse with the expected attributes.""" + defaults = { + "signature_id": signature_id, + "verification_url": f"https://asqav.com/verify/{signature_id}", + "algorithm": "ml-dsa-65", + "compliance_mode": True, + "receipt_type": ASQAV_DEFAULT_RECEIPT_TYPE, + "risk_class": None, + } + resp = Mock(**{**defaults, **overrides}) + return resp + + +def test_agent_name_required(): + with pytest.raises(TypeError, match="agent_name"): + asqav_audit() + + with pytest.raises(ValueError, match="agent_name must be set"): + asqav_audit(agent_name="") + + +def test_get_extra_config(): + decorator = asqav_audit( + agent_name="model-trainer", + secret=secret, + receipt_type="protectmcp:lifecycle", + risk_class="medium", + ) + config = decorator.get_extra_config() + assert config[ASQAV_LINK_TYPE_KEY] == "asqav-audit" + assert config[ASQAV_AGENT_NAME_KEY] == "model-trainer" + assert config[ASQAV_SCOPE_KEY] == "flyte-task-lifecycle" + assert config[ASQAV_RECEIPT_TYPE_KEY] == "protectmcp:lifecycle" + assert config[ASQAV_COMPLIANCE_MODE_KEY] == "True" + assert config[ASQAV_RISK_CLASS_KEY] == "medium" + + +def test_get_extra_config_defaults(): + decorator = asqav_audit(agent_name="trainer", secret=secret) + config = decorator.get_extra_config() + assert config[ASQAV_LINK_TYPE_KEY] == "asqav-audit" + assert config[ASQAV_AGENT_NAME_KEY] == "trainer" + assert config[ASQAV_SCOPE_KEY] == "flyte-task-lifecycle" + assert config[ASQAV_RECEIPT_TYPE_KEY] == ASQAV_DEFAULT_RECEIPT_TYPE + assert config[ASQAV_COMPLIANCE_MODE_KEY] == "True" + assert ASQAV_RISK_CLASS_KEY not in config + + +@task +@asqav_audit(agent_name="model-trainer", secret=secret) +def train_model() -> str: + return "done" + + +@patch("flytekitplugins.asqav.tracking.asqav") +@patch("flytekitplugins.asqav.tracking.Deck") +def test_local_execution(deck_mock, asqav_mock, monkeypatch): + """Verify the decorator signs started + finished receipts on success.""" + monkeypatch.setenv("ASQAV_API_KEY", "local-test-key") + agent_mock = MagicMock() + agent_mock.sign.return_value = _make_signature_response("receipt-started") + asqav_mock.Agent.create.return_value = agent_mock + + result = train_model() + + assert result == "done" + asqav_mock.init.assert_called_once() + asqav_mock.Agent.create.assert_called_once_with(name="model-trainer") + assert agent_mock.sign.call_count == 2 # started + finished + deck_mock.assert_called_once() + + +@patch("flytekitplugins.asqav.tracking.asqav") +@patch("flytekitplugins.asqav.tracking.Deck") +def test_failure_signs_failed_receipt(deck_mock, asqav_mock, monkeypatch): + """Verify the decorator signs a failed receipt when the task raises.""" + monkeypatch.setenv("ASQAV_API_KEY", "local-test-key") + + @task + @asqav_audit(agent_name="model-trainer", secret=secret) + def failing_task(): + raise RuntimeError("model training failed") + + agent_mock = MagicMock() + agent_mock.sign.return_value = _make_signature_response("receipt-failed") + asqav_mock.Agent.create.return_value = agent_mock + + with pytest.raises(RuntimeError, match="model training failed"): + failing_task() + + assert agent_mock.sign.call_count == 2 # started + failed + deck_mock.assert_called_once() + + +@patch("flytekitplugins.asqav.tracking.asqav") +@patch("flytekitplugins.asqav.tracking.FlyteContextManager") +@patch("flytekitplugins.asqav.tracking.Deck") +def test_remote_execution_secret_resolution(deck_mock, manager_mock, asqav_mock): + """Verify secrets are resolved in remote execution.""" + ctx_mock = Mock() + ctx_mock.execution_state.is_local_execution.return_value = False + ctx_mock.user_space_params.secrets.get.return_value = "asqav-key-123" + ctx_mock.user_space_params.execution_id = "exec-001" + manager_mock.current_context.return_value = ctx_mock + + agent_mock = MagicMock() + agent_mock.sign.return_value = _make_signature_response("receipt-remote") + asqav_mock.Agent.create.return_value = agent_mock + + result = train_model() + + assert result == "done" + ctx_mock.user_space_params.secrets.get.assert_called_with(key="asqav-api-key", group="asqav") + asqav_mock.init.assert_called_with(api_key="asqav-key-123") + assert agent_mock.sign.call_count == 2 + + +def get_my_key(): + return "callable-key-456" + + +@task +@asqav_audit(agent_name="callable-agent", secret=get_my_key) +def callable_secret_task() -> int: + return 42 + + +@patch("flytekitplugins.asqav.tracking.asqav") +@patch("flytekitplugins.asqav.tracking.FlyteContextManager") +@patch("flytekitplugins.asqav.tracking.Deck") +def test_callable_secret(deck_mock, manager_mock, asqav_mock): + """Verify a callable secret is called to get the API key.""" + ctx_mock = Mock() + ctx_mock.execution_state.is_local_execution.return_value = False + ctx_mock.user_space_params.execution_id = "exec-002" + manager_mock.current_context.return_value = ctx_mock + + agent_mock = MagicMock() + agent_mock.sign.return_value = _make_signature_response("receipt-callable") + asqav_mock.Agent.create.return_value = agent_mock + + result = callable_secret_task() + + assert result == 42 + asqav_mock.init.assert_called_with(api_key="callable-key-456") + + +@task +@asqav_audit(agent_name="env-agent") +def env_var_task() -> str: + return "from-env" + + +@patch("flytekitplugins.asqav.tracking.asqav") +@patch("flytekitplugins.asqav.tracking.FlyteContextManager") +@patch("flytekitplugins.asqav.tracking.Deck") +def test_env_var_fallback(deck_mock, manager_mock, asqav_mock, monkeypatch): + """Verify the ASQAV_API_KEY env var is used when no Secret or callable is provided.""" + monkeypatch.setenv("ASQAV_API_KEY", "env-key-789") + + ctx_mock = Mock() + ctx_mock.execution_state.is_local_execution.return_value = False + ctx_mock.user_space_params.execution_id = "exec-003" + manager_mock.current_context.return_value = ctx_mock + + agent_mock = MagicMock() + agent_mock.sign.return_value = _make_signature_response("receipt-env") + asqav_mock.Agent.create.return_value = agent_mock + + result = env_var_task() + + assert result == "from-env" + asqav_mock.init.assert_called_with(api_key="env-key-789") + + +def test_no_api_key_error(): + """Verify a clear error is raised when no API key source is available.""" + decorator = asqav_audit(agent_name="no-key-agent") + + with pytest.raises(ValueError, match="No Asqav API key found"): + decorator.execute() diff --git a/uv.lock b/uv.lock index 4ac2ed0629..a04937af82 100644 --- a/uv.lock +++ b/uv.lock @@ -545,16 +545,16 @@ wheels = [ [[package]] name = "flyteidl" -version = "1.16.4" +version = "1.16.7" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "googleapis-common-protos" }, { name = "protobuf" }, { name = "protoc-gen-openapiv2" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b1/b0/5ffd5521c6fb19e1f85ec92122668b6111502b766de402f0ddf177c7a5fe/flyteidl-1.16.4.tar.gz", hash = "sha256:5d2193a414fbf7cd5e27d486ff52c4b3c38e4d0479298fe82838d97c5d9c09fe", size = 130538, upload-time = "2026-02-11T21:12:43.744Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c5/1a/56b7f6aa0a9694b2853a0dbee9d51e2c63915dd722053840aa99059963bf/flyteidl-1.16.7.tar.gz", hash = "sha256:fd77b66a0c911086887d21fe4d2f157ccb65fdad057a63b82cf666006e6344be", size = 131156, upload-time = "2026-05-28T17:37:09.034Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/90/8b/2168a62542ca49ae83c9cc3c83b212d741645cdc063cd43f541eb2ce114b/flyteidl-1.16.4-py3-none-any.whl", hash = "sha256:95028df0265552f2e23ba5b771e889fe34581ba8f93bf01e4657d156daa132e3", size = 221745, upload-time = "2026-02-11T21:12:42.129Z" }, + { url = "https://files.pythonhosted.org/packages/1f/9a/488653e77b19ef505cdaa1ffb7d7600e84dbb89134067c59c3fc93c30818/flyteidl-1.16.7-py3-none-any.whl", hash = "sha256:821818f7ffc73f2931178b4dffc443eb2138475ac6e9535127f8ed46a4ef8222", size = 222439, upload-time = "2026-05-28T17:37:07.342Z" }, ] [[package]] @@ -625,9 +625,9 @@ requires-dist = [ { name = "diskcache", specifier = ">=5.2.1" }, { name = "docker", specifier = ">=4.0.0" }, { name = "docstring-parser", specifier = ">=0.9.0" }, - { name = "flyteidl", specifier = ">=1.16.1,<2.0.0a0" }, - { name = "fsspec", specifier = ">=2023.3.0" }, - { name = "gcsfs", specifier = ">=2023.3.0,!=2024.2.0,!=2025.5.0,!=2025.5.0.post1" }, + { name = "flyteidl", specifier = ">=1.16.7,<2.0.0a0" }, + { name = "fsspec", specifier = ">=2024.9.0" }, + { name = "gcsfs", specifier = ">=2024.9.0,!=2025.5.0,!=2025.5.0.post1" }, { name = "googleapis-common-protos", specifier = ">=1.57" }, { name = "grpcio" }, { name = "grpcio-health-checking", marker = "extra == 'agent'", specifier = "<=1.68.0" }, @@ -656,7 +656,7 @@ requires-dist = [ { name = "requests", specifier = ">=2.18.4" }, { name = "rich" }, { name = "rich-click" }, - { name = "s3fs", specifier = ">=2023.3.0,!=2024.3.1" }, + { name = "s3fs", specifier = ">=2024.9.0" }, { name = "setuptools", specifier = "<82" }, { name = "statsd", specifier = ">=3.0.0" }, { name = "typing-extensions" }, From 3609235785452e6ab68dedd32eed8019dafc2909 Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Wed, 24 Jun 2026 19:03:15 +0530 Subject: [PATCH 3/4] Revert "fix: add depth guard to sync_node_execution to prevent RecursionError" This reverts commit 51ca4171685df2f99a8c3a13223550fe332bc121. Signed-off-by: Abhishek Shinde --- flytekit/remote/remote.py | 48 +------- .../unit/remote/test_recursion_guard.py | 115 ------------------ 2 files changed, 6 insertions(+), 157 deletions(-) delete mode 100644 tests/flytekit/unit/remote/test_recursion_guard.py diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 0120b88f73..041c701ec4 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -2587,17 +2587,10 @@ def _sync_execution( self, execution: FlyteWorkflowExecution, sync_nodes: bool = False, - _depth: int = 0, - _max_depth: int = 50, ) -> FlyteWorkflowExecution: """ Sync a FlyteWorkflowExecution object with its corresponding remote state. """ - if _depth > _max_depth: - raise FlyteAssertion( - f"Nesting depth {_depth} exceeds _max_depth={_max_depth} for execution " - f"{execution.id}. Refusing to recurse further to avoid RecursionError." - ) # Update closure, and then data, because we don't want the execution to finish between when we get the data, # and then for the closure to have is_done to be true. execution._closure = self.client.get_execution(execution.id).closure @@ -2649,9 +2642,7 @@ def _sync_execution( if sync_nodes: node_execs = {} for n in underlying_node_executions: - node_execs[n.id.node_id] = self.sync_node_execution( # noqa - n, node_mapping, _depth=_depth + 1, _max_depth=_max_depth - ) + node_execs[n.id.node_id] = self.sync_node_execution(n, node_mapping) # noqa execution._node_executions = node_execs return self._assign_inputs_and_outputs(execution, execution_data, node_interface) @@ -2659,8 +2650,6 @@ def sync_node_execution( self, execution: FlyteNodeExecution, node_mapping: typing.Dict[str, FlyteNode], - _depth: int = 0, - _max_depth: int = 50, ) -> FlyteNodeExecution: """ Get data backing a node execution. These FlyteNodeExecution objects should've come from Admin with the model @@ -2680,11 +2669,6 @@ def sync_node_execution( The data model is complicated, so ascertaining which of these happened is a bit tricky. That logic is encapsulated in this function. """ - if _depth > _max_depth: - raise FlyteAssertion( - f"Nesting depth {_depth} exceeds _max_depth={_max_depth} for node " - f"{execution.id}. Refusing to recurse further to avoid RecursionError." - ) # For single task execution - the metadata spec node id is missing. In these cases, revert to regular node id node_id = execution.metadata.spec_node_id # This case supports single-task execution compiled workflows. @@ -2728,7 +2712,7 @@ def sync_node_execution( launched_exec = self.fetch_execution( project=launched_exec_id.project, domain=launched_exec_id.domain, name=launched_exec_id.name ) - self._sync_execution(launched_exec, sync_nodes=True, _depth=_depth + 1, _max_depth=_max_depth) + self.sync_execution(launched_exec, sync_nodes=True) if launched_exec.is_done: # The synced underlying execution should've had these populated. execution._inputs = launched_exec.inputs @@ -2759,12 +2743,7 @@ def sync_node_execution( dynamic_flyte_wf = FlyteWorkflow.promote_from_closure(compiled_wf, node_launch_plans) execution._underlying_node_executions = [ - self.sync_node_execution( - FlyteNodeExecution.promote_from_model(cne), - dynamic_flyte_wf._node_map, - _depth=_depth + 1, - _max_depth=_max_depth, - ) + self.sync_node_execution(FlyteNodeExecution.promote_from_model(cne), dynamic_flyte_wf._node_map) for cne in child_node_executions ] execution._task_executions = [ @@ -2778,12 +2757,7 @@ def sync_node_execution( sub_flyte_workflow = execution._node.flyte_entity sub_node_mapping = {n.id: n for n in sub_flyte_workflow.flyte_nodes} execution._underlying_node_executions = [ - self.sync_node_execution( - FlyteNodeExecution.promote_from_model(cne), - sub_node_mapping, - _depth=_depth + 1, - _max_depth=_max_depth, - ) + self.sync_node_execution(FlyteNodeExecution.promote_from_model(cne), sub_node_mapping) for cne in child_node_executions ] execution._interface = sub_flyte_workflow.interface @@ -2804,12 +2778,7 @@ def sync_node_execution( sub_node_mapping[else_node.id] = else_node execution._underlying_node_executions = [ - self.sync_node_execution( - FlyteNodeExecution.promote_from_model(cne), - sub_node_mapping, - _depth=_depth + 1, - _max_depth=_max_depth, - ) + self.sync_node_execution(FlyteNodeExecution.promote_from_model(cne), sub_node_mapping) for cne in child_node_executions ] else: @@ -2882,12 +2851,7 @@ def sync_node_execution( sub_node_mapping[else_node.id] = else_node execution._underlying_node_executions = [ - self.sync_node_execution( - FlyteNodeExecution.promote_from_model(cne), - sub_node_mapping, - _depth=_depth + 1, - _max_depth=_max_depth, - ) + self.sync_node_execution(FlyteNodeExecution.promote_from_model(cne), sub_node_mapping) for cne in child_node_executions ] diff --git a/tests/flytekit/unit/remote/test_recursion_guard.py b/tests/flytekit/unit/remote/test_recursion_guard.py deleted file mode 100644 index 2413d38e5b..0000000000 --- a/tests/flytekit/unit/remote/test_recursion_guard.py +++ /dev/null @@ -1,115 +0,0 @@ -from datetime import datetime, timedelta, timezone -from unittest.mock import MagicMock, patch - -import pytest - -from flytekit.configuration import Config -from flytekit.exceptions.user import FlyteAssertion -from flytekit.models.core.identifier import NodeExecutionIdentifier, WorkflowExecutionIdentifier -from flytekit.models.node_execution import ( - NodeExecutionClosure, - WorkflowNodeMetadata, -) -from flytekit.remote.executions import FlyteNodeExecution, FlyteWorkflowExecution -from flytekit.remote.interface import TypedInterface -from flytekit.remote.remote import FlyteRemote - - -def _mock_launched_exec(): - exec = MagicMock(spec=FlyteWorkflowExecution) - wf = MagicMock() - wf.interface = MagicMock(spec=TypedInterface) - exec._flyte_workflow = wf - exec.is_done = False - exec.inputs = {"a": 1} - exec.outputs = {"b": 2} - return exec - - -def _make_node_execution(node_id: str, execution_name: str, with_workflow_node_metadata: bool = False): - wf_exec_id = WorkflowExecutionIdentifier("p1", "d1", execution_name) - ne_id = NodeExecutionIdentifier(node_id, wf_exec_id) - meta = MagicMock() - meta.is_parent_node = False - meta.is_array = False - meta.spec_node_id = node_id - - wf_node_meta = None - if with_workflow_node_metadata: - wf_node_meta = WorkflowNodeMetadata( - execution_id=WorkflowExecutionIdentifier("p1", "d1", f"launched_{execution_name}") - ) - - closure = NodeExecutionClosure( - phase=0, - started_at=datetime.now(timezone.utc), - duration=timedelta(seconds=1), - workflow_node_metadata=wf_node_meta, - ) - return FlyteNodeExecution(id=ne_id, input_uri="s3://bucket/input", closure=closure, metadata=meta) - - -@pytest.fixture -def remote(): - with patch("flytekit.clients.friendly.SynchronousFlyteClient"): - flyte_remote = FlyteRemote( - config=Config.auto(), - default_project="p1", - default_domain="d1", - ) - flyte_remote._client_initialized = True - flyte_remote._client = MagicMock() - return flyte_remote - - -def test_max_depth_raises_flyte_assertion(remote): - ne = _make_node_execution("n1", "exec1", with_workflow_node_metadata=True) - - launched_exec = _mock_launched_exec() - remote.fetch_execution = MagicMock(return_value=launched_exec) - remote.client.get_node_execution_data = MagicMock( - return_value=MagicMock(dynamic_workflow=None) - ) - - # Don't mock _sync_execution — the guard fires inside it when _depth exceeds _max_depth. - # With _max_depth=0, the initial sync_node_execution enters at _depth=0 (passes guard). - # It reaches the launched LP path, which calls _sync_execution with _depth=1. - # _sync_execution then sees 1 > 0 and raises FlyteAssertion. - with pytest.raises(FlyteAssertion, match="Nesting depth"): - remote.sync_node_execution(ne, {"n1": MagicMock()}, _max_depth=0) - - -def test_reasonable_depth_does_not_raise(remote): - ne = _make_node_execution("n1", "exec1", with_workflow_node_metadata=True) - - launched_exec = _mock_launched_exec() - remote.fetch_execution = MagicMock(return_value=launched_exec) - remote.client.get_node_execution_data = MagicMock( - return_value=MagicMock(dynamic_workflow=None) - ) - remote._sync_execution = MagicMock() - - result = remote.sync_node_execution(ne, {"n1": MagicMock()}) - assert result is ne - - -def test_nested_under_default_limit(remote): - ne = _make_node_execution("n1", "exec1", with_workflow_node_metadata=True) - - launched_exec = _mock_launched_exec() - remote.fetch_execution = MagicMock(return_value=launched_exec) - remote.client.get_node_execution_data = MagicMock( - return_value=MagicMock(dynamic_workflow=None) - ) - remote._sync_execution = MagicMock() - - result = remote.sync_node_execution(ne, {"n1": MagicMock()}, _depth=1, _max_depth=50) - assert result is ne - - -def test_sync_execution_depth_guard(remote): - wf_exec = MagicMock(spec=FlyteWorkflowExecution) - wf_exec.id = WorkflowExecutionIdentifier("p1", "d1", "deep_exec") - - with pytest.raises(FlyteAssertion, match="Nesting depth"): - remote._sync_execution(wf_exec, sync_nodes=True, _depth=51, _max_depth=50) From dd0662314e50315eda984b94d9fe757a72795975 Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Thu, 25 Jun 2026 08:44:36 +0530 Subject: [PATCH 4/4] fix: cache agent, fix secret priority, add fail_closed, add CI Signed-off-by: Abhishek Shinde --- .github/workflows/pythonbuild.yml | 1 + plugins/flytekit-asqav/README.md | 1 + .../flytekitplugins/asqav/tracking.py | 163 ++++++++++-------- .../tests/test_asqav_tracking.py | 108 +++++++++--- 4 files changed, 174 insertions(+), 99 deletions(-) diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index a62d7d0b7e..219cc31257 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -291,6 +291,7 @@ jobs: plugin-names: # Please maintain an alphabetical order in the following list - flytekit-airflow + - flytekit-asqav - flytekit-async-fsspec - flytekit-aws-athena - flytekit-aws-batch diff --git a/plugins/flytekit-asqav/README.md b/plugins/flytekit-asqav/README.md index f11c9fb21a..955d572341 100644 --- a/plugins/flytekit-asqav/README.md +++ b/plugins/flytekit-asqav/README.md @@ -44,6 +44,7 @@ def main() -> float: | `receipt_type` | `str` | ❌ | `protectmcp:lifecycle` | Asqav receipt type | | `risk_class` | `str` | ❌ | `None` | Risk classification (`low`, `medium`, `high`, `unknown`) | | `compliance_mode` | `bool` | ❌ | `True` | Generate IETF-compliant compliance receipts | +| `fail_closed` | `bool` | ❌ | `False` | If `True`, Asqav errors propagate (fail-closed); if `False`, errors are logged and the task continues (fail-open) | ### Secret Resolution (priority order) diff --git a/plugins/flytekit-asqav/flytekitplugins/asqav/tracking.py b/plugins/flytekit-asqav/flytekitplugins/asqav/tracking.py index e95de635ac..3227262b6c 100644 --- a/plugins/flytekit-asqav/flytekitplugins/asqav/tracking.py +++ b/plugins/flytekit-asqav/flytekitplugins/asqav/tracking.py @@ -78,6 +78,7 @@ def __init__( receipt_type: str = ASQAV_DEFAULT_RECEIPT_TYPE, risk_class: Optional[str] = None, compliance_mode: bool = True, + fail_closed: bool = False, **init_kwargs, ): if not agent_name: @@ -89,6 +90,7 @@ def __init__( receipt_type=receipt_type, risk_class=risk_class, compliance_mode=compliance_mode, + fail_closed=fail_closed, **init_kwargs, ) self.agent_name = agent_name @@ -96,20 +98,25 @@ def __init__( self.receipt_type = receipt_type self.risk_class = risk_class self.compliance_mode = compliance_mode + self.fail_closed = fail_closed + self._agent = None def _resolve_api_key(self) -> str: - """Resolve the Asqav API key from the most specific source available.""" - # 1. Callable — highest priority for programmatic injection + """Resolve the Asqav API key from the most controlled source available.""" + # 1. Flyte Secret — highest priority for remote deployment + if isinstance(self.secret, Secret): + ctx = FlyteContextManager.current_context() + try: + return ctx.user_space_params.secrets.get(key=self.secret.key, group=self.secret.group) + except ValueError: + pass # Secret not available in this context (e.g. local execution), fall through + # 2. Callable — for programmatic injection if callable(self.secret): return self.secret() - # 2. Environment variable — convenient for local development + # 3. Environment variable — fallback for local development api_key = os.environ.get(ASQAV_ENV_API_KEY) if api_key: return api_key - # 3. Flyte Secret — for remote execution with secret_requests - if isinstance(self.secret, Secret): - ctx = FlyteContextManager.current_context() - return ctx.user_space_params.secrets.get(key=self.secret.key, group=self.secret.group) raise ValueError( f"No Asqav API key found. Provide a `Secret` or callable to " f"`asqav_audit(secret=...)`, or set the {ASQAV_ENV_API_KEY} env var." @@ -134,29 +141,36 @@ def execute(self, *args, **kwargs): api_key = self._resolve_api_key() - # Initialise the Asqav SDK and create an agent - asqav.init(api_key=api_key) - agent = asqav.Agent.create(name=self.agent_name) - receipts = [] start_time = datetime.datetime.utcnow() - # --- started receipt --- - started_context = { - "event_type": "flyte.task.started", - "execution_id": execution_id, - "agent_name": self.agent_name, - "timestamp": datetime.datetime.utcnow().isoformat() + "Z", - } - started_resp = agent.sign( - action_type="flyte.task.started", - context=started_context, - compliance_mode=self.compliance_mode, - receipt_type=self.receipt_type, - risk_class=self.risk_class, - ) - receipts.append(("started", started_resp)) - logger.info(f"Signed asqav started receipt: {started_resp.signature_id}") + # --- init Asqav agent (cached per process) and sign started receipt --- + agent = None + try: + if self._agent is None: + asqav.init(api_key=api_key) + self._agent = asqav.Agent.create(name=self.agent_name) + agent = self._agent + + started_context = { + "event_type": "flyte.task.started", + "execution_id": execution_id, + "agent_name": self.agent_name, + "timestamp": datetime.datetime.utcnow().isoformat() + "Z", + } + started_resp = agent.sign( + action_type="flyte.task.started", + context=started_context, + compliance_mode=self.compliance_mode, + receipt_type=self.receipt_type, + risk_class=self.risk_class, + ) + receipts.append(("started", started_resp)) + logger.info(f"Signed asqav started receipt: {started_resp.signature_id}") + except Exception as e: + if self.fail_closed: + raise + logger.error(f"Asqav audit trail init/started receipt failed: {e}") # --- task execution --- output = None @@ -167,51 +181,56 @@ def execute(self, *args, **kwargs): exception = e raise finally: - elapsed = (datetime.datetime.utcnow() - start_time).total_seconds() - - if exception is not None: - error_info = f"{type(exception).__name__}: {exception}" - failed_context = { - "event_type": "flyte.task.failed", - "execution_id": execution_id, - "agent_name": self.agent_name, - "duration_seconds": round(elapsed, 3), - "error": error_info, - "timestamp": datetime.datetime.utcnow().isoformat() + "Z", - } - try: - failed_resp = agent.sign( - action_type="flyte.task.failed", - context=failed_context, - compliance_mode=self.compliance_mode, - receipt_type=self.receipt_type, - risk_class=self.risk_class, - ) - receipts.append(("failed", failed_resp)) - logger.info(f"Signed asqav failed receipt: {failed_resp.signature_id}") - except Exception as sign_err: - logger.error(f"Failed to sign asqav receipt for error: {sign_err}") - else: - finished_context = { - "event_type": "flyte.task.finished", - "execution_id": execution_id, - "agent_name": self.agent_name, - "duration_seconds": round(elapsed, 3), - "output_type": type(output).__name__ if output is not None else "None", - "timestamp": datetime.datetime.utcnow().isoformat() + "Z", - } - try: - finished_resp = agent.sign( - action_type="flyte.task.finished", - context=finished_context, - compliance_mode=self.compliance_mode, - receipt_type=self.receipt_type, - risk_class=self.risk_class, - ) - receipts.append(("finished", finished_resp)) - logger.info(f"Signed asqav finished receipt: {finished_resp.signature_id}") - except Exception as sign_err: - logger.error(f"Failed to sign asqav receipt for success: {sign_err}") + if agent is not None: + elapsed = (datetime.datetime.utcnow() - start_time).total_seconds() + + if exception is not None: + error_info = f"{type(exception).__name__}: {exception}" + failed_context = { + "event_type": "flyte.task.failed", + "execution_id": execution_id, + "agent_name": self.agent_name, + "duration_seconds": round(elapsed, 3), + "error": error_info, + "timestamp": datetime.datetime.utcnow().isoformat() + "Z", + } + try: + failed_resp = agent.sign( + action_type="flyte.task.failed", + context=failed_context, + compliance_mode=self.compliance_mode, + receipt_type=self.receipt_type, + risk_class=self.risk_class, + ) + receipts.append(("failed", failed_resp)) + logger.info(f"Signed asqav failed receipt: {failed_resp.signature_id}") + except Exception as sign_err: + if self.fail_closed: + raise + logger.error(f"Failed to sign asqav receipt for error: {sign_err}") + else: + finished_context = { + "event_type": "flyte.task.finished", + "execution_id": execution_id, + "agent_name": self.agent_name, + "duration_seconds": round(elapsed, 3), + "output_type": type(output).__name__ if output is not None else "None", + "timestamp": datetime.datetime.utcnow().isoformat() + "Z", + } + try: + finished_resp = agent.sign( + action_type="flyte.task.finished", + context=finished_context, + compliance_mode=self.compliance_mode, + receipt_type=self.receipt_type, + risk_class=self.risk_class, + ) + receipts.append(("finished", finished_resp)) + logger.info(f"Signed asqav finished receipt: {finished_resp.signature_id}") + except Exception as sign_err: + if self.fail_closed: + raise + logger.error(f"Failed to sign asqav receipt for success: {sign_err}") # Render Flyte Deck with all receipts if receipts: diff --git a/plugins/flytekit-asqav/tests/test_asqav_tracking.py b/plugins/flytekit-asqav/tests/test_asqav_tracking.py index fb7f57d082..e913268bfb 100644 --- a/plugins/flytekit-asqav/tests/test_asqav_tracking.py +++ b/plugins/flytekit-asqav/tests/test_asqav_tracking.py @@ -18,7 +18,6 @@ def _make_signature_response(signature_id: str, **overrides) -> Mock: - """Create a mock SignatureResponse with the expected attributes.""" defaults = { "signature_id": signature_id, "verification_url": f"https://asqav.com/verify/{signature_id}", @@ -66,17 +65,16 @@ def test_get_extra_config_defaults(): assert ASQAV_RISK_CLASS_KEY not in config -@task -@asqav_audit(agent_name="model-trainer", secret=secret) -def train_model() -> str: - return "done" - - @patch("flytekitplugins.asqav.tracking.asqav") @patch("flytekitplugins.asqav.tracking.Deck") def test_local_execution(deck_mock, asqav_mock, monkeypatch): - """Verify the decorator signs started + finished receipts on success.""" monkeypatch.setenv("ASQAV_API_KEY", "local-test-key") + + @task + @asqav_audit(agent_name="model-trainer", secret=secret) + def train_model() -> str: + return "done" + agent_mock = MagicMock() agent_mock.sign.return_value = _make_signature_response("receipt-started") asqav_mock.Agent.create.return_value = agent_mock @@ -86,14 +84,13 @@ def test_local_execution(deck_mock, asqav_mock, monkeypatch): assert result == "done" asqav_mock.init.assert_called_once() asqav_mock.Agent.create.assert_called_once_with(name="model-trainer") - assert agent_mock.sign.call_count == 2 # started + finished + assert agent_mock.sign.call_count == 2 deck_mock.assert_called_once() @patch("flytekitplugins.asqav.tracking.asqav") @patch("flytekitplugins.asqav.tracking.Deck") def test_failure_signs_failed_receipt(deck_mock, asqav_mock, monkeypatch): - """Verify the decorator signs a failed receipt when the task raises.""" monkeypatch.setenv("ASQAV_API_KEY", "local-test-key") @task @@ -108,7 +105,7 @@ def failing_task(): with pytest.raises(RuntimeError, match="model training failed"): failing_task() - assert agent_mock.sign.call_count == 2 # started + failed + assert agent_mock.sign.call_count == 2 deck_mock.assert_called_once() @@ -116,13 +113,17 @@ def failing_task(): @patch("flytekitplugins.asqav.tracking.FlyteContextManager") @patch("flytekitplugins.asqav.tracking.Deck") def test_remote_execution_secret_resolution(deck_mock, manager_mock, asqav_mock): - """Verify secrets are resolved in remote execution.""" ctx_mock = Mock() ctx_mock.execution_state.is_local_execution.return_value = False ctx_mock.user_space_params.secrets.get.return_value = "asqav-key-123" ctx_mock.user_space_params.execution_id = "exec-001" manager_mock.current_context.return_value = ctx_mock + @task + @asqav_audit(agent_name="model-trainer", secret=secret) + def train_model() -> str: + return "done" + agent_mock = MagicMock() agent_mock.sign.return_value = _make_signature_response("receipt-remote") asqav_mock.Agent.create.return_value = agent_mock @@ -139,22 +140,20 @@ def get_my_key(): return "callable-key-456" -@task -@asqav_audit(agent_name="callable-agent", secret=get_my_key) -def callable_secret_task() -> int: - return 42 - - @patch("flytekitplugins.asqav.tracking.asqav") @patch("flytekitplugins.asqav.tracking.FlyteContextManager") @patch("flytekitplugins.asqav.tracking.Deck") def test_callable_secret(deck_mock, manager_mock, asqav_mock): - """Verify a callable secret is called to get the API key.""" ctx_mock = Mock() ctx_mock.execution_state.is_local_execution.return_value = False ctx_mock.user_space_params.execution_id = "exec-002" manager_mock.current_context.return_value = ctx_mock + @task + @asqav_audit(agent_name="callable-agent", secret=get_my_key) + def callable_secret_task() -> int: + return 42 + agent_mock = MagicMock() agent_mock.sign.return_value = _make_signature_response("receipt-callable") asqav_mock.Agent.create.return_value = agent_mock @@ -165,17 +164,10 @@ def test_callable_secret(deck_mock, manager_mock, asqav_mock): asqav_mock.init.assert_called_with(api_key="callable-key-456") -@task -@asqav_audit(agent_name="env-agent") -def env_var_task() -> str: - return "from-env" - - @patch("flytekitplugins.asqav.tracking.asqav") @patch("flytekitplugins.asqav.tracking.FlyteContextManager") @patch("flytekitplugins.asqav.tracking.Deck") def test_env_var_fallback(deck_mock, manager_mock, asqav_mock, monkeypatch): - """Verify the ASQAV_API_KEY env var is used when no Secret or callable is provided.""" monkeypatch.setenv("ASQAV_API_KEY", "env-key-789") ctx_mock = Mock() @@ -183,6 +175,11 @@ def test_env_var_fallback(deck_mock, manager_mock, asqav_mock, monkeypatch): ctx_mock.user_space_params.execution_id = "exec-003" manager_mock.current_context.return_value = ctx_mock + @task + @asqav_audit(agent_name="env-agent") + def env_var_task() -> str: + return "from-env" + agent_mock = MagicMock() agent_mock.sign.return_value = _make_signature_response("receipt-env") asqav_mock.Agent.create.return_value = agent_mock @@ -194,8 +191,65 @@ def test_env_var_fallback(deck_mock, manager_mock, asqav_mock, monkeypatch): def test_no_api_key_error(): - """Verify a clear error is raised when no API key source is available.""" decorator = asqav_audit(agent_name="no-key-agent") with pytest.raises(ValueError, match="No Asqav API key found"): decorator.execute() + + +@patch("flytekitplugins.asqav.tracking.asqav") +@patch("flytekitplugins.asqav.tracking.FlyteContextManager") +@patch("flytekitplugins.asqav.tracking.Deck") +def test_secret_beats_env_var(deck_mock, manager_mock, asqav_mock, monkeypatch): + monkeypatch.setenv("ASQAV_API_KEY", "env-key-should-not-be-used") + + ctx_mock = Mock() + ctx_mock.execution_state.is_local_execution.return_value = False + ctx_mock.user_space_params.secrets.get.return_value = "secret-key-789" + ctx_mock.user_space_params.execution_id = "exec-secret-beats-env" + manager_mock.current_context.return_value = ctx_mock + + @task + @asqav_audit(agent_name="model-trainer", secret=secret) + def train_model() -> str: + return "done" + + agent_mock = MagicMock() + agent_mock.sign.return_value = _make_signature_response("receipt-secret") + asqav_mock.Agent.create.return_value = agent_mock + + result = train_model() + + assert result == "done" + ctx_mock.user_space_params.secrets.get.assert_called_with(key="asqav-api-key", group="asqav") + asqav_mock.init.assert_called_with(api_key="secret-key-789") + + +@patch("flytekitplugins.asqav.tracking.asqav") +@patch("flytekitplugins.asqav.tracking.Deck") +def test_fail_closed_propagates_init_failure(deck_mock, asqav_mock, monkeypatch): + monkeypatch.setenv("ASQAV_API_KEY", "test-key") + asqav_mock.init.side_effect = ConnectionError("Asqav unreachable") + + @task + @asqav_audit(agent_name="fail-closed-agent", fail_closed=True) + def my_task(): + return 1 + + with pytest.raises(ConnectionError, match="Asqav unreachable"): + my_task() + + +@patch("flytekitplugins.asqav.tracking.asqav") +@patch("flytekitplugins.asqav.tracking.Deck") +def test_fail_open_continues_on_init_failure(deck_mock, asqav_mock, monkeypatch): + monkeypatch.setenv("ASQAV_API_KEY", "test-key") + asqav_mock.init.side_effect = ConnectionError("Asqav unreachable") + + @task + @asqav_audit(agent_name="fail-open-agent") + def my_task() -> int: + return 42 + + result = my_task() + assert result == 42