wafer-cli 0.2.14__py3-none-any.whl

wafer/config.py

@@ -0,0 +1,105 @@
+"""Configuration management for Wafer CLI.
+
+Immutable dataclasses for config with TOML parsing.
+"""
+
+import tomllib
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass(frozen=True)
+class WaferEnvironment:
+    """Docker environment configuration. Immutable."""
+
+    name: str
+    docker: str
+    description: str = ""
+
+    def __post_init__(self) -> None:
+        """Validate environment configuration."""
+        assert self.name, "environment name cannot be empty"
+        assert self.docker, "docker image cannot be empty"
+
+
+@dataclass(frozen=True)
+class WaferConfig:
+    """Wafer CLI configuration. Immutable."""
+
+    target: str
+    ssh_key: str
+    environments: dict[str, WaferEnvironment]
+    default_environment: str | None = None
+
+    def __post_init__(self) -> None:
+        """Validate configuration."""
+        assert self.target, "target cannot be empty"
+        assert self.ssh_key, "ssh_key cannot be empty"
+        assert self.environments, "at least one environment must be defined"
+
+        # Validate default_environment exists if specified
+        if self.default_environment:
+            assert (
+                self.default_environment in self.environments
+            ), f"default_environment '{self.default_environment}' not found in environments"
+
+    @classmethod
+    def from_toml(cls, path: Path) -> "WaferConfig":
+        """Parse config from TOML file.
+
+        Args:
+            path: Path to config file
+
+        Returns:
+            WaferConfig instance
+
+        Raises:
+            AssertionError: If config is invalid or missing required fields
+            FileNotFoundError: If config file doesn't exist
+
+        Example config file (~/.wafer/config.toml):
+
+            [default]
+            target = "root@b200:22"
+            ssh_key = "~/.ssh/id_ed25519"
+            environment = "cutlass"  # Optional default
+
+            [environments.cutlass]
+            docker = "nvcr.io/nvidia/cutlass:4.3-devel"
+            description = "CUDA 13 + Cutlass 4.3"
+
+            [environments.pytorch]
+            docker = "pytorch/pytorch:2.5-cuda12.4"
+            description = "PyTorch with CUDA 12.4"
+        """
+        assert path.exists(), f"Config file not found: {path}"
+
+        with open(path, "rb") as f:
+            data = tomllib.load(f)
+
+        # Validate required sections
+        assert "default" in data, "Config must have [default] section"
+        assert "target" in data["default"], "Config must have default.target"
+        assert "ssh_key" in data["default"], "Config must have default.ssh_key"
+
+        # Parse environments
+        environments = {}
+        env_data = data.get("environments", {})
+        assert env_data, "Config must have at least one environment defined"
+
+        for name, env_config in env_data.items():
+            assert isinstance(env_config, dict), f"Environment {name} must be a table/dict"
+            assert "docker" in env_config, f"Environment {name} must have docker image"
+
+            environments[name] = WaferEnvironment(
+                name=name,
+                docker=env_config["docker"],
+                description=env_config.get("description", ""),
+            )
+
+        return cls(
+            target=data["default"]["target"],
+            ssh_key=data["default"]["ssh_key"],
+            environments=environments,
+            default_environment=data["default"].get("environment"),
+        )

wafer/corpus.py

