FlexLock API Reference

FlexLock: A lightweight library for reproducible ML experiments.

exception flexlock.FlexLockBackendError[source]

Bases: FlexLockError

Raised when there is an error with HPC backends (Slurm/PBS).

exception flexlock.FlexLockCacheError[source]

Bases: FlexLockError

Raised when there is an error with smart run caching.

exception flexlock.FlexLockConfigError[source]

Bases: FlexLockError

Raised when there is an error in configuration.

exception flexlock.FlexLockError[source]

Bases: Exception

Base exception for all FlexLock errors.

exception flexlock.FlexLockExecutionError[source]

Bases: FlexLockError

Raised when there is an error during execution.

class flexlock.FlexLockRunner[source]

Bases: object

load_config(args)[source]
run(cli_args=None, base_cfg=None)[source]

Thin layer over Project.submit() — shapes CLI args and dispatches.

exception flexlock.FlexLockSnapshotError[source]

Bases: FlexLockError

Raised when there is an error creating or loading snapshots.

exception flexlock.FlexLockValidationError[source]

Bases: FlexLockError

Raised when validation of inputs fails.

class flexlock.Project(defaults: str | DictConfig | dict | None = None)[source]

Bases: object

check(config=None, *, sweep: List[Dict] | None = None, sweep_target: str | None = None, overrides: dict | List[str] | None = None, merge: str | Path | dict | None = None) → List[dict][source]

Preflight: fully resolve the config (and every sweep item) without touching the filesystem or executing anything.

Mirrors submit’s merge → freeze pipeline, but runs under the freeze stubs so no resolver fires (vinc/now take no mkdir, run_lock/latest read nothing). Every unresolvable interpolation is reported — not just the first — as a dict with item (sweep index or None), full_key, and error.

Returns an empty list when everything resolves.

static collect_results(indices, task_configs, db_path, tag, job_id=None) → list[source]

Build per-item ExecutionResults from the task DB’s terminal state.

Replaces the old blanket status="SUCCESS" — a task that raised is now reported as FAILED with its traceback, so a sweep result set faithfully reflects what happened (issue 2).

exists(cfg: DictConfig, search_dirs: List[str] = None) → bool[source]

Check if a run with the given configuration already exists.

Parameters:
  • cfg – Configuration to check

  • search_dirs – Optional list of directories to search

Returns:

True if matching run exists, False otherwise

get(key: str)[source]

Get a configuration by key from the defaults.

The returned config is self-contained: root-scope ${...} references are frozen into concrete values, while resolver calls (${vinc:}, ${latest:}, ${run_lock:}, etc.) and intra-sub-tree refs are preserved for resolution at submit time. This means the returned config can be pickled, modified, and submitted without losing root context (essential for HPC submission).

Parameters:

key – Dot-path to select a specific node from the defaults.

Returns:

The selected configuration (as DictConfig).

get_result(cfg: DictConfig, search_dirs: List[str] = None) → ExecutionResult[source]

Retrieve results from a previously completed run.

Parameters:
  • cfg – Configuration to match

  • search_dirs – Optional list of directories to search

Returns:

ExecutionResult object with cached results

Raises:

ValueError – If no matching run is found

override_module = None
run_stage(cfg, stage_name=None, smart_run=True, search_dirs=None, **submit_kwargs)[source]

Run a single stage with automatic search_dirs and save_dir propagation.

Parameters:
  • cfg – Stage configuration (DictConfig)

  • stage_name – Name of the stage (inferred from save_dir if None)

  • smart_run – If True, checks for cached runs

  • search_dirs – Directories to search for cached runs (auto-discovered if None)

  • **submit_kwargs – Additional arguments passed to submit()

Returns:

ExecutionResult (or list for sweeps)

save_snapshot(save_dir)[source]

Save the current project defaults as pipeline.yaml.

Parameters:

save_dir – Directory to save the pipeline snapshot to

submit(config: DictConfig | str | None = None, sweep: List[Dict] = None, sweep_target: str = None, sweep_root: str | None = None, n_jobs: int = 1, smart_run: bool = True, search_dirs: List[str] = None, wait: bool = True, pbs_config: str = None, slurm_config: str = None, sweep_dir_suffix: bool = False, match_include: List[str] = None, match_exclude: List[str] = None, isolated: bool = False, force: bool = False, overrides: dict | List[str] | None = None, merge: str | Path | dict | None = None, debug: bool = False, print_config: bool = False, dry_run: bool = False, tag: str | None = None, timeout: int | None = None, note: str | None = None, save_dir_policy: str | None = None, task_record: str | None = None, after=None) → ExecutionResult | List[ExecutionResult] | None[source]

Submit a configuration for execution.

Parameters:
  • config – The configuration to execute. Accepts a DictConfig, a key string (looked up via get()), or None (uses self.defaults as-is).

  • sweep – Optional list of override dicts for parameter sweep.

  • sweep_target – Dot-path inside each task config where the sweep item is merged. When None, items merge at the root of config.

  • n_jobs – Number of parallel workers (for sweeps).

  • smart_run – If True, checks for an existing cached run before executing.

  • search_dirs – Directories to search for cached runs.

  • wait – If True, blocks until completion.

  • slurm_config (pbs_config /) – Path to HPC backend YAML.

  • after – Job IDs (sequence or colon-separated string) that must finish successfully before this Slurm/PBS job starts. Requires an HPC backend. Submitted results expose the scheduler ID as job_id. Unpinned run references bind on the worker; pins bind at submit.

  • sweep_dir_suffix – If True, append _sweep_{i:04d} to each sweep item’s save_dir.

  • match_exclude (match_include /) – Override git path patterns used during smart_run comparison.

  • isolated – If True, run in a spawned subprocess even for a single task (use for GPU stages to confine the CUDA context).

  • force – If True, invalidate the cache marker (run.complete) for this save_dir and re-execute. Outputs and run.lock are preserved; the user function overwrites in place.

  • overrides – Dict ({'lr': 0.01}) or dotlist (['lr=0.01']) merged into config before execution.

  • merge – Path to a YAML file (or a dict) merged into config before execution. overrides is applied after merge.

  • debug – Wrap the user function with the post-mortem debugger so exceptions drop into PDB.

  • print_config – Print the fully-resolved config and return None without executing.

  • dry_run – When using an HPC backend, render the would-be Slurm or PBS submission script, print it (with any validation warnings), and return None without submitting. No-op for local execution.

  • tag – Human-readable label for this sweep’s rows in the shared task DB (e.g. "extract" or "collocate"). Passed through to ParallelExecutor so workers and the status CLI can scope to it. When None a deterministic hash is auto-generated.

  • note – Free-text intent recorded as a top-level note: key in run.lock (sibling of timestamp/config). Never enters the fingerprint, so it can’t perturb caching. For sweeps the note lands on the master run.lock only; sweep items inherit it for display via their .flexlock_marker → master lookup.

  • save_dir_policy – What to do when config.save_dir is already occupied by a previous run (contains run.lock), applied once at submit time, after the smart_run cache check: None/"raise" (default) refuses and raises; "increment" versions the dir (run → run_0000, claimed atomically); "overwrite" deletes the occupied dir’s contents first (never touches a dir without run.lock); "skip" returns the existing complete run’s result without executing; "unsafe" runs in place with no check; "timestamp" appends config.TIMESTAMP_FORMAT. force=True bypasses the default guard (explicit in-place rerun). Replaces the ${vinc:} / ${now:} resolvers. For sweeps the policy applies once to the sweep root; items nest beneath it, and "skip" means resume (reuse complete items, run the rest).

  • task_record – Where each sweep task’s record goes: "dir" (default, or $FLEXLOCK_TASK_RECORD) writes a full run.lock into the task’s save_dir; "db" keeps it in the task DB only (fewer files for very large sweeps). Single runs, including HPC ones, always get a run.lock. Readers find either through flexlock.record.load_record().

Returns:

ExecutionResult (single), List[ExecutionResult] (sweep), or None (print_config=True).

submit_chained(config=None, *, sweep: List[Dict] | None = None, downstream: List[tuple] | None = None, sweep_kwargs: dict | None = None, downstream_kwargs: dict | None = None) → ChainedResult[source]

Run a sweep, then chain a sequence of downstream stages per result.

For each item in sweep, this method submits the base config and then, for each (stage_key, anchor_wiring) in downstream, propagates the sweep result’s attributes into proj.defaults and submits the downstream stage.

anchor_wiring maps an anchor name in proj.defaults to an attribute of the parent ExecutionResult (typically 'save_dir').

Parameters:
  • config – Base config for the sweep (passed to submit()).

  • sweep – List of override dicts for the parameter sweep.

  • downstream – List of (stage_key, anchor_wiring) tuples describing the stages to run per sweep result.

  • sweep_kwargs – Extra kwargs forwarded to submit for the sweep itself (e.g. slurm_config, smart_run).

  • downstream_kwargs – Extra kwargs forwarded to submit for every downstream stage. Defaults to smart_run=False to avoid stale-cache false hits between iterations.

