flexlock 0.8.2__py3-none-any.whl

flexlock/diff.py

@@ -0,0 +1,253 @@
+"""Diff utilities for FlexLock."""
+
+import os
+from typing import Any, Dict, Set
+from omegaconf import OmegaConf, DictConfig, ListConfig
+from loguru import logger
+
+
+class RunDiff:
+    def __init__(
+        self,
+        current: dict,
+        target: dict,
+        ignore_keys: list = None,
+        current_save_dir: str = None,
+        target_save_dir: str = None,
+        match_include: list = None,
+        match_exclude: list = None,
+    ):
+        """
+        Initialize RunDiff for comparing two run snapshots.
+
+        Args:
+            current: Current/proposed run snapshot
+            target: Target/existing run snapshot
+            ignore_keys: Additional keys to ignore during comparison
+            current_save_dir: Save directory of current run (for normalization)
+            target_save_dir: Save directory of target run (for normalization)
+            match_include: Override include patterns for git comparison
+            match_exclude: Override exclude patterns for git comparison
+        """
+        self.current = current
+        self.target = target
+
+        # Keys that legitimately appear inside the user config and are
+        # FlexLock-managed (or explicitly requested) — ignored at *any* depth.
+        self.always_ignore = set(ignore_keys or []) | {"save_dir", "_snapshot_"}
+
+        # Keys FlexLock injects at the snapshot top level (run metadata). These
+        # are only ignored at the top of the config subtree — a user
+        # hyperparameter named e.g. ``time`` or ``system`` nested deeper must
+        # NOT be silently ignored (issue 7), which would cause false cache hits.
+        self.toplevel_ignore = {
+            "timestamp",
+            "note",
+            "system",
+            "job_id",
+            "work_dir",
+            "cwd",
+            "date",
+            "time",
+            "datetime",
+        }
+
+        # Back-compat: some callers/tests read ``ignore_keys``.
+        self.ignore_keys = self.always_ignore | self.toplevel_ignore
+
+        # For value normalization (handling interpolation)
+        self.c_dir = str(current_save_dir) if current_save_dir else None
+        self.t_dir = str(target_save_dir) if target_save_dir else None
+
+        # Override patterns for git comparison (takes priority over snapshot-level patterns)
+        self.match_include = match_include
+        self.match_exclude = match_exclude
+
+        self.diffs = {}
+
+    def _normalize_val(self, val: Any, root_dir: str) -> Any:
+        """
+        If the value is a string and contains the save_dir path,
+        replace it with a placeholder <SAVE_DIR> to allow comparison.
+
+        Args:
+            val: Value to normalize
+            root_dir: Root directory path to replace
+
+        Returns:
+            Normalized value
+        """
+        if not (root_dir and isinstance(val, str)):
+            return val
+        # Prefix-only: replace an exact match or a genuine path prefix, never an
+        # arbitrary substring (issue 8). ``outputs/other`` must not be rewritten
+        # just because save_dir is ``outputs``.
+        if val == root_dir:
+            return "<SAVE_DIR>"
+        if val.startswith(root_dir + os.sep):
+            return "<SAVE_DIR>" + val[len(root_dir):]
+        return val
+
+    def compare_git(self):
+        """Compare git tree hashes, with optional include/exclude filtering."""
+        diff = []
+        c_repos = self.current.get("repos", {})
+        t_repos = self.target.get("repos", {})
+
+        # Iterate the union so a repo present on only one side is flagged
+        # symmetrically (issue 9) instead of being silently ignored.
+        for name in set(c_repos) | set(t_repos):
+            c_info = c_repos.get(name)
+            t_info = t_repos.get(name)
+            if not c_info:
+                diff.append(f"Repo {name} only in target")
+                continue
+            if not t_info:
+                diff.append(f"Repo {name} missing")
+                continue
+
+            # Compare Tree Hashes (Content Identity)
+            if c_info.get("tree") != t_info.get("tree"):
+                # Trees differ — check if RELEVANT files changed
+                # Priority: RunDiff-level override > snapshot-level patterns
+                include = (
+                    self.match_include or c_info.get("include") or t_info.get("include")
+                )
+                exclude = (
+                    self.match_exclude or c_info.get("exclude") or t_info.get("exclude")
+                )
+
+                if include or exclude:
+                    repo_path = c_info.get("path") or t_info.get("path")
+                    if repo_path and self._trees_match_filtered(
+                        repo_path, c_info["tree"], t_info["tree"], include, exclude
+                    ):
+                        continue  # Relevant files unchanged — match
+
+                diff.append(f"Repo {name}: Content changed")
+
+        if diff:
+            self.diffs["git"] = diff
+        return len(diff) == 0
+
+    def _trees_match_filtered(
+        self, repo_path, tree1, tree2, include=None, exclude=None
+    ):
+        """
+        Check if two trees match when filtered by include/exclude patterns.
+
+        Uses git diff-tree with pathspec filtering to compare only relevant files.
+        Returns True if no relevant files differ, False otherwise.
+        """
+        try:
+            from git.repo import Repo as GitRepo
+
+            repo = GitRepo(repo_path, search_parent_directories=True)
+
+            # Build git pathspec: include patterns + :(exclude) patterns
+            pathspec = list(include or [])
+            if exclude:
+                pathspec.extend(f":(exclude){pat}" for pat in exclude)
+
+            args = ["-r", "--name-only", "--no-commit-id", tree1, tree2]
+            if pathspec:
+                args.append("--")
+                args.extend(pathspec)
+
+            output = repo.git.diff_tree(*args).strip()
+            return len(output) == 0  # No relevant files changed
+        except Exception as e:
+            logger.debug(f"Filtered git comparison failed: {e}")
+            return False  # Conservative: treat as different
+
+    def compare_config(self):
+        """
+        Compare configurations with recursive diff, ignoring specified keys
+        and normalizing path values.
+        """
+        c_cfg = self.current.get("config", {})
+        t_cfg = self.target.get("config", {})
+
+        def _recursive_diff(d1, d2, path=""):
+            """Recursively compare two config structures."""
+            diff = []
+
+            # Handle DictConfigs vs Primitives
+            if isinstance(d1, (dict, DictConfig)) and isinstance(
+                d2, (dict, DictConfig)
+            ):
+                # Injected run-metadata keys are ignored only at the top of the
+                # config subtree; user keys are compared at every depth.
+                ignore_here = self.always_ignore
+                if path == "":
+                    ignore_here = ignore_here | self.toplevel_ignore
+
+                all_keys = set(d1.keys()) | set(d2.keys())
+                for k in all_keys:
+                    if k in ignore_here:
+                        continue
+
+                    new_path = f"{path}.{k}" if path else k
+
+                    if k not in d1:
+                        diff.append(f"Missing in current: {new_path}")
+                    elif k not in d2:
+                        diff.append(f"Extra in current: {new_path}")
+                    else:
+                        diff.extend(_recursive_diff(d1[k], d2[k], new_path))
+
+            elif isinstance(d1, (list, tuple, ListConfig)) and isinstance(
+                d2, (list, tuple, ListConfig)
+            ):
+                if len(d1) != len(d2):
+                    diff.append(f"List length mismatch {path}: {len(d1)} vs {len(d2)}")
+                else:
+                    for i, (v1, v2) in enumerate(zip(d1, d2)):
+                        diff.extend(_recursive_diff(v1, v2, f"{path}[{i}]"))
+
+            else:
+                # Value Comparison with Normalization
+                logger.debug(
+                    f"Comparing values at {path}: {d1} vs {d2} after normalization with dirs {self.c_dir}, {self.t_dir}"
+                )
+
+                v1_norm = self._normalize_val(d1, self.c_dir)
+                v2_norm = self._normalize_val(d2, self.t_dir)
+
+                if v1_norm != v2_norm:
+                    diff.append(f"Value mismatch {path}: {v1_norm} != {v2_norm}")
+
+            return diff
+
+        diff = _recursive_diff(c_cfg, t_cfg)
+
+        if diff:
+            self.diffs["config"] = diff
+
+        return len(diff) == 0
+
+    def compare_data(self):
+        """Compare data hashes."""
+        c_data = self.current.get("data", {})
+        t_data = self.target.get("data", {})
+
+        if c_data == t_data:
+            return True
+
+        # Name which keys/hashes differ instead of a bare "Data differs" (issue 9).
+        detail = []
+        for key in sorted(set(c_data) | set(t_data)):
+            cv = c_data.get(key)
+            tv = t_data.get(key)
+            if key not in c_data:
+                detail.append(f"{key}: only in target")
+            elif key not in t_data:
+                detail.append(f"{key}: only in current")
+            elif cv != tv:
+                detail.append(f"{key}: {cv} != {tv}")
+        self.diffs["data"] = detail or ["Data differs"]
+        return False
+
+    def is_match(self):
+        """Check if the current run matches the target run."""
+        return self.compare_git() and self.compare_config() and self.compare_data()

