flexlock 0.8.2__py3-none-any.whl

flexlock/snapshot.py

@@ -0,0 +1,245 @@
+"""Snapshotting utilities for FlexLock."""
+
+from datetime import datetime
+from pathlib import Path
+from omegaconf import OmegaConf
+from .git_utils import create_shadow_snapshot
+from .data_hash import hash_data
+from .load_stage import load_stage_from_path
+from .run_record import RunRecord, COMPLETE_MARKER, COMPLETE_VERSION
+from loguru import logger
+import uuid
+
+
+def write_complete_marker(save_dir: Path, result=None) -> Path:
+    """Write run.complete atomically into ``save_dir``.
+
+    Thin wrapper over :meth:`RunRecord.mark_complete` kept for backward
+    compatibility. ``_find_matching_run`` requires both ``run.lock`` and
+    ``run.complete`` for a cache hit.
+    """
+    return RunRecord(save_dir).mark_complete(result=result)
+
+
+def is_complete(run_dir: Path) -> bool:
+    """Return True if ``run_dir`` has both ``run.lock`` and ``run.complete``."""
+    return RunRecord(run_dir).is_complete
+
+
+class RunTracker:
+    def __init__(self, save_dir, parent_lock=None, note=None):
+        self.save_dir = Path(save_dir)
+        self.parent_lock = Path(parent_lock) if parent_lock else None
+        self.data = {"timestamp": datetime.now().isoformat()}
+
+        # Free-text intent recorded alongside run metadata (sibling of
+        # timestamp/config). It never enters the fingerprint, so it can't
+        # perturb caching or diffs.
+        if note:
+            self.data["note"] = str(note)
+
+        if self.parent_lock:
+            # We record the link, effectively saying "See parent for Git/Env"
+            self.data["parent"] = str(self.parent_lock)
+
+    def record_env(self, repos: dict):
+        # Optimization: If we have a parent, we MIGHT skip recording git
+        # IF we are sure code hasn't changed between main process and worker
+        # (which is true for multiprocessing/slurm typically).
+        if self.parent_lock:
+            return
+
+        self.data["repos"] = {}
+        for name, repo_info in repos.items():
+            path = repo_info["path"]
+            snapshot_data = create_shadow_snapshot(
+                path,
+                ref_name=str(self.save_dir) + f"_{uuid.uuid1().hex}",
+            )
+            # Store metadata for comparison-time filtering
+            if repo_info.get("include"):
+                snapshot_data["include"] = repo_info["include"]
+            if repo_info.get("exclude"):
+                snapshot_data["exclude"] = repo_info["exclude"]
+            if repo_info.get("module"):
+                snapshot_data["module"] = repo_info["module"]
+            snapshot_data["path"] = path
+            self.data["repos"][name] = snapshot_data
+
+    def record_data(self, data_paths: dict):
+        self.data["data"] = {k: hash_data(v) for k, v in data_paths.items()}
+
+    def add_lineage(self, name: str, path: str, info: dict):
+        """Add lineage information from upstream FlexLock runs."""
+        if "lineage" not in self.data:
+            self.data["lineage"] = {}
+
+        self.data["lineage"][name] = {"path": path, "info": info}
+
+    def finalize(self, config):
+        """
+        Prepares the final snapshot dict but does not write it to disk.
+
+        This allows callers to decide where to store the snapshot
+        (file, DB, or both).
+
+        Args:
+            config: OmegaConf configuration object
+
+        Returns:
+            dict: Complete snapshot data structure
+        """
+        self.data["config"] = OmegaConf.to_container(config, resolve=True)
+        return self.data
+
+    def save(self, config):
+        """
+        Finalize and write to disk as run.lock.
+        Now implemented by calling finalize() then writing.
+
+        Returns:
+            dict: The snapshot data that was written
+        """
+        snapshot_data = self.finalize(config)
+
+        # Read the concrete save_dir from the finalized snapshot. save_dir is
+        # baked to a string once at submit time (apply_save_dir_policy), so this
+        # is just a lookup — no resolver fires here.
+        resolved_save_dir = Path(
+            snapshot_data["config"].get("save_dir", str(self.save_dir))
+        )
+
+        # Single on-disk-contract owner performs the atomic write.
+        RunRecord(resolved_save_dir).write_lock(snapshot_data)
+
+        # Remember the resolved dir so mark_complete can write next to run.lock
+        self._resolved_save_dir = resolved_save_dir
+
+        return snapshot_data
+
+    def mark_complete(self, result=None) -> Path:
+        """Write ``run.complete`` next to the previously-written ``run.lock``.
+
+        Call this only after the user function returns successfully. The
+        ``_find_matching_run`` cache check requires both files; an
+        interrupted run leaves ``run.lock`` without ``run.complete`` and is
+        correctly skipped instead of being treated as a stale cache hit.
+        """
+        target = getattr(self, "_resolved_save_dir", self.save_dir)
+        return RunRecord(target).mark_complete(result=result)
+
+
+def snapshot(
+    cfg,
+    repos=None,
+    data=None,
+    prevs=None,
+    parent_lock=None,
+    save_path=None,
+    return_snapshot=False,
+    fingerprint=None,
+    note=None,
+):
+    """
+    Create a snapshot of the current run state.
+
+    Args:
+        cfg: OmegaConf configuration
+        repos: Dict of repository paths to track
+        data: Dict of data paths to hash
+        prevs: List of paths to search for upstream lineage
+        parent_lock: Path to parent run.lock (for delta snapshots)
+        save_path: Custom save directory (overrides cfg.save_dir)
+        return_snapshot: If True, return snapshot dict instead of/in addition to saving
+
+    Returns:
+        dict if return_snapshot=True, else None
+    """
+    if "save_dir" not in cfg:
+        logger.warning("No save_dir specified in config; skipping snapshot.")
+        return None
+
+    # Use custom save_path if provided, otherwise use cfg.save_dir
+    save_dir = Path(save_path) if save_path else Path(cfg.save_dir)
+
+    tracker = RunTracker(save_dir, parent_lock=parent_lock, note=note)
+
+    # Store the precomputed fingerprint (index key) so run.lock/DB snapshots
+    # carry it and `flexlock reindex` can rebuild the index from disk.
+    if fingerprint is not None:
+        tracker.data["fingerprint"] = fingerprint
+
+    # 1. Record Git & Data (Hashing)
+    if repos:
+        tracker.record_env(repos)
+    if data:
+        logger.debug(f"Recording data for snapshot: {data}")
+        tracker.record_data(data)
+
+    # 2. Record Lineage (Automatic Discovery)
+    if prevs:
+        logger.debug(f"Looking for upstream FlexLock runs in: {prevs}")
+
+        def _find_snapshot_dir(start_path: Path) -> tuple[Path, dict] | None:
+            """
+            Find run.lock at or directly in start_path.
+
+            Returns:
+                tuple: (snapshot_dir, snapshot_data) or None
+            """
+            try:
+                p = Path(start_path).resolve()
+            except Exception:
+                return None
+            if p.is_file():
+                p = p.parent
+
+            lock_file = p / "run.lock"
+            if lock_file.exists():
+                try:
+                    data = OmegaConf.to_container(
+                        OmegaConf.load(lock_file), resolve=True
+                    )
+                    logger.debug(f"Found snapshot at: {p}")
+                    return (p, data)
+                except Exception as e:
+                    logger.warning(f"Failed to read run.lock at {p}: {e}")
+
+            return None
+
+        for path_str in prevs:
+            result = _find_snapshot_dir(path_str)
+
+            if result:
+                snapshot_dir, snapshot_data = result
+                # We found an upstream FlexLock run!
+                logger.debug(f"Found upstream FlexLock run at: {snapshot_dir}")
+                try:
+                    # Extract stage info from snapshot data
+                    stage_info = {
+                        "config": snapshot_data.get("config", {}),
+                        "timestamp": snapshot_data.get("timestamp"),
+                        "repos": snapshot_data.get("repos", {}),
+                        "parent": snapshot_data.get("parent"),
+                    }
+
+                    # Add to lineage
+                    tracker.add_lineage(
+                        name=snapshot_dir.name,  # e.g. "run_2023..." or "task_abc123"
+                        path=str(snapshot_dir),
+                        info=stage_info,
+                    )
+                except Exception as e:
+                    logger.warning(
+                        f"Found snapshot at {snapshot_dir} but failed to process: {e}"
+                    )
+
+    # 3. Save and/or return
+    if return_snapshot:
+        # For DB storage: return snapshot without writing file
+        snapshot_data = tracker.finalize(cfg)
+        return snapshot_data
+    else:
+        # Traditional behavior: write to file
+        tracker.save(cfg)
+        return None