Returns:

ChainedResult with sweep (list) and downstream (a list of lists, one per sweep item, each inner list aligned to the order of downstream).

classmethod submit_config(config=None, **kwargs)[source]

Shortcut for one-off submissions without holding a Project instance.

Equivalent to Project().submit(config, **kwargs). Use this when you already have a DictConfig (from py2cfg or otherwise) and don’t need proj.get/proj.exists/proj.defaults plumbing.

See flexlock.submit() for the module-level alias.

submit_pipeline(items: List[List[DictConfig]], *, n_jobs: int = 1, smart_run: bool = False, search_dirs: List[str] = None, slurm_config: str = None, pbs_config: str = None, wait: bool = True, sweep_root: str | None = None, tag: str | None = None, force: bool = False, timeout: int | None = None, note: str | None = None, save_dir_policy: str | None = None, debug: bool = False, dry_run: bool = False, after=None) → List[List[ExecutionResult]] | None[source]

Submit composite pipeline tasks: one task per item, stages in order.

Each item is a list of already-compiled, frozen stage configs (the caller applied freeze_deferred(); deferred resolvers fire on the worker at stage start). One item becomes one composite task in the sweep task DB ({"_stages_": [...]}): the worker runs its stages sequentially, so intra-item ordering holds by construction, while parallelism (local n_jobs, Slurm/PBS array workers) fans out across items.

Composite tasks are self-contained — the base config is never re-merged into them at dequeue time.

Known limitation: all stages of an item run inside one scheduler job / worker slot, so per-stage HPC resource configs are not possible.

Parameters:
  • items – items[i][j] is the j-th stage config of item i.

  • n_jobs – Parallel workers fanning out across items.

  • smart_run – Driver-side cache check — an item whose every stage cache-hits returns CACHED without queueing; partial hits queue the item with stage-level resume (_skip_complete_).

  • search_dirs – Directories searched for cached runs (smart_run).

  • pbs_config (slurm_config /) – Path to HPC backend YAML.

  • wait – Block until completion (HPC). False returns SUBMITTED skeleton results.

  • after – Scheduler jobs that must finish successfully before this job starts. Each submitted stage result exposes the same job_id. Unpinned run references bind at stage start; pins bind at submit.

  • sweep_root – Override the master root hosting the task DB; also opts out of the containment validation.

  • timeout (tag / note /) – As in submit().

  • force – Invalidate every stage’s completion markers and reset the task DB so all items rerun.

  • save_dir_policy – Collision guard applied per stage dir before queueing — raise (default, unless force), overwrite, skip (stage-level resume), unsafe. Naming policies (increment/timestamp) don’t compose with pre-baked cross-stage paths and are rejected.

  • debug – Post-mortem debugger (local in-process execution only).

  • dry_run – With an HPC backend, render the submission script and return None without queueing anything.

Returns:

results[i][j] is the ExecutionResult of item i’s stage j (or None for dry_run).

flexlock.debug_on_fail(fn=None, **legacy_kwargs)[source]

A decorator that provides enhanced debugging on exception.

Features: - Smart frame selection (prefers project code over libraries) - Frame navigation in notebooks (_debug_up, _debug_down, _debug_goto) - Handles C extensions gracefully - PDB post-mortem for scripts - Configurable via environment variables

Environment Variables:
FLEXLOCK_NODEBUG: Set to ‘1’ or ‘true’ to disable this decorator

(turns it into a no-op).

FLEXLOCK_DEBUG_STRATEGY: ‘auto’ (default), ‘pdb’, or ‘inject’.

Note: FLEXLOCK_DEBUG=1 is consumed by flexlock-run and @flexcli to opt into wrapping the user function with this decorator. Setting FLEXLOCK_DEBUG=0 does not disable a manually applied @debug_on_fail — use FLEXLOCK_NODEBUG=1 for that.

Usage:

@debug_on_fail
def my_function():
    ...

In notebooks after exception, the following helpers are injected:

  • _debug_show(): Show all frames

  • _debug_up(): Move to caller

  • _debug_down(): Move toward exception

  • _debug_goto(n): Jump to frame n

flexlock.enqueue_to_file(path: str | Path, cfg_dict: dict) → int[source]

Atomically append cfg_dict to a YAML list file and return the new queue length.

The file is created if it does not exist. Existing content must be a YAML list (or an empty file). The write is atomic: a temporary sibling file is written first, then renamed over the target so partial writes are never visible to concurrent readers.

Parameters:
  • path – Destination YAML file (created if absent).

  • cfg_dict – A plain Python dict (no OmegaConf nodes) to append.

Returns:

Number of items in the queue after appending.

flexlock.flexcli(_func=None, snapshot_config: Dict | None = None, **defaults)[source]

Decorator for FlexLock entry points with progressive framework support.

Parameters:
  • 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): ...
flexlock.get_git_tree_hash(path: str = '.') → str[source]

Gets the current git tree hash for a repository without creating a new commit. This represents the content fingerprint of the repository.

Parameters:

path (str) – The path to the git repository.

Returns:

The tree hash.

Return type:

str

Raises:

FlexLockSnapshotError – if path is not a usable git repository.

flexlock.hash_data(path, match=None, ignore=None, jobs=4, algorithm=<class 'xxhash.xxh64'>, chunk_size=262144, use_cache=True)[source]

Computes a hash for a file or a directory, using an SQLite cache to avoid re-computation.

flexlock.load_python_defaults(import_path: str)[source]

Dynamically import a module or file path to retrieve a variable.

Accepted forms:

  • 'pkg.config.defaults' — dotted module path, variable is the last segment.

  • 'pkg.config:defaults' — module + colon-separated variable name.

  • 'configs/defaults.py:defaults' — file path + colon-separated variable name.

  • 'configs/defaults.py' — bare file path, variable defaults to defaults by convention.

flexlock.load_sweep(*, sweep: list | str | None = None, sweep_file: str | Path | list[str | Path] | None = None, sweep_key: str | None = None, root_cfg: DictConfig | None = None) → list[source]

Load a sweep list from one of several sources.

Exactly one source must be provided. Returns the parsed sweep list. If the source yields a single value or dict, it is wrapped in a list.

Parameters:
  • sweep – A pre-built list, or a CLI-style comma-separated string.

  • sweep_file – Path (or list of paths) to .yaml/.yml, .json, or text files. When a list is given, each file is loaded independently and results are concatenated. A file containing a single dict is treated as one task; a file containing a list expands into multiple tasks.

  • sweep_key – Dotted key into root_cfg whose value is the sweep list.

  • root_cfg – Required when sweep_key is used.

flexlock.log_to_file(path)[source]
flexlock.mlflow_context(save_dir: str | Path, experiment_name: str = 'Default', run_name: str = None, tags: Dict[str, str] = None, log_config: bool = True, log_artifacts: bool = True)[source]

Manages a “Shadow Run” in MLflow linked to a physical directory.

This context manager implements the “Always New + Deprecate Old” strategy, treating the MLflow Run as a temporary view of the persistent physical state (the save_dir).

Behavior: 1. SEARCH: Finds the currently ‘active’ run for this directory. 2. CREATE: Starts a completely new run. 3. INHERIT: Logs existing config/artifacts found in save_dir to the new run. 4. INHERIT TAGS: Copies user tags (e.g., pipeline tags) from previous run. 5. SWAP: Marks new run ‘active’, marks old run ‘deprecated’.

Parameters:
  • save_dir – The directory containing the run data (e.g., outputs/exp/run_01)

  • experiment_name – MLflow experiment name (default: “Default”)

  • run_name – Name for this MLflow run (default: save_dir.name)

  • tags – Additional tags to add to the run (overrides inherited tags)

  • log_config – If True, logs parameters from run.lock file

  • log_artifacts – If True, logs standard artifacts (logs, lock file)

Yields:

mlflow.ActiveRun – The active MLflow run object, or None if MLflow not available

Note

User tags (not starting with ‘flexlock.’ or ‘mlflow.’) are automatically inherited from the previous active run. This is useful for pipeline workflows where you want tags like ‘pipeline_run=xyz’ to persist across stages. Explicitly passing a tag via the tags parameter will override the inherited value.

Example

>>> # Stage 1: Train (creates run with pipeline tag)
>>> with mlflow_context(save_dir, tags={"pipeline_run": "run_001"}):
...     mlflow.log_metric("train_loss", 0.32)
...
>>> # Stage 2: Evaluate (inherits pipeline_run tag automatically)
>>> with mlflow_context(save_dir):  # pipeline_run=run_001 inherited!
...     mlflow.log_metric("test_accuracy", 0.95)
flexlock.parse_sweep_string(sweep_str: str) → list[source]

