tritonparse 0.1.1__py3-none-any.whl

tritonparse/reproducer/orchestrator.py

@@ -0,0 +1,65 @@
+from pathlib import Path
+from typing import Any, Dict
+
+from .ingestion.ndjson import build_context_bundle
+from .param_generator import generate_allocation_snippet, generate_kwargs_dict
+from .prompts.loader import render_prompt
+from .providers.base import LLMProvider
+from .runtime.executor import run_python
+
+
+def _excerpt(s: str, n: int = 160):
+    lines = s.splitlines()
+    return "\n".join(lines[:n])
+
+
+def generate_from_ndjson(
+    ndjson_path: str,
+    provider: LLMProvider,
+    *,
+    launch_index=0,
+    out_py="repro.py",
+    execute=False,
+    retries: int = 0,
+    **gen_kwargs,
+) -> Dict[str, Any]:
+    bundle = build_context_bundle(ndjson_path, launch_index=launch_index)
+    # Augment bundle with pre-generated parameter allocation code to reduce LLM burden
+    allocation_snippet = generate_allocation_snippet(bundle)
+    kwargs_dict = generate_kwargs_dict(bundle)
+    context = {
+        **bundle,
+        "allocation_snippet": allocation_snippet,
+        "kwargs_dict": kwargs_dict,
+    }
+    system_prompt = render_prompt("system.txt", context)
+    user_prompt = render_prompt("generate_one_shot.txt", context)
+
+    code = provider.generate_code(system_prompt, user_prompt, **gen_kwargs)
+    Path(out_py).write_text(code, encoding="utf-8")
+
+    if not execute:
+        return {"path": out_py}
+
+    # Execute and optionally repair
+    rc, out, err = run_python(out_py)
+    attempt = 0
+    while rc != 0 and attempt < retries:
+        attempt += 1
+        # Build repair prompt
+        repair_ctx = {
+            "prev_code_excerpt": _excerpt(code, 200),
+            "error_text": err[-4000:] if err else "(no stderr)",
+        }
+        repair_prompt = render_prompt("repair_loop.txt", repair_ctx)
+        code = provider.generate_code(system_prompt, repair_prompt, **gen_kwargs)
+        Path(out_py).write_text(code, encoding="utf-8")
+        rc, out, err = run_python(out_py)
+
+    return {
+        "path": out_py,
+        "returncode": rc,
+        "stdout": out,
+        "stderr": err,
+        "retries_used": attempt,
+    }

tritonparse/reproducer/param_generator.py

