dstrack 0.1.0__py3-none-any.whl

dstrack/_benchmark/_runner.py

@@ -0,0 +1,196 @@
+"""Runs the snapshot pipeline against a synthetic CSV and times each phase.
+
+The runner performs no output of its own: progress is announced through a
+[BenchmarkObserver][dstrack._benchmark._runner.BenchmarkObserver] and everything
+measured is returned as a
+[BenchmarkRun][dstrack._benchmark._runner.BenchmarkRun], so the same run can be
+driven from a test with no console attached.
+"""
+
+import cProfile
+import time
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Protocol
+
+from dstrack._benchmark._environment import EnvironmentInfo, collect_environment_info
+from dstrack._benchmark._profiling import CallGraph, dstrack_scope
+from dstrack._benchmark._synthetic import SyntheticCsvSpec, write_synthetic_csv
+from dstrack.readers import CsvReader
+from dstrack.snapshot import MetadataBuilder, StatsComputer
+
+_DATASET_NAME = "benchmark-dataset"
+_CREATED_BY = "dstrack-benchmark"
+
+
+@dataclass(frozen=True)
+class BenchmarkResult:
+    """Timings and sizing for one synthetic-CSV snapshot benchmark run.
+
+    Attributes:
+        num_rows: Rows written to the synthetic CSV.
+        num_rows_read: Rows the statistics pass actually read back.
+        num_columns: Columns the reader inferred, including ``id``.
+        csv_size_bytes: Size of the generated CSV on disk.
+        generation_seconds: Time to write the synthetic CSV.
+        metadata_seconds: Time to build snapshot metadata from the reader.
+        stats_seconds: Time for the statistics data pass.
+        schema_hash: Schema hash of the snapshot that was built.
+    """
+
+    num_rows: int
+    num_rows_read: int
+    num_columns: int
+    csv_size_bytes: int
+    generation_seconds: float
+    metadata_seconds: float
+    stats_seconds: float
+    schema_hash: str
+
+    @property
+    def total_snapshot_seconds(self) -> float:
+        """Time to build the snapshot: metadata plus statistics."""
+        return self.metadata_seconds + self.stats_seconds
+
+    @property
+    def rows_per_second(self) -> float:
+        """Rows processed per second of snapshot time."""
+        return self.num_rows_read / self.total_snapshot_seconds
+
+
+@dataclass(frozen=True)
+class BenchmarkRun:
+    """Everything one benchmark run produced, ready to be reported."""
+
+    result: BenchmarkResult
+    environment: EnvironmentInfo
+    call_graph: CallGraph | None
+
+
+class BenchmarkObserver(Protocol):
+    """Receives a benchmark's progress as each phase starts."""
+
+    def generating_csv(self, path: Path, num_rows: int) -> None:
+        """Called before the synthetic CSV is written."""
+        ...
+
+    def building_metadata(self) -> None:
+        """Called before snapshot metadata is built."""
+        ...
+
+    def computing_stats(self) -> None:
+        """Called before the statistics data pass."""
+        ...
+
+
+class SilentObserver:
+    """Observer that reports nothing; the default for programmatic runs."""
+
+    def generating_csv(self, path: Path, num_rows: int) -> None:
+        return None
+
+    def building_metadata(self) -> None:
+        return None
+
+    def computing_stats(self) -> None:
+        return None
+
+
+@dataclass
+class _Elapsed:
+    """Wall-clock seconds a phase took, filled in when its block exits."""
+
+    seconds: float = 0.0
+
+
+class BenchmarkRunner:
+    """Generates a synthetic CSV and times a snapshot being built from it.
+
+    Metadata and statistics are timed separately, and share one profiler so the
+    call tree covers the whole snapshot pipeline.  CSV generation is never
+    profiled: it measures the harness, not dstrack.
+
+    Args:
+        spec: The synthetic dataset to benchmark against.
+        observer: Notified as each phase starts.  Defaults to reporting nothing.
+        profile: Whether to profile dstrack's own methods while building the
+            snapshot.  Adds cProfile overhead to the reported timings.
+    """
+
+    def __init__(
+        self,
+        spec: SyntheticCsvSpec,
+        *,
+        observer: BenchmarkObserver | None = None,
+        profile: bool = True,
+    ) -> None:
+        self._spec = spec
+        self._observer = observer or SilentObserver()
+        self._profiler = cProfile.Profile() if profile else None
+
+    def run(self, csv_path: Path) -> BenchmarkRun:
+        """Generate the dataset at ``csv_path`` and benchmark a snapshot of it."""
+        # clean profiler if already created to avoid misleading profiles on re-runs.
+        if self._profiler is not None:
+            self._profiler.clear()
+
+        # generate data for benchmark
+        self._observer.generating_csv(csv_path, self._spec.num_rows)
+        with self._timed() as generation:
+            write_synthetic_csv(csv_path, self._spec)
+        csv_size_bytes = csv_path.stat().st_size
+
+        reader = CsvReader(csv_path)
+
+        self._observer.building_metadata()
+        with self._timed(profiled=True) as metadata_phase:
+            metadata = MetadataBuilder().build(
+                reader,
+                dataset_name=_DATASET_NAME,
+                dataset_path=csv_path,
+                source_type="file",
+                created_by=_CREATED_BY,
+            )
+
+        self._observer.computing_stats()
+        with self._timed(profiled=True) as stats_phase:
+            stats = StatsComputer().compute(reader)
+
+        result = BenchmarkResult(
+            num_rows=self._spec.num_rows,
+            num_rows_read=stats.num_rows,
+            num_columns=metadata.num_columns,
+            csv_size_bytes=csv_size_bytes,
+            generation_seconds=generation.seconds,
+            metadata_seconds=metadata_phase.seconds,
+            stats_seconds=stats_phase.seconds,
+            schema_hash=metadata.schema_hash,
+        )
+        return BenchmarkRun(
+            result=result,
+            environment=collect_environment_info(),
+            call_graph=self._call_graph(),
+        )
+
+    @contextmanager
+    def _timed(self, *, profiled: bool = False) -> Iterator[_Elapsed]:
+        """Time the block, optionally recording it into the shared profiler."""
+        profiler = self._profiler if profiled else None
+        elapsed = _Elapsed()
+        if profiler is not None:
+            profiler.enable()
+        start = time.perf_counter()
+        try:
+            yield elapsed
+        finally:
+            elapsed.seconds = time.perf_counter() - start
+            if profiler is not None:
+                profiler.disable()
+
+    def _call_graph(self) -> CallGraph | None:
+        """Get a call graph of the snapshot phases, or None if profiling was disabled."""
+        if self._profiler is None:
+            return None
+        return CallGraph.from_profile(self._profiler, dstrack_scope())

