lagora-cli 1.1.0

package/dist/scripts/legato_runner_launchers.py

@@ -0,0 +1,287 @@
+from __future__ import annotations
+
+import ast
+import importlib.util
+import inspect
+import os
+import signal
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any
+
+from legato_runner_env import launcher_env, launcher_import_paths, launcher_working_directory, script_command
+
+CALLABLE_ENTRY_NAMES = ("main", "run", "launch", "lower")
+CALLABLE_ENTRY_PARAMETERS = {"output_type_name", "output_type", "stage", "output_root", "repo_root"}
+LOWERING_ARTIFACT_SUFFIXES = {".mlir", ".ll", ".s", ".asm", ".bin", ".json", ".bc", ".elf", ".o"}
+LEGATO_KERNEL_DECORATORS = {"compile", "primitive"}
+
+
+class LauncherError(RuntimeError):
+    pass
+
+
+class LauncherScriptError(LauncherError):
+    """The submitted kernel ran but failed on its own terms.
+
+    Distinct from LauncherError so callers can report a kernel failure instead
+    of blaming the SDK environment.
+    """
+
+
+ACTIVE_PROCESS: subprocess.Popen[str] | None = None
+
+
+def statically_decorated_kernel_names(kernel_path: Path) -> set[str]:
+    tree = parse_kernel(kernel_path)
+    return {
+        node.name
+        for node in tree.body
+        if isinstance(node, ast.FunctionDef)
+        and any(is_legato_kernel_decorator(decorator) for decorator in node.decorator_list)
+    }
+
+
+def kernel_has_launcher(kernel_path: Path) -> bool:
+    return script_has_main_guard(kernel_path) or bool(callable_entry_names(kernel_path)) or bool(argparse_options(kernel_path))
+
+
+def launcher_runs_in_subprocess(kernel_path: Path) -> bool:
+    """Whether run_launcher will execute this kernel in a child process.
+
+    In-process device state (hart's fsim environment) does not cross that
+    boundary, so callers that configure a device must know which side runs
+    the kernel. Mirrors the dispatch in run_launcher.
+    """
+    return script_has_main_guard(kernel_path) or not callable_entry_names(kernel_path)
+
+
+def run_launcher(kernel_path: Path, stages: list[str], output_root: Path) -> None:
+    options = argparse_options(kernel_path)
+    if script_has_main_guard(kernel_path):
+        run_cli_launcher(kernel_path, stages, output_root, options)
+        return
+    if callable_entry_names(kernel_path):
+        run_callable_launcher(kernel_path, stages, output_root)
+        return
+    run_cli_launcher(kernel_path, stages, output_root, options)
+
+
+def parse_kernel(kernel_path: Path) -> ast.Module:
+    with kernel_path.open("r", encoding="utf-8") as source:
+        return ast.parse(source.read(), filename=str(kernel_path))
+
+
+def is_legato_kernel_decorator(decorator: ast.expr) -> bool:
+    target = decorator.func if isinstance(decorator, ast.Call) else decorator
+    return (
+        isinstance(target, ast.Attribute)
+        and target.attr in LEGATO_KERNEL_DECORATORS
+    ) or (
+        isinstance(target, ast.Name)
+        and target.id in LEGATO_KERNEL_DECORATORS
+    )
+
+
+def script_has_main_guard(kernel_path: Path) -> bool:
+    for node in parse_kernel(kernel_path).body:
+        if isinstance(node, ast.If) and is_main_guard(node.test):
+            return True
+    return False
+
+
+def is_main_guard(test: ast.expr) -> bool:
+    if not isinstance(test, ast.Compare) or len(test.ops) != 1 or len(test.comparators) != 1:
+        return False
+    left = test.left
+    right = test.comparators[0]
+    return (
+        isinstance(test.ops[0], ast.Eq)
+        and isinstance(left, ast.Name)
+        and left.id == "__name__"
+        and isinstance(right, ast.Constant)
+        and right.value == "__main__"
+    )
+
+
+def argparse_options(kernel_path: Path) -> set[str]:
+    options: set[str] = set()
+    for node in ast.walk(parse_kernel(kernel_path)):
+        if not isinstance(node, ast.Call):
+            continue
+        if not isinstance(node.func, ast.Attribute) or node.func.attr != "add_argument":
+            continue
+        for arg in node.args:
+            if isinstance(arg, ast.Constant) and isinstance(arg.value, str) and arg.value.startswith("--"):
+                options.add(arg.value)
+    return options
+
+
+def callable_entry_names(kernel_path: Path) -> set[str]:
+    return {
+        node.name
+        for node in parse_kernel(kernel_path).body
+        if isinstance(node, ast.FunctionDef)
+        and node.name in CALLABLE_ENTRY_NAMES
+        and has_supported_launcher_parameters(node)
+    }
+
+
+def has_supported_launcher_parameters(node: ast.FunctionDef) -> bool:
+    positional = list(node.args.posonlyargs) + list(node.args.args)
+    all_parameters = [arg.arg for arg in positional]
+    all_parameters.extend(arg.arg for arg in node.args.kwonlyargs)
+    if not any(name in CALLABLE_ENTRY_PARAMETERS for name in all_parameters):
+        return False
+    required_count = len(positional) - len(node.args.defaults)
+    required = [arg.arg for arg in positional[:required_count]]
+    required.extend(arg.arg for arg, default in zip(node.args.kwonlyargs, node.args.kw_defaults) if default is None)
+    return all(name in CALLABLE_ENTRY_PARAMETERS for name in required)
+
+
+def run_cli_launcher(kernel_path: Path, stages: list[str], output_root: Path, options: set[str]) -> None:
+    if "--output-type" in options:
+        for stage in stages:
+            run_cli_stage(kernel_path, stage, output_root / stage.lower(), options)
+        return
+    # One process per stage. legato compiles a kernel once per process and keys
+    # what it cached without the output type, so a launcher asked for MLIR,
+    # CORE_IR and BACKEND_IR in one run wrote the first and returned the cached
+    # result for the other two -- three listings collapsed into one, silently.
+    for stage in stages:
+        stage_root = output_root / stage.lower()
+        stage_root.mkdir(parents=True, exist_ok=True)
+        command = script_command(kernel_path)
+        if "--repo-root" in options:
+            command.extend(["--repo-root", str(Path.cwd())])
+        if "--output-root" in options:
+            command.extend(["--output-root", str(stage_root)])
+        if "--stages" in options:
+            command.extend(["--stages", stage])
+        run_command(command, launcher_working_directory(), launcher_env(kernel_path, stage_root, [stage]), stage_root, stage)
+
+
+def run_cli_stage(kernel_path: Path, stage: str, stage_root: Path, options: set[str]) -> None:
+    command = [*script_command(kernel_path), "--output-type", stage]
+    if "--repo-root" in options:
+        command.extend(["--repo-root", str(Path.cwd())])
+    if "--output-root" in options:
+        command.extend(["--output-root", str(stage_root)])
+    run_command(command, launcher_working_directory(), launcher_env(kernel_path, stage_root, [stage]), stage_root, stage)
+
+
+def run_command(command: list[str], cwd: Path, env: dict[str, str], output_root: Path, stage: str) -> None:
+    global ACTIVE_PROCESS
+    output_root.mkdir(parents=True, exist_ok=True)
+    print("[LEGATO_WORKER] launcher_mode=script", flush=True)
+    print(f"[LEGATO_WORKER] stage={stage} status=running", flush=True)
+    print(f"[LEGATO_WORKER] cwd={cwd}", flush=True)
+    print(f"[LEGATO_WORKER] output_path={output_root}", flush=True)
+    print(f"[LEGATO_WORKER] import_paths={env.get('AGORA_LEGATO_IMPORT_PATHS', '')}", flush=True)
+    print(f"[LEGATO_WORKER] command={' '.join(command)}", flush=True)
+    process = subprocess.Popen(command, cwd=cwd, env=env, start_new_session=True)
+    ACTIVE_PROCESS = process
+    try:
+        return_code = process.wait()
+    finally:
+        ACTIVE_PROCESS = None
+    if return_code != 0:
+        raise LauncherScriptError(f"launcher script exited with code {return_code}")
+    require_artifacts(output_root, "launcher script completed with no lowering artifacts")
+    print(f"[LEGATO_WORKER] stage={stage} status=success", flush=True)
+
+
+def run_callable_launcher(kernel_path: Path, stages: list[str], output_root: Path) -> None:
+    module = load_module(kernel_path)
+    entry = first_callable_entry(module)
+    for stage in stages:
+        stage_root = output_root / stage.lower()
+        stage_root.mkdir(parents=True, exist_ok=True)
+        print("[LEGATO_WORKER] launcher_mode=callable", flush=True)
+        print(f"[LEGATO_WORKER] stage={stage} status=running", flush=True)
+        print(f"[LEGATO_WORKER] output_path={stage_root}", flush=True)
+        call_entry(entry, stage, stage_root)
+        require_artifacts(stage_root, "launcher callable completed without generated artifacts")
+        print(f"[LEGATO_WORKER] stage={stage} status=success", flush=True)
+
+
+def load_module(kernel_path: Path) -> Any:
+    spec = importlib.util.spec_from_file_location("agora_launcher_kernel", kernel_path)
+    if spec is None or spec.loader is None:
+        raise LauncherError(f"cannot import launcher module: {kernel_path}")
+    module = importlib.util.module_from_spec(spec)
+    import_paths = launcher_import_paths(kernel_path)
+    for path in reversed(import_paths):
+        sys.path.insert(0, str(path))
+    try:
+        spec.loader.exec_module(module)
+    finally:
+        for path in import_paths:
+            if sys.path and sys.path[0] == str(path):
+                sys.path.pop(0)
+    return module
+
+
+def require_artifacts(output_root: Path, message: str) -> None:
+    artifacts = [
+        path for path in output_root.rglob("*")
+        if path.is_file()
+    ]
+    lowering_artifacts = [
+        path for path in artifacts
+        if path.suffix.lower() in LOWERING_ARTIFACT_SUFFIXES
+    ]
+    if not lowering_artifacts:
+        if artifacts:
+            names = ", ".join(str(path.relative_to(output_root)) for path in artifacts[:5])
+            raise LauncherScriptError(f"{message}; found non-lowering files: {names}")
+        raise LauncherScriptError(message)
+    print(f"[LEGATO_WORKER] generated_artifacts={len(lowering_artifacts)}", flush=True)
+
+
+def first_callable_entry(module: Any) -> Any:
+    for name in CALLABLE_ENTRY_NAMES:
+        value = getattr(module, name, None)
+        if callable(value) and has_supported_signature(value):
+            return value
+    raise LauncherError("no launcher entrypoint found; expected main(), run(), launch(), or lower()")
+
+
+def has_supported_signature(value: Any) -> bool:
+    signature = inspect.signature(value)
+    if not any(name in CALLABLE_ENTRY_PARAMETERS for name in signature.parameters):
+        return False
+    for name, parameter in signature.parameters.items():
+        if name in CALLABLE_ENTRY_PARAMETERS:
+            continue
+        if parameter.default is inspect.Signature.empty:
+            return False
+    return True
+
+
+def call_entry(entry: Any, stage: str, output_root: Path) -> None:
+    signature = inspect.signature(entry)
+    kwargs: dict[str, Any] = {}
+    for name, parameter in signature.parameters.items():
+        if name in {"output_type_name", "output_type", "stage"}:
+            kwargs[name] = stage
+        elif name == "output_root":
+            kwargs[name] = output_root
+        elif name == "repo_root":
+            kwargs[name] = launcher_working_directory()
+        elif parameter.default is inspect.Signature.empty:
+            raise LauncherError(f"launcher entrypoint has unsupported required parameter: {name}")
+    result = entry(**kwargs)
+    if isinstance(result, int) and result != 0:
+        raise LauncherError(f"launcher entrypoint returned {result}")
+
+
+def terminate_active_process(signum: int, _frame: Any) -> None:
+    if ACTIVE_PROCESS is not None and ACTIVE_PROCESS.poll() is None:
+        os.killpg(ACTIVE_PROCESS.pid, signum)
+    raise SystemExit(128 + signum)
+
+
+signal.signal(signal.SIGTERM, terminate_active_process)
+signal.signal(signal.SIGINT, terminate_active_process)

