adam/executor.py

from __future__ import annotations

import importlib
import inspect
import logging
import queue
import subprocess
import sys
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable

from adam.registry import ToolRegistry, ToolSpec


class ToolExecutionError(RuntimeError):
    pass


class ToolCancelled(ToolExecutionError):
    pass


ProgressCallback = Callable[[int, str], None]
LogCallback = Callable[[str], None]
PreviewCallback = Callable[[dict[str, Any]], None]


@dataclass(slots=True)
class ToolContext:
    root: Path
    job_id: str
    tool: ToolSpec
    cancel_event: threading.Event
    run_event: threading.Event
    progress_callback: ProgressCallback
    log_callback: LogCallback
    preview_callback: PreviewCallback = lambda _preview: None
    step_delay: float = 0.2

    def log(self, message: str) -> None:
        self.log_callback(message)

    def progress(self, percent: int, message: str) -> None:
        self.checkpoint()
        self.progress_callback(max(0, min(int(percent), 100)), message)

    def preview(
        self, path: str | Path, *, epoch: int = 0, next_epoch: int = 0,
        prompt: str = "", seed: int | None = None, steps: int = 0,
        kind: str = "training", current: int = 0, total: int = 0,
        image_index: int = 0, image_count: int = 0,
    ) -> None:
        """Publish a trainer-created preview without coupling the UI to a backend."""
        preview = Path(path).expanduser().resolve()
        if not preview.is_file():
            return
        self.preview_callback({
            "path": str(preview), "epoch": max(0, int(epoch)),
            "next_epoch": max(0, int(next_epoch)), "prompt": str(prompt),
            "seed": seed, "steps": max(0, int(steps)), "kind": str(kind),
            "current": max(0, int(current)), "total": max(0, int(total)),
            "image_index": max(0, int(image_index)), "image_count": max(0, int(image_count)),
        })

    def checkpoint(self) -> None:
        if self.cancel_event.is_set():
            raise ToolCancelled("Job cancelled by user.")
        while not self.run_event.wait(timeout=0.15):
            if self.cancel_event.is_set():
                raise ToolCancelled("Job cancelled by user.")

    def wait(self, multiplier: float = 1.0) -> None:
        deadline = time.monotonic() + max(0.0, self.step_delay * multiplier)
        while time.monotonic() < deadline:
            self.checkpoint()
            time.sleep(min(0.05, max(0.0, deadline - time.monotonic())))


