Source code for gyoza.worker.execution_pipeline.steps.complete_oprun

import json
import os
import time

import httpx

from gyoza.client import gyoza_client
from gyoza.worker.execution_pipeline.types import ExecutionContext

_SEND_ATTEMPTS = 10
_SEND_RETRY_SECONDS = 3.0
# Short enough to show in a list or a header; the full log tail is in ``msg``.
_ERROR_MESSAGE_LEN = 500




def complete_oprun(ctx: ExecutionContext) -> ExecutionContext:
    """Send COMPLETED or FAILED event to the master API based on exit code.

    Loads outputs from the pickle file and sends them as payload on success.
    If the server returns 400 (run already terminal) or 409 (the run moved on
    to another attempt), the update is skipped and context is returned
    without raising, so the worker does not treat a reclaimed run as a
    local failure.

    Parameters
    ----------
    ctx : ExecutionContext
        Execution context with exit_code and env_vars set.

    Returns
    -------
    ExecutionContext
        Context with outputs populated on success.
    """
    try:
        if ctx.exit_code == 0:
            ctx.outputs = _load_outputs(ctx.env_vars.get("GYOZA_OUTPUT_PATH", ""))
            _send_event(
                run_id=ctx.oprun_id,
                attempt=ctx.attempt,
                worker_id=ctx.worker_id,
                event_type="COMPLETED",
                msg="Execution completed successfully",
                payload={"outputs": ctx.outputs},
            )
        else:
            lines = ctx.logs.splitlines() if ctx.logs else []
            tail = "\n".join(lines[-1000:])
            last = next((line for line in reversed(lines) if line.strip()), None)
            error = f"Exited with code {ctx.exit_code}"
            if last is not None:
                error = f"{error}: {last.strip()}"
            _send_event(
                run_id=ctx.oprun_id,
                attempt=ctx.attempt,
                worker_id=ctx.worker_id,
                event_type="FAILED",
                msg=f"Execution failed with exit code {ctx.exit_code}:\n{tail}",
                payload={"error_message": error[:_ERROR_MESSAGE_LEN]},
            )
    except httpx.HTTPStatusError as e:
        if e.response.status_code in (400, 409):
            print(
                f"Run {ctx.oprun_id} already finalized on server "
                f"(status {e.response.status_code}), skipping event"
            )
        else:
            raise

    return ctx



def _send_event(**event: object) -> None:
    """Deliver the terminal event across a server that is down for a moment.

    A restart of the server takes seconds; a run that finishes inside them
    would otherwise stay RUNNING until the reaper hands it to another worker.
    """
    for attempt in range(_SEND_ATTEMPTS):
        try:
            gyoza_client.add_event(**event)
            return
        except httpx.TransportError:
            if attempt == _SEND_ATTEMPTS - 1:
                raise
            time.sleep(_SEND_RETRY_SECONDS)


def _load_outputs(output_path: str) -> dict:
    """Load outputs from the pickle file.

    Parameters
    ----------
    output_path : str
        Path to the output pickle file.

    Returns
    -------
    dict
        Loaded outputs, or empty dict if file doesn't exist.
    """
    print(f"Loading outputs from {output_path}")
    if not output_path or not os.path.exists(output_path):
        return {}

    with open(output_path) as f:
        return json.load(f)