@@ -0,0 +1,366 @@
+"""Corpus management for Wafer CLI.
+
+Download and manage documentation corpora for agent filesystem access.
+"""
+
+import shutil
+import tarfile
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+from urllib.parse import urlparse
+
+import httpx
+
+CACHE_DIR = Path.home() / ".cache" / "wafer" / "corpora"
+
+CorpusName = Literal["cuda", "cutlass", "hip", "amd"]
+
+
+@dataclass
+class RepoSource:
+    """A single GitHub repo source within a corpus."""
+
+    repo: str
+    paths: list[str]
+    branch: str = "main"
+
+
+@dataclass
+class CorpusConfig:
+    """Configuration for a downloadable corpus."""
+
+    name: CorpusName
+    description: str
+    source_type: Literal["nvidia_md", "github_repo", "github_multi_repo"]
+    urls: list[str] | None = None
+    repo: str | None = None
+    repo_paths: list[str] | None = None
+    repos: list[RepoSource] | None = None  # For multi-repo corpora
+
+
+CORPORA: dict[CorpusName, CorpusConfig] = {
+    "cuda": CorpusConfig(
+        name="cuda",
+        description="CUDA Programming Guide and Best Practices",
+        source_type="nvidia_md",
+        urls=[
+            "https://docs.nvidia.com/cuda/cuda-programming-guide/index.html",
+            "https://docs.nvidia.com/cuda/cuda-programming-guide/01-introduction/introduction.html",
+            "https://docs.nvidia.com/cuda/cuda-programming-guide/01-introduction/programming-model.html",
+            "https://docs.nvidia.com/cuda/cuda-programming-guide/01-introduction/cuda-platform.html",
+            "https://docs.nvidia.com/cuda/cuda-programming-guide/02-basics/intro-to-cuda-cpp.html",
+            "https://docs.nvidia.com/cuda/cuda-programming-guide/02-basics/understanding-memory.html",
+            "https://docs.nvidia.com/cuda/cuda-programming-guide/02-basics/nvcc.html",
+            "https://docs.nvidia.com/cuda/cuda-programming-guide/03-advanced/advanced-kernel-programming.html",
+            "https://docs.nvidia.com/cuda/cuda-programming-guide/03-advanced/advanced-host-programming.html",
+            "https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/cuda-graphs.html",
+            "https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/stream-ordered-memory-allocation.html",
+            "https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/dynamic-parallelism.html",
+            "https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/virtual-memory-management.html",
+            "https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/cluster-launch-control.html",
+            "https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/graphics-interop.html",
+            "https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/cpp-language-extensions.html",
+            "https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/index.html",
+        ],
+    ),
+    "cutlass": CorpusConfig(
+        name="cutlass",
+        description="CUTLASS and CuTe DSL documentation",
+        source_type="github_repo",
+        repo="NVIDIA/cutlass",
+        repo_paths=["media/docs", "python/cutlass/docs"],
+    ),
+    "hip": CorpusConfig(
+        name="hip",
+        description="HIP programming guide and API reference",
+        source_type="github_repo",
+        repo="ROCm/HIP",
+        repo_paths=["docs"],
+    ),
+    "amd": CorpusConfig(
+        name="amd",
+        description="AMD GPU kernel development (rocWMMA, CK, AITER, rocBLAS, HipKittens, vLLM)",
+        source_type="github_multi_repo",
+        repos=[
+            # rocWMMA - wave matrix multiply-accumulate (WMMA) intrinsics
+            RepoSource(
+                repo="ROCm/rocWMMA",
+                paths=["docs", "samples", "library/include"],
+                branch="develop",
+            ),
+            # Composable Kernel - tile-based GPU programming
+            RepoSource(
+                repo="ROCm/composable_kernel",
+                paths=["docs", "example", "tutorial", "include/ck_tile"],
+                branch="develop",
+            ),
+            # AITER - AMD inference tensor runtime
+            RepoSource(
+                repo="ROCm/aiter",
+                paths=["docs", "aiter/ops"],
+            ),
+            # MIOpen - deep learning primitives (deprecated, use rocm-libraries)
+            RepoSource(
+                repo="ROCm/MIOpen",
+                paths=["docs"],
+                branch="develop_deprecated",
+            ),
+            # rocBLAS - BLAS library (deprecated, use rocm-libraries)
+            RepoSource(
+                repo="ROCm/rocBLAS",
+                paths=["docs"],
+                branch="develop_deprecated",
+            ),
+            # hipBLASLt - lightweight BLAS (deprecated, use rocm-libraries)
+            RepoSource(
+                repo="ROCm/hipBLASLt",
+                paths=["docs"],
+                branch="develop_deprecated",
+            ),
+            # Tensile - GEMM code generator (deprecated, use rocm-libraries)
+            RepoSource(
+                repo="ROCm/Tensile",
+                paths=["docs"],
+                branch="develop_deprecated",
+            ),
+            # HipKittens - high-performance AMD kernels
+            RepoSource(
+                repo="HazyResearch/HipKittens",
+                paths=["docs", "kernels", "include"],
+            ),
+            # vLLM AMD kernels
+            RepoSource(
+                repo="vllm-project/vllm",
+                paths=["csrc/rocm"],
+            ),
+            # SGLang AMD kernels
+            RepoSource(
+                repo="sgl-project/sglang",
+                paths=["3rdparty/amd"],
+            ),
+            # HuggingFace ROCm kernels
+            RepoSource(
+                repo="huggingface/hf-rocm-kernels",
+                paths=["csrc", "hf_rocm_kernels", "docs"],
+            ),
+        ],
+    ),
+}
+
+
+def _corpus_path(name: CorpusName) -> Path:
+    """Get local path for corpus."""
+    return CACHE_DIR / name
+
+
+def _ensure_cache_dir() -> None:
+    """Ensure cache directory exists."""
+    CACHE_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def _url_to_filepath(url: str, base_dir: Path) -> Path:
+    """Convert URL to local filepath preserving structure."""
+    parsed = urlparse(url)
+    path_parts = parsed.path.strip("/").split("/")
+    if path_parts[-1].endswith(".html"):
+        path_parts[-1] = path_parts[-1].replace(".html", ".md")
+    return base_dir / "/".join(path_parts)
+
+
+def _download_nvidia_md(config: CorpusConfig, dest: Path, verbose: bool = True) -> int:
+    """Download NVIDIA docs using .md endpoint."""
+    assert config.urls is not None
+    downloaded = 0
+    with httpx.Client(timeout=30.0, follow_redirects=True) as client:
+        for url in config.urls:
+            md_url = f"{url}.md"
+            filepath = _url_to_filepath(url, dest)
+            filepath.parent.mkdir(parents=True, exist_ok=True)
+            try:
+                resp = client.get(md_url)
+                resp.raise_for_status()
+                filepath.write_text(resp.text)
+                downloaded += 1
+                if verbose:
+                    print(f"  ✓ {filepath.relative_to(dest)}")
+            except httpx.HTTPError as e:
+                if verbose:
+                    print(f"  ✗ {url}: {e}")
+    return downloaded
+
+
+def _extract_matching_files(
+    tar: tarfile.TarFile,
+    repo_paths: list[str],
+    dest: Path,
+    verbose: bool,
+) -> int:
+    """Extract files matching repo_paths from tarball."""
+    downloaded = 0
+    for member in tar.getmembers():
+        if not member.isfile():
+            continue
+        rel_path = "/".join(member.name.split("/")[1:])
+        if not any(rel_path.startswith(rp) for rp in repo_paths):
+            continue
+        target = dest / rel_path
+        target.parent.mkdir(parents=True, exist_ok=True)
+        src = tar.extractfile(member)
+        if src:
+            target.write_bytes(src.read())
+            downloaded += 1
+            if verbose:
+                print(f"  ✓ {rel_path}")
+    return downloaded
+
+
+def _download_single_github_repo(
+    client: httpx.Client,
+    repo: str,
+    repo_paths: list[str],
+    dest: Path,
+    branch: str = "main",
+    verbose: bool = True,
+) -> int:
+    """Download specific paths from a single GitHub repo."""
+    tarball_url = f"https://api.github.com/repos/{repo}/tarball/{branch}"
+    if verbose:
+        print(f"  Fetching {repo}...")
+    resp = client.get(tarball_url)
+    resp.raise_for_status()
+    with tempfile.NamedTemporaryFile(suffix=".tar.gz", delete=False) as tmp:
+        tmp.write(resp.content)
+        tmp_path = Path(tmp.name)
+    try:
+        with tarfile.open(tmp_path, "r:gz") as tar:
+            return _extract_matching_files(tar, repo_paths, dest, verbose)
+    finally:
+        tmp_path.unlink()
+
+
+def _download_github_repo(config: CorpusConfig, dest: Path, verbose: bool = True) -> int:
+    """Download specific paths from GitHub repo."""
+    assert config.repo is not None
+    assert config.repo_paths is not None
+    with httpx.Client(timeout=60.0, follow_redirects=True) as client:
+        return _download_single_github_repo(
+            client, config.repo, config.repo_paths, dest, verbose=verbose
+        )
+
+
+def _download_github_multi_repo(config: CorpusConfig, dest: Path, verbose: bool = True) -> int:
+    """Download specific paths from multiple GitHub repos."""
+    assert config.repos is not None
+    downloaded = 0
+    with httpx.Client(timeout=120.0, follow_redirects=True) as client:
+        for repo_source in config.repos:
+            repo_name = repo_source.repo.split("/")[-1]
+            repo_dest = dest / repo_name
+            repo_dest.mkdir(parents=True, exist_ok=True)
+            try:
+                count = _download_single_github_repo(
+                    client,
+                    repo_source.repo,
+                    repo_source.paths,
+                    repo_dest,
+                    branch=repo_source.branch,
+                    verbose=verbose,
+                )
+                downloaded += count
+            except httpx.HTTPError as e:
+                if verbose:
+                    print(f"  ✗ {repo_source.repo}: {e}")
+    return downloaded
+
+
+def download_corpus(name: CorpusName, force: bool = False, verbose: bool = True) -> Path:
+    """Download a corpus to local cache.
+
+    Args:
+        name: Corpus name
+        force: Re-download even if exists
+        verbose: Print progress
+
+    Returns:
+        Path to downloaded corpus
+
+    Raises:
+        ValueError: If corpus name is unknown
+        httpx.HTTPError: If download fails
+    """
+    if name not in CORPORA:
+        raise ValueError(f"Unknown corpus: {name}. Available: {list(CORPORA.keys())}")
+    config = CORPORA[name]
+    dest = _corpus_path(name)
+    if dest.exists() and not force:
+        if verbose:
+            print(f"Corpus '{name}' already exists at {dest}")
+            print("Use --force to re-download")
+        return dest
+    _ensure_cache_dir()
+    if dest.exists():
+        shutil.rmtree(dest)
+    dest.mkdir(parents=True)
+    if verbose:
+        print(f"Downloading {name}: {config.description}")
+    if config.source_type == "nvidia_md":
+        count = _download_nvidia_md(config, dest, verbose)
+    elif config.source_type == "github_repo":
+        count = _download_github_repo(config, dest, verbose)
+    elif config.source_type == "github_multi_repo":
+        count = _download_github_multi_repo(config, dest, verbose)
+    else:
+        raise ValueError(f"Unknown source type: {config.source_type}")
+    if verbose:
+        print(f"Downloaded {count} files to {dest}")
+    return dest
+
+
+def sync_corpus(name: CorpusName, verbose: bool = True) -> Path:
+    """Sync (re-download) a corpus.
+
+    Args:
+        name: Corpus name
+        verbose: Print progress
+
+    Returns:
+        Path to synced corpus
+    """
+    return download_corpus(name, force=True, verbose=verbose)
+
+
+def list_corpora(verbose: bool = True) -> dict[CorpusName, bool]:
+    """List available corpora and their download status.
+
+    Returns:
+        Dict of corpus name -> is_downloaded
+    """
+    result: dict[CorpusName, bool] = {}
+    for name, config in CORPORA.items():
+        path = _corpus_path(name)
+        exists = path.exists()
+        result[name] = exists
+        if verbose:
+            status = "✓" if exists else " "
+            print(f"[{status}] {name}: {config.description}")
+            if exists:
+                file_count = sum(1 for _ in path.rglob("*") if _.is_file())
+                print(f"    {path} ({file_count} files)")
+    return result
+
+
+def get_corpus_path(name: CorpusName) -> Path | None:
+    """Get path to downloaded corpus, or None if not downloaded.
+
+    Args:
+        name: Corpus name
+
+    Returns:
+        Path if downloaded, None otherwise
+    """
+    if name not in CORPORA:
+        return None
+    path = _corpus_path(name)
+    return path if path.exists() else None