flexlock/diff_cli.py

@@ -0,0 +1,155 @@
+"""CLI for comparing FlexLock snapshots from various sources.
+
+Exit codes: ``0`` match, ``1`` differ, ``2`` error (so the command is usable
+in scripts and CI — the old behaviour always exited 0).
+"""
+
+import argparse
+import json
+import sys
+import yaml
+from pathlib import Path
+from loguru import logger
+from flexlock.diff import RunDiff
+from flexlock.taskdb import get_task_snapshot
+
+
+def load_snapshot_from_dir(dir_path: Path) -> dict:
+    """Load snapshot from a directory (finds run.lock file)."""
+    lock_file = dir_path / "run.lock"
+    if not lock_file.exists():
+        raise FileNotFoundError(f"No run.lock found in {dir_path}")
+
+    with open(lock_file) as f:
+        return yaml.safe_load(f)
+
+
+def load_snapshot_from_db(db_path: Path, task_id: str) -> dict:
+    """Load snapshot from database by task_id."""
+    snapshot = get_task_snapshot(db_path, task_id)
+    if snapshot is None:
+        raise ValueError(f"No snapshot found for task_id '{task_id}' in {db_path}")
+    return snapshot
+
+
+def run_comparison(snap1: dict, snap2: dict) -> "tuple[bool, dict]":
+    """Compare two snapshots, returning ``(is_match, diffs)``.
+
+    The single source of truth for both the text/JSON printers here and
+    ``flexlock why``. ``diffs`` is a JSON-serializable dict of lists keyed by
+    ``git``/``config``/``data`` (only the categories that differ appear).
+    """
+    diff = RunDiff(snap1, snap2)
+    # is_match() runs all three comparisons, populating diff.diffs as a side
+    # effect. Call it first so diffs is complete regardless of short-circuiting.
+    is_match = diff.is_match()
+    return is_match, dict(diff.diffs)
+
+
+def compare_snapshots(snap1: dict, snap2: dict, show_details: bool = False) -> bool:
+    """Compare two snapshots and print a human-readable report."""
+    is_match, diffs = run_comparison(snap1, snap2)
+
+    print("\n=== Snapshot Comparison ===\n")
+    for label, key in (("Git", "git"), ("Config", "config"), ("Data", "data")):
+        section_match = key not in diffs
+        print(f"{label:6s}: {'✓ Match' if section_match else '✗ Differ'}")
+        if not section_match and show_details:
+            for d in diffs.get(key, []):
+                print(f"  - {d}")
+
+    print(f"\nOverall: {'✓ Snapshots Match' if is_match else '✗ Snapshots Differ'}\n")
+    return is_match
+
+
+def main():
+    """CLI entry point for flexlock diff command."""
+    parser = argparse.ArgumentParser(
+        description="Compare FlexLock snapshots from various sources"
+    )
+
+    # Shared options for every subcommand.
+    common = argparse.ArgumentParser(add_help=False)
+    common.add_argument(
+        "--details", action="store_true", help="Show detailed differences (text)"
+    )
+    common.add_argument(
+        "--format", choices=["text", "json"], default="text",
+        help="Output format (default: text)",
+    )
+
+    subparsers = parser.add_subparsers(
+        dest="mode", required=True, help="Comparison mode"
+    )
+
+    # Mode 1: Compare two directories (traditional)
+    dir_parser = subparsers.add_parser(
+        "dirs", parents=[common], help="Compare two directory-based snapshots"
+    )
+    dir_parser.add_argument("dir1", type=Path, help="First directory")
+    dir_parser.add_argument("dir2", type=Path, help="Second directory")
+
+    # Mode 2: Compare two tasks in DB
+    db_parser = subparsers.add_parser(
+        "db", parents=[common], help="Compare two DB-based snapshots"
+    )
+    db_parser.add_argument("db_path", type=Path, help="Path to tasks database")
+    db_parser.add_argument("task_id1", help="First task ID (hash)")
+    db_parser.add_argument("task_id2", help="Second task ID (hash)")
+
+    # Mode 3: Compare directory to DB task
+    mixed_parser = subparsers.add_parser(
+        "mixed", parents=[common], help="Compare directory snapshot to DB snapshot"
+    )
+    mixed_parser.add_argument("dir_path", type=Path, help="Directory path")
+    mixed_parser.add_argument("db_path", type=Path, help="Database path")
+    mixed_parser.add_argument("task_id", help="Task ID in database")
+
+    args = parser.parse_args()
+
+    try:
+        if args.mode == "dirs":
+            if not args.dir1.exists():
+                logger.error(f"Directory not found: {args.dir1}")
+                sys.exit(2)
+            if not args.dir2.exists():
+                logger.error(f"Directory not found: {args.dir2}")
+                sys.exit(2)
+            snap1 = load_snapshot_from_dir(args.dir1)
+            snap2 = load_snapshot_from_dir(args.dir2)
+
+        elif args.mode == "db":
+            if not args.db_path.exists():
+                logger.error(f"Database not found: {args.db_path}")
+                sys.exit(2)
+            snap1 = load_snapshot_from_db(args.db_path, args.task_id1)
+            snap2 = load_snapshot_from_db(args.db_path, args.task_id2)
+
+        elif args.mode == "mixed":
+            if not args.dir_path.exists():
+                logger.error(f"Directory not found: {args.dir_path}")
+                sys.exit(2)
+            if not args.db_path.exists():
+                logger.error(f"Database not found: {args.db_path}")
+                sys.exit(2)
+            snap1 = load_snapshot_from_dir(args.dir_path)
+            snap2 = load_snapshot_from_db(args.db_path, args.task_id)
+
+        if args.format == "json":
+            is_match, diffs = run_comparison(snap1, snap2)
+            print(json.dumps({"match": is_match, "diffs": diffs}, indent=2, default=str))
+        else:
+            is_match = compare_snapshots(snap1, snap2, args.details)
+
+        # 0 match / 1 differ.
+        sys.exit(0 if is_match else 1)
+
+    except SystemExit:
+        raise
+    except Exception as e:
+        logger.error(f"Comparison failed: {e}")
+        sys.exit(2)
+
+
+if __name__ == "__main__":
+    main()

