wafer-core 0.1.26__py3-none-any.whl → 0.1.28__py3-none-any.whl

wafer_core/lib/trace_compare/same_kernel_analyzer.py

@@ -0,0 +1,119 @@
+"""Same kernel analysis - comparing identical kernel names across platforms.
+
+Identifies kernels where AMD and NVIDIA use the same kernel name/pattern
+and compares their performance directly.
+"""
+
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import Any
+
+from .aligner import KernelPair, LayerAlignment
+
+
+@dataclass
+class SameKernelComparison:
+    """Comparison of identical kernels across platforms."""
+
+    layer: int
+    kernel_name: str
+    operation: str
+    amd_avg_us: float
+    nvidia_avg_us: float
+    ratio: float
+    gap_us: float
+    amd_count: int
+    nvidia_count: int
+
+
+@dataclass
+class SameKernelAnalysis:
+    """Complete same kernel analysis result."""
+
+    kernels: list[SameKernelComparison] = field(default_factory=list)
+    summary: dict[str, Any] = field(default_factory=dict)
+
+
+def analyze_same_kernels(
+    layer_alignments: list[LayerAlignment],
+) -> SameKernelAnalysis:
+    """Find and compare kernels with identical names across platforms.
+
+    Args:
+        layer_alignments: List of aligned layers
+
+    Returns:
+        SameKernelAnalysis with comparisons
+    """
+    same_kernels: list[SameKernelComparison] = []
+
+    for layer_alignment in layer_alignments:
+        for pair in layer_alignment.kernel_pairs:
+            if pair.is_same_kernel and pair.amd_kernel and pair.nvidia_kernel:
+                same_kernels.append(
+                    SameKernelComparison(
+                        layer=layer_alignment.layer,
+                        kernel_name=pair.amd_kernel,
+                        operation=pair.operation,
+                        amd_avg_us=pair.amd_avg_us,
+                        nvidia_avg_us=pair.nvidia_avg_us,
+                        ratio=pair.ratio,
+                        gap_us=pair.gap_us,
+                        amd_count=pair.amd_count,
+                        nvidia_count=pair.nvidia_count,
+                    )
+                )
+
+    if same_kernels:
+        ratios = [k.ratio for k in same_kernels if k.ratio != float("inf")]
+        avg_ratio = sum(ratios) / len(ratios) if ratios else 1.0
+        amd_faster = sum(1 for k in same_kernels if k.ratio < 1.0)
+        nvidia_faster = sum(1 for k in same_kernels if k.ratio > 1.0)
+    else:
+        avg_ratio = 1.0
+        amd_faster = 0
+        nvidia_faster = 0
+
+    return SameKernelAnalysis(
+        kernels=same_kernels,
+        summary={
+            "total_same_kernels": len(same_kernels),
+            "avg_ratio": avg_ratio,
+            "kernels_where_amd_faster": amd_faster,
+            "kernels_where_nvidia_faster": nvidia_faster,
+        },
+    )
+
+
+def analyze_same_kernels_from_alignment(
+    layer_alignments: list[LayerAlignment],
+) -> dict[str, Any]:
+    """Analyze same kernels from alignment data (for API compatibility).
+
+    Args:
+        layer_alignments: List of aligned layers
+
+    Returns:
+        Dictionary with same kernel analysis results
+    """
+    analysis = analyze_same_kernels(layer_alignments)
+
+    kernels = [
+        {
+            "layer": k.layer,
+            "kernel_name": k.kernel_name,
+            "operation": k.operation,
+            "amd_avg_us": k.amd_avg_us,
+            "nvidia_avg_us": k.nvidia_avg_us,
+            "ratio": k.ratio,
+            "gap_us": k.gap_us,
+            "amd_count": k.amd_count,
+            "nvidia_count": k.nvidia_count,
+        }
+        for k in analysis.kernels
+    ]
+
+    return {
+        "kernels": kernels,
+        "summary": analysis.summary,
+    }

wafer_core/lib/trace_compare/warnings.py

