zrb 1.5.8__py3-none-any.whl → 1.5.10__py3-none-any.whl

zrb/task/base/execution.py

@@ -0,0 +1,274 @@
+import asyncio
+from typing import TYPE_CHECKING, Any
+
+from zrb.context.any_context import AnyContext
+from zrb.session.any_session import AnySession
+from zrb.util.attr import get_bool_attr
+from zrb.util.run import run_async
+from zrb.xcom.xcom import Xcom
+
+if TYPE_CHECKING:
+    from zrb.task.base_task import BaseTask
+
+
+async def execute_task_chain(task: "BaseTask", session: AnySession):
+    """
+    Executes the task and its downstream successors if conditions are met.
+    """
+    if session.is_terminated or not session.is_allowed_to_run(task):
+        return
+    result = await execute_task_action(task, session)
+    # Get next tasks
+    nexts = session.get_next_tasks(task)
+    if session.is_terminated or len(nexts) == 0:
+        return result
+    # Run next tasks asynchronously
+    next_coros = [run_async(next_task.exec_chain(session)) for next_task in nexts]
+    # Wait for the next tasks to complete. The result of the current task is returned.
+    await asyncio.gather(*next_coros)
+    return result
+
+
+async def execute_task_action(task: "BaseTask", session: AnySession):
+    """
+    Executes a single task's action, handling conditions and readiness checks.
+    """
+    ctx = task.get_ctx(session)
+    if not session.is_allowed_to_run(task):
+        # Task is not allowed to run, skip it for now.
+        # This will be triggered later if dependencies are met.
+        ctx.log_info("Not allowed to run")
+        return
+    if not check_execute_condition(task, session):
+        # Skip the task based on its execute_condition
+        ctx.log_info("Marked as skipped (condition false)")
+        session.get_task_status(task).mark_as_skipped()
+        return
+    # Wait for task to be ready (handles action execution, readiness checks)
+    return await run_async(execute_action_until_ready(task, session))
+
+
+def check_execute_condition(task: "BaseTask", session: AnySession) -> bool:
+    """
+    Evaluates the task's execute_condition attribute.
+    """
+    ctx = task.get_ctx(session)
+    execute_condition_attr = getattr(task, "_execute_condition", True)
+    return get_bool_attr(ctx, execute_condition_attr, True, auto_render=True)
+
+
+async def execute_action_until_ready(task: "BaseTask", session: AnySession):
+    """
+    Manages the execution of the task's action, coordinating with readiness checks.
+    """
+    ctx = task.get_ctx(session)
+    readiness_checks = task.readiness_checks
+    readiness_check_delay = getattr(task, "_readiness_check_delay", 0.5)
+    monitor_readiness = getattr(task, "_monitor_readiness", False)
+
+    if not readiness_checks:  # Simplified check for empty list
+        ctx.log_info("No readiness checks")
+        # Task has no readiness check, execute action directly
+        result = await run_async(execute_action_with_retry(task, session))
+        # Mark ready only if the action completed successfully (not failed/cancelled)
+        if session.get_task_status(task).is_completed:
+            ctx.log_info("Marked as ready")
+            session.get_task_status(task).mark_as_ready()
+        return result
+
+    # Start the task action and readiness checks concurrently
+    ctx.log_info("Starting action and readiness checks")
+    action_coro = asyncio.create_task(
+        run_async(execute_action_with_retry(task, session))
+    )
+
+    await asyncio.sleep(readiness_check_delay)
+
+    readiness_check_coros = [
+        run_async(check.exec_chain(session)) for check in readiness_checks
+    ]
+
+    # Wait primarily for readiness checks to complete
+    ctx.log_info("Waiting for readiness checks")
+    readiness_passed = False
+    try:
+        # Gather results, but primarily interested in completion/errors
+        await asyncio.gather(*readiness_check_coros)
+        # Check if all readiness tasks actually completed successfully
+        all_readiness_completed = all(
+            session.get_task_status(check).is_completed for check in readiness_checks
+        )
+        if all_readiness_completed:
+            ctx.log_info("Readiness checks completed successfully")
+            readiness_passed = True
+            # Mark task as ready only if checks passed and action didn't fail during checks
+            if not session.get_task_status(task).is_failed:
+                ctx.log_info("Marked as ready")
+                session.get_task_status(task).mark_as_ready()
+        else:
+            ctx.log_warning(
+                "One or more readiness checks did not complete successfully."
+            )
+
+    except Exception as e:
+        ctx.log_error(f"Readiness check failed with exception: {e}")
+        # If readiness checks fail with an exception, the task is not ready.
+        # The action_coro might still be running or have failed.
+        # execute_action_with_retry handles marking the main task status.
+
+    # Defer the main action coroutine; it will be awaited later if needed
+    session.defer_action(task, action_coro)
+
+    # Start monitoring only if readiness passed and monitoring is enabled
+    if readiness_passed and monitor_readiness:
+        # Import dynamically to avoid circular dependency if monitoring imports execution
+        from zrb.task.base.monitoring import monitor_task_readiness
+
+        monitor_coro = asyncio.create_task(
+            run_async(monitor_task_readiness(task, session, action_coro))
+        )
+        session.defer_monitoring(task, monitor_coro)
+
+    # The result here is primarily about readiness check completion.
+    # The actual task result is handled by the deferred action_coro.
+    return None
+
+
+async def execute_action_with_retry(task: "BaseTask", session: AnySession) -> Any:
+    """
+    Executes the task's core action (`_exec_action`) with retry logic,
+    handling success (triggering successors) and failure (triggering fallbacks).
+    """
+    ctx = task.get_ctx(session)
+    retries = getattr(task, "_retries", 2)
+    retry_period = getattr(task, "_retry_period", 0)
+    max_attempt = retries + 1
+    ctx.set_max_attempt(max_attempt)
+
+    for attempt in range(max_attempt):
+        ctx.set_attempt(attempt + 1)
+        if attempt > 0:
+            ctx.log_info(f"Retrying in {retry_period}s...")
+            await asyncio.sleep(retry_period)
+
+        try:
+            ctx.log_info("Marked as started")
+            session.get_task_status(task).mark_as_started()
+
+            # Execute the underlying action (which might be overridden in subclasses)
+            # We call the task's _exec_action method directly here.
+            result = await run_async(task._exec_action(ctx))
+
+            ctx.log_info("Marked as completed")
+            session.get_task_status(task).mark_as_completed()
+
+            # Store result in XCom
+            task_xcom: Xcom = ctx.xcom.get(task.name)
+            task_xcom.push(result)
+
+            # Skip fallbacks and execute successors on success
+            skip_fallbacks(task, session)
+            await run_async(execute_successors(task, session))
+            return result
+
+        except (asyncio.CancelledError, KeyboardInterrupt):
+            ctx.log_warning("Task cancelled or interrupted")
+            session.get_task_status(task).mark_as_failed()  # Mark as failed on cancel
+            # Do not trigger fallbacks/successors on cancellation
+            return  # Or re-raise? Depends on desired cancellation behavior
+
+        except BaseException as e:
+            ctx.log_error(f"Attempt {attempt + 1}/{max_attempt} failed: {e}")
+            session.get_task_status(
+                task
+            ).mark_as_failed()  # Mark failed for this attempt
+
+            if attempt < max_attempt - 1:
+                # More retries available
+                continue
+            else:
+                # Final attempt failed
+                ctx.log_error("Marked as permanently failed")
+                session.get_task_status(task).mark_as_permanently_failed()
+                # Skip successors and execute fallbacks on permanent failure
+                skip_successors(task, session)
+                await run_async(execute_fallbacks(task, session))
+                raise e  # Re-raise the exception after handling fallbacks
+
+
+async def run_default_action(task: "BaseTask", ctx: AnyContext) -> Any:
+    """
+    Executes the specific action defined by the '_action' attribute for BaseTask.
+    This is the default implementation called by BaseTask._exec_action.
+    Subclasses like LLMTask override _exec_action with their own logic.
+    """
+    action = getattr(task, "_action", None)
+    if action is None:
+        ctx.log_debug("No action defined for this task.")
+        return None
+    if isinstance(action, str):
+        # Render f-string action
+        rendered_action = ctx.render(action)
+        ctx.log_debug(f"Rendered action string: {rendered_action}")
+        # Assuming string actions are meant to be returned as is.
+        # If they need execution (e.g., shell commands), that logic would go here.
+        return rendered_action
+    elif callable(action):
+        # Execute callable action
+        ctx.log_debug(f"Executing callable action: {action.__name__}")
+        return await run_async(action(ctx))
+    else:
+        ctx.log_warning(f"Unsupported action type: {type(action)}")
+        return None
+
+
+async def execute_successors(task: "BaseTask", session: "AnySession"):
+    """Executes all successor tasks."""
+    ctx = task.get_ctx(session)
+    successors = task.successors
+    if successors:
+        ctx.log_info(f"Executing {len(successors)} successor(s)")
+        successor_coros = [
+            run_async(successor.exec_chain(session)) for successor in successors
+        ]
+        await asyncio.gather(*successor_coros)
+    else:
+        ctx.log_debug("No successors to execute.")
+
+
+def skip_successors(task: "BaseTask", session: AnySession):
+    """Marks all successor tasks as skipped."""
+    ctx = task.get_ctx(session)
+    successors = task.successors
+    if successors:
+        ctx.log_info(f"Skipping {len(successors)} successor(s)")
+        for successor in successors:
+            # Check if already skipped to avoid redundant logging/state changes
+            if not session.get_task_status(successor).is_skipped:
+                session.get_task_status(successor).mark_as_skipped()
+
+
+async def execute_fallbacks(task: "BaseTask", session: AnySession):
+    """Executes all fallback tasks."""
+    ctx = task.get_ctx(session)
+    fallbacks = task.fallbacks
+    if fallbacks:
+        ctx.log_info(f"Executing {len(fallbacks)} fallback(s)")
+        fallback_coros = [
+            run_async(fallback.exec_chain(session)) for fallback in fallbacks
+        ]
+        await asyncio.gather(*fallback_coros)
+    else:
+        ctx.log_debug("No fallbacks to execute.")
+
+
+def skip_fallbacks(task: "BaseTask", session: AnySession):
+    """Marks all fallback tasks as skipped."""
+    ctx = task.get_ctx(session)
+    fallbacks = task.fallbacks
+    if fallbacks:
+        ctx.log_info(f"Skipping {len(fallbacks)} fallback(s)")
+        for fallback in fallbacks:
+            # Check if already skipped
+            if not session.get_task_status(fallback).is_skipped:
+                session.get_task_status(fallback).mark_as_skipped()

