cotlab 0.8.0__py3-none-any.whl

cotlab/backends/vllm_backend.py

@@ -0,0 +1,227 @@
+"""vLLM backend for high-throughput inference.
+
+Supports multiple platforms:
+- CUDA (NVIDIA GPUs) - standard vLLM
+- ROCm (AMD GPUs) - via ROCm Docker or HIP
+- Metal (Apple Silicon) - via vllm-metal plugin
+"""
+
+import os
+import platform as plat
+from typing import List, Optional
+
+import torch
+
+from ..core.base import GenerationOutput
+from ..core.registry import Registry
+from .base import InferenceBackend
+
+# Fix CUDA forking issue: vLLM's V1 engine uses multiprocessing, but when
+# Hydra/PyTorch have already initialized CUDA, forking fails. Setting spawn
+# method creates fresh Python processes instead.
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+
+def _is_apple_silicon() -> bool:
+    """Detect if running on Apple Silicon."""
+    return plat.system() == "Darwin" and plat.processor() == "arm"
+
+
+def _detect_platform() -> str:
+    """Detect the GPU platform.
+
+    Returns:
+        "metal" for Apple Silicon
+        "cuda" for NVIDIA/AMD GPUs (ROCm uses CUDA-compatible API)
+        "cpu" if no GPU available
+    """
+    if _is_apple_silicon():
+        return "metal"
+    if torch.cuda.is_available():
+        return "cuda"
+    return "cpu"
+
+
+@Registry.register_backend("vllm")
+class VLLMBackend(InferenceBackend):
+    """
+    High-throughput inference backend using vLLM.
+
+    Best for:
+    - Large-scale experiments (1000+ samples)
+    - Batch inference
+    - When activation access is not needed
+
+    Platforms:
+    - CUDA: Standard vLLM (pip install vllm)
+    - ROCm: vLLM in ROCm Docker container
+    - Metal: vLLM-Metal plugin (install via curl script)
+
+    Note:
+        Does NOT support activation extraction or patching.
+    """
+
+    def __init__(
+        self,
+        tensor_parallel_size: int = 1,
+        dtype: str = "bfloat16",
+        max_model_len: int | None = None,
+        trust_remote_code: bool = True,
+        quantization: str | None = None,
+        gpu_memory_utilization: float = 0.9,
+        enforce_eager: bool = False,
+        limit_mm_per_prompt: dict | str | None = None,
+        platform: str = "auto",
+        **kwargs,
+    ):
+        self.tensor_parallel_size = tensor_parallel_size
+        self.dtype = dtype
+        self.max_model_len = max_model_len
+        self.trust_remote_code = trust_remote_code
+        self.quantization = quantization
+        self.gpu_memory_utilization = gpu_memory_utilization
+        self.enforce_eager = enforce_eager
+        self.limit_mm_per_prompt = limit_mm_per_prompt
+        self._model = None
+        self._model_name: Optional[str] = None
+
+        # Platform detection
+        self._platform = _detect_platform() if platform == "auto" else platform
+        self._setup_platform_env()
+
+    def _setup_platform_env(self) -> None:
+        """Configure environment variables for the detected platform."""
+        if self._platform == "metal":
+            # vLLM-Metal configuration
+            os.environ.setdefault("VLLM_METAL_MEMORY_FRACTION", str(self.gpu_memory_utilization))
+            os.environ.setdefault("VLLM_METAL_USE_MLX", "1")
+            os.environ.setdefault("VLLM_METAL_BLOCK_SIZE", "16")
+            print("  Platform: Apple Silicon (Metal/MLX)")
+        elif self._platform == "cuda":
+            print("  Platform: CUDA")
+        else:
+            print(f"  Platform: {self._platform}")
+
+    def load_model(self, model_name: str, **kwargs) -> None:
+        """Load model with vLLM."""
+        try:
+            from vllm import LLM
+        except ImportError:
+            if self._platform == "metal":
+                raise ImportError(
+                    "vLLM not found. On Apple Silicon, install vllm-metal:\n"
+                    "curl -fsSL https://raw.githubusercontent.com/vllm-project/vllm-metal/main/install.sh | bash"
+                )
+            raise ImportError("vLLM not installed. Run: pip install vllm")
+
+        # Build LLM kwargs
+        llm_kwargs = {
+            "model": model_name,
+            "tensor_parallel_size": self.tensor_parallel_size,
+            "dtype": self.dtype,
+            "trust_remote_code": self.trust_remote_code,
+            "enforce_eager": self.enforce_eager,
+            **kwargs,
+        }
+
+        # gpu_memory_utilization only applies to CUDA/ROCm
+        if self._platform != "metal":
+            llm_kwargs["gpu_memory_utilization"] = self.gpu_memory_utilization
+
+        if self.max_model_len is not None:
+            llm_kwargs["max_model_len"] = self.max_model_len
+
+        if self.quantization is not None:
+            llm_kwargs["quantization"] = self.quantization
+
+        if self.limit_mm_per_prompt is not None:
+            llm_kwargs["limit_mm_per_prompt"] = self.limit_mm_per_prompt
+
+        print(f"DEBUG: VLLM args: {llm_kwargs}")
+        self._model = LLM(**llm_kwargs)
+        self._model_name = model_name
+
+    def generate(
+        self,
+        prompt: str,
+        max_new_tokens: int = 512,
+        temperature: float = 0.7,
+        top_p: float = 0.9,
+        system_prompt: Optional[str] = None,
+        **kwargs,
+    ) -> GenerationOutput:
+        """Generate from a single prompt."""
+        outputs = self.generate_batch(
+            [prompt],
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            system_prompt=system_prompt,
+            **kwargs,
+        )
+        return outputs[0]
+
+    def generate_batch(
+        self,
+        prompts: List[str],
+        max_new_tokens: int = 512,
+        temperature: float = 0.7,
+        top_p: float = 0.9,
+        system_prompt: Optional[str] = None,
+        **kwargs,
+    ) -> List[GenerationOutput]:
+        """Generate from multiple prompts efficiently."""
+        from vllm import SamplingParams
+
+        if self._model is None:
+            raise RuntimeError("Model not loaded. Call load_model() first.")
+
+        prompts = self._apply_system_prompt(prompts, system_prompt)
+
+        sampling_params = SamplingParams(
+            max_tokens=max_new_tokens, temperature=temperature, top_p=top_p, **kwargs
+        )
+
+        outputs = self._model.generate(prompts, sampling_params)
+
+        return [
+            GenerationOutput(
+                text=output.outputs[0].text,
+                tokens=list(output.outputs[0].token_ids),
+                logprobs=None,  # Can be enabled in SamplingParams if needed
+            )
+            for output in outputs
+        ]
+
+    @staticmethod
+    def _apply_system_prompt(prompts: List[str], system_prompt: Optional[str]) -> List[str]:
+        if not system_prompt:
+            return prompts
+        system_prompt = system_prompt.strip()
+        if not system_prompt:
+            return prompts
+        return [f"{system_prompt}\n\n{prompt}" for prompt in prompts]
+
+    @property
+    def platform(self) -> str:
+        """Return the detected platform (cuda, metal, cpu)."""
+        return self._platform
+
+    @property
+    def supports_activations(self) -> bool:
+        """vLLM optimizes away intermediate activations."""
+        return False
+
+    @property
+    def model_name(self) -> Optional[str]:
+        return self._model_name
+
+    def unload(self) -> None:
+        """Free GPU memory."""
+        if self._model is not None:
+            del self._model
+            self._model = None
+
+        # Platform-specific cleanup
+        if self._platform != "metal" and torch.cuda.is_available():
+            torch.cuda.empty_cache()

