vatsrahul1001 commented on code in PR #69003: URL: https://github.com/apache/airflow/pull/69003#discussion_r3519817716
########## providers/anthropic/src/airflow/providers/anthropic/operators/anthropic.py: ########## @@ -0,0 +1,200 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import time +from collections.abc import Sequence +from datetime import timedelta +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.anthropic.exceptions import AnthropicBatchJobError, AnthropicBatchTimeout +from airflow.providers.anthropic.hooks.anthropic import AnthropicHook, evaluate_batch_counts +from airflow.providers.anthropic.triggers.anthropic import AnthropicBatchTrigger +from airflow.providers.common.compat.sdk import BaseOperator, conf + +if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context + + +class AnthropicBatchOperator(BaseOperator): + """ + Submit an Anthropic Message Batch and wait for it to complete. + + Message Batches process many ``messages.create`` requests asynchronously at 50% of + standard cost; most complete within an hour (24h SLA). This operator submits the + batch and, in deferrable mode, releases the worker slot while a trigger polls for + completion. + + The operator returns the **batch ID only** — never the results. Pull results with + :meth:`~airflow.providers.anthropic.hooks.anthropic.AnthropicHook.stream_batch_results` + and persist them to object storage; results can be very large and must not be pushed + to XCom. Results are retained for 29 days after the batch is created. + + .. note:: + A retry re-submits a brand-new batch. Prefer ``retries=0`` on this task (the + submitted ``batch_id`` is pushed to XCom under key ``batch_id`` immediately, so + a crashed run never loses track of an in-flight batch). + + .. seealso:: + For more information, take a look at the guide: + :ref:`howto/operator:AnthropicBatchOperator` + + :param requests: A list of ``{"custom_id": str, "params": {...}}`` dicts, where + ``params`` is a ``messages.create`` payload (``model``, ``max_tokens``, ``messages``, ...). + :param conn_id: The Anthropic connection ID to use. + :param deferrable: Run the operator in deferrable mode. + :param poll_interval: Seconds between status checks, in both the synchronous and + deferrable paths. + :param timeout: Seconds to wait for the batch to reach a terminal status. Defaults to + 24 hours (the Message Batches SLA). In deferrable mode this also bounds the + deferral; set ``execution_timeout`` only if you want a shorter hard cap (note a + shorter ``execution_timeout`` preempts the graceful cancel-on-timeout path). + :param wait_for_completion: Whether to wait for the batch to complete. If ``False``, + the operator returns the batch ID immediately after submission. + :param fail_on_partial_error: If ``True``, fail the task when any request errored or + expired. Defaults to ``False`` (succeed and log a warning so the successful + results are not discarded). + """ + + template_fields: Sequence[str] = ("requests",) + + def __init__( + self, + requests: list[dict[str, Any]], + conn_id: str = AnthropicHook.default_conn_name, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + poll_interval: float = 60, + timeout: float = 24 * 60 * 60, + wait_for_completion: bool = True, + fail_on_partial_error: bool = False, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.requests = requests + self.conn_id = conn_id + self.deferrable = deferrable + self.poll_interval = poll_interval + self.timeout = timeout + self.wait_for_completion = wait_for_completion + self.fail_on_partial_error = fail_on_partial_error + self.batch_id: str | None = None + + @cached_property + def hook(self) -> AnthropicHook: + """Return an instance of the AnthropicHook.""" + return AnthropicHook(conn_id=self.conn_id) + + def execute(self, context: Context) -> str | None: + if not self.requests: + raise ValueError("AnthropicBatchOperator requires at least one request; got an empty list.") + batch = self.hook.create_batch(self.requests) + self.batch_id = batch.id + # Push immediately so a crash between submit and completion never loses the batch. + context["ti"].xcom_push(key="batch_id", value=batch.id) + self.log.info("Submitted Anthropic Message Batch %s (%d requests)", batch.id, len(self.requests)) + + if not self.wait_for_completion: + return self.batch_id + + if self.deferrable: + self.defer( + # Backstop the deferral slightly beyond the trigger's own end_time so the + # trigger's clean "timeout" event (which cancels the batch) wins over a + # generic AirflowTaskTimeout. A user-set execution_timeout still applies + # as a shorter hard cap. + timeout=self.execution_timeout or timedelta(seconds=self.timeout + self.poll_interval + 60), + trigger=AnthropicBatchTrigger( + conn_id=self.conn_id, + batch_id=self.batch_id, + poll_interval=self.poll_interval, + end_time=time.time() + self.timeout, + ), + method_name="execute_complete", + ) + + self.log.info("Waiting for batch %s to complete", self.batch_id) + try: + batch = self.hook.wait_for_batch( + self.batch_id, wait_seconds=self.poll_interval, timeout=self.timeout + ) + except AnthropicBatchTimeout: + # Mirror the deferrable execute_complete: tear down the still-running batch + # before the task fails, so a sync timeout does not leave it billing. + self.log.warning("Batch %s timed out; requesting cancellation.", self.batch_id) + self._cancel_batch_quietly() + raise + counts = batch.request_counts + self._apply_policy(counts.canceled, counts.errored, counts.expired, counts.succeeded) + return self.batch_id + + def execute_complete(self, context: Context, event: Any = None) -> str: + """ + Resume after the trigger fires. + + The deferred task is a fresh instance, so the batch ID is read from the event, + not ``self.batch_id``. + """ + self.batch_id = event["batch_id"] + status = event["status"] + if status == "timeout": + self.log.warning("Batch %s timed out; requesting cancellation.", self.batch_id) + self._cancel_batch_quietly() + raise AnthropicBatchTimeout(event["message"]) + if status == "error": Review Comment: Non-blocking (resource/billing leak on the error path): the deferrable `timeout` branch tears down the remote resource (`_cancel_batch_quietly()` here, `_archive_session()` in the agent operator), but the `error` branch just raises: ```python if status == "error": raise AnthropicBatchJobError(event["message"]) # batch still running / billing ``` The trigger yields `status="error"` not only for a genuinely terminal batch, but also when polling gives up while the batch is *still in progress* — after `MAX_CONSECUTIVE_POLL_FAILURES` transient errors, or when a poll happens to raise right at the deadline (`triggers/anthropic.py`): ```python if consecutive_failures >= MAX_CONSECUTIVE_POLL_FAILURES or time.time() > self.end_time: yield TriggerEvent({"status": "error", ...}) ``` In that case the batch keeps running (and billing) with nothing cancelling it — the opposite of the documented cancel-on-timeout guarantee. Cancelling on the `error` path too (best-effort, as the timeout branch does) would close the gap: ```python if status == "error": self._cancel_batch_quietly() raise AnthropicBatchJobError(event["message"]) ``` Same shape in `operators/agent.py` (the `error` branch should `_archive_session()`), and the synchronous paths have a related gap — `execute()` only catches `AnthropicBatchTimeout`/`AnthropicAgentSessionTimeout`, so a non-timeout SDK error (5xx, auth expiry) propagates without cancelling/archiving. This is adjacent to the trigger-`on_kill` suggestion already raised, but a distinct path (normal error-resume rather than a killed deferred task). -- 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]
