llama-tornado

#!/usr/bin/env python3
"""
llama-tornado: GPU-accelerated Java LLM runner with TornadoVM
Run LLM models using either OpenCL or PTX backends.
"""

import argparse
import glob
import os
import subprocess
import sys
import time
import platform
from pathlib import Path
from typing import List, Optional, Dict, Any
from enum import Enum


class Backend(Enum):
    OPENCL = "opencl"
    PTX = "ptx"


class LlamaRunner:
    """Main class for managing LLM execution with GPU acceleration."""

    def __init__(self):
        self.java_home = os.environ.get("JAVA_HOME")
        self.tornado_sdk = os.environ.get("TORNADOVM_HOME")
        self.llama_root = os.environ.get("LLAMA_ROOT")

        if not all([self.java_home, self.tornado_sdk, self.llama_root]):
            print("Error: Required environment variables not set")
            print("Please ensure JAVA_HOME, TORNADOVM_HOME, and LLAMA_ROOT are defined")
            print("Note: check set_path in root dir -> source set_path")
            sys.exit(1)

    def _validate_paths(self):
        """Validate that required paths exist."""
        paths_to_check = {
            "JAVA_HOME": self.java_home,
            "TORNADOVM_HOME": self.tornado_sdk,
            "LLAMA_ROOT": self.llama_root,
        }

        for name, path in paths_to_check.items():
            if not Path(path).exists():
                print(f"Error: {name} path does not exist: {path}")
                sys.exit(1)

    @staticmethod
    def module_path_colon_sep(paths: List[str]) -> str:
        """Return OS-specific separator for Java module paths."""
        return ";".join(paths) if platform.system() == "Windows" else ":".join(paths)

    def _build_base_command(self, args: argparse.Namespace) -> List[str]:
        """Build the base Java command with JVM options."""
        cmd = [
            f"{self.java_home}/bin/java",
            "-server",
            "-XX:+UnlockExperimentalVMOptions",
            "-XX:+EnableJVMCI",
            f"-Xms{args.heap_min}",
            f"-Xmx{args.heap_max}",
            "--enable-preview",
            f"-Djava.library.path={self.tornado_sdk}/lib",
            "-Djdk.module.showModuleResolution=false",
            "--module-path",
            self.module_path_colon_sep([".", f"{self.tornado_sdk}/share/java/tornado"]),
        ]

        # TornadoVM configuration
        tornado_config = [
            "-Dtornado.load.api.implementation=uk.ac.manchester.tornado.runtime.tasks.TornadoTaskGraph",
            "-Dtornado.load.runtime.implementation=uk.ac.manchester.tornado.runtime.TornadoCoreRuntime",
            "-Dtornado.load.tornado.implementation=uk.ac.manchester.tornado.runtime.common.Tornado",
            "-Dtornado.load.annotation.implementation=uk.ac.manchester.tornado.annotation.ASMClassVisitor",
            "-Dtornado.load.annotation.parallel=uk.ac.manchester.tornado.api.annotations.Parallel",
            "-Dtornado.tvm.maxbytecodesize=65536"
        ]
        cmd.extend(tornado_config)

        # GPU options
        if args.use_gpu:
            cmd.append("-Duse.tornadovm=true")

        if args.verbose_init:
            cmd.append("-Dllama.EnableTimingForTornadoVMInit=true")

        # Debug options
        debug_config = []

        if args.debug:
            debug_config.extend(
                [
                    "-Dtornado.debug=true",
                    "-Dtornado.threadInfo=True"
                    if args.threads
                    else "-Dtornado.threadInfo=false",
                ]
            )
        else:
            debug_config.extend(
                [
                    "-Dtornado.threadInfo=True"
                    if args.threads
                    else "-Dtornado.threadInfo=false",
                    "-Dtornado.debug=false",
                ]
            )

        # Additional debug options
        debug_config.extend(
            [
                "-Dtornado.fullDebug=True"
                if args.full_dump
                else "-Dtornado.fullDebug=false",
                "-Dtornado.printKernel=True"
                if args.print_kernel
                else "-Dtornado.printKernel=false",
                "-Dtornado.print.bytecodes=True"
                if args.print_bytecodes
                else "-Dtornado.print.bytecodes=false",
            ]
        )

        cmd.extend(debug_config)

        # Additional TornadoVM settings
        tornado_runtime_config = [
            f"-Dtornado.device.memory={args.gpu_memory}",
            f"-Dtornado.profiler={str(args.profiler).lower()}",
            "-Dtornado.log.profiler=false",
            f"-Dtornado.profiler.dump.dir={args.profiler_dump_dir}",
            "-Dtornado.enable.fastMathOptimizations=true",
            "-Dtornado.enable.mathOptimizations=false",
            "-Dtornado.enable.nativeFunctions=true",
            "-Dtornado.loop.interchange=true",
            f"-Dtornado.eventpool.maxwaitevents={args.max_wait_events}",
        ]
        cmd.extend(tornado_runtime_config)

        # Backend-specific configuration
        if args.backend == Backend.OPENCL:
            # OpenCL specific flags
            cmd.append(f"-Dtornado.opencl.compiler.flags={args.opencl_flags}")

        # Module configuration - varies by backend
        module_config = [
            f"--upgrade-module-path",
            f"{self.tornado_sdk}/share/java/graalJars",
            f"@{self.tornado_sdk}/etc/exportLists/common-exports",
        ]
        # Add backend-specific exports and modules
        if args.backend == Backend.OPENCL:
            module_config.extend(
                [
                    f"@{self.tornado_sdk}/etc/exportLists/opencl-exports",
                    "--add-modules",
                    "ALL-SYSTEM,jdk.incubator.vector,tornado.runtime,tornado.annotation,tornado.drivers.common,tornado.drivers.opencl",
                ]
            )
        elif args.backend == Backend.PTX:
            module_config.extend(
                [
                    f"@{self.tornado_sdk}/etc/exportLists/ptx-exports",
                    "--add-modules",
                    "ALL-SYSTEM,jdk.incubator.vector,tornado.runtime,tornado.annotation,tornado.drivers.common,tornado.drivers.ptx",
                ]
            )

        module_config.extend(
            [
                "-cp",
                self._find_llama_jar(),
                "org.beehive.gpullama3.LlamaApp",
            ]
        )
        cmd.extend(module_config)

        return cmd

    def _find_llama_jar(self) -> str:
        """Find the LLaMA JAR file automatically using glob pattern."""
        target_dir = f"{self.llama_root}/target"
        jar_pattern = f"{target_dir}/gpu-llama3-*-SNAPSHOT.jar"
        jar_files = glob.glob(jar_pattern)

        if not jar_files:
            # Fallback: try any gpu-llama3 jar
            jar_pattern = f"{target_dir}/gpu-llama3-*.jar"
            jar_files = glob.glob(jar_pattern)

        if not jar_files:
            raise FileNotFoundError(f"No gpu-llama3 JAR file found in {target_dir}")

        # Sort to get the latest version if multiple exist
        jar_files.sort(reverse=True)
        return jar_files[0]

    def _add_llama_args(self, cmd: List[str], args: argparse.Namespace) -> List[str]:
        """Add LLaMA-specific arguments to the command."""
        llama_args = [
            "-m",
            args.model_path,
            "--temperature",
            str(args.temperature),
            "--top-p",
            str(args.top_p),
            "--seed",
            str(args.seed),
            "--max-tokens",
            str(args.max_tokens),
            "--stream",
            str(args.stream).lower(),
            "--echo",
            str(args.echo).lower(),
        ]

        if args.prompt:
            llama_args.extend(["-p", args.prompt])

        if args.system_prompt:
            llama_args.extend(["-sp", args.system_prompt])

        if args.interactive:
            llama_args.append("--interactive")
        elif args.instruct:
            llama_args.append("--instruct")

        return cmd + llama_args

    def run(self, args: argparse.Namespace) -> int:
        """Execute the LLaMA model with the specified arguments."""
        self._validate_paths()

        # Build the complete command
        cmd = self._build_base_command(args)
        cmd = self._add_llama_args(cmd, args)

        # Print command if requested (before verbose output)
        if args.show_command:
            print("Full Java command:")
            print("-" * 80)

            # Create a properly formatted command for easy copy-paste
            escaped_cmd = []
            for arg in cmd:
                # Escape arguments that contain spaces or special characters
                if " " in arg or '"' in arg or "'" in arg:
                    escaped_cmd.append(f'"{arg}"')
                else:
                    escaped_cmd.append(arg)

            # Print as a continuous line that can be easily copied
            print(" ".join(escaped_cmd))
            print("-" * 80)
            print()

            # If user only wants to see the command without executing
            if not args.execute_after_show:
                print("Command built successfully. Exiting without execution.")
                print(
                    "Use --execute-after-show to run the command after displaying it."
                )
                return 0

        if args.verbose:
            print("Executing command:")
            for arg in cmd:
                print(f"  {arg}")
            print()

        # Execute the command
        try:
            result = subprocess.run(cmd, check=True)
            return result.returncode
        except subprocess.CalledProcessError as e:
            print(f"Error: Command failed with return code {e.returncode}")
            return e.returncode
        except KeyboardInterrupt:
            print("\nOperation cancelled by user")
            return 130
        except Exception as e:
            print(f"Error: {e}")
            return 1