zrb/task/base/lifecycle.py

@@ -0,0 +1,182 @@
+import asyncio
+from typing import Any
+
+from zrb.context.shared_context import SharedContext
+from zrb.session.any_session import AnySession
+from zrb.session.session import Session
+from zrb.task.any_task import AnyTask
+from zrb.task.base.context import fill_shared_context_envs, fill_shared_context_inputs
+from zrb.util.run import run_async
+
+
+async def run_and_cleanup(
+    task: AnyTask,
+    session: AnySession | None = None,
+    str_kwargs: dict[str, str] = {},
+) -> Any:
+    """
+    Wrapper for async_run that ensures session termination and cleanup of
+    other concurrent asyncio tasks. This is the main entry point for `task.run()`.
+    """
+    # Ensure a session exists
+    if session is None:
+        session = Session(shared_ctx=SharedContext())
+
+    # Create the main task execution coroutine
+    main_task_coro = asyncio.create_task(run_task_async(task, session, str_kwargs))
+
+    try:
+        result = await main_task_coro
+        return result
+    except (asyncio.CancelledError, KeyboardInterrupt) as e:
+        ctx = task.get_ctx(session)  # Get context for logging
+        ctx.log_warning(f"Run cancelled/interrupted: {e}")
+        raise  # Re-raise to propagate
+    finally:
+        # Ensure session termination if it exists and wasn't terminated by the run
+        if session and not session.is_terminated:
+            ctx = task.get_ctx(session)  # Get context for logging
+            ctx.log_info("Terminating session after run completion/error.")
+            session.terminate()
+
+        # Clean up other potentially running asyncio tasks (excluding the main one)
+        # Be cautious with blanket cancellation if other background tasks are expected
+        pending = [
+            t for t in asyncio.all_tasks() if t is not main_task_coro and not t.done()
+        ]
+        if pending:
+            ctx = task.get_ctx(session)  # Get context for logging
+            ctx.log_debug(f"Cleaning up {len(pending)} pending asyncio tasks...")
+            for t in pending:
+                t.cancel()
+            try:
+                # Give cancelled tasks a moment to process cancellation
+                await asyncio.wait(pending, timeout=1.0)
+            except asyncio.CancelledError:
+                # Expected if tasks handle cancellation promptly
+                pass
+            except Exception as cleanup_exc:
+                # Log errors during cleanup if necessary
+                ctx.log_warning(f"Error during task cleanup: {cleanup_exc}")
+
+
+async def run_task_async(
+    task: AnyTask,
+    session: AnySession | None = None,
+    str_kwargs: dict[str, str] = {},
+) -> Any:
+    """
+    Asynchronous entry point for running a task (`task.async_run()`).
+    Sets up the session and initiates the root task execution chain.
+    """
+    if session is None:
+        session = Session(shared_ctx=SharedContext())
+
+    # Populate shared context with inputs and environment variables
+    fill_shared_context_inputs(task, session.shared_ctx, str_kwargs)
+    fill_shared_context_envs(session.shared_ctx)  # Inject OS env vars
+
+    # Start the execution chain from the root tasks
+    result = await task.exec_root_tasks(session)
+    return result
+
+
+async def execute_root_tasks(task: AnyTask, session: AnySession):
+    """
+    Identifies and executes the root tasks required for the main task,
+    manages session state logging, and handles overall execution flow.
+    """
+    session.set_main_task(task)
+    session.state_logger.write(session.as_state_log())  # Initial state log
+    ctx = task.get_ctx(session)  # Get context early for logging
+
+    log_state_task = None
+    try:
+        # Start background state logging
+        log_state_task = asyncio.create_task(log_session_state(task, session))
+
+        # Identify root tasks allowed to run
+        root_tasks = [
+            t for t in session.get_root_tasks(task) if session.is_allowed_to_run(t)
+        ]
+
+        if not root_tasks:
+            ctx.log_info("No root tasks to execute for this task.")
+            # If the main task itself should run even with no explicit roots?
+            # Current logic seems to imply if no roots, nothing runs.
+            session.terminate()  # Terminate if nothing to run
+            return None
+
+        ctx.log_info(f"Executing {len(root_tasks)} root task(s)")
+        root_task_coros = [
+            # Assuming exec_chain exists on AnyTask (it's abstract)
+            run_async(root_task.exec_chain(session))
+            for root_task in root_tasks
+        ]
+
+        # Wait for all root chains to complete
+        await asyncio.gather(*root_task_coros)
+
+        # Wait for any deferred actions (like long-running task bodies)
+        ctx.log_info("Waiting for deferred actions...")
+        await session.wait_deferred()
+        ctx.log_info("Deferred actions complete.")
+
+        # Final termination and logging
+        session.terminate()
+        if log_state_task and not log_state_task.done():
+            await log_state_task  # Ensure final state is logged
+        ctx.log_info("Session finished.")
+        return session.final_result
+
+    except IndexError:
+        # This might occur if get_root_tasks fails unexpectedly
+        ctx.log_error(
+            "IndexError during root task execution, potentially session issue."
+        )
+        session.terminate()  # Ensure termination on error
+        return None
+    except (asyncio.CancelledError, KeyboardInterrupt):
+        ctx.log_warning("Session execution cancelled or interrupted.")
+        # Session termination happens in finally block
+        return None  # Indicate abnormal termination
+    finally:
+        # Ensure termination and final state logging regardless of outcome
+        if not session.is_terminated:
+            session.terminate()
+        # Ensure the state logger task is awaited/cancelled properly
+        if log_state_task:
+            if not log_state_task.done():
+                log_state_task.cancel()
+                try:
+                    await log_state_task
+                except asyncio.CancelledError:
+                    pass  # Expected cancellation
+            # Log final state after ensuring logger task is finished
+            session.state_logger.write(session.as_state_log())
+        else:
+            # Log final state even if logger task didn't start
+            session.state_logger.write(session.as_state_log())
+
+        ctx.log_debug(f"Final session state: {session}")  # Log final session details
+
+
+async def log_session_state(task: AnyTask, session: AnySession):
+    """
+    Periodically logs the session state until the session is terminated.
+    """
+    try:
+        while not session.is_terminated:
+            session.state_logger.write(session.as_state_log())
+            await asyncio.sleep(0.1)  # Log interval
+        # Log one final time after termination signal
+        session.state_logger.write(session.as_state_log())
+    except asyncio.CancelledError:
+        # Log final state on cancellation too
+        session.state_logger.write(session.as_state_log())
+        ctx = task.get_ctx(session)
+        ctx.log_debug("Session state logger cancelled.")
+    except Exception as e:
+        # Log any unexpected errors in the logger itself
+        ctx = task.get_ctx(session)
+        ctx.log_error(f"Error in session state logger: {e}")