@@ -0,0 +1,99 @@
+"""Warning detection and reporting for trace analysis.
+
+Detects issues with trace data quality and provides actionable suggestions.
+"""
+
+from dataclasses import dataclass
+from typing import Literal
+
+
+@dataclass(frozen=True)
+class TraceWarning:
+    """A warning about trace data quality or analysis limitations."""
+
+    code: str  # e.g., "NO_PHASE_ANNOTATIONS", "NO_PYTHON_STACKS"
+    severity: Literal["info", "warning", "error"]
+    message: str
+    suggestion: str
+
+
+def detect_warnings(
+    events: list[dict],
+    kernel_names: list[str],
+    phases: list[dict] | None = None,
+    layers_detected: int = 0,
+    total_kernels: int = 0,
+) -> list[TraceWarning]:
+    """Detect warnings from trace data.
+    
+    Args:
+        events: All trace events
+        kernel_names: List of all kernel names
+        phases: Optional list of phase events (for checking phase annotations)
+        layers_detected: Number of layers detected
+        total_kernels: Total number of kernels
+        
+    Returns:
+        List of warnings
+    """
+    warnings: list[TraceWarning] = []
+    
+    # Check for phase annotations
+    has_phase_annotations = any(
+        ev.get("cat") == "user_annotation" and ev.get("name", "").startswith("execute_context")
+        for ev in events
+    )
+    
+    if not has_phase_annotations:
+        warnings.append(
+            TraceWarning(
+                code="NO_PHASE_ANNOTATIONS",
+                severity="warning",
+                message="No vLLM phase annotations found. Phase analysis (prefill/decode) will be unavailable.",
+                suggestion="Ensure you're using vLLM v1.0+ with profiling enabled. Re-profile with torch.profiler.profile() to capture phase markers.",
+            )
+        )
+    
+    # Check for Python stack traces
+    has_python_stacks = any(
+        ev.get("cat") == "python_function"
+        for ev in events
+    )
+    
+    if not has_python_stacks:
+        warnings.append(
+            TraceWarning(
+                code="NO_PYTHON_STACKS",
+                severity="info",
+                message="No Python stack traces available. CPU→kernel mapping will be limited.",
+                suggestion="Re-profile with with_stack=True: torch.profiler.profile(with_stack=True) for better CPU operator identification.",
+            )
+        )
+    
+    # Check for high percentage of unknown kernels
+    if total_kernels > 0:
+        unknown_count = sum(1 for name in kernel_names if "unknown" in name.lower() or name == "Other")
+        unknown_percentage = (unknown_count / total_kernels) * 100
+        
+        if unknown_percentage > 20:
+            warnings.append(
+                TraceWarning(
+                    code="HIGH_UNKNOWN_KERNELS",
+                    severity="warning",
+                    message=f"{unknown_percentage:.1f}% of kernels are classified as 'Unknown'. Kernel registry may be outdated.",
+                    suggestion="Update kernel pattern registry or report unrecognized kernel patterns for support.",
+                )
+            )
+    
+    # Check for layer detection failure
+    if layers_detected == 0 and total_kernels > 100:
+        warnings.append(
+            TraceWarning(
+                code="LAYER_DETECTION_FAILED",
+                severity="warning",
+                message="No transformer layers detected. Layer-wise analysis unavailable.",
+                suggestion="This may indicate a non-transformer model (e.g., SSM/Mamba) or insufficient correlation data. Check model architecture.",
+            )
+        )
+    
+    return warnings

wafer_core/targets/__init__.py

