Surface Fuser and Pipeline UI as script entrypoints (#36)

* Surface Fuser and Pipeline UI as script entrypoints

* Cleanup

* Remove TODOs

* lint

Author: Jack-Khuu

README.md

@@ -99,8 +99,8 @@ More knobs live in `triton_kernel_agent/agent.py` and `Fuser/config.py`.
 
 - **UIs** — interactive runs with Gradio frontends:
   - Triton KernelAgent UI: `kernel-agent` or `python scripts/triton_ui.py`
-  - Fuser orchestration UI: `python -m Fuser.fuser_ui`
-  - Full pipeline UI: `python -m Fuser.pipeline_ui`
+  - Fuser orchestration UI: `fuser-ui` or `python scripts/fuser_ui`
+  - Full pipeline UI: `pipeline-ui` or `python scripts/pipeline_ui`
 
 ## Component Details
 

pyproject.toml

@@ -43,8 +43,10 @@ dev = [
 ]
 
 [project.scripts]
+fuser-ui = "scripts.fuser_ui:main"
 kernel-agent = "scripts.triton_ui:main"
 list-models = "scripts.list_models:main"
+pipeline-ui = "scripts.pipeline_ui:main"
 
 [project.urls]
 "Homepage" = "https://github.com/pytorch-labs/KernelAgent"

Fuser/fuser_ui.py renamed to scripts/fuser_ui.py

@@ -19,7 +19,6 @@
 import argparse
 import ast
 import os
-import sys
 import tarfile
 import time
 import traceback
@@ -38,19 +37,9 @@
 except Exception:  # pragma: no cover
     extract_subgraphs_to_json = None  # type: ignore
 
-# Support both package and script execution contexts
-if __package__ is None or __package__ == "":
-    PACKAGE_ROOT = Path(__file__).resolve().parent
-    REPO_ROOT = PACKAGE_ROOT.parent
-    if str(REPO_ROOT) not in sys.path:
-        sys.path.insert(0, str(REPO_ROOT))
-    from Fuser.config import OrchestratorConfig, new_run_id
-    from Fuser.orchestrator import Orchestrator
-    from Fuser.paths import ensure_abs_regular_file, make_run_dirs, PathSafetyError
-else:
-    from .config import OrchestratorConfig, new_run_id
-    from .orchestrator import Orchestrator
-    from .paths import ensure_abs_regular_file, make_run_dirs, PathSafetyError
+from Fuser.config import OrchestratorConfig, new_run_id
+from Fuser.orchestrator import Orchestrator
+from Fuser.paths import ensure_abs_regular_file, make_run_dirs, PathSafetyError
 
 
 @dataclass
@@ -327,6 +316,7 @@ def __init__(self) -> None:
             repo_root / "external" / "KernelBench" / "KernelBench",
             Path.cwd() / "external" / "KernelBench" / "KernelBench",
             Path.cwd() / "KernelBench" / "KernelBench",
+            Path.cwd().parent / "KernelBench" / "KernelBench",
         ]
         seen: set[str] = set()
         collected: list[tuple[str, str]] = []

Fuser/pipeline_ui.py renamed to scripts/pipeline_ui.py

@@ -29,22 +29,9 @@
 import gradio as gr
 from dotenv import load_dotenv
 
-try:
-    # Support both package and script execution contexts (match Fuser/fuser_ui.py pattern)
-    if __package__ is None or __package__ == "":
-        PACKAGE_ROOT = Path(__file__).resolve().parent
-        REPO_ROOT = PACKAGE_ROOT.parent
-        if str(REPO_ROOT) not in sys.path:
-            sys.path.insert(0, str(REPO_ROOT))
-        from Fuser.pipeline import run_pipeline
-        from Fuser.auto_agent import AutoKernelRouter
-        from Fuser.code_extractor import extract_single_python_file
-    else:
-        from .pipeline import run_pipeline
-        from .auto_agent import AutoKernelRouter
-        from .code_extractor import extract_single_python_file
-except Exception:
-    raise
+from Fuser.pipeline import run_pipeline
+from Fuser.auto_agent import AutoKernelRouter
+from Fuser.code_extractor import extract_single_python_file
 from triton_kernel_agent.providers.models import (
     get_model_provider,
     MODEL_NAME_TO_CONFIG,
@@ -416,9 +403,21 @@ def run_pipeline_ui(
 class PipelineUI:
     def __init__(self) -> None:
         load_dotenv()
-        self.problem_choices = _list_kernelbench_problems(
-            Path.cwd() / "external" / "KernelBench" / "KernelBench"
-        )
+        candidate_roots = [
+            Path.cwd() / "external" / "KernelBench" / "KernelBench",
+            Path.cwd() / "KernelBench" / "KernelBench",
+            Path.cwd().parent / "KernelBench" / "KernelBench",
+        ]
+        seen: set[str] = set()
+        collected: list[tuple[str, str]] = []
+        for base in candidate_roots:
+            print(base, file=sys.stderr)
+            for label, abspath in _list_kernelbench_problems(base):
+                if abspath not in seen:
+                    collected.append((label, abspath))
+                    seen.add(abspath)
+        self.problem_choices = collected
+
         control_flow_path = (
             Path(__file__).resolve().parent.parent.parent
             / "external"