"""Run container step - executes container and streams logs."""
import codecs
import json
import logging
import os
import threading
from typing import Any
from uuid import uuid4
import httpx
import docker
from gyoza.client import gyoza_client
from gyoza.worker.execution_pipeline.types import ExecutionContext
from gyoza.worker.worker_context import worker_context
# Prefixes emitted by GyozaLogger: "[GYOZA] [<TYPE>] <content>"
_GYOZA_EVENT_PREFIX = "[GYOZA]"
_INFO_EVENT_PREFIX = f"{_GYOZA_EVENT_PREFIX} [INFO]"
_PROGRESS_EVENT_PREFIX = f"{_GYOZA_EVENT_PREFIX} [PROGRESS]"
_RESERVATION_EVENT_PREFIX = f"{_GYOZA_EVENT_PREFIX} [RESERVATION]"
_logger = logging.getLogger(__name__)
# How often a running container checks whether the server gave up on it.
_CANCEL_POLL_SECONDS = float(os.getenv("GYOZA_CANCEL_POLL", "5"))
[docs]
class RunContainer:
"""Run the container and stream logs.
Runs the container attached, streaming logs in real-time.
Logs compatible with the gyoza events are parsed and sent to master.
"""
def __init__(self) -> None:
self._client = docker.from_env()
def __call__(self, ctx: ExecutionContext) -> ExecutionContext:
"""Execute the step.
Parameters
----------
ctx : ExecutionContext
Execution context with image, env_vars, and volumes configured.
Returns
-------
ExecutionContext
Context with logs and exit_code populated.
"""
container = self._create_container(ctx)
done = threading.Event()
watcher = threading.Thread(
target=self._stop_when_abandoned, args=(container, ctx, done), daemon=True
)
watcher.start()
try:
ctx.logs, ctx.exit_code = self._run_and_stream(container, ctx)
finally:
done.set()
return ctx
def _stop_when_abandoned(
self,
container: docker.models.containers.Container,
ctx: ExecutionContext,
done: threading.Event,
) -> None:
"""Stop the container once the run is cancelled or moved to another attempt.
The log loop only notices between chunks, so a container that prints
nothing would otherwise run to the end after a cancel or a retry.
Parameters
----------
container : Container
The container executing this attempt.
ctx : ExecutionContext
Execution context with ``oprun_id`` and ``attempt``.
done : threading.Event
Set when the container has finished.
"""
while not done.wait(_CANCEL_POLL_SECONDS):
try:
run = gyoza_client.get_run(ctx.oprun_id)
except Exception: # noqa: BLE001
continue # server unreachable: keep running, check again later
superseded = (
ctx.attempt is not None and run.get("current_attempt") != ctx.attempt
)
if run.get("state") == "CANCELLED" or superseded:
print(f"Run {ctx.oprun_id} abandoned on server, stopping container")
try:
container.stop()
except Exception as error: # noqa: BLE001
print(f"Failed to stop container {container.name}: {error}")
return
def _create_container(
self, ctx: ExecutionContext
) -> docker.models.containers.Container:
"""Create the docker container.
Parameters
----------
ctx : ExecutionContext
Execution context with image, env_vars, and volumes.
Returns
-------
Container
Created container ready to start.
"""
create_kwargs = {
"image": ctx.image,
"environment": ctx.env_vars,
"volumes": ctx.volumes,
"detach": True,
"name": f"gyoza-{ctx.oprun_id}-{uuid4().hex[:6]}",
"labels": {
"gyoza.managed": "true",
"gyoza.worker_id": worker_context.worker_id,
"gyoza.oprun_id": ctx.oprun_id,
},
}
gpu_id = ctx.constraints.get("gpu_id")
if gpu_id is not None:
create_kwargs["device_requests"] = [
docker.types.DeviceRequest(
capabilities=[["gpu"]],
device_ids=[str(gpu_id)],
)
]
container = self._client.containers.create(**create_kwargs)
return container
def _run_and_stream(
self, container: docker.models.containers.Container, ctx: ExecutionContext
) -> tuple[str, int | None]:
"""Run the container and stream logs.
Parameters
----------
container : Container
Docker container to run.
ctx : ExecutionContext
Execution context; its run id and attempt tag the events sent.
Returns
-------
tuple[str, int | None]
Tuple of (logs, exit_code).
"""
logs_lines: list[str] = []
log_line_buffer = ""
decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
container.start()
for chunk in container.logs(stream=True, follow=True):
if worker_context.is_run_cancelled(ctx.oprun_id, ctx.attempt):
container.stop()
break
log_line_buffer += decoder.decode(chunk)
while "\n" in log_line_buffer:
line, log_line_buffer = log_line_buffer.split("\n", 1)
logs_lines.append(line)
if line.startswith(_GYOZA_EVENT_PREFIX):
self._send_log_event(ctx, line)
log_line_buffer += decoder.decode(b"", final=True)
if log_line_buffer:
logs_lines.append(log_line_buffer)
if log_line_buffer.startswith(_GYOZA_EVENT_PREFIX):
self._send_log_event(ctx, log_line_buffer)
result = container.wait()
exit_code = result.get("StatusCode")
self._remove(container)
return "\n".join(logs_lines), exit_code
def _remove(self, container: docker.models.containers.Container) -> None:
"""Discard the finished container without risking the run's result.
A busy daemon can time out here long after the op produced its
outputs; failing the run then reports success as a failure. The
gyoza.worker_id label lets the next boot collect whatever is left.
"""
try:
container.remove()
except Exception as error: # noqa: BLE001
print(f"Failed to remove container {container.name}: {error}")
def _send_log_event(self, ctx: ExecutionContext, line: str) -> None:
"""Parse a [GYOZA] log line and forward it as the correct event type.
Parameters
----------
ctx : ExecutionContext
Execution context of the run the line came from.
line : str
Raw log line emitted by GyozaLogger.
"""
print(f"Sending event: {line}")
if line.startswith(_RESERVATION_EVENT_PREFIX):
try:
reservation = json.loads(line[len(_RESERVATION_EVENT_PREFIX) :].strip())
if not isinstance(reservation, dict):
raise ValueError("Reservation must be an object")
except ValueError as error:
_logger.warning(
"Resource update for %s was not confirmed: %s", ctx.oprun_id, error
)
return
response = self._add_event(
ctx,
"RESOURCES_UPDATED",
"Resource reservation updated",
payload={"reservation": reservation},
)
if isinstance(response, dict):
worker_context.update_from_server(ctx.oprun_id, ctx.attempt, response)
elif line.startswith(_PROGRESS_EVENT_PREFIX):
raw = line[len(_PROGRESS_EVENT_PREFIX) :].strip().rstrip("%")
try:
value = float(raw)
except ValueError:
return
if not value.is_integer():
return
self._add_event(ctx, "PROGRESS", int(value))
elif line.startswith(_INFO_EVENT_PREFIX):
msg = line[len(_INFO_EVENT_PREFIX) :].strip()
self._add_event(ctx, "INFO", msg)
else:
self._add_event(ctx, "INFO", line)
def _add_event(
self,
ctx: ExecutionContext,
event_type: str,
msg: str | int,
payload: dict[str, Any] | None = None,
) -> dict[str, Any] | None:
"""Send an event to master, retrying 4 times; a lost event must not
fail the run.
Parameters
----------
ctx : ExecutionContext
Execution context of the run the event belongs to.
event_type : str
Gyoza event type (``"PROGRESS"`` or ``"INFO"``).
msg : str | int
Event payload.
"""
for attempt in range(5):
try:
return gyoza_client.add_event(
run_id=ctx.oprun_id,
event_type=event_type,
msg=msg,
attempt=ctx.attempt,
worker_id=ctx.worker_id,
payload=payload,
)
except httpx.HTTPError as error:
if (
isinstance(error, httpx.HTTPStatusError)
and error.response.status_code == 409
):
return None # GET reconciliation determines whether it moved on
if attempt < 4:
continue
if event_type == "RESOURCES_UPDATED":
_logger.warning(
"Resource update for %s was not confirmed: %s",
ctx.oprun_id,
error,
)
else:
print(f"Failed to send {event_type} event, dropping it: {error}")
return None
run_container = RunContainer()