flexlock/exceptions.py

@@ -0,0 +1,50 @@
+"""Custom exceptions for FlexLock."""
+
+
+class FlexLockError(Exception):
+    """Base exception for all FlexLock errors."""
+
+    pass
+
+
+class FlexLockConfigError(FlexLockError):
+    """Raised when there is an error in configuration."""
+
+    pass
+
+
+class FlexLockExecutionError(FlexLockError):
+    """Raised when there is an error during execution."""
+
+    pass
+
+
+class FlexLockSnapshotError(FlexLockError):
+    """Raised when there is an error creating or loading snapshots."""
+
+    pass
+
+
+class FlexLockValidationError(FlexLockError):
+    """Raised when validation of inputs fails."""
+
+    pass
+
+
+class FlexLockCacheError(FlexLockError):
+    """Raised when there is an error with smart run caching."""
+
+    pass
+
+
+class FlexLockBackendError(FlexLockError):
+    """Raised when there is an error with HPC backends (Slurm/PBS)."""
+
+    pass
+
+
+class UnresolvedInterpolationError(FlexLockConfigError):
+    """Raised when a sub-node interpolation references a key that is not present
+    in the sub-tree and cannot be resolved against the root config either."""
+
+    pass

flexlock/export.py

@@ -0,0 +1,134 @@
+"""Export utilities for FlexLock - extract snapshots from database to files."""
+
+import argparse
+import sys
+import json
+import yaml
+import tempfile
+import os
+from pathlib import Path
+from loguru import logger
+from flexlock.taskdb import get_task_snapshot, list_task_snapshots
+
+
+def export_task(db_path: Path, task_id: str, output_dir: Path) -> None:
+    """
+    Export a single task snapshot from database to a standalone directory.
+
+    Args:
+        db_path: Path to SQLite database
+        task_id: Hash of the task to export
+        output_dir: Directory to write the exported snapshot
+
+    Raises:
+        ValueError: If task not found in database
+    """
+    snapshot_data = get_task_snapshot(db_path, task_id)
+    if not snapshot_data:
+        raise ValueError(f"Task {task_id} not found in {db_path}")
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Atomic write of run.lock file
+    with tempfile.NamedTemporaryFile("w", dir=output_dir, delete=False) as tf:
+        yaml.dump(snapshot_data, tf, sort_keys=False)
+        tmp_name = tf.name
+    os.replace(tmp_name, output_dir / "run.lock")
+
+    logger.info(f"Exported task {task_id} to {output_dir}")
+
+
+def export_all_tasks(db_path: Path, output_base_dir: Path, status: str = None) -> None:
+    """
+    Export all tasks (or filtered by status) from database to separate directories.
+
+    Args:
+        db_path: Path to SQLite database
+        output_base_dir: Base directory for exports (each task gets a subdirectory)
+        status: Optional filter by status (pending, running, done, failed)
+    """
+    tasks = list_task_snapshots(db_path, status)
+
+    if not tasks:
+        logger.warning(
+            f"No tasks found in {db_path}"
+            + (f" with status={status}" if status else "")
+        )
+        return
+
+    output_base_dir.mkdir(parents=True, exist_ok=True)
+
+    for task_id, snapshot_data, task_status in tasks:
+        if snapshot_data:
+            task_output_dir = output_base_dir / f"task_{task_id[:8]}"
+            task_output_dir.mkdir(parents=True, exist_ok=True)
+
+            # Atomic write
+            with tempfile.NamedTemporaryFile(
+                "w", dir=task_output_dir, delete=False
+            ) as tf:
+                yaml.dump(snapshot_data, tf, sort_keys=False)
+                tmp_name = tf.name
+            os.replace(tmp_name, task_output_dir / "run.lock")
+
+            logger.info(
+                f"Exported task {task_id[:8]} (status={task_status}) to {task_output_dir}"
+            )
+
+    logger.info(f"Exported {len(tasks)} tasks to {output_base_dir}")
+
+
+def main():
+    """CLI entry point for flexlock export command."""
+    parser = argparse.ArgumentParser(
+        description="Export FlexLock task snapshots from database to files"
+    )
+
+    parser.add_argument(
+        "--db",
+        type=Path,
+        required=True,
+        help="Path to tasks database (e.g., outputs/sweep/tasks.db)",
+    )
+
+    parser.add_argument(
+        "--task", help="Task ID to export (hash). If not specified, exports all tasks."
+    )
+
+    parser.add_argument(
+        "--out",
+        type=Path,
+        required=True,
+        help="Output directory for exported snapshot(s)",
+    )
+
+    parser.add_argument(
+        "--status",
+        choices=["pending", "running", "done", "failed"],
+        help="Filter tasks by status (only used when exporting all tasks)",
+    )
+
+    args = parser.parse_args()
+
+    # Validate database exists
+    if not args.db.exists():
+        logger.error(f"Database not found: {args.db}")
+        sys.exit(1)
+
+    try:
+        if args.task:
+            # Export single task
+            export_task(args.db, args.task, args.out)
+        else:
+            # Export all tasks
+            export_all_tasks(args.db, args.out, args.status)
+
+        logger.success("Export completed successfully")
+
+    except Exception as e:
+        logger.error(f"Export failed: {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()