def load_env_from_script():
    if "LLAMA_ROOT" in os.environ and os.environ["LLAMA_ROOT"]:
        return

    system = platform.system()

    if system == "Windows":
        # Call set_paths.cmd and capture output as environment
        result = subprocess.run(
            ["cmd.exe", "/c", "set_paths.cmd && set"],
            capture_output=True,
            text=True,
            shell=False,
        )
        if result.returncode != 0:
            print("Failed to run set_paths.cmd")
            sys.exit(1)

        # Parse environment variables from output
        for line in result.stdout.splitlines():
            if "=" in line:
                key, value = line.strip().split("=", 1)
                os.environ[key] = value

    elif system in ("Linux", "Darwin"):
        # Source the set_paths file and capture env
        command = ["bash", "-c", "source ./set_paths && env"]
        result = subprocess.run(command, capture_output=True, text=True)
        if result.returncode != 0:
            print("Failed to source set_paths")
            sys.exit(1)

        for line in result.stdout.splitlines():
            if "=" in line:
                key, value = line.strip().split("=", 1)
                os.environ[key] = value
    else:
        print(f"Unsupported OS: {system}")
        sys.exit(1)


def create_parser() -> argparse.ArgumentParser:
    """Create and configure the argument parser."""
    parser = argparse.ArgumentParser(
        prog="llama-tornado",
        description="GPU-accelerated LLM runner using TornadoVM",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    # Required arguments
    parser.add_argument(
        "--model",
        dest="model_path",
        required=True,
        help="Path to the LLM gguf file (e.g., Llama-3.2-1B-Instruct-Q8_0.gguf)",
    )

    # LLM arguments
    llm_group = parser.add_argument_group("LLaMA Configuration")
    llm_group.add_argument("--prompt", help="Input prompt for the model")
    llm_group.add_argument("-sp", "--system-prompt", help="System prompt for the model")
    llm_group.add_argument(
        "--temperature",
        type=float,
        default=0.1,
        help="Sampling temperature (0.0 to 2.0)",
    )
    llm_group.add_argument(
        "--top-p", type=float, default=0.95, help="Top-p sampling parameter"
    )
    llm_group.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Random seed (default: current timestamp)",
    )
    llm_group.add_argument(
        "-n",
        "--max-tokens",
        type=int,
        default=512,
        help="Maximum number of tokens to generate",
    )
    llm_group.add_argument(
        "--stream", type=bool, default=True, help="Enable streaming output"
    )
    llm_group.add_argument(
        "--echo", type=bool, default=False, help="Echo the input prompt"
    )
    llm_group.add_argument(
        "--suffix", help="Suffix for fill-in-the-middle request (Codestral)"
    )

    # Mode selection
    mode_group = parser.add_argument_group("Mode Selection")
    mode_group.add_argument(
        "-i", "--interactive", action="store_true", help="Run in interactive/chat mode"
    )
    mode_group.add_argument(
        "--instruct",
        action="store_true",
        default=True,
        help="Run in instruction mode (default)",
    )

    # Hardware configuration
    hw_group = parser.add_argument_group("Hardware Configuration")
    hw_group.add_argument(
        "--gpu", dest="use_gpu", action="store_true", help="Enable GPU acceleration"
    )
    hw_group.add_argument(
        "--opencl",
        dest="backend",
        action="store_const",
        const=Backend.OPENCL,
        help="Use OpenCL backend (default)",
    )
    hw_group.add_argument(
        "--ptx",
        dest="backend",
        action="store_const",
        const=Backend.PTX,
        help="Use PTX/CUDA backend",
    )
    hw_group.add_argument("--gpu-memory", default="14GB", help="GPU memory allocation")
    hw_group.add_argument("--heap-min", default="20g", help="Minimum JVM heap size")
    hw_group.add_argument("--heap-max", default="20g", help="Maximum JVM heap size")

    # Debug and profiling
    debug_group = parser.add_argument_group("Debug and Profiling")
    debug_group.add_argument("--debug", action="store_true", help="Enable debug output")
    debug_group.add_argument(
        "--profiler", action="store_true", help="Enable TornadoVM profiler"
    )
    debug_group.add_argument(
        "--profiler-dump-dir",
        default=None,
        help="Directory for profiler output",
    )

    # TornadoVM Execution Verbose options
    verbose_group = parser.add_argument_group("TornadoVM Execution Verbose")
    verbose_group.add_argument(
        "--print-bytecodes",
        dest="print_bytecodes",
        action="store_true",
        help="Print bytecodes (tornado.print.bytecodes=true)",
    )
    verbose_group.add_argument(
        "--print-threads",
        dest="threads",
        action="store_true",
        help="Print thread information (tornado.threadInfo=true)",
    )
    verbose_group.add_argument(
        "--print-kernel",
        dest="print_kernel",
        action="store_true",
        help="Print kernel information (tornado.printKernel=true)",
    )
    verbose_group.add_argument(
        "--full-dump",
        dest="full_dump",
        action="store_true",
        help="Enable full debug dump (tornado.fullDebug=true)",
    )
    verbose_group.add_argument(
        "--verbose-init",
        dest="verbose_init",
        action="store_true",
        help="Enable timers for TornadoVM initialization (llama.EnableTimingForTornadoVMInit=true)",
    )

    # Command display options
    command_group = parser.add_argument_group("Command Display Options")
    command_group.add_argument(
        "--show-command",
        action="store_true",
        help="Display the full Java command that will be executed",
    )
    command_group.add_argument(
        "--execute-after-show",
        action="store_true",
        help="Execute the command after showing it (use with --show-command)",
    )

    # Advanced options
    advanced_group = parser.add_argument_group("Advanced Options")
    advanced_group.add_argument(
        "--opencl-flags",
        default="-cl-denorms-are-zero -cl-no-signed-zeros -cl-finite-math-only",
        help="OpenCL compiler flags",
    )
    advanced_group.add_argument(
        "--max-wait-events",
        type=int,
        default=32000,
        help="Maximum wait events for TornadoVM event pool",
    )
    advanced_group.add_argument(
        "--verbose", "-v", action="store_true", help="Verbose output"
    )

    return parser


def main():
    """Main entry point."""
    load_env_from_script()
    parser = create_parser()
    args = parser.parse_args()

    # Set default profiler log path relative to LLAMA_ROOT
    if args.profiler_dump_dir is None:
        llama_root = os.environ.get("LLAMA_ROOT")
        args.profiler_dump_dir = os.path.join(llama_root, "profiler-log.json")

    # Set default seed if not provided
    if args.seed is None:
        args.seed = int(time.time())

    # Set default backend to OpenCL if not specified
    if not hasattr(args, "backend") or args.backend is None:
        args.backend = Backend.OPENCL

    # Handle mode selection logic
    if args.interactive:
        args.instruct = False

    # Create and run the LLaMA runner
    runner = LlamaRunner()
    return runner.run(args)


if __name__ == "__main__":
    sys.exit(main())