Parse a comma-separated CLI sweep value into a list.

Examples:

parse_sweep_string("1,2,3")              # → [1, 2, 3]
parse_sweep_string("lr=0.1,lr=0.2")      # → [{'lr': 0.1}, {'lr': 0.2}]
parse_sweep_string('"a,b",c')            # → ['a,b', 'c']
flexlock.py2cfg(obj, /, *pos, **overrides)[source]

Generates a default configuration dict from a function or class signature. Supports nested py2cfg calls and handles decorated functions.

flexlock.select_and_freeze_root_refs(root_cfg: DictConfig, key: str | None) → DictConfig[source]

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.

Parameters:
  • 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.

flexlock.snapshot(cfg, repos=None, data=None, prevs=None, parent_lock=None, save_path=None, return_snapshot=False, fingerprint=None, note=None, meta=None)[source]

Create a snapshot of the current run state.

Parameters:
  • 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

  • meta – Extra top-level keys for the record (e.g. a sweep master’s task_record). Never part of the fingerprint.

Returns:

dict if return_snapshot=True, else None

flexlock.submit(config=None, **kwargs)[source]

Submit a configuration without instantiating a Project.

Sugar over Project().submit(config, **kwargs) — use when you have a ready DictConfig and don’t need the multi-stage pipeline plumbing (proj.get, proj.defaults, etc.). All Project.submit kwargs (sweep, slurm_config, smart_run, overrides, …) are forwarded as-is.

Example:

from flexlock import submit, py2cfg
cfg = py2cfg(train, lr=0.01, save_dir='outputs/train')
result = submit(cfg, slurm_config='configs/slurm_gpu.yaml')

Base classes for FlexLock backend integrations.

class flexlock.backends.base.Backend[source]

Bases: ABC

Abstract base class for FlexLock job submission backends.

abstractmethod environment() → JobEnvironment[source]

Returns a JobEnvironment object providing environment-specific variables.

abstractmethod submit(fn, *args, **kwargs) → Job[source]

Submits a single function for execution.

class flexlock.backends.base.Job[source]

Bases: object

Abstract base class representing a submitted job.

property job_id: str

The unique identifier for the job.

class flexlock.backends.base.JobEnvironment[source]

Bases: object

Abstract base class providing environment information for a running job.

property global_rank: int

The global rank of the current process within the job (0-indexed).

property world_size: int

The total number of processes in the job.

PBS backend for FlexLock parallel execution.

class flexlock.backends.pbs.PBSBackend(folder: Path, startup_lines: list[str], configure_logging: bool = True, configure_name: bool = True, python_exe: str = 'python', after=None)[source]

Bases: Backend

Implements the FlexLock backend for PBS (Portable Batch System) job submission.

TERMINAL_STATES = frozenset({'C', 'F'})
cancel_job(job_id: str) → bool[source]

Cancel a PBS job.

Parameters:

job_id – PBS job identifier

Returns:

True if cancellation succeeded, False otherwise

check_status(job_id: str) → str[source]

Check the status of a PBS job.

Returns:

‘Q’ (queued), ‘R’ (running), ‘C’ (completed), ‘E’ (exiting), ‘H’ (held), or ‘unknown’

Return type:

Status string

environment()[source]

Returns a JobEnvironment object providing PBS-specific environment variables.

is_terminal(job_id: str) → bool[source]

Return True only when PBS confirms the job has ended.

Conservative on purpose (see SlurmBackend.is_terminal): ambiguous reads return False so a healthy job is never reconciled away. The controller additionally debounces across several polls, which guards against qstat transiently failing (check_status maps a not-found job to "C").

render_script(pickled_path=None) → str[source]

Render the submission script without submitting or pickling.

submit(fn, *args, **kwargs)[source]

Submits a single function for execution as a PBS job.

wait_for_job(job_id: str, timeout=None, poll_interval=5) → bool[source]

Wait for a PBS job to complete.

Parameters:
  • job_id – PBS job identifier

  • timeout – Maximum time to wait in seconds (None for no timeout)

  • poll_interval – Time between status checks in seconds

Returns:

True if job completed successfully, False otherwise

class flexlock.backends.pbs.PBSJob(job_id, backend=None)[source]

Bases: Job

Represents a PBS job.

cancel()[source]

Cancel the job.

property job_id

The unique identifier for the job.

status()[source]

Get current job status.

wait(timeout=None, poll_interval=5)[source]

Wait for job to complete.

Slurm backend for FlexLock parallel execution.

class flexlock.backends.slurm.SlurmBackend(folder: Path, startup_lines: list[str], configure_logging: bool = True, python_exe='python', after=None)[source]

Bases: Backend

Implements the FlexLock backend for Slurm job submission.

TERMINAL_STATES = frozenset({'BOOT_FAIL', 'CANCELLED', 'COMPLETED', 'DEADLINE', 'FAILED', 'NODE_FAIL', 'OUT_OF_MEMORY', 'PREEMPTED', 'REVOKED', 'SPECIAL_EXIT', 'TIMEOUT'})
cancel_job(job_id: str) → bool[source]

Cancel a Slurm job.

Parameters:

job_id – Slurm job identifier

Returns:

True if cancellation succeeded, False otherwise

check_status(job_id: str) → str[source]

Check the status of a Slurm job.

Returns:

‘PENDING’, ‘RUNNING’, ‘COMPLETED’, ‘FAILED’, ‘CANCELLED’, or ‘unknown’

Return type:

Status string

environment()[source]

Returns a JobEnvironment object providing Slurm-specific environment variables.

is_terminal(job_id: str) → bool[source]

Return True only when the scheduler confirms the job has ended.

Used by the controller to decide whether tasks still marked running/pending have been orphaned by a dead job. This is deliberately conservative: unknown (e.g. squeue doesn’t list a just-submitted job yet and sacct hasn’t recorded it) and all active states return False, so a transient/racy status read never causes a healthy job to be reconciled away.

render_script(pickled_path: Path | str | None = None) → str[source]

Render the would-be Slurm script without submitting.

pickled_path defaults to a placeholder path so users can preview the script (e.g. via dry_run=True) before any pickling happens.

submit(fn, *args, **kwargs)[source]

Submits a single function for execution as a Slurm job.

wait_for_job(job_id: str, timeout=None, poll_interval=5) → bool[source]

Wait for a Slurm job to complete.

Parameters:
  • job_id – Slurm job identifier

  • timeout – Maximum time to wait in seconds (None for no timeout)

  • poll_interval – Time between status checks in seconds

Returns:

True if job completed successfully, False otherwise

class flexlock.backends.slurm.SlurmJob(job_id, backend=None)[source]

Bases: Job

Represents a Slurm job.

cancel()[source]

Cancel the job.

property job_id

The unique identifier for the job.

status()[source]

Get current job status.

wait(timeout=None, poll_interval=5)[source]

Wait for job to complete.

flexlock.backends.slurm.validate_slurm_script(script: str, *, expects_gpu: bool = False) → list[str][source]

Return a list of human-readable warning strings for a Slurm script.

Empty list means no concerns. Each warning is independent — callers decide whether to log or print. Validation is intentionally lenient (warn, never fail) because power-user setups may intentionally omit pieces this function expects.

Shared context for a FlexLock run.

Data hashing utilities for FlexLock.

flexlock.data_hash.dirhash(path, match=None, ignore=None, jobs=1, algorithm=<built-in function openssl_md5>, chunk_size=65536)[source]

Computes a hash of the directory content using pathlib.glob for filtering.

flexlock.data_hash.hash_data(path, match=None, ignore=None, jobs=4, algorithm=<class 'xxhash.xxh64'>, chunk_size=262144, use_cache=True)[source]

Computes a hash for a file or a directory, using an SQLite cache to avoid re-computation.

Debugging utilities for FlexLock.

flexlock.debug.debug_on_fail(fn=None, **legacy_kwargs)[source]

A decorator that provides enhanced debugging on exception.

Features: - Smart frame selection (prefers project code over libraries) - Frame navigation in notebooks (_debug_up, _debug_down, _debug_goto) - Handles C extensions gracefully - PDB post-mortem for scripts - Configurable via environment variables

Environment Variables:
FLEXLOCK_NODEBUG: Set to ‘1’ or ‘true’ to disable this decorator

(turns it into a no-op).

FLEXLOCK_DEBUG_STRATEGY: ‘auto’ (default), ‘pdb’, or ‘inject’.

Note: FLEXLOCK_DEBUG=1 is consumed by flexlock-run and @flexcli to opt into wrapping the user function with this decorator. Setting FLEXLOCK_DEBUG=0 does not disable a manually applied @debug_on_fail — use FLEXLOCK_NODEBUG=1 for that.

Usage:

@debug_on_fail
def my_function():
    ...

In notebooks after exception, the following helpers are injected:

  • _debug_show(): Show all frames

  • _debug_up(): Move to caller

  • _debug_down(): Move toward exception

  • _debug_goto(n): Jump to frame n