zrb/task/base/monitoring.py

@@ -0,0 +1,134 @@
+import asyncio
+
+from zrb.session.any_session import AnySession
+from zrb.task.base.execution import execute_action_with_retry
+from zrb.task.base_task import BaseTask
+from zrb.util.run import run_async
+from zrb.xcom.xcom import Xcom
+
+
+async def monitor_task_readiness(
+    task: BaseTask, session: AnySession, action_coro: asyncio.Task
+):
+    """
+    Monitors the readiness of a task after its initial execution.
+    If readiness checks fail beyond a threshold, it cancels the original action,
+    resets the task status, and re-executes the action.
+    """
+    ctx = task.get_ctx(session)
+    readiness_checks = task.readiness_checks
+    readiness_check_period = getattr(task, "_readiness_check_period", 5.0)
+    readiness_failure_threshold = getattr(task, "_readiness_failure_threshold", 1)
+    readiness_timeout = getattr(task, "_readiness_timeout", 60)
+
+    if not readiness_checks:
+        ctx.log_debug("No readiness checks defined, monitoring is not applicable.")
+        return
+
+    failure_count = 0
+    ctx.log_info("Starting readiness monitoring...")
+
+    while not session.is_terminated:
+        await asyncio.sleep(readiness_check_period)
+
+        if session.is_terminated:
+            break  # Exit loop if session terminated during sleep
+
+        if failure_count < readiness_failure_threshold:
+            ctx.log_info("Performing periodic readiness check...")
+            # Reset status and XCom for readiness check tasks before re-running
+            for check in readiness_checks:
+                session.get_task_status(check).reset_history()
+                session.get_task_status(check).reset()
+                # Clear previous XCom data for the check task if needed
+                check_xcom: Xcom = ctx.xcom.get(check.name)
+                check_xcom.clear()
+
+            readiness_check_coros = [
+                run_async(check.exec_chain(session)) for check in readiness_checks
+            ]
+
+            try:
+                # Wait for checks with a timeout
+                await asyncio.wait_for(
+                    asyncio.gather(*readiness_check_coros),
+                    timeout=readiness_timeout,
+                )
+                # Check if all checks actually completed successfully
+                all_checks_completed = all(
+                    session.get_task_status(check).is_completed
+                    for check in readiness_checks
+                )
+                if all_checks_completed:
+                    ctx.log_info("Readiness check OK.")
+                    failure_count = 0  # Reset failure count on success
+                    continue  # Continue monitoring
+                else:
+                    ctx.log_warning(
+                        "Periodic readiness check failed (tasks did not complete)."
+                    )
+                    failure_count += 1
+
+            except asyncio.TimeoutError:
+                failure_count += 1
+                ctx.log_warning(
+                    f"Readiness check timed out ({readiness_timeout}s). "
+                    f"Failure count: {failure_count}/{readiness_failure_threshold}"
+                )
+                # Ensure check tasks are marked as failed on timeout
+                for check in readiness_checks:
+                    if not session.get_task_status(check).is_finished:
+                        session.get_task_status(check).mark_as_failed()
+
+            except (asyncio.CancelledError, KeyboardInterrupt):
+                ctx.log_info("Monitoring cancelled or interrupted.")
+                break  # Exit monitoring loop
+
+            except Exception as e:
+                failure_count += 1
+                ctx.log_error(
+                    f"Readiness check failed with exception: {e}. "
+                    f"Failure count: {failure_count}"
+                )
+                # Mark checks as failed
+                for check in readiness_checks:
+                    if not session.get_task_status(check).is_finished:
+                        session.get_task_status(check).mark_as_failed()
+
+        # If failure threshold is reached
+        if failure_count >= readiness_failure_threshold:
+            ctx.log_warning(
+                f"Readiness failure threshold ({readiness_failure_threshold}) reached."
+            )
+
+            # Cancel the original running action if it's still running
+            if action_coro and not action_coro.done():
+                ctx.log_info("Cancelling original task action...")
+                action_coro.cancel()
+                try:
+                    await action_coro  # Allow cancellation to process
+                except asyncio.CancelledError:
+                    ctx.log_info("Original task action cancelled.")
+                except Exception as e:
+                    ctx.log_warning(f"Error during original action cancellation: {e}")
+
+            # Reset the main task status
+            ctx.log_info("Resetting task status.")
+            session.get_task_status(task).reset()
+
+            # Re-execute the action (with retries)
+            ctx.log_info("Re-executing task action...")
+            # Import dynamically to avoid circular dependency
+            new_action_coro = asyncio.create_task(
+                run_async(execute_action_with_retry(task, session))
+            )
+            # Defer the new action coroutine
+            session.defer_action(task, new_action_coro)
+            # Update the reference for the next monitoring cycle
+            action_coro = new_action_coro
+
+            # Reset failure count after attempting restart
+            failure_count = 0
+            ctx.log_info("Continuing monitoring...")
+
+    ctx.log_info("Stopping readiness monitoring.")

