flexlock 0.8.2__py3-none-any.whl

flexlock/fingerprint.py

@@ -0,0 +1,129 @@
+"""Pure, stable run fingerprints.
+
+A *fingerprint* is a deterministic digest over the inputs that define a run:
+
+  * the resolved config, with ``save_dir`` **prefix-normalized** so two runs
+    that differ only in where they write share a fingerprint;
+  * per-repo git **tree hashes** (via :func:`create_shadow_tree` — no commit,
+    no ref side effects), restricted to each repo's ``include``/``exclude``
+    pathspec so the "relevant files unchanged" notion becomes plain digest
+    equality rather than a special case in the matcher;
+  * data hashes for tracked data paths.
+
+The digest is the single key used by the project-wide index (see ``index.py``)
+to decide cache hits, so it must be stable across processes and machines for
+identical inputs.
+"""
+
+import hashlib
+import json
+import os
+from typing import Any, Mapping, Optional
+
+from omegaconf import OmegaConf, DictConfig
+
+from .git_utils import create_shadow_tree
+from .data_hash import hash_data
+
+# Placeholder substituted for the run's own save_dir so location doesn't leak
+# into the fingerprint.
+SAVE_DIR_PLACEHOLDER = "<SAVE_DIR>"
+
+# Keys that describe *tracking intent* rather than run inputs; excluded from the
+# config portion of the digest.
+_TRACKING_KEYS = ("_snapshot_",)
+
+
+def _to_container(cfg: Any) -> Any:
+    """Resolve an OmegaConf config to a plain container; pass dicts through."""
+    if isinstance(cfg, DictConfig):
+        return OmegaConf.to_container(cfg, resolve=True)
+    if OmegaConf.is_config(cfg):
+        return OmegaConf.to_container(cfg, resolve=True)
+    return cfg
+
+
+def _normalize_paths(value: Any, save_dir: Optional[str]) -> Any:
+    """Replace the run's save_dir *prefix* with a stable placeholder.
+
+    Only an exact match or a genuine path prefix (``save_dir`` + separator) is
+    rewritten — never an arbitrary substring — so unrelated paths that merely
+    start with the same characters are left untouched.
+    """
+    if not save_dir:
+        return value
+    if isinstance(value, str):
+        if value == save_dir:
+            return SAVE_DIR_PLACEHOLDER
+        if value.startswith(save_dir + os.sep):
+            return SAVE_DIR_PLACEHOLDER + value[len(save_dir):]
+        return value
+    if isinstance(value, dict):
+        return {k: _normalize_paths(v, save_dir) for k, v in value.items()}
+    if isinstance(value, list):
+        return [_normalize_paths(v, save_dir) for v in value]
+    return value
+
+
+def canonical_config(cfg: Any) -> Any:
+    """Return the config portion of the fingerprint: resolved, save_dir-
+    normalized, and stripped of tracking-only keys."""
+    container = _to_container(cfg)
+    if not isinstance(container, dict):
+        return container
+
+    save_dir = container.get("save_dir")
+    # Drop the run's own save_dir and tracking keys, then normalize any nested
+    # references to save_dir into the placeholder.
+    stripped = {
+        k: v
+        for k, v in container.items()
+        if k != "save_dir" and k not in _TRACKING_KEYS
+    }
+    return _normalize_paths(stripped, save_dir if isinstance(save_dir, str) else None)
+
+
+def _repo_tree(repo_info: Mapping[str, Any]) -> str:
+    """Tree hash for one tracked repo, restricted to its include/exclude."""
+    result = create_shadow_tree(
+        repo_info["path"],
+        include=repo_info.get("include"),
+        exclude=repo_info.get("exclude"),
+    )
+    return result["tree"]
+
+
+def fingerprint(
+    cfg: Any,
+    repos: Optional[Mapping[str, Mapping[str, Any]]] = None,
+    data: Optional[Mapping[str, str]] = None,
+    use_cache: bool = True,
+) -> str:
+    """Compute the stable fingerprint digest for a run.
+
+    Args:
+        cfg: The run config (OmegaConf or plain dict). Resolved and
+            save_dir-normalized before hashing.
+        repos: ``{name: {path, include?, exclude?, module?}}`` — per-repo git
+            identity. Tree hashes are computed via ``create_shadow_tree``.
+        data: ``{name: path}`` data inputs to hash.
+        use_cache: Passed through to :func:`hash_data` for data paths.
+
+    Returns:
+        A hex sha256 digest.
+    """
+    parts: dict[str, Any] = {"config": canonical_config(cfg)}
+
+    if repos:
+        parts["repos"] = {
+            name: _repo_tree(info) for name, info in sorted(repos.items())
+        }
+
+    if data:
+        parts["data"] = {
+            name: hash_data(path, use_cache=use_cache)
+            for name, path in sorted(data.items())
+        }
+
+    blob = json.dumps(parts, sort_keys=True, default=str)
+    return hashlib.sha256(blob.encode("utf-8")).hexdigest()

