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

"""Start OpRun step - sends STARTED event to master."""

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




def start_oprun(ctx: ExecutionContext) -> ExecutionContext:
    """Send STARTED event to the master API.

    Parameters
    ----------
    ctx : ExecutionContext
        Execution context with oprun_id.

    Returns
    -------
    ExecutionContext
        Unchanged context.
    """
    gyoza_client.add_event(
        run_id=ctx.oprun_id,
        event_type="STARTED",
        msg="Execution started",
        payload={
            "worker_id": worker_context.worker_id,
            "gpu_id": ctx.constraints.get("gpu_id"),
        },
    )

    return ctx