package/dist/scripts/legato_runner_script_wrapper.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import inspect
+import os
+import runpy
+import sys
+import types
+from pathlib import Path
+from typing import Any
+
+
+def prepare_bertha(ctx: Any) -> Any:
+    import legato.model.bertha as bertha_model
+
+    return bertha_model.Bertha(ctx, "default", 32, False, 8, 128 * pow(1024, 3))
+
+
+def setup_ha_home(anchor: Path) -> Path:
+    return anchor
+
+
+def install_compat_modules() -> None:
+    install_module("kernels", None)
+    install_module("kernels._common", {"prepare_bertha": prepare_bertha})
+    install_module("utils", None)
+    install_module("utils.legato_cache", {"setup_ha_home": setup_ha_home})
+
+
+def install_module(name: str, members: dict[str, Any] | None) -> None:
+    if name in sys.modules:
+        return
+    module = types.ModuleType(name)
+    if members is None:
+        module.__path__ = []  # type: ignore[attr-defined]
+    else:
+        for key, value in members.items():
+            setattr(module, key, value)
+    sys.modules[name] = module
+
+
+def install_legato_session_redirect() -> None:
+    output_root = os.environ.get("AGORA_LEGATO_OUTPUT_ROOT", "").strip()
+    stages = [stage.strip() for stage in os.environ.get("AGORA_LEGATO_STAGES", "").split(",") if stage.strip()]
+    launch_value = os.environ.get("AGORA_LEGATO_SESSION_LAUNCH", "").strip().lower()
+    launch_override = session_launch_override(launch_value)
+    if not output_root and not stages and launch_override is None:
+        return
+    import legato
+
+    original_session = legato.session
+    effective_output_path = str(Path(output_root).resolve()) if output_root else None
+    effective_output_type = getattr(legato.OutputType, stages[0]) if len(stages) == 1 else None
+    accepts_launch = session_accepts_launch(original_session)
+    output_type_position = session_parameter_position(original_session, "output_type")
+    output_path_position = session_parameter_position(original_session, "output_path")
+    launch_position = session_parameter_position(original_session, "launch")
+
+    def redirected_session(*args: Any, **kwargs: Any) -> Any:
+        next_args = list(args)
+        next_kwargs = dict(kwargs)
+        if effective_output_type is not None:
+            original_output_type = session_argument(next_args, next_kwargs, "output_type", output_type_position)
+            if original_output_type != effective_output_type:
+                print(f"[LEGATO_WORKER] session_output_type_effective={stages[0]}", flush=True)
+            replace_session_argument(
+                next_args,
+                next_kwargs,
+                "output_type",
+                output_type_position,
+                effective_output_type,
+            )
+        if effective_output_path is not None:
+            original_output_path = session_argument(next_args, next_kwargs, "output_path", output_path_position)
+            if original_output_path is not None:
+                original_text = str(original_output_path)
+                if original_text != effective_output_path:
+                    print(f"[LEGATO_WORKER] session_output_path_original={original_text}", flush=True)
+                    print(f"[LEGATO_WORKER] session_output_path_effective={effective_output_path}", flush=True)
+            else:
+                print(f"[LEGATO_WORKER] session_output_path_injected={effective_output_path}", flush=True)
+            replace_session_argument(
+                next_args,
+                next_kwargs,
+                "output_path",
+                output_path_position,
+                effective_output_path,
+            )
+        if launch_override is not None and accepts_launch:
+            original_launch = session_argument(next_args, next_kwargs, "launch", launch_position)
+            if original_launch != launch_override:
+                print(f"[LEGATO_WORKER] session_launch_effective={str(launch_override).lower()}", flush=True)
+            replace_session_argument(next_args, next_kwargs, "launch", launch_position, launch_override)
+        return original_session(*next_args, **next_kwargs)
+
+    legato.session = redirected_session
+
+
+def session_launch_override(value: str) -> bool | None:
+    if value in {"0", "false", "no", "off"}:
+        return False
+    if value in {"1", "true", "yes", "on"}:
+        return True
+    return None
+
+
+def session_accepts_launch(session_factory: Any) -> bool:
+    try:
+        signature = inspect.signature(session_factory)
+    except (TypeError, ValueError):
+        return False
+    return "launch" in signature.parameters or any(
+        parameter.kind is inspect.Parameter.VAR_KEYWORD
+        for parameter in signature.parameters.values()
+    )
+
+
+def session_parameter_position(session_factory: Any, name: str) -> int | None:
+    try:
+        signature = inspect.signature(session_factory)
+    except (TypeError, ValueError):
+        return None
+    positional = [
+        parameter_name
+        for parameter_name, parameter in signature.parameters.items()
+        if parameter.kind in {inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD}
+    ]
+    return positional.index(name) if name in positional else None
+
+
+def session_argument(
+    args: list[Any],
+    kwargs: dict[str, Any],
+    name: str,
+    position: int | None,
+) -> Any:
+    if name in kwargs:
+        return kwargs[name]
+    return args[position] if position is not None and position < len(args) else None
+
+
+def replace_session_argument(
+    args: list[Any],
+    kwargs: dict[str, Any],
+    name: str,
+    position: int | None,
+    value: Any,
+) -> None:
+    if position is not None and position < len(args):
+        args[position] = value
+        kwargs.pop(name, None)
+        return
+    kwargs[name] = value
+
+
+def configure_fsim_device() -> None:
+    """Open the fsim device in this process when the runner delegated it here.
+
+    hart keeps device state in the C++ runtime of the process that calls
+    set_fsim_env, so a script launcher has to open its own.
+    """
+    if os.environ.get("AGORA_FSIM_CONFIGURE_IN_CHILD", "").strip() != "1":
+        return
+    import hart
+
+    device_version = getattr(hart.HaDeviceVersion, os.environ["AGORA_FSIM_DEVICE_VERSION"])
+    compute_mode = getattr(hart.HaFSimComputeMode, os.environ["AGORA_FSIM_COMPUTE_MODE"])
+    hart.set_fsim_env(
+        device_version,
+        compute_mode,
+        os.environ["AGORA_FSIM_LOG_PATH"],
+        os.environ.get("AGORA_FSIM_LOG_LEVEL", "trace"),
+    )
+    if hart.get_device_count() == 0:
+        raise RuntimeError("no fsim devices available in launcher process")
+    hart.set_device(0)
+    print("[LEGATO_WORKER] fsim_device=ready", flush=True)
+
+
+def main() -> None:
+    if len(sys.argv) < 2:
+        raise SystemExit("usage: legato_runner_script_wrapper.py <script> [args...]")
+    install_compat_modules()
+    import_paths = [path for path in os.environ.get("AGORA_LEGATO_IMPORT_PATHS", "").split(os.pathsep) if path]
+    sys.path[:0] = [path for path in import_paths if path not in sys.path]
+    configure_fsim_device()
+    install_legato_session_redirect()
+    script = sys.argv[1]
+    sys.argv = [script, *sys.argv[2:]]
+    runpy.run_path(script, run_name="__main__")
+
+
+if __name__ == "__main__":
+    main()