flexlock 0.8.2__py3-none-any.whl

flexlock/config.py

@@ -0,0 +1,121 @@
+"""Configuration and constants for FlexLock.
+
+This module defines default values and environment variable overrides.
+All environment variables are prefixed with FLEXLOCK_.
+"""
+
+import os
+
+
+def get_env_int(key: str, default: int) -> int:
+    """Get integer from environment variable."""
+    value = os.environ.get(key)
+    if value is None:
+        return default
+    try:
+        return int(value)
+    except ValueError:
+        return default
+
+
+def get_env_float(key: str, default: float) -> float:
+    """Get float from environment variable."""
+    value = os.environ.get(key)
+    if value is None:
+        return default
+    try:
+        return float(value)
+    except ValueError:
+        return default
+
+
+def get_env_bool(key: str, default: bool) -> bool:
+    """Get boolean from environment variable."""
+    value = os.environ.get(key, "").lower()
+    if value in ("1", "true", "yes", "on"):
+        return True
+    elif value in ("0", "false", "no", "off"):
+        return False
+    return default
+
+
+# ==================== Parallel Execution ====================
+
+# Poll interval for checking task status (seconds)
+POLL_INTERVAL = get_env_int("FLEXLOCK_POLL_INTERVAL", 10)
+
+# Log frequency for progress updates (seconds)
+LOG_FREQUENCY = get_env_int("FLEXLOCK_LOG_FREQUENCY", 15)
+
+# Default number of parallel jobs
+DEFAULT_N_JOBS = get_env_int("FLEXLOCK_DEFAULT_N_JOBS", 1)
+
+# Default timeout for waiting on HPC jobs (seconds, None = no timeout)
+_timeout_str = os.environ.get("FLEXLOCK_DEFAULT_TIMEOUT")
+DEFAULT_TIMEOUT = int(_timeout_str) if _timeout_str else 3600
+
+
+# ==================== Smart Run / Caching ====================
+
+# Whether to warn when smart_run=True but search_dirs=None
+WARN_SMART_RUN_NO_SEARCH_DIRS = get_env_bool("FLEXLOCK_WARN_SMART_RUN", True)
+
+
+# ==================== Timestamp Format ====================
+
+# Standard timestamp format for directory names
+# ISO 8601 compatible: YYYY-MM-DDTHH-MM-SS
+TIMESTAMP_FORMAT = os.environ.get("FLEXLOCK_TIMESTAMP_FORMAT", "%Y-%m-%dT%H-%M-%S")
+
+# Timestamp format for snapshots (ISO 8601 with microseconds)
+SNAPSHOT_TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%S.%f"
+
+
+# ==================== Database ====================
+
+# Database filename suffix
+DB_FILENAME_SUFFIX = ".tasks.db"
+
+# Attempts for DB write transactions when SQLite reports "database is locked"
+# (jittered exponential backoff between attempts, on top of busy_timeout).
+DB_RETRY_ATTEMPTS = get_env_int("FLEXLOCK_DB_RETRY_ATTEMPTS", 8)
+
+# Cap on the backoff delay between retry attempts (seconds)
+DB_RETRY_MAX_BACKOFF = get_env_float("FLEXLOCK_DB_RETRY_MAX_BACKOFF", 5.0)
+
+# Number of tasks a worker claims per DB transaction. Raise this (e.g. 16-64)
+# when running hundreds of workers against one DB on a shared filesystem:
+# claim and finish write-lock traffic both drop by this factor. Trade-off: a
+# worker that dies strands up to this many claimed tasks until the controller
+# reconciles orphans / --reclaim resets them.
+CLAIM_BATCH = get_env_int("FLEXLOCK_CLAIM_BATCH", 1)
+
+
+# ==================== Pre-flight checks ====================
+
+# When set, a worker verifies it can actually create a CUDA context before
+# claiming any task. Catches unhealthy GPU nodes (Fabric Manager down,
+# driver/kernel mismatch → cudaErrorSystemNotReady/802) up front with a clear
+# log, instead of an opaque traceback deep inside the user's training.
+PREFLIGHT_CUDA = get_env_bool("FLEXLOCK_PREFLIGHT_CUDA", False)
+
+# Default database filename pattern
+DB_FILENAME_PATTERN = "run.lock{suffix}"
+
+
+# ==================== Logging ====================
+
+# Whether to configure logging by default
+CONFIGURE_LOGGING = get_env_bool("FLEXLOCK_CONFIGURE_LOGGING", True)
+
+
+# ==================== Debug Mode ====================
+
+# Debug mode enabled
+DEBUG = get_env_bool("FLEXLOCK_DEBUG", False)
+
+# Debug mode disabled (takes precedence)
+NODEBUG = get_env_bool("FLEXLOCK_NODEBUG", False)
+
+# Debug strategy: 'auto', 'pdb', or 'inject'
+DEBUG_STRATEGY = os.environ.get("FLEXLOCK_DEBUG_STRATEGY", "auto")