@@ -0,0 +1,142 @@
+"""Parameter generator: produce deterministic allocation code from a bundle.
+
+This module reduces LLM burden by emitting Python code that:
+- selects a device
+- seeds RNG
+- allocates tensors with the exact shape/dtype/device/stride
+- prepares scalar/constexpr kwargs
+
+The generated code is intended to be inserted into the final repro script.
+"""
+
+import json
+from typing import Any, Dict, List, Optional
+
+
+def _torch_dtype_expr(dtype: str) -> str:
+    mapping = {
+        "float16": "torch.float16",
+        "bfloat16": "torch.bfloat16",
+        "float32": "torch.float32",
+        "float": "torch.float32",
+        "float64": "torch.float64",
+        "half": "torch.float16",
+        "bf16": "torch.bfloat16",
+        "fp16": "torch.float16",
+        "fp32": "torch.float32",
+        "fp64": "torch.float64",
+        "int8": "torch.int8",
+        "int16": "torch.int16",
+        "int32": "torch.int32",
+        "int64": "torch.int64",
+        "long": "torch.int64",
+        "bool": "torch.bool",
+    }
+    return mapping.get(str(dtype).lower(), "torch.float32")
+
+
+def _compute_storage_numel(shape: List[int], stride: Optional[List[int]]) -> int:
+    if not shape:
+        return 1
+    if not stride:
+        # contiguous default
+        numel = 1
+        for s in shape:
+            numel *= int(s)
+        return numel
+    # minimal storage size (in elements) to support the given logical shape/stride
+    max_index = 0
+    for dim, (sz, st) in enumerate(zip(shape, stride)):
+        if sz <= 0:
+            continue
+        max_index = max(max_index, (int(sz) - 1) * int(st))
+    return int(max_index) + 1
+
+
+def _emit_tensor_alloc(name: str, spec: Dict[str, Any]) -> str:
+    shape = spec.get("shape") or []
+    dtype = _torch_dtype_expr(spec.get("dtype"))
+    device = spec.get("device") or "cuda:0"
+    stride = spec.get("stride")
+
+    # ensure ints
+    shape = [int(s) for s in shape]
+    if stride is not None:
+        stride_list = [int(x) for x in stride]
+    else:
+        stride_list = None
+
+    lines: List[str] = []
+    # allocate backing storage
+    storage_numel = _compute_storage_numel(shape, stride_list)
+    lines.append(
+        f"# {name}: shape={shape}, dtype={dtype}, device={device}, stride={stride_list}"
+    )
+    lines.append(
+        f"_storage_{name} = torch.empty(({storage_numel},), dtype={dtype}, device=device)"
+    )
+    if stride_list:
+        # Create an as_strided view over the 1D storage
+        sizes_expr = str(tuple(shape))
+        strides_expr = str(tuple(stride_list))
+        lines.append(
+            f"{name} = _storage_{name}.as_strided(size={sizes_expr}, stride={strides_expr})"
+        )
+    else:
+        # contiguous allocation
+        size_expr = str(tuple(shape))
+        lines.append(f"{name} = torch.empty({size_expr}, dtype={dtype}, device=device)")
+    return "\n".join(lines)
+
+
+def _emit_scalar(name: str, spec: Dict[str, Any]) -> str:
+    value = spec.get("value")
+    # Preserve JSON-serializable value as-is
+    return f"{name} = {json.dumps(value)}"
+
+
+def generate_allocation_snippet(bundle: Dict[str, Any]) -> str:
+    """Generate a self-contained code snippet that:
+    - imports torch
+    - sets device
+    - seeds RNG
+    - allocates tensors and defines scalars for all args
+    Returns Python source as a string.
+    """
+    tensor_args: Dict[str, Any] = bundle.get("tensor_args", {}) or {}
+    args_all: Dict[str, Any] = bundle.get("args", {}) or {}
+
+    # Pick device from any tensor arg, fallback to cuda:0
+    device = "cuda:0"
+    for spec in tensor_args.values():
+        dev = spec.get("device")
+        if dev:
+            device = str(dev)
+            break
+
+    lines: List[str] = []
+    lines.append("import torch")
+    lines.append(f"device = '{device}'")
+    lines.append("torch.manual_seed(0)")
+    lines.append("if torch.cuda.is_available(): torch.cuda.manual_seed_all(0)")
+    lines.append("")
+
+    # Emit tensors first for names with type==tensor in args_all
+    for name, spec in args_all.items():
+        if isinstance(spec, dict) and spec.get("type") == "tensor":
+            lines.append(_emit_tensor_alloc(name, spec))
+            lines.append("")
+
+    # Emit non-tensor scalars next
+    for name, spec in args_all.items():
+        if not isinstance(spec, dict) or spec.get("type") == "tensor":
+            continue
+        lines.append(_emit_scalar(name, spec))
+    return "\n".join(lines)
+
+
+def generate_kwargs_dict(bundle: Dict[str, Any]) -> Dict[str, Any]:
+    """Return a kwargs dict derived from bundle['launch']['kwargs'] suitable for kernel call."""
+    launch = bundle.get("launch", {}) or {}
+    kwargs = launch.get("kwargs", {}) or {}
+    return kwargs

tritonparse/reproducer/prompts/__init__.py

@@ -0,0 +1 @@
+__all__ = []

tritonparse/reproducer/prompts/loader.py