flexlock/flexcli.py

@@ -0,0 +1,163 @@
+"""Configuration decorator for FlexLock with progressive framework support."""
+
+import sys
+import os
+import functools
+from loguru import logger
+import inspect
+from typing import Dict, Optional
+from .runner import FlexLockRunner
+from .utils import py2cfg, instantiate
+from .debug import debug_on_fail
+from .snapshot import snapshot
+from . import config
+
+
+def _is_jupyter_or_interactive():
+    """
+    Check if running in Jupyter/IPython interactive environment.
+
+    Returns True for:
+    - Jupyter notebooks
+    - Jupyter console
+    - VSCode interactive windows
+    - IPython REPL
+    """
+    try:
+        from IPython import get_ipython
+
+        ipython = get_ipython()
+        if ipython is not None:
+            return True
+    except ImportError:
+        pass
+
+    # Also check for python -i
+    if hasattr(sys, "ps1"):
+        return True
+
+    return False
+
+
+def _should_use_cli_mode():
+    """
+    Determine if we should parse CLI arguments or execute directly.
+
+    Returns False if:
+    - sys.argv[0] is ipykernel_launcher.py (Jupyter kernel startup)
+    - sys.argv has kernel arguments like --f=... or -f ...
+
+    Returns True if:
+    - Normal CLI execution
+    - IPython %run command (sys.argv[0] is the script name)
+    - Script execution
+
+    This allows:
+    - %run file.py -o param=2 → Parse CLI args ✓
+    - train() in Jupyter → Don't parse ✗
+    - ipykernel_launcher.py --f=kernel.json → Don't parse ✗
+    """
+    if len(sys.argv) == 0:
+        # No arguments, likely interactive
+        logger.debug("No sys.argv detected, assuming interactive mode")
+        return False
+
+    # Check sys.argv[0] for kernel launcher
+    if sys.argv[0].endswith("ipykernel_launcher.py"):
+        logger.debug(
+            f"Detected ipykernel_launcher in sys.argv[0], skipping CLI parsing"
+        )
+        return False
+
+    # Check for kernel connection file arguments
+    for arg in sys.argv[1:]:
+        if "--f=" in arg or arg.startswith("-f"):
+            # Check if next arg or same arg contains .json (connection file)
+            if ".json" in arg or (
+                sys.argv.index(arg) + 1 < len(sys.argv)
+                and ".json" in sys.argv[sys.argv.index(arg) + 1]
+            ):
+                logger.debug(f"Detected Jupyter kernel connection file argument: {arg}")
+                return False
+
+    # If we get here, either:
+    # - Normal script execution
+    # - %run command (sys.argv[0] is script name, args look normal)
+    # Both should parse CLI
+    logger.debug(f"CLI mode enabled. sys.argv: {sys.argv}")
+    return True
+
+
+def flexcli(_func=None, snapshot_config: Optional[Dict] = None, **defaults):
+    """
+    Decorator for FlexLock entry points with progressive framework support.
+
+    Args:
+        snapshot_config: Dict defining repos, data, and lineage tracking.
+                         Injects '_snapshot_' into the config.
+        **defaults: Default values for the function arguments.
+
+    Usage::
+
+        @flexcli
+        def main(param=10, db="postgres"): ...
+
+        @flexcli(param=10)
+        def main(param, db="postgres"): ...
+
+        @flexcli(
+            data_path="./data/mnist",
+            snapshot_config={
+                "repos": {"main": "."},
+                "data": {"input_dataset": "${data_path}"}
+            }
+        )
+        def train(data_path, save_dir=None): ...
+    """
+
+    def decorator(fn):
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            if len(args) > 0 or len(kwargs) > 0:
+                # Direct call with arguments - execute immediately
+                # Apply debug wrapper if enabled
+                if config.get_env_bool("FLEXLOCK_DEBUG", False):
+                    return debug_on_fail(fn)(*args, **kwargs)
+                else:
+                    return fn(*args, **kwargs)
+
+            # Called with no arguments - check if we should parse CLI
+            # 1. Build Base Config from signature + defaults
+            base_cfg = py2cfg(fn, **defaults)
+
+            # 2. Inject Snapshot Config if provided
+            if snapshot_config:
+                base_cfg["_snapshot_"] = snapshot_config
+
+            # 3. Determine execution mode
+            if _should_use_cli_mode():
+                # CLI mode (normal script or %run): Launch Runner to parse arguments
+                runner = FlexLockRunner()
+                return runner.run(base_cfg=base_cfg)
+            else:
+                # Interactive mode (Jupyter kernel): Execute directly with defaults
+                # This happens when train() is called in Jupyter without %run
+                logger.info(
+                    "Interactive mode detected: Executing with defaults (no CLI parsing)"
+                )
+
+                # Execute with debug wrapper if enabled
+                if config.get_env_bool("FLEXLOCK_DEBUG", False):
+                    return debug_on_fail(fn)(**defaults)
+                else:
+                    return fn(**defaults)
+
+        wrapper._original_fn = fn
+        wrapper._defaults = defaults  # Store metadata
+        return wrapper
+
+    # Support both @flexcli and @flexcli(...) syntax
+    if _func is None:
+        return decorator
+    else:
+        return decorator(_func)

