mcpbr 0.4.15__py3-none-any.whl → 0.5.0__py3-none-any.whl

mcpbr/dataset_streaming.py

@@ -0,0 +1,491 @@
+"""Memory-efficient large dataset handling for benchmark evaluations.
+
+This module provides streaming and chunked loading of large HuggingFace datasets,
+enabling benchmark runs on datasets that would otherwise exceed available memory.
+It includes memory monitoring, automatic chunking under memory pressure, and
+iterator-based APIs compatible with existing benchmark ``load_tasks`` patterns.
+
+Key components:
+- ``MemoryMonitor``: Tracks RSS and available memory, detects memory pressure.
+- ``ChunkedLoader``: Iterates over a HuggingFace dataset in configurable chunks.
+- ``StreamingDataset``: High-level API that yields tasks lazily with memory awareness.
+- ``DatasetStats``: Summary statistics for a streaming load session.
+"""
+
+import logging
+import os
+import time
+from collections.abc import Iterator
+from dataclasses import dataclass
+from typing import Any
+
+from datasets import load_dataset, load_dataset_builder
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# DatasetStats
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class DatasetStats:
+    """Summary statistics for a streaming dataset load session.
+
+    Attributes:
+        total_loaded: Number of individual task items yielded so far.
+        peak_memory_mb: Peak RSS observed during loading (in megabytes).
+        chunks_processed: Number of chunks fetched from the underlying loader.
+        load_time_seconds: Wall-clock seconds elapsed during loading.
+    """
+
+    total_loaded: int = 0
+    peak_memory_mb: float = 0.0
+    chunks_processed: int = 0
+    load_time_seconds: float = 0.0
+
+
+# ---------------------------------------------------------------------------
+# MemoryMonitor
+# ---------------------------------------------------------------------------
+
+
+class MemoryMonitor:
+    """Lightweight monitor for process and system memory usage.
+
+    Uses ``psutil`` when available, falling back to reading ``/proc/self/status``
+    and ``/proc/meminfo`` on Linux. On platforms where neither is available the
+    methods return ``0.0`` and memory-pressure detection is disabled.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the memory monitor and detect available backends."""
+        self._has_psutil = False
+        try:
+            import psutil  # noqa: F401
+
+            self._has_psutil = True
+        except ImportError:
+            pass
+
+    # -- public API ---------------------------------------------------------
+
+    def get_memory_usage_mb(self) -> float:
+        """Return the current Resident Set Size (RSS) in megabytes.
+
+        Returns:
+            RSS in MB, or ``0.0`` if measurement is unavailable.
+        """
+        if self._has_psutil:
+            return self._rss_via_psutil()
+        return self._rss_via_proc()
+
+    def get_available_memory_mb(self) -> float:
+        """Return available system memory in megabytes.
+
+        Returns:
+            Available memory in MB, or ``0.0`` if measurement is unavailable.
+        """
+        if self._has_psutil:
+            return self._available_via_psutil()
+        return self._available_via_proc()
+
+    def is_memory_pressure(self, threshold_pct: float = 80.0) -> bool:
+        """Check whether system memory usage exceeds a threshold.
+
+        Args:
+            threshold_pct: Percentage (0--100) of total memory above which
+                the system is considered under pressure.
+
+        Returns:
+            ``True`` if memory usage exceeds *threshold_pct*, ``False``
+            otherwise or if measurement is unavailable.
+        """
+        if self._has_psutil:
+            return self._pressure_via_psutil(threshold_pct)
+        return self._pressure_via_proc(threshold_pct)
+
+    # -- psutil backend -----------------------------------------------------
+
+    def _rss_via_psutil(self) -> float:
+        """Get RSS using psutil."""
+        try:
+            import psutil
+
+            process = psutil.Process(os.getpid())
+            return process.memory_info().rss / (1024 * 1024)
+        except Exception:
+            return 0.0
+
+    def _available_via_psutil(self) -> float:
+        """Get available system memory using psutil."""
+        try:
+            import psutil
+
+            return psutil.virtual_memory().available / (1024 * 1024)
+        except Exception:
+            return 0.0
+
+    def _pressure_via_psutil(self, threshold_pct: float) -> bool:
+        """Check memory pressure using psutil."""
+        try:
+            import psutil
+
+            return psutil.virtual_memory().percent >= threshold_pct
+        except Exception:
+            return False
+
+    # -- /proc fallback -----------------------------------------------------
+
+    @staticmethod
+    def _rss_via_proc() -> float:
+        """Get RSS by parsing ``/proc/self/status``."""
+        try:
+            with open("/proc/self/status") as fh:
+                for line in fh:
+                    if line.startswith("VmRSS:"):
+                        # Value is in kB
+                        return int(line.split()[1]) / 1024
+        except (OSError, ValueError, IndexError):
+            pass
+        return 0.0
+
+    @staticmethod
+    def _available_via_proc() -> float:
+        """Get available memory by parsing ``/proc/meminfo``."""
+        try:
+            with open("/proc/meminfo") as fh:
+                for line in fh:
+                    if line.startswith("MemAvailable:"):
+                        return int(line.split()[1]) / 1024
+        except (OSError, ValueError, IndexError):
+            pass
+        return 0.0
+
+    @staticmethod
+    def _pressure_via_proc(threshold_pct: float) -> bool:
+        """Check memory pressure using ``/proc/meminfo``."""
+        try:
+            mem_total = 0.0
+            mem_available = 0.0
+            with open("/proc/meminfo") as fh:
+                for line in fh:
+                    if line.startswith("MemTotal:"):
+                        mem_total = int(line.split()[1]) / 1024
+                    elif line.startswith("MemAvailable:"):
+                        mem_available = int(line.split()[1]) / 1024
+            if mem_total > 0:
+                used_pct = ((mem_total - mem_available) / mem_total) * 100
+                return used_pct >= threshold_pct
+        except (OSError, ValueError, IndexError):
+            pass
+        return False
+
+
+# ---------------------------------------------------------------------------
+# ChunkedLoader
+# ---------------------------------------------------------------------------
+
+
+class ChunkedLoader:
+    """Iterate over a HuggingFace dataset in fixed-size chunks.
+
+    Each iteration yields a ``list[dict]`` containing up to *chunk_size*
+    records. When the HuggingFace ``datasets`` library supports it, the
+    dataset is loaded with ``streaming=True`` to avoid downloading the
+    entire dataset at once.
+
+    Args:
+        dataset_name: HuggingFace dataset identifier (e.g. ``"openai_humaneval"``).
+        split: Dataset split to load (default ``"test"``).
+        chunk_size: Maximum number of records per chunk.
+        subset: Optional dataset subset / configuration name.
+    """
+
+    def __init__(
+        self,
+        dataset_name: str,
+        split: str = "test",
+        chunk_size: int = 1000,
+        subset: str | None = None,
+    ) -> None:
+        self.dataset_name = dataset_name
+        self.split = split
+        self.chunk_size = chunk_size
+        self.subset = subset
+        self._total_items: int | None = None
+
+    # -- public API ---------------------------------------------------------
+
+    def __len__(self) -> int:
+        """Return the total number of items in the dataset.
+
+        This may trigger a metadata fetch the first time it is called.
+
+        Returns:
+            Total number of items, or ``0`` if the count cannot be determined.
+        """
+        if self._total_items is None:
+            self._total_items = self._fetch_total_items()
+        return self._total_items
+
+    def __iter__(self) -> Iterator[list[dict[str, Any]]]:
+        """Yield successive chunks of dataset records.
+
+        Yields:
+            Lists of up to *chunk_size* task dictionaries.
+        """
+        dataset_iter = self._load_dataset_streaming()
+
+        chunk: list[dict[str, Any]] = []
+        for item in dataset_iter:
+            chunk.append(dict(item))
+            if len(chunk) >= self.chunk_size:
+                yield chunk
+                chunk = []
+
+        # Yield any remaining items
+        if chunk:
+            yield chunk
+
+    # -- internal helpers ---------------------------------------------------
+
+    def _load_dataset_streaming(self) -> Any:
+        """Load the dataset, preferring streaming mode.
+
+        Returns:
+            An iterable of dataset records (either a streaming
+            ``IterableDataset`` or a regular ``Dataset``).
+        """
+        load_kwargs: dict[str, Any] = {}
+        if self.subset is not None:
+            load_kwargs["name"] = self.subset
+
+        # Try streaming first for memory efficiency
+        try:
+            ds = load_dataset(
+                self.dataset_name,
+                split=self.split,
+                streaming=True,
+                **load_kwargs,
+            )
+            logger.info(
+                "Loaded dataset %s (split=%s) in streaming mode",
+                self.dataset_name,
+                self.split,
+            )
+            return ds
+        except Exception:
+            logger.debug(
+                "Streaming not supported for %s; falling back to full load",
+                self.dataset_name,
+            )
+
+        # Fallback: load the full dataset into memory
+        ds = load_dataset(
+            self.dataset_name,
+            split=self.split,
+            **load_kwargs,
+        )
+        logger.info(
+            "Loaded dataset %s (split=%s) fully into memory (%d items)",
+            self.dataset_name,
+            self.split,
+            len(ds),
+        )
+        return ds
+
+    def _fetch_total_items(self) -> int:
+        """Fetch the total item count from dataset metadata.
+
+        Returns:
+            The number of items, or ``0`` if it cannot be determined.
+        """
+        try:
+            load_kwargs: dict[str, Any] = {}
+            if self.subset is not None:
+                load_kwargs["name"] = self.subset
+
+            builder = load_dataset_builder(self.dataset_name, **load_kwargs)
+            info = builder.info
+            if info.splits and self.split in info.splits:
+                return info.splits[self.split].num_examples
+        except Exception:
+            logger.debug(
+                "Could not determine total items for %s/%s",
+                self.dataset_name,
+                self.split,
+            )
+        return 0
+
+
+# ---------------------------------------------------------------------------
+# StreamingDataset
+# ---------------------------------------------------------------------------
+
+
+class StreamingDataset:
+    """High-level memory-aware dataset loader.
+
+    Wraps :class:`ChunkedLoader` with automatic memory monitoring and
+    adaptive chunk sizing. Provides an iterator-based ``load_tasks`` method
+    compatible with the existing :class:`~mcpbr.benchmarks.base.Benchmark`
+    protocol (callers can materialise with ``list(...)`` when needed).
+
+    Args:
+        dataset_name: HuggingFace dataset identifier.
+        split: Dataset split (default ``"test"``).
+        max_memory_mb: Optional soft memory cap. When the process RSS exceeds
+            this value the chunk size is halved to reduce pressure.
+    """
+
+    # Default and minimum chunk sizes
+    _DEFAULT_CHUNK_SIZE = 1000
+    _MIN_CHUNK_SIZE = 50
+
+    def __init__(
+        self,
+        dataset_name: str,
+        split: str = "test",
+        max_memory_mb: float | None = None,
+    ) -> None:
+        self.dataset_name = dataset_name
+        self.split = split
+        self.max_memory_mb = max_memory_mb
+
+        self._monitor = MemoryMonitor()
+        self._stats = DatasetStats()
+        self._chunk_size = self._DEFAULT_CHUNK_SIZE
+        self._start_time: float | None = None
+
+    # -- public API ---------------------------------------------------------
+
+    def load_tasks(
+        self,
+        sample_size: int | None = None,
+        task_ids: list[str] | None = None,
+    ) -> Iterator[dict[str, Any]]:
+        """Lazily yield task dictionaries from the dataset.
+
+        Args:
+            sample_size: Maximum number of tasks to yield (``None`` for all).
+            task_ids: If provided, only yield tasks whose ``instance_id`` or
+                ``task_id`` is in this set.
+
+        Yields:
+            Individual task dictionaries.
+        """
+        self._start_time = time.monotonic()
+        self._stats = DatasetStats()
+
+        if sample_size is not None and sample_size <= 0:
+            self._stats.load_time_seconds = time.monotonic() - self._start_time
+            return
+
+        task_id_set: set[str] | None = set(task_ids) if task_ids else None
+
+        loader = ChunkedLoader(
+            dataset_name=self.dataset_name,
+            split=self.split,
+            chunk_size=self._chunk_size,
+        )
+
+        yielded = 0
+
+        for chunk in loader:
+            self._stats.chunks_processed += 1
+
+            # Adapt chunk size under memory pressure
+            self._maybe_adapt_chunk_size(loader)
+
+            for item in chunk:
+                # Apply task_id filter
+                if task_id_set is not None:
+                    item_id = item.get("instance_id") or item.get("task_id")
+                    if item_id not in task_id_set:
+                        continue
+
+                self._stats.total_loaded += 1
+                yielded += 1
+
+                # Track peak memory
+                current_mb = self._monitor.get_memory_usage_mb()
+                if current_mb > self._stats.peak_memory_mb:
+                    self._stats.peak_memory_mb = current_mb
+
+                yield item
+
+                if sample_size is not None and yielded >= sample_size:
+                    self._stats.load_time_seconds = time.monotonic() - self._start_time
+                    return
+
+        self._stats.load_time_seconds = time.monotonic() - self._start_time
+
+    def get_stats(self) -> DatasetStats:
+        """Return statistics collected during the most recent ``load_tasks`` call.
+
+        Returns:
+            A :class:`DatasetStats` instance with current metrics.
+        """
+        # Update load_time if still in progress
+        if self._start_time is not None and self._stats.load_time_seconds == 0.0:
+            self._stats.load_time_seconds = time.monotonic() - self._start_time
+        return self._stats
+
+    # -- internal helpers ---------------------------------------------------
+
+    def _maybe_adapt_chunk_size(self, loader: ChunkedLoader) -> None:
+        """Reduce the chunk size if memory pressure is detected.
+
+        Args:
+            loader: The active :class:`ChunkedLoader` whose chunk size will be
+                updated in place.
+        """
+        under_pressure = False
+
+        if self.max_memory_mb is not None:
+            current_mb = self._monitor.get_memory_usage_mb()
+            if current_mb > self.max_memory_mb:
+                under_pressure = True
+                logger.warning(
+                    "RSS %.1f MB exceeds max_memory_mb %.1f MB; reducing chunk size",
+                    current_mb,
+                    self.max_memory_mb,
+                )
+
+        if not under_pressure and self._monitor.is_memory_pressure():
+            under_pressure = True
+            logger.warning("System memory pressure detected; reducing chunk size")
+
+        if under_pressure and loader.chunk_size > self._MIN_CHUNK_SIZE:
+            new_size = max(loader.chunk_size // 2, self._MIN_CHUNK_SIZE)
+            logger.info("Chunk size reduced from %d to %d", loader.chunk_size, new_size)
+            loader.chunk_size = new_size
+
+
+# ---------------------------------------------------------------------------
+# Convenience helpers
+# ---------------------------------------------------------------------------
+
+
+def get_memory_usage_mb() -> float:
+    """Return current process RSS in megabytes.
+
+    Convenience wrapper around :meth:`MemoryMonitor.get_memory_usage_mb`.
+
+    Returns:
+        RSS in MB, or ``0.0`` if measurement is unavailable.
+    """
+    return MemoryMonitor().get_memory_usage_mb()
+
+
+def get_available_memory_mb() -> float:
+    """Return available system memory in megabytes.
+
+    Convenience wrapper around :meth:`MemoryMonitor.get_available_memory_mb`.
+
+    Returns:
+        Available memory in MB, or ``0.0`` if measurement is unavailable.
+    """
+    return MemoryMonitor().get_available_memory_mb()

mcpbr/dataset_versioning.py

@@ -0,0 +1,222 @@
+"""Dataset versioning for reproducible benchmark evaluations.
+
+This module provides utilities to pin and track HuggingFace dataset versions,
+ensuring that benchmark runs can be reproduced with the exact same data.
+Version information includes dataset revision hashes, download timestamps,
+and optional checksums for data integrity verification.
+"""
+
+import hashlib
+import json
+import logging
+from dataclasses import asdict, dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from datasets import Dataset, load_dataset
+from huggingface_hub import dataset_info
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class DatasetVersion:
+    """Pinned version information for a HuggingFace dataset.
+
+    Attributes:
+        dataset_id: HuggingFace dataset identifier (e.g., 'SWE-bench/SWE-bench_Lite').
+        revision: Git revision hash of the dataset (None for latest).
+        download_date: ISO 8601 timestamp of when the version was pinned.
+        checksum: Optional SHA256 checksum of the dataset content for integrity verification.
+    """
+
+    dataset_id: str
+    revision: str | None
+    download_date: str
+    checksum: str | None
+
+
+def pin_dataset_version(
+    dataset_id: str,
+    revision: str | None = None,
+) -> DatasetVersion:
+    """Record the current version of a HuggingFace dataset.
+
+    Fetches dataset metadata from the HuggingFace Hub to determine the
+    current revision. If a specific revision is provided, it is used directly.
+
+    Args:
+        dataset_id: HuggingFace dataset identifier (e.g., 'SWE-bench/SWE-bench_Lite').
+        revision: Specific git revision to pin. If None, the latest revision is fetched.
+
+    Returns:
+        DatasetVersion with the pinned revision and metadata.
+
+    Raises:
+        Exception: If the dataset cannot be found or accessed on the HuggingFace Hub.
+    """
+    info = dataset_info(dataset_id, revision=revision)
+    resolved_revision = info.sha
+
+    # Compute a checksum from the dataset card and file metadata for integrity
+    checksum_data = f"{dataset_id}:{resolved_revision}"
+    if info.siblings:
+        file_names = sorted(s.rfilename for s in info.siblings)
+        checksum_data += ":" + ",".join(file_names)
+    checksum = hashlib.sha256(checksum_data.encode()).hexdigest()
+
+    download_date = datetime.now(timezone.utc).isoformat()
+
+    version = DatasetVersion(
+        dataset_id=dataset_id,
+        revision=resolved_revision,
+        download_date=download_date,
+        checksum=checksum,
+    )
+
+    logger.info(
+        "Pinned dataset %s at revision %s",
+        dataset_id,
+        resolved_revision,
+    )
+
+    return version
+
+
+def load_dataset_pinned(
+    dataset_id: str,
+    version: DatasetVersion | None = None,
+    **kwargs: Any,
+) -> Dataset:
+    """Load a HuggingFace dataset using a pinned version for reproducibility.
+
+    Wraps the standard ``datasets.load_dataset`` call, injecting the pinned
+    revision so that the exact same data snapshot is used across runs.
+
+    Args:
+        dataset_id: HuggingFace dataset identifier.
+        version: Pinned version to use. If None, loads the latest version.
+        **kwargs: Additional keyword arguments passed to ``datasets.load_dataset``
+            (e.g., split, name, streaming).
+
+    Returns:
+        The loaded HuggingFace Dataset.
+    """
+    revision = None
+    if version is not None:
+        revision = version.revision
+        logger.info(
+            "Loading dataset %s at pinned revision %s (pinned on %s)",
+            dataset_id,
+            revision,
+            version.download_date,
+        )
+    else:
+        logger.info("Loading dataset %s at latest revision", dataset_id)
+
+    return load_dataset(dataset_id, revision=revision, **kwargs)
+
+
+def save_version_manifest(
+    versions: dict[str, DatasetVersion],
+    path: Path,
+) -> None:
+    """Save dataset version pins to a JSON manifest file.
+
+    The manifest file records all pinned dataset versions so they can be
+    shared across team members or CI environments for reproducible runs.
+
+    Args:
+        versions: Mapping of dataset identifiers to their pinned versions.
+        path: File path to write the JSON manifest.
+    """
+    manifest: dict[str, Any] = {
+        "format_version": "1.0",
+        "created_at": datetime.now(timezone.utc).isoformat(),
+        "datasets": {},
+    }
+
+    for dataset_id, version in versions.items():
+        manifest["datasets"][dataset_id] = asdict(version)
+
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    with open(path, "w") as f:
+        json.dump(manifest, f, indent=2)
+
+    logger.info("Saved version manifest with %d datasets to %s", len(versions), path)
+
+
+def load_version_manifest(path: Path) -> dict[str, DatasetVersion]:
+    """Load pinned dataset versions from a JSON manifest file.
+
+    Args:
+        path: File path to the JSON manifest.
+
+    Returns:
+        Mapping of dataset identifiers to their pinned versions.
+
+    Raises:
+        FileNotFoundError: If the manifest file does not exist.
+        json.JSONDecodeError: If the manifest file contains invalid JSON.
+        KeyError: If the manifest is missing required fields.
+    """
+    with open(path) as f:
+        manifest = json.load(f)
+
+    versions: dict[str, DatasetVersion] = {}
+    datasets_data = manifest.get("datasets", {})
+
+    for dataset_id, version_data in datasets_data.items():
+        versions[dataset_id] = DatasetVersion(
+            dataset_id=version_data["dataset_id"],
+            revision=version_data.get("revision"),
+            download_date=version_data["download_date"],
+            checksum=version_data.get("checksum"),
+        )
+
+    logger.info("Loaded version manifest with %d datasets from %s", len(versions), path)
+
+    return versions
+
+
+def get_dataset_info(dataset_id: str) -> dict[str, Any]:
+    """Get metadata about a HuggingFace dataset.
+
+    Retrieves information such as the latest revision, description,
+    file listing, and other Hub metadata.
+
+    Args:
+        dataset_id: HuggingFace dataset identifier.
+
+    Returns:
+        Dictionary containing dataset metadata with keys:
+            - dataset_id: The dataset identifier.
+            - latest_revision: The current HEAD revision hash.
+            - description: Dataset description text.
+            - tags: List of dataset tags.
+            - downloads: Number of downloads.
+            - last_modified: Last modification timestamp.
+            - files: List of files in the dataset repository.
+
+    Raises:
+        Exception: If the dataset cannot be found or accessed on the HuggingFace Hub.
+    """
+    info = dataset_info(dataset_id)
+
+    files: list[str] = []
+    if info.siblings:
+        files = [s.rfilename for s in info.siblings]
+
+    result: dict[str, Any] = {
+        "dataset_id": dataset_id,
+        "latest_revision": info.sha,
+        "description": info.description or "",
+        "tags": list(info.tags) if info.tags else [],
+        "downloads": info.downloads if info.downloads is not None else 0,
+        "last_modified": info.last_modified.isoformat() if info.last_modified else None,
+        "files": files,
+    }
+
+    return result