Source code for gyoza.worker.runner

"""Worker main loop.

Entry point called by ``gyoza worker start``. One blocking cycle: report
to the server, ask it for work and let it hold the request until there is
some, dispatch each op run to a thread pool.
"""

import os
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any

import docker
from gyoza.cache import sweep
from gyoza.cache.store import CacheStore
from gyoza.client import gyoza_client
from gyoza.worker import shared_runtime
from gyoza.worker.execution_pipeline import ExecutionContext, execution_pipeline
from gyoza.worker.worker_context import worker_context

_CLAIM_WAIT = float(os.getenv("GYOZA_CLAIM_WAIT", "25"))
# A server that answers at once, whether it cannot hold the request or is
# failing, must not turn the loop into a hot one.
_MIN_CYCLE_SECONDS = 1.0




def start() -> None:
    """Start the worker loop.

    Blocks until interrupted with Ctrl-C or a signal. Each cycle:

    1. Sends a heartbeat to the server (registers/updates the worker).
    2. Claims op runs, waiting on the server until some are allocated.
    3. Dispatches each claimed run to a background thread.
    4. Refreshes the server state of the running ops and sweeps the cache.
    """
    print(f"Starting worker: {worker_context.worker_id}")
    print(f"Resources: {worker_context.resources.to_dict()}")

    _remove_orphan_containers()

    executor = ThreadPoolExecutor()
    cache_sweeper = sweep.Scheduler(CacheStore(Path(shared_runtime.cache_dir())))

    try:
        while True:
            started = time.monotonic()
            _cycle(executor, cache_sweeper)
            time.sleep(max(0.0, _MIN_CYCLE_SECONDS - (time.monotonic() - started)))
    except KeyboardInterrupt:
        print("\nShutting down worker...")
    finally:
        executor.shutdown(wait=True)
        print("Worker stopped")



def _cycle(executor: ThreadPoolExecutor, cache_sweeper: sweep.Scheduler) -> None:
    try:
        gyoza_client.heartbeat(**worker_context.to_heartbeat_payload())
        print(f"Heartbeat sent. Running ops: {len(worker_context.running_ops)}")

        ops = gyoza_client.claim_ops(worker_context.worker_id, wait=_CLAIM_WAIT)
        for op in ops:
            print(f"Dispatching task: {op['id']}")
            executor.submit(_run_task, op)

        _refresh_running_ops_states()
    except Exception as e:
        print(f"Error in main loop: {e}")

    _sweep_cache(cache_sweeper)


def _remove_orphan_containers() -> None:
    """Remove op containers this worker left behind on a previous life.

    A worker that dies mid-run leaves its containers running: they keep a GPU
    busy while the run is reclaimed and executed elsewhere. Nothing this worker
    labelled can be in use at start-up, so every match is an orphan.
    """
    try:
        client = docker.from_env()
        orphans = client.containers.list(
            all=True,
            filters={"label": f"gyoza.worker_id={worker_context.worker_id}"},
        )
        for container in orphans:
            container.remove(force=True)
        if orphans:
            print(f"Removed {len(orphans)} orphan container(s)")
    except Exception as e:
        print(f"Could not remove orphan containers: {e}")


def _sweep_cache(scheduler: sweep.Scheduler) -> None:
    try:
        report = scheduler.tick(busy=bool(worker_context.running_ops))
    except Exception as e:
        print(f"Cache sweep failed: {e}")
        return
    if report and report.removed:
        print(
            f"Cache sweep removed {len(report.removed)} file(s), "
            f"freed {report.freed_bytes / sweep.GB:.1f} GB, "
            f"kept {report.kept_bytes / sweep.GB:.1f} GB"
        )


def _run_task(op: dict[str, Any]) -> None:
    """Execute a single op run in a worker thread.

    Parameters
    ----------
    op : dict[str, Any]
        Op run data from ``claim_ops`` (``id``, ``image``, ``inputs``,
        ``constraints``).
    """
    oprun_id = op["id"]
    image = op["image"]
    inputs = op.get("inputs", {})
    outputs = op.get("outputs", {})
    constraints = op.get("constraints", {})

    attempt = op.get("attempt")
    worker_context.add_running_op(
        oprun_id,
        image,
        inputs,
        constraints,
        attempt=attempt,
        reservation=op.get("reservation"),
    )

    try:
        ctx = ExecutionContext(
            oprun_id=oprun_id,
            attempt=attempt,
            worker_id=worker_context.worker_id,
            image=image,
            constraints=constraints,
            inputs=inputs,
            outputs=outputs,
        )
        execution_pipeline.run(ctx)
        print(f"Task {oprun_id} finished with exit code {ctx.exit_code}")
    except Exception as e:
        print(f"Task {oprun_id} failed: {e}")
    finally:
        worker_context.remove_running_op(oprun_id, attempt)


def _refresh_running_ops_states() -> None:
    """Fetch current state from server for each running op and update context."""
    for op in list(worker_context.running_ops):
        try:
            run_data = gyoza_client.get_run(op.id)
            worker_context.update_from_server(op.id, op.attempt, run_data)
        except Exception:
            pass