@@ -0,0 +1,18 @@
+import json
+from pathlib import Path
+from typing import Any, Dict
+
+PROMPTS_DIR = Path(__file__).parent
+
+
+def render_prompt(name: str, context: Dict[str, Any]) -> str:
+    text = (PROMPTS_DIR / name).read_text(encoding="utf-8")
+    # very simple {{key}} replacement for top-level keys; JSON for dicts
+    for k, v in context.items():
+        token = "{{ " + k + " }}"
+        if token in text:
+            if isinstance(v, (dict, list)):
+                text = text.replace(token, json.dumps(v, ensure_ascii=False, indent=2))
+            else:
+                text = text.replace(token, str(v))
+    return text

tritonparse/reproducer/providers/__init__.py

@@ -0,0 +1 @@
+__all__ = []

tritonparse/reproducer/providers/base.py

@@ -0,0 +1,14 @@
+from typing import Any, Dict, List, Optional, Protocol
+
+
+class LLMProvider(Protocol):
+    def generate_code(
+        self,
+        system_prompt: str,
+        user_prompt: str,
+        *,
+        temperature: float = 0.2,
+        max_tokens: int = 8192,
+        stop: Optional[List[str]] = None,
+        extra: Optional[Dict[str, Any]] = None,
+    ) -> str: ...

tritonparse/reproducer/providers/gemini.py

@@ -0,0 +1,47 @@
+import os
+import re
+from typing import Any, Dict, List, Optional
+
+from google.genai import Client
+
+
+def _extract_python_block(s: str) -> str:
+    m = re.search(r"""```python\s+(.*?)```""", s, flags=re.S)
+    return m.group(1).strip() if m else ""
+
+
+class GeminiProvider:
+    def __init__(
+        self, project: str, location: str = "us-central1", model: str = "gemini-2.5-pro"
+    ):
+        # Expect GOOGLE_APPLICATIONS_CREDENTIALS to be set
+        if not os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
+            raise EnvironmentError("GOOGLE_APPLICATION_CREDENTIALS not set.")
+        self.client = Client(vertexai=True, project=project, location=location)
+        self.model = model
+
+    def generate_code(
+        self,
+        system_prompt: str,
+        user_prompt: str,
+        *,
+        temperature: float = 0.2,
+        max_tokens: int = 8192,
+        stop: Optional[List[str]] = None,
+        extra: Optional[Dict[str, Any]] = None,
+    ) -> str:
+        # Gemini doesn't have a 'system' role in this SDK, prepend system to user
+        full_prompt = f"{system_prompt.strip()}\n\n---\n\n{user_prompt.strip()}"
+        resp = self.client.models.generate_content(
+            model=self.model,
+            contents=full_prompt,
+            config={
+                "temperature": temperature,
+                "max_output_tokens": max_tokens,
+            },
+        )
+        text = getattr(resp, "text", "") or ""
+        code = _extract_python_block(text) or text
+        if not code.strip():
+            raise RuntimeError(f"Empty response from Gemini. Raw: {text[:2000]}")
+        return code

tritonparse/reproducer/runtime/__init__.py

@@ -0,0 +1 @@
+__all__ = []

tritonparse/reproducer/runtime/executor.py