dstrack/_benchmark/_synthetic.py

@@ -0,0 +1,158 @@
+"""Synthetic CSV generation for the benchmark.
+
+The dataset is described by a
+[SyntheticCsvSpec][dstrack._benchmark._synthetic.SyntheticCsvSpec] value object
+and written by
+[write_synthetic_csv()][dstrack._benchmark._synthetic.write_synthetic_csv].  Each
+generated dtype is one
+[SyntheticColumn][dstrack._benchmark._synthetic.SyntheticColumn] whose
+[generate][dstrack._benchmark._synthetic.SyntheticColumn.generate] callable
+renders a single cell, so supporting a new dtype means adding one cell factory
+and one entry to ``_COLUMN_KINDS`` rather than threading another count through
+every layer.
+"""
+
+import csv
+import math
+import random
+from collections.abc import Callable
+from dataclasses import dataclass
+from datetime import UTC, datetime, timedelta
+from pathlib import Path
+from typing import TypeAlias
+
+# Renders one cell as the text that lands in the CSV.  Cells are always
+# generated from the spec's seeded RNG, so a spec plus a seed is reproducible.
+CellFactory: TypeAlias = Callable[[random.Random], str]
+
+_WORD_POOL: tuple[str, ...] = (
+    "alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "hotel",
+    "india", "juliet", "kilo", "lima", "mike", "november", "oscar", "papa",
+    "quebec", "romeo", "sierra", "tango", "uniform", "victor", "whiskey",
+)  # fmt: skip
+
+_DATETIME_START = datetime(2020, 1, 1, tzinfo=UTC)
+_DATETIME_END = datetime(2026, 1, 1, tzinfo=UTC)
+
+
+def _numeric_cell(rng: random.Random) -> str:
+    return f"{rng.gauss(100, 25):.4f}"
+
+
+def _string_cell(rng: random.Random) -> str:
+    return " ".join(rng.choices(_WORD_POOL, k=rng.randint(1, 4)))
+
+
+def _datetime_cell(rng: random.Random) -> str:
+    span = (_DATETIME_END - _DATETIME_START).total_seconds()
+    return (_DATETIME_START + timedelta(seconds=rng.uniform(0, span))).isoformat()
+
+
+def _bool_cell(rng: random.Random) -> str:
+    return rng.choice(["True", "False"])
+
+
+@dataclass(frozen=True)
+class SyntheticColumn:
+    """One generated column: its header name and how to render a cell.
+
+    Attributes:
+        name: Column header written to the CSV.
+        generate: Renders a single non-null cell from the seeded RNG.
+    """
+
+    name: str
+    generate: CellFactory
+
+
+# (name prefix, cell factory) per dtype dstrack's CSV reader infers, in the
+# order the columns appear in the file.
+_COLUMN_KINDS: tuple[tuple[str, CellFactory], ...] = (
+    ("numeric", _numeric_cell),
+    ("string", _string_cell),
+    ("datetime", _datetime_cell),
+    ("bool", _bool_cell),
+)
+
+
+@dataclass(frozen=True)
+class SyntheticCsvSpec:
+    """Describes the synthetic dataset a benchmark run should be measured on.
+
+    Attributes:
+        num_rows: Number of data rows to generate.
+        num_numeric_cols: Number of ``float64`` columns.
+        num_string_cols: Number of ``string`` columns.
+        num_datetime_cols: Number of ``datetime64`` columns.
+        num_bool_cols: Number of ``bool`` columns.
+        null_rate: Fraction of non-id cells left empty (null).
+        seed: Seed for the random generator, for reproducible datasets.
+    """
+
+    num_rows: int = 200_000
+    num_numeric_cols: int = 5
+    num_string_cols: int = 5
+    num_datetime_cols: int = 2
+    num_bool_cols: int = 2
+    null_rate: float = 0.02
+    seed: int = 42
+
+    def __post_init__(self) -> None:
+        """Reject counts and a null rate that would silently corrupt the dataset.
+
+        Negative counts and an out-of-range ``null_rate`` have no meaningful
+        dataset to generate, so they are refused here rather than producing a
+        subtly wrong CSV further downstream.
+        """
+        for name in (
+            "num_rows",
+            "num_numeric_cols",
+            "num_string_cols",
+            "num_datetime_cols",
+            "num_bool_cols",
+        ):
+            value: int = getattr(self, name)
+            if value < 0:
+                raise ValueError(f"{name} must be >= 0, got {value}")
+        if not (math.isfinite(self.null_rate) and 0.0 <= self.null_rate <= 1.0):
+            raise ValueError(f"null_rate must be in [0, 1], got {self.null_rate!r}")
+
+    def columns(self) -> list[SyntheticColumn]:
+        """Return the nullable, randomly generated columns, in file order.
+
+        The leading ``id`` column is not included: it is written from the row
+        index and is never null.
+        """
+        counts = (
+            self.num_numeric_cols,
+            self.num_string_cols,
+            self.num_datetime_cols,
+            self.num_bool_cols,
+        )
+        return [
+            SyntheticColumn(f"{prefix}_{i}", generate)
+            for (prefix, generate), count in zip(_COLUMN_KINDS, counts, strict=True)
+            for i in range(count)
+        ]
+
+
+def write_synthetic_csv(path: Path, spec: SyntheticCsvSpec) -> None:
+    """Write the dataset described by ``spec`` to ``path``, overwriting it.
+
+    Args:
+        path: Destination CSV path; overwritten if it already exists.
+        spec: The dataset to generate.
+    """
+    rng = random.Random(spec.seed)
+    columns = spec.columns()
+
+    with path.open("w", newline="", encoding="utf-8") as fh:
+        writer = csv.writer(fh)
+        writer.writerow(["id", *(column.name for column in columns)])
+        for row_id in range(spec.num_rows):
+            row = [str(row_id)]
+            row.extend(
+                "" if rng.random() < spec.null_rate else column.generate(rng)
+                for column in columns
+            )
+            writer.writerow(row)