flexlock/context.py

@@ -0,0 +1,7 @@
+"""Shared context for a FlexLock run."""
+
+from contextvars import ContextVar
+
+# This context will hold information gathered during a run, like the git commit hash.
+# It is initialized with an empty dictionary.
+run_context: ContextVar[dict] = ContextVar("run_context", default={})

flexlock/data_hash.py

@@ -0,0 +1,310 @@
+"""Data hashing utilities for FlexLock."""
+
+import os
+import sqlite3
+import threading
+from pathlib import Path
+import xxhash
+import hashlib
+from joblib import Parallel, delayed
+from contextlib import contextmanager
+
+from . import config
+
+# Hash-format version. Bump when the on-disk hash algorithm changes so stale
+# rows in an old cache file are never returned as valid hits. v2: dirhash now
+# folds each file's relative path into the digest (previously content-only).
+HASH_VERSION = 2
+
+# --- Cache Configuration ---
+CACHE_DIR = (
+    Path(
+        os.environ.get(
+            "FLEXLOCK_CACHE", os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")
+        )
+    )
+    / "flexlock"
+)
+CACHE_DB = CACHE_DIR / f"hashes_v{HASH_VERSION}.db"
+CACHE_DIR.mkdir(parents=True, exist_ok=True)
+DEFAULT_DIR_FILE_LIMIT = os.environ.get("FLEXLOCK_DIR_FILE_LIMIT", 1000)
+
+# Thread-local storage for database connections
+_thread_local_conns = threading.local()
+
+
+@contextmanager
+def _get_db():
+    """
+    A thread-safe context manager for SQLite database connections.
+
+    This function maintains a cache of connections per thread. A new connection
+    is created for each unique database path and reused for subsequent calls
+    with the same path within that thread.
+    """
+    # Use the absolute path as a reliable key for the connections dictionary.
+    db_path_str = str(CACHE_DB.resolve())
+
+    # Initialize the connections dictionary for the current thread if it doesn't exist.
+    if not hasattr(_thread_local_conns, "conns"):
+        _thread_local_conns.conns = {}
+
+    # Check if a connection for this specific db_path already exists in the thread's cache.
+    if db_path_str not in _thread_local_conns.conns:
+        # If not, create a new connection and add it to the cache.
+        try:
+            c = sqlite3.connect(db_path_str, check_same_thread=False)
+            # Set PRAGMA for better performance and concurrency.
+            c.execute("PRAGMA journal_mode=WAL")
+            c.execute("PRAGMA busy_timeout=15000")
+            c.execute(
+                "PRAGMA foreign_keys=ON"
+            )  # Good practice to enforce foreign key constraints
+            c.execute(
+                """
+                CREATE TABLE IF NOT EXISTS cache (
+                    path TEXT PRIMARY KEY,
+                    mtime REAL,
+                    file_count INTEGER,
+                    latest_mtime REAL,
+                    hash TEXT,
+                    is_dir INTEGER
+                )
+                """
+            )
+            _thread_local_conns.conns[db_path_str] = c
+        except sqlite3.Error as e:
+            print(f"Error connecting to database {db_path_str}: {e}")
+            raise
+
+    # Yield the connection from the thread's cache.
+    try:
+        yield _thread_local_conns.conns[db_path_str]
+    finally:
+        # Don't close the connection since we're caching it for reuse
+        pass
+
+
+def _hash_file_content(path):
+    """Hashes a single file using XXHash."""
+    hasher = xxhash.xxh64()
+    with open(path, "rb") as f:
+        while chunk := f.read(65536):
+            hasher.update(chunk)
+    return hasher.hexdigest()
+
+
+def _get_dir_stats(path: Path, limit: int):
+    """
+    Walks a directory to get the file count and the latest modification time.
+
+    If the file count exceeds the limit, it returns (limit + 1, 0) to signal
+    that the directory is "large".
+    """
+    count = 0
+    latest_mtime = path.stat().st_mtime
+
+    for root, _, files in os.walk(path):
+        count += len(files)
+        if count > limit:
+            return count, 0  # Exceeded limit, fallback mode
+
+        for name in files:
+            filepath = Path(root) / name
+            try:
+                mtime = filepath.stat().st_mtime
+                if mtime > latest_mtime:
+                    latest_mtime = mtime
+            except OSError:
+                # File might be a broken symlink, etc.
+                pass
+    return count, latest_mtime
+
+
+def dirhash(
+    path, match=None, ignore=None, jobs=1, algorithm=hashlib.md5, chunk_size=65536
+):
+    """
+    Computes a hash of the directory content using pathlib.glob for filtering.
+    """
+    base_path = Path(path)
+    if not base_path.is_dir():
+        raise ValueError(f"'{path}' is not a valid directory.")
+
+    # Default match to all files recursively if not specified
+    match_pattern = match if match is not None else "**/*"
+    if isinstance(match_pattern, str):
+        match_pattern = [match_pattern]
+
+    ignore_patterns = ignore if ignore is not None else []
+    if isinstance(ignore_patterns, str):
+        ignore_patterns = [ignore_patterns]
+
+    import glob
+    from pathlib import Path as p_path
+
+    # Find all files matching patterns
+    files_to_hash = []
+    for pattern in match_pattern:
+        files_to_hash.extend(base_path.glob(pattern))
+
+    files_to_hash = [f for f in files_to_hash if f.is_file()]
+
+    # Apply ignore patterns
+    final_files = []
+    for f in files_to_hash:
+        should_ignore = False
+        for pattern in ignore_patterns:
+            if f.match(pattern):
+                should_ignore = True
+                break
+        if not should_ignore:
+            final_files.append(f)
+
+    if not final_files:
+        return algorithm().hexdigest()
+
+    def _hash_file(filepath, algorithm, chunk_size):
+        """Hashes a single file."""
+        hasher = algorithm()
+        with open(filepath, "rb") as f:
+            while True:
+                data = f.read(chunk_size)
+                if not data:
+                    break
+                hasher.update(data)
+        return hasher.hexdigest()
+
+    content_hashes = Parallel(n_jobs=jobs)(
+        delayed(_hash_file)(str(f), algorithm, chunk_size) for f in final_files
+    )
+
+    # Fold the file's path into the digest so that two directories with
+    # identical file *contents* under different names hash differently, and so
+    # that swapping contents between two paths changes the hash.
+    entries = [
+        f"{f.relative_to(base_path).as_posix()}\0{h}"
+        for f, h in zip(final_files, content_hashes)
+    ]
+
+    final_hasher = algorithm()
+    for entry in sorted(entries):
+        final_hasher.update(entry.encode("utf-8"))
+
+    return final_hasher.hexdigest()
+
+
+def hash_data(
+    path,
+    match=None,
+    ignore=None,
+    jobs=4,
+    algorithm=xxhash.xxh64,
+    chunk_size=2**18,
+    use_cache=True,
+):
+    """
+    Computes a hash for a file or a directory, using an SQLite cache to avoid re-computation.
+    """
+    path = Path(path).resolve()
+    # Honour the ``use_cache`` argument; the env var can only *force* caching
+    # off (it never re-enables a cache the caller explicitly disabled).
+    use_cache = use_cache and not config.get_env_bool("FLEXLOCK_NO_CACHE", False)
+    dir_file_limit = int(
+        os.environ.get("FLEXLOCK_CACHE_DIR_FILE_LIMIT", DEFAULT_DIR_FILE_LIMIT)
+    )
+
+    if use_cache:
+        with _get_db() as conn:
+            cursor = conn.cursor()
+
+            if path.is_file():
+                # Check cache for file
+                cursor.execute(
+                    "SELECT hash, mtime FROM cache WHERE path=? AND is_dir=0",
+                    (str(path),),
+                )
+                row = cursor.fetchone()
+
+                if row:
+                    cached_hash, cached_mtime = row
+                    current_mtime = path.stat().st_mtime
+                    if cached_mtime == current_mtime:
+                        return cached_hash
+            elif path.is_dir():
+                # Check cache for directory
+                cursor.execute(
+                    "SELECT hash, mtime, file_count, latest_mtime FROM cache WHERE path=? AND is_dir=1",
+                    (str(path),),
+                )
+                row = cursor.fetchone()
+
+                if row:
+                    (
+                        cached_hash,
+                        cached_mtime,
+                        cached_file_count,
+                        cached_latest_mtime,
+                    ) = row
+                    file_count, latest_mtime = _get_dir_stats(path, dir_file_limit)
+
+                    if file_count > dir_file_limit:
+                        # Large directory fallback
+                        current_mtime = path.stat().st_mtime
+                        if cached_mtime == current_mtime:
+                            return cached_hash
+                    else:
+                        if (
+                            cached_file_count == file_count
+                            and cached_latest_mtime == latest_mtime
+                        ):
+                            return cached_hash
+
+    # If not in cache or cache is invalid/disabled, compute the hash
+    if not path.exists():
+        raise FileNotFoundError(f"The specified path does not exist: {path}")
+
+    new_hash = None
+    if path.is_file():
+        new_hash = _hash_file_content(path)
+    elif path.is_dir():
+        new_hash = dirhash(
+            path,
+            match=match,
+            ignore=ignore,
+            jobs=jobs,
+            algorithm=algorithm,
+            chunk_size=chunk_size,
+        )
+
+    if new_hash is None:
+        raise ValueError(f"Could not compute hash for path: {path}")
+
+    # Update and save the cache if enabled
+    if use_cache:
+        with _get_db() as conn:
+            cursor = conn.cursor()
+            if path.is_file():
+                mtime = path.stat().st_mtime
+                cursor.execute(
+                    "INSERT OR REPLACE INTO cache VALUES (?, ?, NULL, NULL, ?, 0)",
+                    (str(path), mtime, new_hash),
+                )
+            elif path.is_dir():
+                file_count, latest_mtime = _get_dir_stats(path, dir_file_limit)
+                mtime = path.stat().st_mtime
+                if file_count > dir_file_limit:
+                    # For large directories, use just the directory's mtime
+                    cursor.execute(
+                        "INSERT OR REPLACE INTO cache VALUES (?, ?, NULL, NULL, ?, 1)",
+                        (str(path), mtime, new_hash),
+                    )
+                else:
+                    # For smaller directories, cache more detailed stats
+                    cursor.execute(
+                        "INSERT OR REPLACE INTO cache VALUES (?, ?, ?, ?, ?, 1)",
+                        (str(path), mtime, file_count, latest_mtime, new_hash),
+                    )
+            conn.commit()
+
+    return new_hash