@@ -0,0 +1,13 @@
+import subprocess
+import sys
+
+
+def run_python(path: str, timeout: int = 60):
+    p = subprocess.Popen(
+        [sys.executable, path],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+    out, err = p.communicate(timeout=timeout)
+    return p.returncode, out, err

tritonparse/reproducer/utils/io.py

@@ -0,0 +1,6 @@
+from pathlib import Path
+
+
+def write_text(path: str, content: str, *, encoding="utf-8"):
+    Path(path).parent.mkdir(parents=True, exist_ok=True)
+    Path(path).write_text(content, encoding=encoding)

tritonparse/shared_vars.py

@@ -0,0 +1,9 @@
+# We'd like to sperate structured logging module and tritonparse module as much as possible. So, put the shared variables here.
+import os
+
+# The compilation information will be stored to /logs/DEFAULT_TRACE_FILE_PREFIX by default
+# unless other flags disable or set another store. Add USER to avoid permission issues in shared servers.
+DEFAULT_TRACE_FILE_PREFIX = (
+    f"dedicated_log_triton_trace_{os.getenv('USER', 'unknown')}_"
+)
+DEFAULT_TRACE_FILE_PREFIX_WITHOUT_USER = "dedicated_log_triton_trace_"

tritonparse/source_type.py

@@ -0,0 +1,56 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+from enum import Enum
+from pathlib import Path
+from typing import Tuple
+
+
+class SourceType(str, Enum):
+    """Enumeration of supported source types for OSS only."""
+
+    LOCAL = "local"
+    LOCAL_FILE = "local_file"
+
+    @classmethod
+    def _missing_(cls, value: object) -> "SourceType":
+        """
+        Handle unknown source types by raising a ValueError.
+
+        Args:
+            value: The unknown value that was attempted to be used as a SourceType
+
+        Returns:
+            Never returns, always raises ValueError
+        """
+        valid_types = [e.value for e in cls]
+        raise ValueError(
+            f"Invalid source type '{value}'. Valid types are: {', '.join(valid_types)}"
+        )
+
+
+class Source:
+    """Represents a source of logs to parse."""
+
+    def __init__(self, source_str: str, verbose: bool = False):
+        """
+        Initialize a Source object by parsing the source string.
+
+        Args:
+            source_str: String representing the source
+            verbose: Whether to print verbose information
+        """
+        self.source_str = source_str
+        self.verbose = verbose
+        self.type, self.value = self._parse_source()
+
+    def _parse_source(self) -> Tuple[SourceType, str]:
+        # Check if it's a local path
+        path = Path(self.source_str)
+        if path.is_dir():
+            return SourceType.LOCAL, str(path.absolute())
+        elif path.is_file():
+            return SourceType.LOCAL_FILE, str(path.absolute())
+        else:
+            raise ValueError(
+                f"Source '{self.source_str}' is not a valid directory or file"
+            )

tritonparse/sourcemap_utils.py

@@ -0,0 +1,72 @@
+from typing import Any, Dict, List
+
+
+def get_file_extension(filename: str) -> str:
+    """
+    Get the file extension from a given filename or return the filename itself if it has no extension.
+
+    Args:
+        filename (str): The filename or file extension.
+
+    Returns:
+        str: The file extension or the filename itself if no extension is present.
+    """
+    # Split the filename by '.' and return the last part if it exists
+    parts = filename.split(".")
+    return parts[-1] if len(parts) > 1 else filename
+
+
+def _flatten_dict(
+    d: Dict[str, Any], parent_key: str = "", sep: str = "."
+) -> Dict[str, Any]:
+    """
+    Flattens a nested dictionary.
+    """
+    items = []
+    for k, v in d.items():
+        new_key = parent_key + sep + k if parent_key else k
+        if isinstance(v, dict):
+            items.extend(_flatten_dict(v, new_key, sep=sep).items())
+        else:
+            items.append((new_key, v))
+    return dict(items)
+
+
+def _unflatten_dict(d: Dict[str, Any], sep: str = ".") -> Dict[str, Any]:
+    """
+    Unflattens a dictionary with delimited keys.
+    """
+    result = {}
+    for key, value in d.items():
+        parts = key.split(sep)
+        d_ref = result
+        for part in parts[:-1]:
+            if part not in d_ref:
+                d_ref[part] = {}
+            d_ref = d_ref[part]
+        d_ref[parts[-1]] = value
+    return result
+
+
+def _to_ranges(indices: List[int]) -> List[Dict[str, int]]:
+    """
+    Converts a sorted list of indices into a list of continuous ranges.
+    e.g., [0, 1, 2, 5, 6, 8] -> [{'start': 0, 'end': 2}, {'start': 5, 'end': 6}, {'start': 8, 'end': 8}]
+    """
+    if not indices:
+        return []
+
+    indices = sorted(indices)
+    ranges = []
+    start = indices[0]
+    end = indices[0]
+
+    for i in range(1, len(indices)):
+        if indices[i] == end + 1:
+            end = indices[i]
+        else:
+            ranges.append({"start": start, "end": end})
+            start = end = indices[i]
+
+    ranges.append({"start": start, "end": end})
+    return ranges