Configuration decorator for FlexLock with progressive framework support.

flexlock.flexcli.flexcli(_func=None, snapshot_config: Dict | None = None, **defaults)[source]

Decorator for FlexLock entry points with progressive framework support.

Parameters:
  • 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): ...

Source code versioning utilities for FlexLock.

flexlock.git_utils.code_drift(repos: dict, since: float, modules=None) → dict[source]

Loaded source files that differ from the recorded tree.

A run records its code as a git tree at snapshot time, but Python loads a module when it is first imported: an edit made while the job waited in the queue, or before a lazy import, runs code the tree does not describe. This checks every module in modules (default sys.modules) whose file lies in a recorded repo and was modified after since (epoch seconds of the snapshot); a file counts as drifted when its content differs from the tree’s blob, or when it is a new, non-ignored file absent from the tree.

Parameters:
  • repos – The repos section of a run record ({name: {"path": ..., "tree": ...}}).

  • since – Snapshot time; files with an older mtime are skipped unread.

  • modules – Iterable of module objects (defaults to sys.modules).

Returns:

{repo name: [repo-relative paths]} for repos with drift; empty when the loaded code matches the recorded tree.

flexlock.git_utils.create_shadow_snapshot(repo_path: str = '.', ignore_patterns: list | None = None, ref_name: str | None = None) → dict[source]

Creates a Shadow Commit. Returns: {commit_hash, tree_hash, is_dirty}

flexlock.git_utils.create_shadow_tree(repo_path: str = '.', include: list | None = None, exclude: list | None = None) → dict[source]

Compute a content tree hash for a working tree without side effects.