dstrack/_cli.py

@@ -0,0 +1,105 @@
+import logging
+import tempfile
+from pathlib import Path
+from typing import Annotated
+
+import typer
+
+from dstrack import __version__, console
+from dstrack._log import log
+from dstrack._store_template import STORE_TEMPLATE, materialize
+from dstrack._track import track
+from dstrack.errors import StoreInitError
+from dstrack.utils import get_invocation_path
+
+_log = logging.getLogger(__name__)
+
+app = typer.Typer(
+    no_args_is_help=True,
+    epilog="Made with :heart:",
+    suggest_commands=True,
+    pretty_exceptions_show_locals=False,
+    help="Main CLI interface to `dstrack`'s capabilities.",
+)
+
+
+def init_local_store() -> Path:
+    """Initializes local store structure.
+
+    Builds [STORE_TEMPLATE][dstrack._store_template.STORE_TEMPLATE] in a
+    temporary directory next to the destination, and only moves it into place
+    once every file and directory in the template has been created successfully.
+    The destination is therefore never touched by a failed or partial build.
+
+    Returns:
+        Path where the local store was created.
+
+    Raises:
+        FileExistsError: If the local store path already exists.
+        StoreInitError: If building the store from its template fails. The
+            temporary build directory is cleaned up before raising.
+    """
+    invocation_path = get_invocation_path()
+    store_path = invocation_path / STORE_TEMPLATE.name
+
+    if store_path.is_dir():
+        raise FileExistsError(f"Local store path already exists: {store_path}")
+
+    with tempfile.TemporaryDirectory(
+        prefix=f"{STORE_TEMPLATE.name}-tmp-", dir=invocation_path
+    ) as tmp_dir:
+        try:
+            built_path = materialize(STORE_TEMPLATE, Path(tmp_dir))
+        except Exception as e:
+            raise StoreInitError(
+                f"Failed to initialize local store at {store_path}"
+            ) from e
+
+        console.info(f"Generating local store structure at {store_path}.")
+        built_path.rename(store_path)
+
+    return store_path
+
+
+@app.command(
+    help="Initialize local store before starting to track any dataset. "
+    "The local store `.dstrack/` is created at the location where this "
+    "command is called."
+)
+def init(
+    allow_exists: Annotated[
+        bool, typer.Option(help="Do not fail if `.dstrack/` already exists.")
+    ] = False,
+) -> None:
+    try:
+        store_path = init_local_store()
+        console.success(f"Finished creating local store: {store_path}")
+    except FileExistsError as e:
+        if not allow_exists:
+            console.error(str(e))
+            raise
+        _log.warning(f"{e}")
+        console.warning(f"{e}")
+    except StoreInitError as e:
+        _log.error(f"Failed to create local store. {e}")
+        console.error(str(e))
+        raise
+
+
+app.command(help="Compute a snapshot of a dataset and store it in the local store.")(
+    track
+)
+
+app.command(
+    help="Show the snapshot history of a tracked dataset, given its dataset id "
+    "or its path."
+)(log)
+
+
+@app.command(help="Print package version.")
+def version() -> None:
+    print(__version__)
+
+
+if __name__ == "__main__":
+    app()