flexlock/freeze.py

@@ -0,0 +1,291 @@
+"""Partial-freeze of OmegaConf interpolations when detaching a sub-node.
+
+``select_and_freeze_root_refs`` selects a sub-node out of a root config and
+makes it *self-contained* — pickle-able, mergeable, and safe to hand to a
+worker — by resolving the references that point *outside* the sub-tree while
+preserving the ones that must stay live.
+
+OmegaConf resolves interpolations lazily against the config's root at access
+time, so a detached node loses every ``${...}`` that pointed to the root. But
+plain ``to_container(resolve=True)`` is too blunt: it would also fire resolver
+calls (``${vinc:}``, ``${latest:}``, ``${run_lock:}``) that must run *later*,
+on the worker, at submit time. So this module walks the ``${...}`` grammar and
+decides, per interpolation, whether to freeze or preserve it:
+
+- **Simple refs outside the sub-tree** (``${root_anchor}``, ``${a.b}``) are
+  resolved against the root by OmegaConf itself — following multi-hop chains,
+  mixed strings (``${pipeline_dir}/split``), and nested resolver-arg refs — and
+  baked to concrete values. Delegating to OmegaConf (rather than hand-rolling a
+  chain-follower) is what makes the historical ``_resolve_in_root`` class of
+  bugs structurally impossible.
+- **Resolver calls** (``${name:args}``) are preserved as frozen call strings
+  (their arguments are pre-resolved by OmegaConf under the freeze stubs), so
+  they fire later at submit/stage time.
+- **Relative refs** (``${.foo}``) and **refs inside the sub-tree** are preserved
+  verbatim.
+
+The only public entry points are :func:`select_and_freeze_root_refs` and
+:func:`freeze_deferred`; everything else is a private helper for the scan.
+"""
+
+from omegaconf import OmegaConf, DictConfig, ListConfig
+
+# Sentinel: an external ref that OmegaConf could not resolve from the root.
+_UNRESOLVED = object()
+
+
+def freeze_deferred(cfg: "DictConfig | ListConfig") -> "DictConfig | ListConfig":
+    """Eager-resolve every interpolation except the deferred resolvers.
+
+    Deep-copies ``cfg`` (detaching it from any parent), then resolves it under
+    :func:`flexlock.resolvers.deferred_stubbed`. OmegaConf resolves all simple
+    refs, cross-tree refs, multi-hop chains, mixed strings, and relative refs
+    itself; the deferred resolvers (``run_lock``/``latest``) collapse to
+    self-contained call strings that fire later at stage start. The result is
+    a plain config of concrete values plus frozen deferred call strings, safe
+    to serialize with ``resolve=False``.
+    """
+    from .resolvers import deferred_stubbed
+
+    detached = OmegaConf.create(
+        OmegaConf.to_container(cfg, resolve=False, throw_on_missing=False)
+    )
+    with deferred_stubbed():
+        OmegaConf.resolve(detached)
+    return detached
+
+
+def select_and_freeze_root_refs(root_cfg: DictConfig, key: str | None) -> DictConfig:
+    """Select a sub-node from `root_cfg` and freeze its root-scope references.
+
+    The returned config is self-contained: it can be pickled, round-tripped
+    through ``OmegaConf.to_container``/``OmegaConf.create``, and merged with
+    other configs without losing context.
+
+    Three kinds of interpolations are handled distinctly:
+
+    - **Simple refs to keys outside the sub-tree** (e.g. ``${root_anchor}``)
+      are resolved against ``root_cfg`` (by OmegaConf) and substituted as
+      concrete values. If the target is itself a resolver call
+      (``${vinc:}``/``${run_lock:}``), the frozen call string is substituted so
+      the resolver still fires later.
+    - **Resolver calls** (``${name:args}``) are preserved as frozen call
+      strings, with any nested simple refs in their args resolved.
+    - **Relative refs** (``${.foo}``) and **refs to keys inside the sub-tree**
+      are preserved verbatim.
+
+    Args:
+        root_cfg: The full root configuration.
+        key: Dot-path to the sub-node. If ``None``, returns ``root_cfg``
+            unchanged.
+
+    Returns:
+        A new ``DictConfig`` (or the same root if ``key is None``).
+
+    Raises:
+        KeyError: If ``key`` is not found in ``root_cfg``.
+        UnresolvedInterpolationError: If the sub-node references a key that
+            exists in neither the sub-tree nor the root.
+    """
+    if key is None:
+        return root_cfg
+
+    sub_node = OmegaConf.select(root_cfg, key, throw_on_missing=False)
+    if sub_node is None:
+        raise KeyError(f"Key '{key}' not found in config")
+
+    if not isinstance(sub_node, (DictConfig, ListConfig)):
+        return sub_node
+
+    from .resolvers import frozen_resolvers
+
+    sub_raw = OmegaConf.to_container(sub_node, resolve=False, throw_on_missing=False)
+    # A detached working copy of the whole root. External refs are resolved
+    # against it by OmegaConf; the sub-tree stays attached for context. Stub
+    # every resolver so ${name:args} calls freeze to call strings instead of
+    # firing at selection time.
+    work = OmegaConf.create(
+        OmegaConf.to_container(root_cfg, resolve=False, throw_on_missing=False)
+    )
+    with frozen_resolvers():
+        transformed = _freeze_walk(sub_raw, sub_raw, work)
+    return OmegaConf.create(transformed)
+
+
+def _freeze_walk(value, sub_raw, work):
+    if isinstance(value, dict):
+        return {k: _freeze_walk(v, sub_raw, work) for k, v in value.items()}
+    if isinstance(value, list):
+        return [_freeze_walk(v, sub_raw, work) for v in value]
+    if isinstance(value, str) and "${" in value:
+        whole = _whole_string_interp(value)
+        if whole is not None:
+            return _freeze_whole_string(whole, sub_raw, work, fallback_str=value)
+        return _freeze_embedded(value, sub_raw, work)
+    return value
+
+
+# --- string-level scan helpers ------------------------------------------------
+
+
+def _whole_string_interp(s: str) -> "str | None":
+    """If ``s`` is exactly ``${...}`` (one balanced block, nothing else), return
+    the inner expression. Otherwise return None."""
+    if not s.startswith("${"):
+        return None
+    found = _find_balanced_interp(s, 0)
+    if found is None:
+        return None
+    start, end = found
+    if start == 0 and end == len(s):
+        return s[2:-1]
+    return None
+
+
+def _find_balanced_interp(s: str, start: int) -> "tuple[int, int] | None":
+    """Find the next ``${...}`` block in ``s`` starting at ``start``. Returns
+    ``(open_idx, close_idx_exclusive)`` or ``None``. Handles nested ``${...}``."""
+    i = s.find("${", start)
+    if i == -1:
+        return None
+    depth = 0
+    j = i
+    while j < len(s):
+        if s[j : j + 2] == "${":
+            depth += 1
+            j += 2
+        elif s[j] == "}":
+            depth -= 1
+            j += 1
+            if depth == 0:
+                return (i, j)
+        else:
+            j += 1
+    return None
+
+
+def _find_top_level_colon(inner: str) -> "int | None":
+    """In the content between ``${`` and ``}``, find the first ``:`` not nested
+    inside a ``${...}`` block. Returns the index, or ``None`` if no colon."""
+    depth = 0
+    for i, ch in enumerate(inner):
+        if inner[i : i + 2] == "${":
+            depth += 1
+        elif ch == "}" and depth > 0:
+            depth -= 1
+        elif ch == ":" and depth == 0:
+            return i
+    return None
+
+
+def _path_exists(d, dotted: str) -> bool:
+    cur = d
+    for p in dotted.split("."):
+        if isinstance(cur, dict) and p in cur:
+            cur = cur[p]
+        else:
+            return False
+    return True
+
+
+def _get_raw(d, dotted: str):
+    cur = d
+    for p in dotted.split("."):
+        if isinstance(cur, dict) and p in cur:
+            cur = cur[p]
+        else:
+            return None
+    return cur
+
+
+# --- freezing --------------------------------------------------------------
+
+
+def _resolve_external(ref: str, work):
+    """Resolve a simple ref against the root via OmegaConf.
+
+    OmegaConf follows the full chain (multi-hop, mixed strings, nested
+    resolver-arg refs). Under the freeze stubs, a target that is itself a
+    resolver call comes back as a frozen call string. Returns ``_UNRESOLVED``
+    when the ref is absent from the root.
+    """
+    val = OmegaConf.select(work, ref, throw_on_missing=False, default=_UNRESOLVED)
+    if isinstance(val, (DictConfig, ListConfig)):
+        return OmegaConf.to_container(val, resolve=False)
+    return val
+
+
+def _is_internal(ref: str, sub_raw) -> bool:
+    """A simple ref is *internal* (preserve verbatim) when it resolves within
+    the sub-tree and isn't self-shadowing (its own value isn't ``${ref}``)."""
+    if not _path_exists(sub_raw, ref):
+        return False
+    sub_val = _get_raw(sub_raw, ref)
+    return not (isinstance(sub_val, str) and f"${{{ref}}}" in sub_val)
+
+
+def _unresolved_error(ref: str):
+    from .exceptions import UnresolvedInterpolationError
+
+    first = ref.split(".")[0]
+    raise UnresolvedInterpolationError(
+        f"Interpolation ${{{ref}}} could not be resolved: '{first}' not "
+        f"found in the sub-tree or root config. Set it via overrides= or "
+        f"OmegaConf.update(proj.defaults, '{first}', ...)."
+    )
+
+
+def _freeze_whole_string(inner: str, sub_raw, work, fallback_str: str):
+    """Freeze a whole-string interpolation ``${inner}``. Preserves the target's
+    native type for simple external refs (e.g. ``${batch_size}`` → int)."""
+    if _find_top_level_colon(inner) is not None:
+        # Resolver call — preserve as a frozen call string.
+        return "${" + _freeze_resolver_args(inner, sub_raw, work) + "}"
+    if inner.startswith("."):
+        return fallback_str  # relative ref — resolved later at access site
+    if _is_internal(inner, sub_raw):
+        return fallback_str  # intra-sub-tree ref — resolves after detachment
+    val = _resolve_external(inner, work)
+    if val is _UNRESOLVED:
+        _unresolved_error(inner)
+    return val
+
+
+def _freeze_embedded(s: str, sub_raw, work) -> str:
+    """Freeze every ``${...}`` block embedded in a larger string."""
+    out = []
+    pos = 0
+    while pos < len(s):
+        found = _find_balanced_interp(s, pos)
+        if found is None:
+            out.append(s[pos:])
+            break
+        start, end = found
+        out.append(s[pos:start])
+        inner = s[start + 2 : end - 1]
+        out.append(_freeze_one_block(inner, sub_raw, work))
+        pos = end
+    return "".join(out)
+
+
+def _freeze_one_block(inner: str, sub_raw, work) -> str:
+    """Freeze one ``${inner}`` block, returning the ``${...}`` text to embed."""
+    if _find_top_level_colon(inner) is not None:
+        return "${" + _freeze_resolver_args(inner, sub_raw, work) + "}"
+    if inner.startswith("."):
+        return "${" + inner + "}"  # relative ref — preserved
+    if _is_internal(inner, sub_raw):
+        return "${" + inner + "}"  # intra-sub-tree ref — preserved
+    val = _resolve_external(inner, work)
+    if val is _UNRESOLVED:
+        _unresolved_error(inner)
+    return str(val) if val is not None else "null"
+
+
+def _freeze_resolver_args(inner: str, sub_raw, work) -> str:
+    """Freeze the ``name:args`` body of a resolver call — the call is preserved,
+    nested simple refs in the args are frozen."""
+    colon = _find_top_level_colon(inner)
+    name = inner[:colon]
+    args = inner[colon + 1 :]
+    return name + ":" + _freeze_embedded(args, sub_raw, work)