@@ -1,5 +1,43 @@
-"""Re-export targets from utils for convenience."""
+"""Target system: specs (config) + targets (live resources) + reconciliation.
 
+New API (preferred):
+    from wafer_core.targets import Target, TargetSpec, ReconcileResult
+    from wafer_core.targets.providers import get_provider
+    from wafer_core.targets.reconcile import reconcile
+    from wafer_core.targets.spec_store import load_spec, list_spec_names
+    from wafer_core.targets.state_cache import get_binding_hints
+
+Legacy API (still works, will be deprecated):
+    from wafer_core.targets import RunPodTarget, runpod_ssh_context, ...
+"""
+
+# ── New types ────────────────────────────────────────────────────────────────
+from wafer_core.targets.digitalocean import (
+    DigitalOceanError,
+    DigitalOceanSSHInfo,
+    cleanup_all_droplets,
+    cleanup_digitalocean_target,
+    digitalocean_ssh_context,
+    get_droplet_state,
+    list_running_droplets,
+)
+from wafer_core.targets.runpod import (
+    RunPodError,
+    RunPodSSHInfo,
+    cleanup_all_pods,
+    cleanup_target,
+    get_pod_state,
+    list_running_pods,
+    runpod_ssh_context,
+)
+from wafer_core.targets.types import (
+    ReconcileResult,
+    Target,
+    TargetProvider,
+    TargetSpec,
+)
+
+# ── Legacy re-exports (unchanged, for backwards compatibility) ───────────────
 from wafer_core.utils.kernel_utils.targets import (
     BaremetalTarget,
     DigitalOceanTarget,
@@ -18,26 +56,14 @@ from wafer_core.utils.kernel_utils.targets import (
     select_target_for_operation,
     target_to_deployment_config,
 )
-from wafer_core.targets.runpod import (
-    RunPodError,
-    RunPodSSHInfo,
-    cleanup_all_pods,
-    cleanup_target,
-    get_pod_state,
-    list_running_pods,
-    runpod_ssh_context,
-)
-from wafer_core.targets.digitalocean import (
-    DigitalOceanError,
-    DigitalOceanSSHInfo,
-    cleanup_all_droplets,
-    cleanup_digitalocean_target,
-    digitalocean_ssh_context,
-    get_droplet_state,
-    list_running_droplets,
-)
 
 __all__ = [
+    # New API
+    "Target",
+    "TargetSpec",
+    "TargetProvider",
+    "ReconcileResult",
+    # Legacy: target config types
     "BaremetalTarget",
     "VMTarget",
     "ModalTarget",
@@ -54,7 +80,7 @@ __all__ = [
     "check_target_available",
     "find_free_gpu",
     "run_operation_on_target",
-    # RunPod provisioning
+    # Legacy: RunPod provisioning
     "RunPodError",
     "RunPodSSHInfo",
     "runpod_ssh_context",
@@ -62,7 +88,7 @@ __all__ = [
     "cleanup_all_pods",
     "list_running_pods",
     "get_pod_state",
-    # DigitalOcean provisioning
+    # Legacy: DigitalOcean provisioning
     "DigitalOceanError",
     "DigitalOceanSSHInfo",
     "digitalocean_ssh_context",

wafer_core/targets/pool.py

@@ -0,0 +1,181 @@
+"""Pool queries: filter live targets by GPU type, provider, and labels.
+
+A pool is a predicate over live targets, not a hardcoded list.
+Pool queries are defined in ~/.wafer/config.toml:
+
+    [pools.mi300x]
+    gpu_type = "MI300X"
+
+    [pools.mi300x-rocm7]
+    gpu_type = "MI300X"
+    labels.rocm_version = "7.0.2"
+
+    [pools.runpod-only]
+    provider = "runpod"
+
+Matching: a target matches a pool query if all specified fields match.
+Fields not specified in the query are ignored (match anything).
+Label matching is AND — all required labels must be present and equal.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from wafer_core.targets.types import Target
+
+WAFER_DIR = Path.home() / ".wafer"
+CONFIG_FILE = WAFER_DIR / "config.toml"
+
+# Fields on PoolQuery that map directly to Target fields
+_TARGET_FIELDS = ("gpu_type", "provider", "status")
+
+
+@dataclass(frozen=True)
+class PoolQuery:
+    """Predicate for filtering live targets.
+
+    All specified fields must match (AND semantics).
+    None means "don't care" for that field.
+    """
+
+    gpu_type: str | None = None
+    provider: str | None = None
+    status: str | None = "running"
+    labels: dict[str, str] = field(default_factory=dict)
+
+
+def match_targets(query: PoolQuery, targets: list[Target]) -> list[Target]:
+    """Filter targets that satisfy the pool query. Pure function."""
+    matched = []
+    for target in targets:
+        if not _matches(query, target):
+            continue
+        matched.append(target)
+    return matched
+
+
+def _matches(query: PoolQuery, target: Target) -> bool:
+    """Check if a single target satisfies the query."""
+    if query.gpu_type is not None and target.gpu_type != query.gpu_type:
+        return False
+    if query.provider is not None and target.provider != query.provider:
+        return False
+    if query.status is not None and target.status != query.status:
+        return False
+
+    # All required labels must be present and equal
+    for key, value in query.labels.items():
+        if target.labels.get(key) != value:
+            return False
+
+    return True
+
+
+def load_pool_query(name: str) -> PoolQuery:
+    """Load a pool query from ~/.wafer/config.toml.
+
+    Raises KeyError if the pool is not defined.
+    """
+    pools = _load_pools_section()
+    if name not in pools:
+        available = ", ".join(sorted(pools)) if pools else "(none)"
+        raise KeyError(f"Pool {name!r} not found. Available: {available}")
+
+    raw = pools[name]
+    assert isinstance(raw, dict), f"Pool {name!r} must be a table, got {type(raw).__name__}"
+
+    labels_raw = raw.get("labels", {})
+    assert isinstance(labels_raw, dict), (
+        f"Pool {name!r} labels must be a table, got {type(labels_raw).__name__}"
+    )
+
+    return PoolQuery(
+        gpu_type=raw.get("gpu_type"),
+        provider=raw.get("provider"),
+        status=raw.get("status", "running"),
+        labels={str(k): str(v) for k, v in labels_raw.items()},
+    )
+
+
+def list_pool_names() -> list[str]:
+    """List all pool names from config.toml."""
+    pools = _load_pools_section()
+    return sorted(pools.keys())
+
+
+def is_query_pool(name: str) -> bool:
+    """Check if a pool is defined as a PoolQuery (new format) vs target list (old format).
+
+    Old format: [pools.name] targets = ["t1", "t2"]
+    New format: [pools.name] gpu_type = "MI300X"
+
+    Returns False if pool doesn't exist or is old format.
+    """
+    pools = _load_pools_section()
+    if name not in pools:
+        return False
+    raw = pools[name]
+    if not isinstance(raw, dict):
+        return False
+    # Old format has a "targets" key with a list of names
+    return "targets" not in raw
+
+
+async def resolve_pool(name: str) -> list[Target]:
+    """Resolve a pool query to live targets.
+
+    Queries all cloud providers, hydrates cached labels, filters by pool query.
+    Returns matching Target objects sorted by resource_id for determinism.
+
+    Raises KeyError if pool not found.
+    """
+    from dataclasses import replace
+
+    from wafer_core.targets.providers import get_all_cloud_providers
+    from wafer_core.targets.state_cache import load_all_labels
+    from wafer_core.targets.types import TargetProvider
+
+    import trio
+
+    query = load_pool_query(name)
+
+    # Fetch all live targets
+    all_targets: list[Target] = []
+
+    async def _fetch(prov_impl: TargetProvider, results: list[Target]) -> None:
+        try:
+            targets = await prov_impl.list_targets()
+            results.extend(targets)
+        except Exception:
+            pass  # Skip providers that fail (missing API key, etc.)
+
+    async with trio.open_nursery() as nursery:
+        for _, prov_impl in get_all_cloud_providers():
+            nursery.start_soon(_fetch, prov_impl, all_targets)
+
+    # Hydrate labels from cache
+    cached_labels = load_all_labels()
+    all_targets = [
+        replace(t, labels=cached_labels[t.resource_id])
+        if t.resource_id in cached_labels
+        else t
+        for t in all_targets
+    ]
+
+    # Filter and sort
+    matched = match_targets(query, all_targets)
+    matched.sort(key=lambda t: t.resource_id)
+    return matched
+
+
+def _load_pools_section() -> dict:
+    """Read the [pools] section from config.toml. Returns empty dict if missing."""
+    if not CONFIG_FILE.exists():
+        return {}
+
+    import tomllib
+
+    data = tomllib.loads(CONFIG_FILE.read_text())
+    return data.get("pools", {})

wafer_core/targets/probe.py

@@ -0,0 +1,113 @@
+"""SSH probe: detect software labels on a live target.
+
+Runs a Python script on the target via SSH that reports installed
+software versions. Returns a flat dict[str, str] of labels.
+
+Only called at provision time or manually via `wafer targets probe`.
+Results are cached in target_state.json — probe is never implicit.
+
+Uses subprocess ssh (not asyncssh) to match existing codebase patterns.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import subprocess
+
+logger = logging.getLogger(__name__)
+
+# Probe script runs on the target machine via SSH.
+# Prints a JSON dict to stdout. Must work with stock Python 3.10+.
+_PROBE_SCRIPT = r"""
+import json, shutil, subprocess, sys
+
+def probe():
+    result = {}
+
+    # Python version
+    result["python_version"] = ".".join(map(str, sys.version_info[:2]))
+
+    # ROCm version from filesystem
+    try:
+        with open("/opt/rocm/.info/version") as f:
+            result["rocm_version"] = f.read().strip().split("-")[0]
+    except Exception:
+        pass
+
+    # CUDA version from nvcc
+    nvcc = shutil.which("nvcc")
+    if nvcc:
+        try:
+            out = subprocess.check_output([nvcc, "--version"], text=True)
+            for line in out.split("\n"):
+                if "release" in line.lower():
+                    parts = line.split("release")
+                    if len(parts) > 1:
+                        result["cuda_version"] = parts[1].split(",")[0].strip()
+                    break
+        except Exception:
+            pass
+
+    # PyTorch version
+    try:
+        import torch
+        result["pytorch_version"] = torch.__version__.split("+")[0]
+    except ImportError:
+        pass
+
+    # Triton version
+    try:
+        import triton
+        result["triton_version"] = triton.__version__
+    except ImportError:
+        pass
+
+    print(json.dumps(result))
+
+probe()
+"""
+
+
+def probe_target_labels(
+    host: str,
+    port: int,
+    username: str,
+    ssh_key_path: str | None = None,
+    timeout: int = 60,
+) -> dict[str, str]:
+    """SSH into a target and probe installed software. Returns labels dict.
+
+    Raises on SSH failure — caller decides how to handle.
+    """
+    ssh_args = [
+        "ssh",
+        "-p", str(port),
+        "-o", "StrictHostKeyChecking=no",
+        "-o", "UserKnownHostsFile=/dev/null",
+        "-o", "LogLevel=ERROR",
+        "-o", "ConnectTimeout=10",
+    ]
+    if ssh_key_path:
+        ssh_args.extend(["-i", ssh_key_path])
+
+    ssh_args.append(f"{username}@{host}")
+    ssh_args.append("python3")
+
+    result = subprocess.run(
+        ssh_args,
+        input=_PROBE_SCRIPT,
+        capture_output=True,
+        text=True,
+        timeout=timeout,
+    )
+
+    if result.returncode != 0:
+        stderr = result.stderr.strip()
+        raise RuntimeError(f"Probe failed (exit {result.returncode}): {stderr}")
+
+    stdout = result.stdout.strip()
+    labels = json.loads(stdout)
+    assert isinstance(labels, dict), f"Probe returned {type(labels).__name__}, expected dict"
+
+    return {str(k): str(v) for k, v in labels.items()}

wafer_core/targets/providers/__init__.py

@@ -0,0 +1,46 @@
+"""Provider registry for GPU resource management.
+
+Each provider implements TargetProvider protocol: list, get, provision, terminate.
+"""
+
+from __future__ import annotations
+
+from wafer_core.targets.providers.baremetal import BaremetalProvider
+from wafer_core.targets.providers.digitalocean import DigitalOceanProvider
+from wafer_core.targets.providers.runpod import RunPodProvider
+from wafer_core.targets.types import TargetProvider
+
+_PROVIDERS: dict[str, type] = {
+    "runpod": RunPodProvider,
+    "digitalocean": DigitalOceanProvider,
+    "baremetal": BaremetalProvider,
+}
+
+
+def get_provider(name: str) -> TargetProvider:
+    """Get a provider instance by name.
+
+    Raises KeyError if provider is not registered.
+    """
+    cls = _PROVIDERS.get(name)
+    if cls is None:
+        raise KeyError(f"Unknown provider: {name!r}. Available: {', '.join(sorted(_PROVIDERS))}")
+    return cls()
+
+
+def get_all_cloud_providers() -> list[tuple[str, TargetProvider]]:
+    """Get all cloud providers that can list remote resources.
+
+    Excludes baremetal (no remote API to query).
+    Returns list of (name, provider) tuples.
+    """
+    return [(name, cls()) for name, cls in _PROVIDERS.items() if name != "baremetal"]
+
+
+__all__ = [
+    "BaremetalProvider",
+    "DigitalOceanProvider",
+    "RunPodProvider",
+    "get_all_cloud_providers",
+    "get_provider",
+]