ResInsight/GrpcInterface/Python/rips/taskmaestro_helper/run.py

"""Execute a taskmaestro workflow against a running ResInsight.

The helper:
1. Connects back to ResInsight via the gRPC port passed by the caller.
2. Reads input.yaml and walks the dict, replacing every
   `{__resinsight_ref__: <type>, ...id...}` map with an `ObjectModel`-style
   `{"value": <live rips object>}` so taskmaestro's Pydantic validation
   accepts it as a wrapped object.
3. Builds the workflow via taskmaestro's loader and runs it.
4. Streams progress events as newline-delimited JSON to stdout for the
   ResInsight log dialog to render.
"""

from __future__ import annotations

import json
import sys
from pathlib import Path
from typing import Any, Callable

import yaml

from .refs import REF_MARKER


PROGRESS_PREFIX = "@@RI_WORKFLOW_EVENT@@"


def _emit_progress(
    run_id: str, task: str, state: str, error: str = "", item: str | None = None
) -> None:
    payload = {"event": "task_state", "run_id": run_id, "task": task, "state": state}
    if item is not None:
        # One item of a mapped task (`map:`); the task itself reports its own states
        payload["event"] = "map_item_state"
        payload["item"] = item
    if error:
        payload["error"] = error
    sys.stdout.write(PROGRESS_PREFIX + json.dumps(payload) + "\n")
    sys.stdout.flush()


class ProgressHook:
    """Report taskmaestro's resolved task-instance lifecycle to ResInsight."""

    def __init__(self, run_id: str) -> None:
        self.run_id = run_id

    def on_task_start(self, job: Any, task: Any) -> None:
        _emit_progress(self.run_id, task.name, "running")

    def on_task_complete(self, job: Any, task: Any, output: Any) -> None:
        _emit_progress(self.run_id, task.name, "completed")

    def on_task_fail(self, job: Any, task: Any, error: Exception) -> None:
        _emit_progress(self.run_id, task.name, "failed", str(error))

    def on_map_item_start(self, job: Any, task: Any, key: str) -> None:
        _emit_progress(self.run_id, task.name, "running", item=key)

    def on_map_item_complete(self, job: Any, task: Any, key: str, output: Any) -> None:
        _emit_progress(self.run_id, task.name, "completed", item=key)

    def on_map_item_fail(self, job: Any, task: Any, key: str, error: Exception) -> None:
        _emit_progress(self.run_id, task.name, "failed", str(error), item=key)


def _emit(event: str, **fields: Any) -> None:
    sys.stdout.write(json.dumps({"event": event, **fields}) + "\n")
    sys.stdout.flush()


def resolve_refs(data: Any, resolver: Callable[[str, dict[str, Any]], Any]) -> Any:
    """Recursively replace `{__resinsight_ref__: ..., ...}` maps in `data`.

    The replacement is a `{"value": <resolved>}` dict, so Pydantic's
    ObjectModel validator can accept it directly.
    """
    if isinstance(data, dict):
        if REF_MARKER in data:
            ri_type = data[REF_MARKER]
            extras = {k: v for k, v in data.items() if k != REF_MARKER}
            return {"value": resolver(ri_type, extras)}
        return {k: resolve_refs(v, resolver) for k, v in data.items()}
    if isinstance(data, list):
        return [resolve_refs(v, resolver) for v in data]
    return data


def make_rips_resolver(rips_instance: Any) -> Callable[[str, dict[str, Any]], Any]:
    """Resolve a `__resinsight_ref__` map to a live `rips` object."""

    def resolve(ri_type: str, extras: dict[str, Any]) -> Any:
        project = rips_instance.project
        if ri_type == "EclipseCase":
            return project.case(case_id=extras["case_id"])
        if ri_type == "WellPath":
            return project.well_path_by_name(extras["well_path_name"])
        if ri_type == "View":
            for v in project.views():
                if getattr(v, "id", None) == extras["view_id"]:
                    return v
            raise LookupError(f"View id {extras['view_id']} not found")
        raise LookupError(f"Unknown resinsight_type: {ri_type}")

    return resolve


def _load_workflow(workflow_dir: Path | None, registered_id: str | None) -> Any:
    if registered_id is not None:
        from taskmaestro.discovery import get_registered_workflow

        return get_registered_workflow(registered_id)

    from taskmaestro.yaml_config import _load_workflow_only

    assert workflow_dir is not None
    workflow, _ = _load_workflow_only(workflow_dir / "workflow.yaml")
    return workflow


def run_workflow(
    workflow_dir: Path | None,
    input_path: Path,
    grpc_port: int,
    run_id: str = "",
    registered_id: str | None = None,
) -> int:
    """Run a YAML workflow folder, or a registered workflow by its id."""
    if workflow_dir is not None:
        workflow_dir_str = str(workflow_dir)
        if workflow_dir_str not in sys.path:
            sys.path.insert(0, workflow_dir_str)

    import rips

    instance = rips.Instance.find(start_port=grpc_port, end_port=grpc_port + 1)
    if instance is None:
        _emit("error", message=f"Could not connect to ResInsight on port {grpc_port}")
        return 1
    _emit("connected", port=grpc_port)

    raw_input = yaml.safe_load(input_path.read_text()) or {}
    if not isinstance(raw_input, dict):
        _emit("error", message="input.yaml must contain a mapping")
        return 1

    resolver = make_rips_resolver(instance)
    try:
        resolved = resolve_refs(raw_input, resolver)
    except Exception as exc:
        _emit("error", message=f"Failed to resolve references: {exc}")
        return 1

    from taskmaestro import EmptyConfig, ExecutionContext, Job, JobConfiguration, Runner

    try:
        workflow = _load_workflow(workflow_dir, registered_id)
    except Exception as exc:
        _emit("error", message=f"Failed to load workflow: {exc}")
        return 1

    job_config = JobConfiguration(resolved)
    job: Job[Any] = Job(
        workflow=workflow, config=EmptyConfig(), job_configuration=job_config
    )

    _emit("starting", workflow=workflow.name)
    try:
        result = Runner(hooks=[ProgressHook(run_id)]).run(job, ctx=ExecutionContext())
    except Exception as exc:
        _emit("error", message=str(exc))
        return 1

    _emit("finished", status=str(result.status), error=str(result.error or ""))
    return 0 if result.error is None else 1


def main(argv: list[str]) -> int:
    import argparse

    parser = argparse.ArgumentParser(
        prog="run", description="Run a taskmaestro workflow."
    )
    parser.add_argument("workflow_dir", type=Path, nargs="?")
    parser.add_argument(
        "--registered", metavar="ID", help="Run a registered workflow by its id"
    )
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--grpc-port", type=int, required=True)
    parser.add_argument("--run-id", default="")
    args = parser.parse_args(argv)
    if (args.workflow_dir is None) == (args.registered is None):
        parser.error("specify either a workflow folder or --registered <id>")
    return run_workflow(
        args.workflow_dir.resolve() if args.workflow_dir is not None else None,
        args.input.resolve(),
        args.grpc_port,
        args.run_id,
        registered_id=args.registered,
    )