class ToolExecutor:
    def __init__(
        self,
        root: Path,
        registry: ToolRegistry,
        logger: logging.Logger,
        *,
        step_delay: float = 0.2,
    ) -> None:
        self.root = root.resolve()
        self.registry = registry
        self.logger = logger
        self.step_delay = step_delay

    def execute(
        self,
        tool_id: str,
        arguments: dict[str, Any],
        *,
        job_id: str,
        cancel_event: threading.Event,
        run_event: threading.Event,
        progress_callback: ProgressCallback,
        log_callback: LogCallback,
        preview_callback: PreviewCallback | None = None,
    ) -> dict[str, Any]:
        spec = self.registry.get(tool_id)
        self._validate_arguments(spec, arguments)
        context = ToolContext(
            root=self.root,
            job_id=job_id,
            tool=spec,
            cancel_event=cancel_event,
            run_event=run_event,
            progress_callback=progress_callback,
            log_callback=log_callback,
            preview_callback=preview_callback or (lambda _preview: None),
            step_delay=self.step_delay,
        )
        backend_type = str(spec.backend.get("type", "")).lower()
        self.logger.info("Job %s starting tool %s", job_id, spec.id)

        if backend_type == "python":
            result = self._execute_python(spec, context, arguments)
        elif backend_type == "script":
            result = self._execute_script(spec, context, arguments)
        else:
            raise ToolExecutionError(f"{spec.name} has no configured backend.")

        self.logger.info("Job %s completed tool %s", job_id, spec.id)
        return result or {}

    def _validate_arguments(self, spec: ToolSpec, arguments: dict[str, Any]) -> None:
        if not isinstance(arguments, dict):
            raise ToolExecutionError("Tool arguments must be a dictionary.")
        unknown = set(arguments) - set(spec.arguments)
        if unknown:
            raise ToolExecutionError(
                f"{spec.name} received unsupported arguments: "
                f"{', '.join(sorted(unknown))}"
            )
        missing = set(spec.required_arguments) - set(arguments)
        if missing:
            raise ToolExecutionError(
                f"{spec.name} is missing required arguments: "
                f"{', '.join(sorted(missing))}"
            )
        for key, value in arguments.items():
            if isinstance(value, str):
                if "\x00" in value or len(value) > 2_000:
                    raise ToolExecutionError(f"Unsafe value supplied for {key}.")
            elif not isinstance(value, (int, float, bool, list, dict, type(None))):
                raise ToolExecutionError(f"Unsupported value type supplied for {key}.")

    def _execute_python(
        self,
        spec: ToolSpec,
        context: ToolContext,
        arguments: dict[str, Any],
    ) -> dict[str, Any]:
        module_name = str(spec.backend.get("module", ""))
        function_name = str(spec.backend.get("function", ""))
        if not module_name or not function_name:
            raise ToolExecutionError(f"{spec.name} has an incomplete Python backend.")
        try:
            module = importlib.import_module(module_name)
            function = getattr(module, function_name)
        except (ImportError, AttributeError) as exc:
            raise ToolExecutionError(
                f"Could not load backend for {spec.name}: {exc}"
            ) from exc
        if not callable(function):
            raise ToolExecutionError(f"Backend for {spec.name} is not callable.")

        signature = inspect.signature(function)
        try:
            signature.bind(context, **arguments)
        except TypeError as exc:
            raise ToolExecutionError(
                f"Arguments do not match {spec.name}'s backend: {exc}"
            ) from exc

        result = function(context, **arguments)
        if result is None:
            return {}
        if not isinstance(result, dict):
            raise ToolExecutionError(f"{spec.name} must return a dictionary or None.")
        return result

    def _execute_script(
        self,
        spec: ToolSpec,
        context: ToolContext,
        arguments: dict[str, Any],
    ) -> dict[str, Any]:
        raw_path = str(spec.backend.get("path", ""))
        if not raw_path:
            raise ToolExecutionError(f"{spec.name} has no script path.")
        script = Path(raw_path).expanduser()
        if not script.is_absolute():
            script = self.root / script
        script = script.resolve()
        if not script.is_file() or script.suffix.lower() != ".py":
            raise ToolExecutionError(f"Python script was not found: {script}")

        command = [sys.executable, str(script)]
        for key in spec.arguments:
            if key not in arguments:
                continue
            value = arguments[key]
            command.extend([f"--{key.replace('_', '-')}", str(value)])

        context.log(f"Launching registered script: {script.name}")
        process = subprocess.Popen(
            command,
            cwd=str(script.parent),
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            encoding="utf-8",
            errors="replace",
            shell=False,
        )
        output_queue: queue.Queue[str | None] = queue.Queue()

        def collect_output() -> None:
            assert process.stdout is not None
            for line in process.stdout:
                output_queue.put(line.rstrip())
            output_queue.put(None)

        reader = threading.Thread(target=collect_output, daemon=True)
        reader.start()
        process_controller = None
        try:
            import psutil

            process_controller = psutil.Process(process.pid)
        except Exception:
            process_controller = None
        suspended = False

        def stop_process_tree() -> None:
            """Stop the script and any workers it launched."""
            if process_controller is not None:
                try:
                    descendants = process_controller.children(recursive=True)
                    for child in descendants:
                        try:
                            child.terminate()
                        except Exception:
                            pass
                    process_controller.terminate()
                    try:
                        import psutil

                        _gone, alive = psutil.wait_procs(descendants, timeout=2)
                        for child in alive:
                            try:
                                child.kill()
                            except Exception:
                                pass
                    except Exception:
                        pass
                    return
                except Exception:
                    pass
            process.terminate()

        try:
            while True:
                if context.cancel_event.is_set():
                    raise ToolCancelled("Job cancelled by user.")
                if not context.run_event.is_set():
                    if process_controller is not None and not suspended:
                        try:
                            process_controller.suspend()
                            suspended = True
                        except Exception:
                            process_controller = None
                    context.checkpoint()
                    if process_controller is not None and suspended:
                        try:
                            process_controller.resume()
                        except Exception:
                            pass
                        suspended = False
                try:
                    line = output_queue.get(timeout=0.08)
                    if line is not None and line:
                        context.log(line)
                except queue.Empty:
                    pass
                if process.poll() is not None and output_queue.empty():
                    break
        except ToolCancelled:
            if process_controller is not None and suspended:
                try:
                    process_controller.resume()
                except Exception:
                    pass
            stop_process_tree()
            try:
                process.wait(timeout=3)
            except subprocess.TimeoutExpired:
                process.kill()
            raise

        if process.returncode != 0:
            raise ToolExecutionError(
                f"{spec.name} exited with code {process.returncode}."
            )
        context.progress(100, f"{spec.name} completed")
        return {}