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

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 run_async(execute_root_tasks(task, 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}")