cotlab/cli.py

@@ -0,0 +1,83 @@
+"""CLI utilities for CoTLab."""
+
+import shutil
+from pathlib import Path
+
+import click
+
+
+@click.group()
+def cli():
+    """CoTLab command-line utilities."""
+    pass
+
+
+@cli.command()
+@click.argument("model_name")
+@click.option(
+    "--backend",
+    default="vllm",
+    type=click.Choice(["vllm", "transformers"]),
+    help="Backend: vllm or transformers",
+)
+@click.option("--output", "-o", help="Output path (default: conf/model/<safe_name>.yaml)")
+def template(model_name: str, backend: str, output: str):
+    """Generate model config from template.
+    Use this when you want to pre-create a model config file before the first run.
+    This is optional: CoTLab can also auto-generate a model config at runtime when
+    you run with `model=org/repo-id`.
+
+    Examples:
+
+        cotlab-template meta-llama/Llama-3.1-8B
+
+        cotlab-template google/gemma-3-12b --backend transformers
+
+        cotlab-template mistralai/Mistral-7B-v0.1 -o conf/model/mistral7b.yaml
+    """
+    # Get project root (where this file is located)
+    cli_file = Path(__file__)
+    project_root = cli_file.parent.parent.parent
+
+    # Get template
+    template_path = project_root / f"conf/model/_base/{backend}_default.yaml"
+
+    if not template_path.exists():
+        click.echo(f"Template not found: {template_path}", err=True)
+        raise click.Abort()
+
+    # Generate output path
+    if not output:
+        # Convert model name to safe filename
+        # meta-llama/Llama-3.1-8B -> meta_llama_llama_3_1_8b
+        safe_name = model_name.replace("/", "_").replace("-", "_").replace(".", "_").lower()
+        output = f"conf/model/{safe_name}.yaml"
+
+    output_path = Path(output)
+
+    # Check if exists
+    if output_path.exists():
+        if not click.confirm(f"{output} already exists. Overwrite?"):
+            click.echo("Aborted.")
+            raise click.Abort()
+
+    # Ensure parent directory exists
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Copy template
+    shutil.copy(template_path, output_path)
+
+    # Replace placeholder
+    content = output_path.read_text()
+    content = content.replace("huggingface/model-name", model_name)
+    output_path.write_text(content)
+
+    # Success message
+    click.echo(f"Created: {output}")
+    click.echo(f"Edit {output} to customize parameters")
+    click.echo("\nUsage:")
+    click.echo(f"  python -m cotlab.main model={output_path.stem}")
+
+
+if __name__ == "__main__":
+    cli()

cotlab/core/__init__.py

@@ -0,0 +1,34 @@
+"""Core module - base classes and configuration."""
+
+from .base import (
+    BaseExperiment,
+    BasePromptStrategy,
+    ExperimentResult,
+    GenerationOutput,
+    StructuredOutputMixin,
+)
+from .config import (
+    BackendConfig,
+    Config,
+    DatasetConfig,
+    ExperimentConfig,
+    ModelConfig,
+    PromptConfig,
+)
+from .registry import Registry, create_component
+
+__all__ = [
+    "GenerationOutput",
+    "ExperimentResult",
+    "BasePromptStrategy",
+    "BaseExperiment",
+    "StructuredOutputMixin",
+    "BackendConfig",
+    "ModelConfig",
+    "PromptConfig",
+    "DatasetConfig",
+    "ExperimentConfig",
+    "Config",
+    "Registry",
+    "create_component",
+]