flexlock/status_cli.py

@@ -0,0 +1,299 @@
+"""CLI for monitoring FlexLock task database status."""
+
+import argparse
+import sys
+import time
+from pathlib import Path
+from loguru import logger
+from flexlock.taskdb import (
+    get_status_counts,
+    get_status_counts_by_tag,
+    get_failed_tasks,
+    get_all_tasks,
+)
+
+
+def print_status_summary(db_path: Path, tags=None, label: str = ""):
+    """Print summary of task statuses, optionally scoped to ``tags``."""
+    status_counts = get_status_counts(db_path, tags=tags)
+
+    pending = status_counts.get("pending", 0)
+    running = status_counts.get("running", 0)
+    done = status_counts.get("done", 0)
+    failed = status_counts.get("failed", 0)
+    interrupted = status_counts.get("interrupted", 0)
+    total = pending + running + done + failed + interrupted
+
+    header = f"Task Status Summary{label}"
+    print("\n" + "=" * 60)
+    print(header)
+    print("=" * 60)
+    print(f"Pending:      {pending:>5}")
+    print(f"Running:      {running:>5}")
+    print(f"Done:         {done:>5}")
+    print(f"Failed:       {failed:>5}")
+    print(f"Interrupted:  {interrupted:>5}")
+    print("-" * 60)
+    print(f"Total:        {total:>5}")
+
+    if total > 0:
+        terminal = done + failed + interrupted
+        progress = terminal / total * 100
+        print(f"Progress:     {progress:>5.1f}% ({terminal}/{total} completed)")
+
+        if pending == 0 and running == 0:
+            issues = failed + interrupted
+            if issues:
+                label_parts = []
+                if failed:
+                    label_parts.append(f"{failed} failed")
+                if interrupted:
+                    label_parts.append(f"{interrupted} interrupted")
+                print(f"\nStatus:     ⚠️  Finished with issues ({', '.join(label_parts)})")
+            else:
+                print(f"\nStatus:     ✓  All tasks completed successfully")
+        elif pending == 0 and running > 0:
+            print(f"\nStatus:     ⚠️  {running} task(s) running with no pending — possible zombie(s)")
+        else:
+            print(f"\nStatus:     ⏳ In progress")
+
+    print("=" * 60 + "\n")
+
+
+def print_tags_breakdown(db_path: Path):
+    """Print a per-tag status breakdown table."""
+    by_tag = get_status_counts_by_tag(db_path)
+    if not by_tag:
+        print("No tasks found in DB.\n")
+        return
+
+    print("\n" + "=" * 60)
+    print("Per-Tag Status Breakdown")
+    print("=" * 60)
+
+    for tag, counts in sorted(by_tag.items(), key=lambda x: (x[0] is None, x[0])):
+        tag_label = repr(tag) if tag is not None else "(untagged)"
+        pending = counts.get("pending", 0)
+        running = counts.get("running", 0)
+        done = counts.get("done", 0)
+        failed = counts.get("failed", 0)
+        interrupted = counts.get("interrupted", 0)
+        total = pending + running + done + failed + interrupted
+        terminal = done + failed + interrupted
+        pct = terminal / total * 100 if total else 0
+
+        print(f"\nTag: {tag_label}  ({total} total, {pct:.1f}% done)")
+        print(f"  pending={pending}  running={running}  done={done}  "
+              f"failed={failed}  interrupted={interrupted}")
+
+    print("\n" + "=" * 60 + "\n")
+
+
+def print_failed_tasks(db_path: Path, verbose: bool = False, tags=None):
+    """Print details of failed and interrupted tasks."""
+    failed_tasks = get_failed_tasks(db_path, tags=tags)
+
+    if not failed_tasks:
+        print("No failed or interrupted tasks found.\n")
+        return
+
+    print(f"\nFailed / Interrupted Tasks ({len(failed_tasks)})")
+    print("=" * 60)
+
+    for i, task_info in enumerate(failed_tasks, 1):
+        status = task_info.get("status", "failed")
+        print(f"\nTask #{i}  [{status.upper()}]")
+        print("-" * 60)
+
+        task = task_info["task"]
+        print("Config:")
+        for key, value in task.items():
+            if not key.startswith("_"):
+                print(f"  {key}: {value}")
+
+        error = task_info["error"]
+        if error:
+            print(f"\nError:")
+            if len(error) > 500 and not verbose:
+                print(f"  {error[:500]}...")
+                print(f"  (Use --verbose to see full error)")
+            else:
+                for line in error.split("\n"):
+                    print(f"  {line}")
+
+        if task_info["node"]:
+            print(f"\nNode: {task_info['node']}")
+        if task_info["ts_start"]:
+            print(f"Started:  {task_info['ts_start']}")
+        if task_info["ts_end"]:
+            print(f"Finished: {task_info['ts_end']}")
+
+    print("\n" + "=" * 60 + "\n")
+
+
+def print_all_tasks(db_path: Path, status_filter: str = None, tags=None):
+    """Print all tasks, optionally filtered by status and/or tag."""
+    tasks = get_all_tasks(db_path, status=status_filter, tags=tags)
+
+    filter_desc = f" ({status_filter})" if status_filter else ""
+    print(f"\nAll Tasks{filter_desc} ({len(tasks)})")
+    print("=" * 60)
+
+    if not tasks:
+        print("No tasks found.\n")
+        return
+
+    print(f"{'Status':<13} {'Task ID':<12} {'Node':<15} {'Timestamp'}")
+    print("-" * 60)
+
+    for task_info in tasks:
+        status = task_info["status"]
+        task_id = task_info["task_id"][:10]
+        node = task_info["node"] or "N/A"
+        node = node[:13] if len(node) > 13 else node
+        timestamp = task_info["ts_end"] or task_info["ts_start"] or "N/A"
+        if timestamp != "N/A":
+            timestamp = timestamp[:19]
+
+        print(f"{status:<13} {task_id:<12} {node:<15} {timestamp}")
+
+    print("=" * 60 + "\n")
+
+
+def watch_status(db_path: Path, interval: int = 10, tags=None):
+    """Watch task status in real-time, optionally scoped to ``tags``."""
+    print("Watching task status (Ctrl+C to stop)...\n")
+
+    try:
+        while True:
+            print("\033[2J\033[H", end="")
+
+            from datetime import datetime
+            print(f"Last updated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
+
+            print_status_summary(db_path, tags=tags)
+
+            status_counts = get_status_counts(db_path, tags=tags)
+            pending = status_counts.get("pending", 0)
+            running = status_counts.get("running", 0)
+
+            if pending == 0 and running == 0:
+                print("All tasks reached a terminal state. Exiting watch mode.\n")
+                break
+
+            print(f"Refreshing in {interval}s... (Ctrl+C to stop)")
+            time.sleep(interval)
+
+    except KeyboardInterrupt:
+        print("\n\nWatch mode stopped.\n")
+
+
+def main():
+    """CLI entry point for flexlock-status command."""
+    parser = argparse.ArgumentParser(
+        description="Monitor FlexLock task database status",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Show status summary
+  flexlock-status outputs/sweep/run.lock.tasks.db
+
+  # Per-tag breakdown (useful when multiple sweeps share a DB)
+  flexlock-status outputs/sweep/run.lock.tasks.db --tags
+
+  # Scope the entire view to one tag
+  flexlock-status outputs/sweep/run.lock.tasks.db --tag extract
+
+  # Show failed / interrupted tasks
+  flexlock-status outputs/sweep/run.lock.tasks.db --failed
+
+  # Show all tasks
+  flexlock-status outputs/sweep/run.lock.tasks.db --all
+
+  # Filter by status
+  flexlock-status outputs/sweep/run.lock.tasks.db --all --status interrupted
+
+  # Watch status in real-time (scoped to a tag)
+  flexlock-status outputs/sweep/run.lock.tasks.db --watch --tag collocate
+        """,
+    )
+
+    parser.add_argument(
+        "db_path", type=Path, help="Path to task database (run.lock.tasks.db)"
+    )
+    parser.add_argument(
+        "--failed", action="store_true", help="Show details of failed / interrupted tasks"
+    )
+    parser.add_argument("--all", action="store_true", help="Show all tasks")
+    parser.add_argument(
+        "--status",
+        choices=["pending", "running", "done", "failed", "interrupted"],
+        help="Filter tasks by status (use with --all)",
+    )
+    parser.add_argument(
+        "--watch",
+        action="store_true",
+        help="Watch status in real-time (updates every 10s)",
+    )
+    parser.add_argument(
+        "--interval",
+        type=int,
+        default=10,
+        help="Update interval for watch mode in seconds (default: 10)",
+    )
+    parser.add_argument(
+        "--verbose", "-v", action="store_true",
+        help="Show verbose output (full error messages)",
+    )
+    parser.add_argument(
+        "--tag",
+        metavar="NAME",
+        help="Scope the entire view (summary, --failed, --all, --watch) to a "
+             "single tag. Useful to monitor one sweep inside a shared DB.",
+    )
+    parser.add_argument(
+        "--tags",
+        action="store_true",
+        dest="show_tags",
+        help="Print a per-tag status breakdown before the global summary. "
+             "Useful to discover which tags exist and what to pass to "
+             "flexlock-worker --tags.",
+    )
+
+    args = parser.parse_args()
+
+    if not args.db_path.exists():
+        logger.error(f"Database not found: {args.db_path}")
+        sys.exit(1)
+
+    # --tag NAME scopes everything; convert to a list for the taskdb API.
+    scope_tags = [args.tag] if args.tag else None
+    tag_label = f" [tag={args.tag!r}]" if args.tag else ""
+
+    try:
+        if args.watch:
+            watch_status(args.db_path, args.interval, tags=scope_tags)
+        elif args.failed:
+            if args.show_tags:
+                print_tags_breakdown(args.db_path)
+            print_status_summary(args.db_path, tags=scope_tags, label=tag_label)
+            print_failed_tasks(args.db_path, verbose=args.verbose, tags=scope_tags)
+        elif args.all:
+            if args.show_tags:
+                print_tags_breakdown(args.db_path)
+            print_status_summary(args.db_path, tags=scope_tags, label=tag_label)
+            print_all_tasks(args.db_path, status_filter=args.status, tags=scope_tags)
+        else:
+            if args.show_tags:
+                print_tags_breakdown(args.db_path)
+            print_status_summary(args.db_path, tags=scope_tags, label=tag_label)
+
+    except Exception as e:
+        logger.error(f"Error reading database: {e}")
+        import traceback
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()