zrb/task/base/operators.py

@@ -0,0 +1,41 @@
+# No specific imports needed from typing for these changes
+from zrb.task.any_task import AnyTask
+
+
+def handle_rshift(
+    left_task: AnyTask, right_operand: AnyTask | list[AnyTask]
+) -> AnyTask | list[AnyTask]:
+    """
+    Implements the >> operator logic: left_task becomes an upstream for right_operand.
+    Modifies the right_operand(s) by calling append_upstream.
+    Returns the right_operand.
+    """
+    try:
+        if isinstance(right_operand, list):
+            for task in right_operand:
+                # Assuming append_upstream exists and handles duplicates
+                task.append_upstream(left_task)
+        else:
+            # Assuming right_operand is a single AnyTask
+            right_operand.append_upstream(left_task)
+        return right_operand
+    except Exception as e:
+        # Catch potential errors during append_upstream or type issues
+        raise ValueError(f"Invalid operation {left_task} >> {right_operand}: {e}")
+
+
+def handle_lshift(
+    left_task: AnyTask, right_operand: AnyTask | list[AnyTask]
+) -> AnyTask:
+    """
+    Implements the << operator logic: right_operand becomes an upstream for left_task.
+    Modifies the left_task by calling append_upstream.
+    Returns the left_task.
+    """
+    try:
+        # Assuming append_upstream exists and handles single or list input
+        left_task.append_upstream(right_operand)
+        return left_task
+    except Exception as e:
+        # Catch potential errors during append_upstream or type issues
+        raise ValueError(f"Invalid operation {left_task} << {right_operand}: {e}")