Unlike create_shadow_snapshot(), this stages into a throwaway shadow index and runs ``write-tree`` only — it creates no commit object and no refs/flexlock/runs/* ref. It is therefore safe to call on every fingerprint check (smart-run) without accumulating objects/refs in .git.

Parameters:
  • repo_path – Path inside the repository.

  • include – Optional pathspec(s); when given, only these paths are staged (the fingerprint is restricted to the “relevant” subtree, so an include-match becomes plain tree-hash equality). Defaults to all tracked + untracked files.

  • exclude – Optional pathspec(s) removed from the staged index.

Returns:

{"tree": <hash>, "is_dirty": <bool>}.

Return type:

dict

flexlock.git_utils.get_git_commit(path: str = '.') → str[source]

Gets the current commit hash for a git repository without creating a new commit.

Parameters:

path (str) – The path to the git repository.

Returns:

The commit hash.

Return type:

str

Raises:

FlexLockSnapshotError – if path is not a usable git repository.

flexlock.git_utils.get_git_tree_hash(path: str = '.') → str[source]

Gets the current git tree hash for a repository without creating a new commit. This represents the content fingerprint of the repository.

Parameters:

path (str) – The path to the git repository.

Returns:

The tree hash.

Return type:

str

Raises:

FlexLockSnapshotError – if path is not a usable git repository.

flexlock.git_utils.sanitize_ref_name(name: str) → str[source]

Sanitize a string to be a valid git ref name.

flexlock.git_utils.shadow_index(repo: Repo)[source]

Context manager for Git Plumbing operations without touching user index.

Diff utilities for FlexLock.

class flexlock.diff.RunDiff(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)[source]

Bases: object

compare_config()[source]

Compare configurations with recursive diff, ignoring specified keys and normalizing path values.

compare_data()[source]

Compare data hashes.

compare_env()[source]

Compare environment lockfile hashes.

A target without an env section predates environment hashing; it is treated as a match so existing caches keep working (the index backfill then records the run under its env-aware fingerprint).

compare_git()[source]

Compare git tree hashes, with optional include/exclude filtering.

is_match()[source]

Check if the current run matches the target run.

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).

flexlock.diff_cli.compare_snapshots(snap1: dict, snap2: dict, show_details: bool = False) → bool[source]

Compare two snapshots and print a human-readable report.

flexlock.diff_cli.load_snapshot_from_db(db_path: Path, task_id: str) → dict[source]

Load snapshot from database by task_id.

flexlock.diff_cli.load_snapshot_from_dir(dir_path: Path) → dict[source]

Load a run’s record (run.lock, or its task DB row for sweep tasks).

flexlock.diff_cli.main()[source]

CLI entry point for flexlock diff command.

flexlock.diff_cli.run_comparison(snap1: dict, snap2: dict) → tuple[bool, dict][source]

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/env (only the categories that differ appear).

CLI entry point for FlexLock runner.

flexlock.run_cli.main()[source]

CLI entry point for flexlock-run command.

This is a standalone CLI that provides direct access to FlexLockRunner without requiring a Python script with @flexcli decorator.

Usage:

flexlock-run –config experiment.yml flexlock-run –defaults module.config:defaults –config exp.yml flexlock-run –config exp.yml –select experiments.exp1 flexlock-run –sweep-file sweep.yaml –n_jobs 4

Export utilities for FlexLock - extract snapshots from database to files.

flexlock.export.export_all_tasks(db_path: Path, output_base_dir: Path, status: str = None) → None[source]

Export all tasks (or filtered by status) from database to separate directories.

Parameters:
  • 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)

flexlock.export.export_task(db_path: Path, task_id: str, output_dir: Path) → None[source]

Export a single task snapshot from database to a standalone directory.

Parameters:
  • 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

flexlock.export.main()[source]

CLI entry point for flexlock export command.

Utility for loading data from a previous FlexLock stage.

flexlock.load_stage.load_stage_from_path(path: str) → dict[source]

Loads a stage from a given path and returns its flattened data, including all its ancestors.

Parameters:

path (str) – The path to the directory of the previous stage.

Returns:

A flattened, deduplicated dictionary of all ancestor runs.

Return type:

dict

MLflow integration for FlexLock.

flexlock.mlflow.mlflow_context(save_dir: str | Path, experiment_name: str = 'Default', run_name: str = None, tags: Dict[str, str] = None, log_config: bool = True, log_artifacts: bool = True)[source]

Manages a “Shadow Run” in MLflow linked to a physical directory.

This context manager implements the “Always New + Deprecate Old” strategy, treating the MLflow Run as a temporary view of the persistent physical state (the save_dir).

Behavior: 1. SEARCH: Finds the currently ‘active’ run for this directory. 2. CREATE: Starts a completely new run. 3. INHERIT: Logs existing config/artifacts found in save_dir to the new run. 4. INHERIT TAGS: Copies user tags (e.g., pipeline tags) from previous run. 5. SWAP: Marks new run ‘active’, marks old run ‘deprecated’.

Parameters:
  • save_dir – The directory containing the run data (e.g., outputs/exp/run_01)

  • experiment_name – MLflow experiment name (default: “Default”)

  • run_name – Name for this MLflow run (default: save_dir.name)

  • tags – Additional tags to add to the run (overrides inherited tags)

  • log_config – If True, logs parameters from run.lock file

  • log_artifacts – If True, logs standard artifacts (logs, lock file)

Yields:

mlflow.ActiveRun – The active MLflow run object, or None if MLflow not available

Note

User tags (not starting with ‘flexlock.’ or ‘mlflow.’) are automatically inherited from the previous active run. This is useful for pipeline workflows where you want tags like ‘pipeline_run=xyz’ to persist across stages. Explicitly passing a tag via the tags parameter will override the inherited value.

Example

>>> # Stage 1: Train (creates run with pipeline tag)
>>> with mlflow_context(save_dir, tags={"pipeline_run": "run_001"}):
...     mlflow.log_metric("train_loss", 0.32)
...
>>> # Stage 2: Evaluate (inherits pipeline_run tag automatically)
>>> with mlflow_context(save_dir):  # pipeline_run=run_001 inherited!
...     mlflow.log_metric("test_accuracy", 0.95)

Manages parallel task execution for FlexLock.

class flexlock.parallel.ParallelExecutor(func, tasks: List[Any], task_target: str | None, cfg: DictConfig, n_jobs: int = 1, slurm_config: str | None = None, pbs_config: str | None = None, local_workers: int | None = None, isolated: bool = False, tag: str | None = None, note: str | None = None, task_record: str | None = None, after=None)[source]

Bases: object

Manages the execution of tasks in parallel using a centralized task queue.

This class orchestrates task execution across different backends (local, Slurm, PBS) by interacting with a SQLite database for task management. It supports dynamic task distribution (pull model) and result aggregation.

run(wait: bool = True, timeout: int = None, poll_interval: int = None)[source]

Execute tasks via backend or locally.

Parameters:
  • wait – If True, blocks until all tasks complete (default: True for local, optional for HPC)

  • timeout – Maximum time to wait in seconds (None = no timeout, applies only if wait=True)

  • poll_interval – How often to check task status in seconds (defaults to config.POLL_INTERVAL, applies only if wait=True)

Returns:

True if tasks completed successfully (or not waiting), False on timeout/failure

Return type:

bool

flexlock.parallel.load_tasks(tasks: str, tasks_key: str, cfg: DictConfig) → List[Any][source]

Load tasks from a file or from the config.

OmegaConf resolvers for FlexLock.

flexlock.resolvers.defer_unpinned_runs(enabled=True)[source]

Preserve unpinned run calls while still resolving explicit pins.

flexlock.resolvers.deferred_stubbed(defer_runs=False)[source]

Stub only the deferred resolvers (run_lock/latest).

Used by the merge-before-resolve sweep freeze, where vinc/now should still fire (they are being retired via save_dir_policy).

flexlock.resolvers.frozen_resolvers()[source]

Stub every flexlock resolver, so all ${name:args} calls are preserved as frozen call strings during the submit-time node freeze.

flexlock.resolvers.latest_resolver(path_glob: str, default=<object object>) → str[source]

OmegaConf resolver that returns the latest path matching the given pattern.

Raises FileNotFoundError when nothing matches, unless a default is supplied (${latest:pattern,fallback}), in which case the default is returned. Returning the unmatched pattern silently — the old behaviour — produced a bogus path that failed far downstream.

flexlock.resolvers.now_resolver(fmt: str = None) → str[source]

OmegaConf resolver that returns the current time as a formatted string.

Deprecated: prefer submit(..., save_dir_policy="timestamp") over embedding ${now:} in save_dir. See docs/resolution_simplification_plan.md.

Parameters:

fmt – Format string for strftime (defaults to config.TIMESTAMP_FORMAT)

flexlock.resolvers.register_resolvers()[source]

Registers the flexlock resolvers with OmegaConf.

flexlock.resolvers.resolve_deferred(cfg: DictConfig) → DictConfig[source]

Fire the deferred resolvers once, at stage start, and return a detached config of concrete values.

This is the single point where run_lock/latest (and any other interpolations left after the submit-time freeze) are evaluated with the real resolvers. Called on the worker right after the task override is merged (and on the local single-run path). Re-wrapping as a plain config detaches it from any parent so downstream config.copy() cannot re-fire anything.

flexlock.resolvers.run_lock_resolver(run_dir: str, key: str, default=<object object>)[source]

OmegaConf resolver that reads a field from an upstream run.lock.

Navigates a dot-separated key path into the run.lock YAML and returns the value with its native type (int, float, bool, str, list, dict).

Usage in configs:

${run_lock:path/to/run_dir,config.datamodule.stats_file} ${run_lock:${run_dir},config.lit_module.regression_checkpoint_path,null}

flexlock.resolvers.run_resolver(address, select=None, pin=None, strict=None) → str[source]

${run:<preset>[,<select>[,<pin>[,strict]]]} → newest complete run dir.

See flexlock.presets.resolve_run(). Bound by freeze_run_refs at submit, or at worker stage start for unpinned scheduler-chain references.

flexlock.resolvers.vinc_resolver(path: str, fmt: str = '_{i:04d}') → str[source]

OmegaConf resolver that finds the highest existing version of a folder/file and returns the next versioned path as a string.

The value is a pure function of the current filesystem state (highest existing version + 1) so that multiple references to the same ${vinc:} within one submit() resolve idempotently to the same directory — the counter only advances once the previous run’s directory actually exists on disk.

Concurrency note: because the claim is not taken here, two submits racing from separate processes can compute the same next version and collide.

Deprecated: prefer submit(..., save_dir_policy="increment"), which takes the atomic mkdir(exist_ok=False) claim at the right layer and so fixes the race this resolver could not. See docs/resolution_simplification_plan.md.

Snapshotting utilities for FlexLock.

class flexlock.snapshot.RunTracker(save_dir, parent_lock=None, note=None)[source]

Bases: object

add_lineage(name: str, path: str, info: dict)[source]

Add lineage information from upstream FlexLock runs.

finalize(config)[source]

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).

Parameters:

config – OmegaConf configuration object

Returns:

Complete snapshot data structure

Return type:

dict

mark_complete(result=None) → Path[source]

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.

record_data(data_paths: dict)[source]
record_env(repos: dict)[source]
save(config)[source]

Finalize and write to disk as run.lock. Now implemented by calling finalize() then writing.

Returns:

The snapshot data that was written

Return type:

dict

flexlock.snapshot.is_complete(run_dir: Path) → bool[source]

Return True if run_dir has both run.lock and run.complete.

flexlock.snapshot.record_code_drift(save_dir) → dict[source]

Append code_drift to save_dir/run.lock if loaded code changed.

Call when the run ends (success or failure). Compares the modules loaded in this process against the tree recorded in run.lock; see flexlock.git_utils.code_drift(). Never raises.

flexlock.snapshot.snapshot(cfg, repos=None, data=None, prevs=None, parent_lock=None, save_path=None, return_snapshot=False, fingerprint=None, note=None, meta=None)[source]

Create a snapshot of the current run state.

Parameters:
  • 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

  • meta – Extra top-level keys for the record (e.g. a sweep master’s task_record). Never part of the fingerprint.

Returns:

dict if return_snapshot=True, else None

flexlock.snapshot.write_complete_marker(save_dir: Path, result=None) → Path[source]

Write run.complete atomically into save_dir.

Thin wrapper over RunRecord.mark_complete() kept for backward compatibility. _find_matching_run requires both run.lock and run.complete for a cache hit.

Runner for FlexLock experiments.

class flexlock.runner.FlexLockRunner[source]

Bases: object

load_config(args)[source]
run(cli_args=None, base_cfg=None)[source]

Thin layer over Project.submit() — shapes CLI args and dispatches.

CLI for monitoring FlexLock task database status.

flexlock.status_cli.main()[source]

CLI entry point for flexlock-status command.

flexlock.status_cli.print_all_tasks(db_path: Path, status_filter: str = None, tags=None)[source]

Print all tasks, optionally filtered by status and/or tag.

flexlock.status_cli.print_failed_tasks(db_path: Path, verbose: bool = False, tags=None)[source]

Print details of failed and interrupted tasks.

flexlock.status_cli.print_status_summary(db_path: Path, tags=None, label: str = '')[source]

Print summary of task statuses, optionally scoped to tags.

flexlock.status_cli.print_tags_breakdown(db_path: Path)[source]

Print a per-tag status breakdown table.

flexlock.status_cli.watch_status(db_path: Path, interval: int = 10, tags=None)[source]

Watch task status in real-time, optionally scoped to tags.

Python API for FlexLock.

class flexlock.api.ChainedResult(sweep: List[ExecutionResult], downstream: List[List[ExecutionResult]])[source]

Bases: object

Return type for Project.submit_chained().

sweep

List of sweep ExecutionResult objects.

downstream

List of lists; downstream[i][j] is the result of the j-th downstream stage for the i-th sweep item.

class flexlock.api.ExecutionResult(save_dir: str, status: str | Status, result: Any = None, cfg: DictConfig = None, error: str | None = None, job_id: str | None = None)[source]

Bases: object

Typed result object from a run.

The function’s return payload is kept in result (alias metrics); its keys are still reachable as attributes (res.accuracy) via __getattr__ — but only when they don’t collide with a real attribute, so a payload key named status/get no longer clobbers the object’s own API (issue 6).

get(key, default=None)[source]

Dict-like get method.

property is_success: bool
property metrics

The return payload (alias for result).

raise_on_failure() → ExecutionResult[source]

Raise if the run did not succeed; return self otherwise (chainable).

class flexlock.api.Project(defaults: str | DictConfig | dict | None = None)[source]

Bases: object

check(config=None, *, sweep: List[Dict] | None = None, sweep_target: str | None = None, overrides: dict | List[str] | None = None, merge: str | Path | dict | None = None) → List[dict][source]

Preflight: fully resolve the config (and every sweep item) without touching the filesystem or executing anything.

Mirrors submit’s merge → freeze pipeline, but runs under the freeze stubs so no resolver fires (vinc/now take no mkdir, run_lock/latest read nothing). Every unresolvable interpolation is reported — not just the first — as a dict with item (sweep index or None), full_key, and error.

Returns an empty list when everything resolves.

static collect_results(indices, task_configs, db_path, tag, job_id=None) → list[source]

Build per-item ExecutionResults from the task DB’s terminal state.

Replaces the old blanket status="SUCCESS" — a task that raised is now reported as FAILED with its traceback, so a sweep result set faithfully reflects what happened (issue 2).

exists(cfg: DictConfig, search_dirs: List[str] = None) → bool[source]

Check if a run with the given configuration already exists.

Parameters:
  • cfg – Configuration to check

  • search_dirs – Optional list of directories to search

Returns:

True if matching run exists, False otherwise

get(key: str)[source]

Get a configuration by key from the defaults.

The returned config is self-contained: root-scope ${...} references are frozen into concrete values, while resolver calls (${vinc:}, ${latest:}, ${run_lock:}, etc.) and intra-sub-tree refs are preserved for resolution at submit time. This means the returned config can be pickled, modified, and submitted without losing root context (essential for HPC submission).

Parameters:

key – Dot-path to select a specific node from the defaults.

Returns:

The selected configuration (as DictConfig).

get_result(cfg: DictConfig, search_dirs: List[str] = None) → ExecutionResult[source]

Retrieve results from a previously completed run.

Parameters:
  • cfg – Configuration to match

  • search_dirs – Optional list of directories to search

Returns:

ExecutionResult object with cached results

Raises:

ValueError – If no matching run is found

override_module = None
run_stage(cfg, stage_name=None, smart_run=True, search_dirs=None, **submit_kwargs)[source]

Run a single stage with automatic search_dirs and save_dir propagation.

Parameters:
  • cfg – Stage configuration (DictConfig)

  • stage_name – Name of the stage (inferred from save_dir if None)

  • smart_run – If True, checks for cached runs

  • search_dirs – Directories to search for cached runs (auto-discovered if None)

  • **submit_kwargs – Additional arguments passed to submit()

Returns:

ExecutionResult (or list for sweeps)

save_snapshot(save_dir)[source]

Save the current project defaults as pipeline.yaml.

Parameters:

save_dir – Directory to save the pipeline snapshot to

submit(config: DictConfig | str | None = None, sweep: List[Dict] = None, sweep_target: str = None, sweep_root: str | None = None, n_jobs: int = 1, smart_run: bool = True, search_dirs: List[str] = None, wait: bool = True, pbs_config: str = None, slurm_config: str = None, sweep_dir_suffix: bool = False, match_include: List[str] = None, match_exclude: List[str] = None, isolated: bool = False, force: bool = False, overrides: dict | List[str] | None = None, merge: str | Path | dict | None = None, debug: bool = False, print_config: bool = False, dry_run: bool = False, tag: str | None = None, timeout: int | None = None, note: str | None = None, save_dir_policy: str | None = None, task_record: str | None = None, after=None) → ExecutionResult | List[ExecutionResult] | None[source]

Submit a configuration for execution.

Parameters:
  • config – The configuration to execute. Accepts a DictConfig, a key string (looked up via get()), or None (uses self.defaults as-is).

  • sweep – Optional list of override dicts for parameter sweep.

  • sweep_target – Dot-path inside each task config where the sweep item is merged. When None, items merge at the root of config.

  • n_jobs – Number of parallel workers (for sweeps).

  • smart_run – If True, checks for an existing cached run before executing.

  • search_dirs – Directories to search for cached runs.

  • wait – If True, blocks until completion.

  • slurm_config (pbs_config /) – Path to HPC backend YAML.

  • after – Job IDs (sequence or colon-separated string) that must finish successfully before this Slurm/PBS job starts. Requires an HPC backend. Submitted results expose the scheduler ID as job_id. Unpinned run references bind on the worker; pins bind at submit.

  • sweep_dir_suffix – If True, append _sweep_{i:04d} to each sweep item’s save_dir.

  • match_exclude (match_include /) – Override git path patterns used during smart_run comparison.

  • isolated – If True, run in a spawned subprocess even for a single task (use for GPU stages to confine the CUDA context).

  • force – If True, invalidate the cache marker (run.complete) for this save_dir and re-execute. Outputs and run.lock are preserved; the user function overwrites in place.

  • overrides – Dict ({'lr': 0.01}) or dotlist (['lr=0.01']) merged into config before execution.

  • merge – Path to a YAML file (or a dict) merged into config before execution. overrides is applied after merge.

  • debug – Wrap the user function with the post-mortem debugger so exceptions drop into PDB.

  • print_config – Print the fully-resolved config and return None without executing.

  • dry_run – When using an HPC backend, render the would-be Slurm or PBS submission script, print it (with any validation warnings), and return None without submitting. No-op for local execution.

  • tag – Human-readable label for this sweep’s rows in the shared task DB (e.g. "extract" or "collocate"). Passed through to ParallelExecutor so workers and the status CLI can scope to it. When None a deterministic hash is auto-generated.

  • note – Free-text intent recorded as a top-level note: key in run.lock (sibling of timestamp/config). Never enters the fingerprint, so it can’t perturb caching. For sweeps the note lands on the master run.lock only; sweep items inherit it for display via their .flexlock_marker → master lookup.

  • save_dir_policy – What to do when config.save_dir is already occupied by a previous run (contains run.lock), applied once at submit time, after the smart_run cache check: None/"raise" (default) refuses and raises; "increment" versions the dir (run → run_0000, claimed atomically); "overwrite" deletes the occupied dir’s contents first (never touches a dir without run.lock); "skip" returns the existing complete run’s result without executing; "unsafe" runs in place with no check; "timestamp" appends config.TIMESTAMP_FORMAT. force=True bypasses the default guard (explicit in-place rerun). Replaces the ${vinc:} / ${now:} resolvers. For sweeps the policy applies once to the sweep root; items nest beneath it, and "skip" means resume (reuse complete items, run the rest).

  • task_record – Where each sweep task’s record goes: "dir" (default, or $FLEXLOCK_TASK_RECORD) writes a full run.lock into the task’s save_dir; "db" keeps it in the task DB only (fewer files for very large sweeps). Single runs, including HPC ones, always get a run.lock. Readers find either through flexlock.record.load_record().

Returns:

ExecutionResult (single), List[ExecutionResult] (sweep), or None (print_config=True).

submit_chained(config=None, *, sweep: List[Dict] | None = None, downstream: List[tuple] | None = None, sweep_kwargs: dict | None = None, downstream_kwargs: dict | None = None) → ChainedResult[source]

Run a sweep, then chain a sequence of downstream stages per result.

For each item in sweep, this method submits the base config and then, for each (stage_key, anchor_wiring) in downstream, propagates the sweep result’s attributes into proj.defaults and submits the downstream stage.

anchor_wiring maps an anchor name in proj.defaults to an attribute of the parent ExecutionResult (typically 'save_dir').

Parameters:
  • config – Base config for the sweep (passed to submit()).

  • sweep – List of override dicts for the parameter sweep.

  • downstream – List of (stage_key, anchor_wiring) tuples describing the stages to run per sweep result.

  • sweep_kwargs – Extra kwargs forwarded to submit for the sweep itself (e.g. slurm_config, smart_run).

  • downstream_kwargs – Extra kwargs forwarded to submit for every downstream stage. Defaults to smart_run=False to avoid stale-cache false hits between iterations.

Returns:

ChainedResult with sweep (list) and downstream (a list of lists, one per sweep item, each inner list aligned to the order of downstream).

classmethod submit_config(config=None, **kwargs)[source]

Shortcut for one-off submissions without holding a Project instance.

Equivalent to Project().submit(config, **kwargs). Use this when you already have a DictConfig (from py2cfg or otherwise) and don’t need proj.get/proj.exists/proj.defaults plumbing.

See flexlock.submit() for the module-level alias.

submit_pipeline(items: List[List[DictConfig]], *, n_jobs: int = 1, smart_run: bool = False, search_dirs: List[str] = None, slurm_config: str = None, pbs_config: str = None, wait: bool = True, sweep_root: str | None = None, tag: str | None = None, force: bool = False, timeout: int | None = None, note: str | None = None, save_dir_policy: str | None = None, debug: bool = False, dry_run: bool = False, after=None) → List[List[ExecutionResult]] | None[source]

Submit composite pipeline tasks: one task per item, stages in order.

Each item is a list of already-compiled, frozen stage configs (the caller applied freeze_deferred(); deferred resolvers fire on the worker at stage start). One item becomes one composite task in the sweep task DB ({"_stages_": [...]}): the worker runs its stages sequentially, so intra-item ordering holds by construction, while parallelism (local n_jobs, Slurm/PBS array workers) fans out across items.

Composite tasks are self-contained — the base config is never re-merged into them at dequeue time.

Known limitation: all stages of an item run inside one scheduler job / worker slot, so per-stage HPC resource configs are not possible.

Parameters:
  • items – items[i][j] is the j-th stage config of item i.

  • n_jobs – Parallel workers fanning out across items.

  • smart_run – Driver-side cache check — an item whose every stage cache-hits returns CACHED without queueing; partial hits queue the item with stage-level resume (_skip_complete_).

  • search_dirs – Directories searched for cached runs (smart_run).

  • pbs_config (slurm_config /) – Path to HPC backend YAML.

  • wait – Block until completion (HPC). False returns SUBMITTED skeleton results.

  • after – Scheduler jobs that must finish successfully before this job starts. Each submitted stage result exposes the same job_id. Unpinned run references bind at stage start; pins bind at submit.

  • sweep_root – Override the master root hosting the task DB; also opts out of the containment validation.

  • timeout (tag / note /) – As in submit().

  • force – Invalidate every stage’s completion markers and reset the task DB so all items rerun.

  • save_dir_policy – Collision guard applied per stage dir before queueing — raise (default, unless force), overwrite, skip (stage-level resume), unsafe. Naming policies (increment/timestamp) don’t compose with pre-baked cross-stage paths and are rejected.

  • debug – Post-mortem debugger (local in-process execution only).

  • dry_run – With an HPC backend, render the submission script and return None without queueing anything.

Returns:

results[i][j] is the ExecutionResult of item i’s stage j (or None for dry_run).

class flexlock.api.Status(value)[source]

Bases: str, Enum

Terminal state of a run. A str subclass so status == "SUCCESS" and f-string formatting keep working for existing callers.

CACHED = 'CACHED'
FAILED = 'FAILED'
INTERRUPTED = 'INTERRUPTED'
SKIPPED = 'SKIPPED'
SUBMITTED = 'SUBMITTED'
SUCCESS = 'SUCCESS'

The on-disk contract for a single run directory.

RunRecord is the one owner of the files that define a run on disk — run.lock (the snapshot written before execution), run.complete (the success marker), and results.json (the return payload). Every writer — the serial Project.submit path, the sweep worker, and the legacy write_complete_marker helper — routes through it, so the driver and worker are structurally identical and there is a single place to hook the fingerprint index (see index.py, wired in Phase 2.2).

All writes are atomic (temp file + os.replace) so a concurrent reader never observes a partial file.

class flexlock.run_record.RunRecord(save_dir)[source]

Bases: object

Encapsulates one run directory’s files and lifecycle.

clear_error() → None[source]

Remove a stale run.error if present (no-op otherwise).

property complete_path: Path
property error_path: Path
property is_complete: bool
is_complete_task() → bool[source]

Completeness for a sweep-task dir, which carries run.complete and results.json but no run.lock (its snapshot lives in the task DB).

load() → dict | None[source]

Return the parsed run.lock snapshot, or None if absent.

load_error() → dict | None[source]

Return the parsed run.error payload, or None if absent.

load_results() → Any | None[source]

Return the parsed results.json payload, or None if absent.

property lock_path: Path
mark_complete(result: Any = None) → Path[source]

Write run.complete after the user function returns successfully.

A cache hit requires both run.lock and run.complete; an interrupted run leaves run.lock alone and is correctly skipped. This is the single place the fingerprint index is refreshed (2.2).

property results_path: Path
property status: str
write_error(exc: BaseException, *, task_id=None, node=None, stage=None) → Path | None[source]

Record a run.error JSON sidecar describing a failed run.

Capturing the failure must never mask the original exception, so the entire body is defensive: any problem writing the record is logged and swallowed, and the caller re-raises the user’s exception unchanged.

write_lock(snapshot_data: dict) → Path[source]

Write run.lock (the snapshot) atomically before execution.

write_results(result: Any) → Path[source]

Write results.json with the run’s return payload.

flexlock.run_record.is_placeholder_lock(lock: Any) → bool[source]

True for a master-dir placeholder: config ⊆ {save_dir, _snapshot_}.

flexlock.run_record.load_lock_data(run_dir) → dict | None[source]

Side-effect-free effective run.lock of run_dir (or None).

Parses run.lock with yaml.safe_load (no resolvers fire). When it is a master placeholder and the dir’s task DB holds the snapshot of a task that ran in this very dir (single-task HPC/isolated submission), returns that snapshot merged with the master’s run-level metadata instead. Sweep masters (tasks in sub-dirs) keep their placeholder: there is no single config to show for them.

flexlock.run_record.materialize_lock(run_dir, backup: bool = True) → bool[source]

Rewrite a placeholder run.lock with its effective (task-DB) content.

One-off repair for run dirs created before workers finalized run.lock. Keeps the placeholder as run.lock.placeholder.bak when backup. Returns True when the file was rewritten.

flexlock.run_record.merge_master_into_task_snapshot(task_snap: dict, master: dict | None, run_dir) → dict[source]

Build the run.lock of a task that ran in its master dir.

Starts from the task snapshot (resolved config, fingerprint, lineage) and inherits run-level metadata the worker snapshot does not carry (note, and repos/data when the task recorded none). A parent pointer to the dir’s own run.lock is dropped (it would be self-referential).

SQLite-based task database for FlexLock parallel execution.

flexlock.taskdb.claim_next_task(db_path: Path, node: str, job_id: str | None = None, tags=None) → Any | None[source]

Claims the next available pending task and marks it as running.

Single-task convenience wrapper over claim_next_tasks.

flexlock.taskdb.claim_next_tasks(db_path: Path, node: str, job_id: str | None = None, tags=None, n: int = 1) → List[Any][source]

Claims up to n pending tasks in one write transaction.

Batching is the main lever against DB contention at scale: with many workers, one claim transaction per batch instead of per task divides the write-lock traffic by n. Returns the claimed tasks (possibly fewer than n, empty when nothing is pending).

node and job_id record which worker/scheduler job owns the tasks, so the controller can later reconcile orphaned tasks against the OS/ scheduler’s view of worker liveness (no self-reported heartbeat needed).

tags restricts the claim to rows whose tag matches (None claims any row, preserving backwards-compatible behaviour for untagged DBs).

flexlock.taskdb.dump_to_yaml(db_path: Path, yaml_path: Path, tags=None) → None[source]

Dumps all terminal (done/failed/interrupted) tasks and their results to a YAML file.

tags restricts the dump to the specified sweep’s rows so each sweep writes only its own results.

flexlock.taskdb.finish_task(db_path: Path, task: Any, error: str | None = None, result: Any | None = None, status: str | None = None) → None[source]

Marks a task as finished and records its result or error.

status defaults to 'done' on success, 'failed' on error. Pass status='interrupted' when the worker was killed mid-task.

flexlock.taskdb.finish_tasks(db_path: Path, entries: List[dict]) → None[source]

Marks a batch of tasks finished in one write transaction.

Each entry is a dict with key task (required) and optional error, result, status — same semantics as finish_task — and snapshot, which replaces the task’s stored snapshot (used to append end-of-run facts such as code_drift without an extra transaction). Workers buffer per-task outcomes and flush them here once per claimed batch.

flexlock.taskdb.get_all_tasks(db_path: Path, status: str = None, tags=None) → list[source]

Get all tasks, optionally filtered by status and/or tag.

Parameters:
  • db_path – Path to database

  • status – Optional status filter (‘pending’, ‘running’, ‘done’, ‘failed’, ‘interrupted’)

  • tags – Optional tag filter (str or list of str); None returns all rows.

Returns:

List of dicts with task details

Return type:

list

flexlock.taskdb.get_failed_tasks(db_path: Path, tags=None) → list[source]

Get details of failed and interrupted tasks.

tags restricts results to the specified sweep’s rows.

Returns:

List of dicts with task info, error, and timestamps

Return type:

list

flexlock.taskdb.get_status_counts(db_path: Path, tags=None) → dict[source]

Get counts of tasks by status.

Returns a dict with keys for every status that has at least one task, always including ‘pending’, ‘running’, ‘done’, ‘failed’, ‘interrupted’ (defaulting to 0 when absent).

tags restricts the counts to the specified sweep’s rows.

flexlock.taskdb.get_status_counts_by_tag(db_path: Path) → dict[source]

Get per-tag status counts for all rows in the DB.

Returns {tag: {status: count, ...}} where tag may be None for legacy untagged rows. Each inner dict is zero-filled for the five standard statuses.

flexlock.taskdb.get_task_snapshot(db_path: Path, task_id: str) → dict | None[source]

Retrieves the snapshot for a specific task from the database.

Parameters:
  • db_path – Path to SQLite database

  • task_id – Hash of the task (from _hash_task)

Returns:

Snapshot data, or None if not found

Return type:

dict

flexlock.taskdb.get_task_status(db_path: Path, task_id: str) → dict | None[source]

Return a single task’s status row, or None if not found.

Thin SELECT used by the query layer to classify a sweep-task dir from its .flexlock_marker without loading the whole DB.

flexlock.taskdb.list_task_snapshots(db_path: Path, status: str = None) → List[tuple][source]

Lists all tasks with their snapshots.

Parameters:
  • db_path – Path to SQLite database

  • status – Optional filter by status (pending, running, done, failed)

Returns:

(task_id, snapshot_dict, status)

Return type:

List of tuples

flexlock.taskdb.mark_orphans_interrupted(db_path: Path, job_id: str | None = None, tags=None) → int[source]

Mark still-running tasks as interrupted.

Called by the controller once it knows the owning worker(s) are dead (local processes joined, or the HPC job left the scheduler’s active set). tags limits the reconcile to this sweep’s rows so a parent pipeline row in the same DB is never touched. Returns the number of tasks marked.

flexlock.taskdb.pending_count(db_path: Path, tags=None) → int[source]

Returns the number of pending tasks in the database.

tags restricts the count to rows matching the given tag(s).

flexlock.taskdb.queue_tasks(db_path: Path, tasks: List[Any], tag: str | None = None) → None[source]

Adds a list of tasks to the database if they don’t already exist.

tag scopes these rows to a particular sweep so concurrent sweeps (or a parent pipeline job) sharing the same DB don’t interfere with each other.

Note: task_id is a content hash of the task dict. If two sweeps queue an identical task dict they share one row; INSERT OR IGNORE keeps the first tag. That only happens when the work is truly identical, which is acceptable.

flexlock.taskdb.reclaim_running(db_path: Path, job_id: str | None = None, tags=None) → int[source]

Reset running tasks back to pending so they can be re-claimed.

Used by flexlock-worker --reclaim to recover tasks stranded by a previous worker that died without cleanup. tags limits the reclaim to the specified sweep. Returns the number reset.

flexlock.taskdb.reset_db(db_path: Path) → None[source]

Delete a task DB (and its WAL/SHM files) so its tasks are re-queued.

Closes this thread’s cached connection first: otherwise _conn keeps writing to the unlinked inode and the path is never recreated.

flexlock.taskdb.update_task_snapshot(db_path: Path, task_id: str, snapshot_data: dict) → None[source]

Updates the snapshot column for a specific task.

Parameters:
  • db_path – Path to SQLite database

  • task_id – Hash of the task (from _hash_task)

  • snapshot_data – Complete snapshot dictionary

Utility functions for FlexLock.

class flexlock.utils.Swap(name: str, value: Any)[source]

Bases: object

A config value that replaces its node instead of merging into it.

Produced by expand_swaps() for @name values; consumed by merge_task_into_cfg() and apply_overrides(). Replacement (as in Hydra config groups) is what swapping needs: merging a new model into the old one would keep the old model’s extra keys.

flexlock.utils.apply_overrides(cfg: DictConfig, overrides, module=None) → None[source]

Apply dotlist overrides in order, with key=@name swaps.

Plain key=value entries merge as before (consecutive ones in one from_dotlist); key=@name replaces the node at key with the config named name (an attribute of module, the -d module, or a fully qualified pkg.mod.attr). Order matters: model=@big model.width=16 swaps, then tweaks the new model. key=@@x sets the literal @x (a bare leading @ was never valid: YAML rejects it). A dict is accepted too ({"model": "@big", "lr": 0.1}).

flexlock.utils.collect_target_include_patterns(cfg, repo_path=None)[source]

Recursively walk a config, collect all _target_ values, resolve them to source file paths relative to the repo root. Returns a list of relative file paths suitable for use as match_include patterns in smart_run.

External targets (site-packages, builtins) are silently skipped.

flexlock.utils.collect_task_repos(cfg, tasks, task_to=None) → Dict[source]

Repos to record in a sweep master’s run.lock, including the tasks’ own.

extract_tracking_info(cfg) only sees the base config. When _target_ (or _snapshot_.repos) lives in the tasks instead (a single HPC run, or --sweep-file of full configs), the base yields nothing, and since workers record tasks as deltas against the master, no code state would be recorded at all. Tasks are grouped by their _target_/_snapshot_ so the cost is one merge per distinct target, not per task.

flexlock.utils.defaults_module(import_path: str | None)[source]

The module a -d / Project(defaults=...) string points into.

flexlock.utils.enqueue_to_file(path: str | Path, cfg_dict: dict) → int[source]

Atomically append cfg_dict to a YAML list file and return the new queue length.

The file is created if it does not exist. Existing content must be a YAML list (or an empty file). The write is atomic: a temporary sibling file is written first, then renamed over the target so partial writes are never visible to concurrent readers.

Parameters:
  • path – Destination YAML file (created if absent).

  • cfg_dict – A plain Python dict (no OmegaConf nodes) to append.

Returns:

Number of items in the queue after appending.

flexlock.utils.expand_swaps(value: Any, module=None) → Any[source]

Turn "@name" strings (at any depth) into Swap markers.

"@@x" is the escape for the literal string "@x".

flexlock.utils.extract_tracking_info(cfg) → Tuple[Dict, Dict, List][source]

Extract tracking information from config’s _snapshot_ field.

This is a unified function used across FlexLock to extract repository tracking, data hashing, and lineage information from configurations.

If no repos.main is specified but the config has a _target_, automatically detects the git repository of the target function’s source file.

Parameters:

cfg – OmegaConf DictConfig with optional _snapshot_ field

Returns:

(repos, data, prevs) where:
  • repos: dict of {name: {“path”: str, “include”: list|None, “exclude”: list|None}}

  • data: dict of {name: path} for data files/dirs to hash

  • prevs: list of paths to check for upstream FlexLock runs

Return type:

tuple

Raises:

FlexLockConfigError – If _snapshot_ contains invalid keys like ‘repo’ (singular)

Examples

>>> cfg = OmegaConf.create({
...     '_snapshot_': {
...         'repos': {'main': '.'},
...         'data': {'input': 'data/train.csv'}
...     }
... })
>>> repos, data, prevs = extract_tracking_info(cfg)
>>> repos
{'main': {'path': '.'}}
flexlock.utils.instantiate(config, *args, **kwargs)[source]

Recursively instantiate objects defined in dictionaries with a “_target_” key.

Parameters:
  • config – The configuration dictionary (or value).

  • *args – Additional arguments to pass to the root object.

  • **kwargs – Additional arguments to pass to the root object.

flexlock.utils.load_python_defaults(import_path: str)[source]

Dynamically import a module or file path to retrieve a variable.

Accepted forms:

  • 'pkg.config.defaults' — dotted module path, variable is the last segment.

  • 'pkg.config:defaults' — module + colon-separated variable name.

  • 'configs/defaults.py:defaults' — file path + colon-separated variable name.

  • 'configs/defaults.py' — bare file path, variable defaults to defaults by convention.

flexlock.utils.load_sweep(*, sweep: list | str | None = None, sweep_file: str | Path | list[str | Path] | None = None, sweep_key: str | None = None, root_cfg: DictConfig | None = None) → list[source]

Load a sweep list from one of several sources.

Exactly one source must be provided. Returns the parsed sweep list. If the source yields a single value or dict, it is wrapped in a list.

Parameters:
  • sweep – A pre-built list, or a CLI-style comma-separated string.

  • sweep_file – Path (or list of paths) to .yaml/.yml, .json, or text files. When a list is given, each file is loaded independently and results are concatenated. A file containing a single dict is treated as one task; a file containing a list expands into multiple tasks.

  • sweep_key – Dotted key into root_cfg whose value is the sweep list.

  • root_cfg – Required when sweep_key is used.

flexlock.utils.log_to_file(path)[source]
flexlock.utils.merge_task_into_cfg(cfg: DictConfig, task: Any, task_to: str | None) → DictConfig[source]

Merge a task into the config (@name swaps replace their node).

flexlock.utils.parse_sweep_string(sweep_str: str) → list[source]

Parse a comma-separated CLI sweep value into a list.

Examples:

parse_sweep_string("1,2,3")              # → [1, 2, 3]
parse_sweep_string("lr=0.1,lr=0.2")      # → [{'lr': 0.1}, {'lr': 0.2}]
parse_sweep_string('"a,b",c')            # → ['a,b', 'c']
flexlock.utils.py2cfg(obj, /, *pos, **overrides)[source]

Generates a default configuration dict from a function or class signature. Supports nested py2cfg calls and handles decorated functions.

flexlock.utils.resolve_module_to_repo_path(module_name: str) → str[source]

Resolve a Python module name to its containing git repository’s working directory.

flexlock.utils.resolve_swap(name: str, module=None) → Any[source]

Value of @name: an attribute of module, else pkg.mod.attr.

flexlock.utils.to_dictconfig(incfg)[source]

Convert various config formats (dataclass, dict, class instance, DictConfig) into a DictConfig object. Warns if dataclass fields are missing type annotations.

Worker process for executing FlexLock tasks.

The worker is deliberately single-threaded while the user function runs: no background heartbeat thread or process. A live background thread would turn any fork the user’s code performs (PyTorch DataLoader workers, DDP, …) into a multi-threaded fork and corrupt the CUDA runtime in the child (cudaErrorSystemNotReady / “Unexpected error from cudaGetDeviceCount()”). Liveness is instead reconciled by the controller against the OS/scheduler’s authoritative view of worker state (see ParallelExecutor).

flexlock.worker.worker_loop(func, cfg, task_to: str, db_path, tags=None)[source]

Claim and execute tasks from the task database until none remain.

tags restricts which rows this worker claims (and the idle-check pending count). None (default) claims any row, matching legacy untagged DBs and the flexlock-worker CLI without --tags.