calkit-python 0.42.0__tar.gz → 0.42.1__tar.gz

{calkit_python-0.42.0 → calkit_python-0.42.1}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: calkit-python
-Version: 0.42.0
+Version: 0.42.1
 Summary: Reproducibility simplified.
 Project-URL: Homepage, https://calkit.org
 Project-URL: Issues, https://github.com/calkit/calkit/issues
@@ -132,6 +132,11 @@ without the cognitive overhead.
 - Extensions for doing all of the above graphically in
   [JupyterLab](https://docs.calkit.org/jupyterlab) and
   [VS Code](https://marketplace.visualstudio.com/items?itemName=Calkit.calkit-vscode).
+- A [browser extension](https://docs.calkit.org/browser-ext) for collecting references
+  directly to BibTeX (optionally synced with Zotero),
+  viewing DVC-stored files on GitHub,
+  and syncing figures and results with Overleaf directly in Chrome,
+  Microsoft Edge, and more.
 
 <!-- END INCLUDE -->
 

{calkit_python-0.42.0 → calkit_python-0.42.1}/README.md

@@ -73,6 +73,11 @@ without the cognitive overhead.
 - Extensions for doing all of the above graphically in
   [JupyterLab](https://docs.calkit.org/jupyterlab) and
   [VS Code](https://marketplace.visualstudio.com/items?itemName=Calkit.calkit-vscode).
+- A [browser extension](https://docs.calkit.org/browser-ext) for collecting references
+  directly to BibTeX (optionally synced with Zotero),
+  viewing DVC-stored files on GitHub,
+  and syncing figures and results with Overleaf directly in Chrome,
+  Microsoft Edge, and more.
 
 <!-- END INCLUDE -->
 

{calkit_python-0.42.0 → calkit_python-0.42.1}/calkit/cli/check.py

@@ -1203,10 +1203,13 @@ def check_venv(
                 else os.path.join(wdir or ".", prefix)
             )
             if os.path.isdir(prefix_full_path):
+                # This can fail if the environment is in use, e.g., on
+                # Windows, where files can't be removed while open, in which
+                # case we keep it and fall back to rebuilding from the spec
                 shutil.rmtree(prefix_full_path)
             create_venv()
             pip_install_and_freeze(dep_file_txt)
-        except subprocess.CalledProcessError:
+        except (subprocess.CalledProcessError, OSError):
             warn(
                 f"Failed to create environment from lock file ({reqs_to_use}); "
                 f"attempting rebuild from input file {path}"

{calkit_python-0.42.0 → calkit_python-0.42.1}/calkit/cli/main/core.py

@@ -2440,6 +2440,26 @@ def manual_step(
     typer.echo("Done")
 
 
+def _env_prefix_exists(
+    env: dict, envs: dict, env_name: str, wdir: str | None
+) -> bool:
+    """Check that a venv-like environment's prefix exists on disk.
+
+    The environment check cache only hashes a prefix that's pinned in
+    ``calkit.yaml``, so a virtual environment living at its default location
+    could have been deleted since the last check without invalidating the
+    cache, in which case it still needs to be rebuilt.
+    """
+    if env.get("kind") not in ["uv-venv", "venv"]:
+        return True
+    if env.get("prefix") or "path" not in env:
+        return True
+    prefix = calkit.environments.get_default_venv_prefix(
+        envs, env["path"], env_name
+    )
+    return os.path.isdir(os.path.join(wdir or ".", prefix))
+
+
 @app.command(
     name="xenv|runenv",
     help="Execute a command in an environment.",
@@ -2529,6 +2549,38 @@ def run_in_env(
         docker_wdir = posixpath.join(docker_wdir, wdir)
     shell = env.get("shell", "sh")
     platform = env.get("platform")
+
+    def save_env_check_cache() -> None:
+        """Record a successful check so repeat calls can skip it."""
+        try:
+            calkit.environments.save_cache(
+                env_name=env_name, env=env, wdir=wdir, success=True
+            )
+        except Exception as e:
+            if verbose:
+                typer.echo(f"Failed to save environment check cache: {e}")
+
+    # This command is typically called once per pipeline stage, so checking
+    # the environment every time is wasteful, and on Windows can fail
+    # outright, e.g., if a package installed into the environment provides
+    # the executable that's currently running. Consult the same cache
+    # ``calkit run`` uses and only check if something relevant has changed.
+    if (
+        not no_check
+        and env["kind"] not in calkit.environments.KINDS_NO_CHECK
+        and calkit.environments.check_cache(
+            env_name=env_name, env=env, wdir=wdir
+        )
+        and _env_prefix_exists(
+            env=env, envs=envs, env_name=env_name, wdir=wdir
+        )
+    ):
+        if verbose:
+            typer.echo(
+                f"Skipping check for environment '{env_name}' since it was "
+                "recently checked and nothing has changed"
+            )
+        no_check = True
     if env["kind"] == "docker":
         if "image" not in env:
             raise_error("Image must be defined for Docker environments")
@@ -2555,6 +2607,7 @@ def run_in_env(
                 args=env.get("args", []),
                 quiet=not verbose,
             )
+            save_env_check_cache()
         docker_cmd = [
             "docker",
             "run",
@@ -2630,6 +2683,7 @@ def run_in_env(
                 relaxed=relaxed_check,
                 quiet=not verbose,
             )
+            save_env_check_cache()
         # TODO: Prefix should only be in the env file or calkit.yaml, not both?
         prefix = env.get("prefix")
         # Conda is often not on the PATH (especially on Windows), so search
@@ -2708,6 +2762,7 @@ def run_in_env(
                 quiet=True,
                 verbose=verbose,
             )
+            save_env_check_cache()
         # Now run the command
         cmd = f"{activate_cmd} && {shell_cmd} && deactivate"  # type: ignore
         if verbose:
@@ -2719,6 +2774,7 @@ def run_in_env(
     elif env["kind"] == "julia":
         if not no_check:
             check_environment(env_name=env_name, verbose=verbose)
+            save_env_check_cache()
         env_path = env.get("path")
         if env_path is None:
             raise_error(
@@ -2884,6 +2940,7 @@ def run_in_env(
         assert isinstance(env_path, str)
         if not no_check:
             check_renv(env_path=env_path, verbose=verbose)
+            save_env_check_cache()
         # For renv, we need to run from the renv project directory so renv
         # properly initializes the library, but the script needs to run
         # from its original working directory
@@ -2935,6 +2992,7 @@ def run_in_env(
             )
         if not no_check:
             check_nix_env(env=env, verbose=verbose)
+            save_env_check_cache()
         env_dir = os.path.dirname(os.path.abspath(env_path)) or "."
         flake_ref = f"path:{env_dir}"
         shell_name = env.get("shell")
@@ -2966,6 +3024,7 @@ def run_in_env(
                     env=env, env_name=env_name, as_posix=False
                 ),  # type: ignore
             )
+            save_env_check_cache()
         image_name = calkit.matlab.get_docker_image_name(
             ck_info=ck_info,
             env_name=env_name,

{calkit_python-0.42.0 → calkit_python-0.42.1}/calkit/cli/new.py

@@ -1572,6 +1572,18 @@ def new_slurm_env(
             ),
         ),
     ] = [],
+    max_concurrent_jobs: Annotated[
+        int | None,
+        typer.Option(
+            "--max-concurrent-jobs",
+            help=(
+                "Maximum number of this project's jobs allowed in the "
+                "queue at once, or 0 for no limit. Submissions beyond this "
+                "wait for a slot, so an iterated stage does not take over a "
+                "shared cluster's queue. Unlimited by default."
+            ),
+        ),
+    ] = None,
     description: Annotated[
         str | None, typer.Option("--description", help="Description.")
     ] = None,
@@ -1612,6 +1624,12 @@ def new_slurm_env(
         env["default_options"] = normalized_default_options  # type: ignore
     if normalized_default_setup:
         env["default_setup"] = normalized_default_setup  # type: ignore
+    if max_concurrent_jobs is not None:
+        if max_concurrent_jobs < 0:
+            raise_error("--max-concurrent-jobs cannot be negative")
+        # 0 means no limit, the same as it does for `calkit update`.
+        if max_concurrent_jobs:
+            env["max_concurrent_jobs"] = max_concurrent_jobs  # type: ignore
     if description is not None:
         env["description"] = description
     envs[name] = env
@@ -1661,6 +1679,18 @@ def new_pbs_env(
             ),
         ),
     ] = [],
+    max_concurrent_jobs: Annotated[
+        int | None,
+        typer.Option(
+            "--max-concurrent-jobs",
+            help=(
+                "Maximum number of this project's jobs allowed in the "
+                "queue at once, or 0 for no limit. Submissions beyond this "
+                "wait for a slot, so an iterated stage does not take over a "
+                "shared cluster's queue. Unlimited by default."
+            ),
+        ),
+    ] = None,
     description: Annotated[
         str | None, typer.Option("--description", help="Description.")
     ] = None,
@@ -1703,6 +1733,12 @@ def new_pbs_env(
         env["default_options"] = normalized_default_options  # type: ignore
     if normalized_default_setup:
         env["default_setup"] = normalized_default_setup  # type: ignore
+    if max_concurrent_jobs is not None:
+        if max_concurrent_jobs < 0:
+            raise_error("--max-concurrent-jobs cannot be negative")
+        # 0 means no limit, the same as it does for `calkit update`.
+        if max_concurrent_jobs:
+            env["max_concurrent_jobs"] = max_concurrent_jobs  # type: ignore
     if description is not None:
         env["description"] = description
     envs[name] = env

{calkit_python-0.42.0 → calkit_python-0.42.1}/calkit/cli/scheduler.py

@@ -8,12 +8,15 @@ Registered as ``scheduler|sch``.
 
 from __future__ import annotations
 
+import contextlib
 import os
+import random
 import re
 import shlex
 import shutil
 import signal
 import socket
+import sqlite3
 import subprocess
 import sys
 import time
@@ -44,6 +47,23 @@ LOCAL_DIR = os.path.join(".calkit", "local")
 # (e.g. `scheduler queue`) never see a half-written file.
 JOBS_DB_PATH = os.path.join(LOCAL_DIR, "scheduler-jobs.db")
 MOCK_DIR = os.path.join(LOCAL_DIR, "mock-scheduler")
+# Serializes the check-for-a-free-slot-then-submit sequence across the
+# parallel batch processes that fan out an iterated stage. Without it they
+# would all see the same free slot and submit together, which is exactly what
+# an environment's max_concurrent_jobs exists to prevent.
+QUEUE_LOCK_PATH = os.path.join(LOCAL_DIR, "scheduler-queue-lock.db")
+# Seconds SQLite waits for the lock before raising; retried past this, so it
+# only bounds how long a single attempt blocks.
+QUEUE_LOCK_TIMEOUT = 60
+# How often a job waiting for a queue slot rechecks. Long enough not to hammer
+# the scheduler with status queries, short enough that a freed slot is taken
+# promptly relative to the runtime of a job worth queueing.
+QUEUE_SLOT_POLL_INTERVAL = 10.0
+MOCK_QUEUE_SLOT_POLL_INTERVAL = 0.25
+# Fraction by which each waiter randomizes its poll interval. A capped sweep
+# leaves one waiting process per unsubmitted case, and without jitter they all
+# wake, take the lock, and query the scheduler in lockstep.
+QUEUE_SLOT_POLL_JITTER = 0.25
 # When set, scheduler commands run jobs on the local host instead of
 # dispatching to a real SLURM/PBS install (see _mock_enabled).
 MOCK_ENV_VAR = "CALKIT_MOCK_SCHEDULER"
@@ -434,6 +454,190 @@ def _is_active(kind: str, job_id: str) -> bool:
     return _poll_job(kind, job_id)[0]
 
 
+def _active_job_ids(kind: str, job_ids: list[str]) -> set[str]:
+    """Return which of ``job_ids`` are still queued or running.
+
+    Answers for the whole set in one scheduler query where possible. Counting
+    queue occupancy means asking about every job this project has submitted,
+    and doing that one ``squeue`` call per job---from every process waiting
+    for a slot, on every poll---would put more load on the scheduler than the
+    submissions we are trying to pace.
+    """
+    if not job_ids:
+        return set()
+    if not _mock_enabled() and kind == "slurm":
+        # Query our own jobs rather than passing `--jobs <ids>`: squeue errors
+        # out when any listed id is unknown, and ids do go unknown once the
+        # scheduler purges a finished job, which is the common case here.
+        p = subprocess.run(
+            ["squeue", "--me", "-h", "-o", "%i"],
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        if p.returncode == 0:
+            wanted = set(job_ids)
+            # Array and step ids appear as `123_4` and `123.batch`; both
+            # belong to the job we recorded as `123`.
+            active = {
+                line.strip().split("_")[0].split(".")[0]
+                for line in p.stdout.splitlines()
+                if line.strip()
+            }
+            return wanted & active
+        # squeue failed---fall through and ask about each job individually
+        # rather than reporting an empty queue we have not confirmed.
+    return {job_id for job_id in job_ids if _is_active(kind, job_id)}
+
+
+def _count_queued_jobs(environment: str, exclude: str) -> int:
+    """Count this project's jobs sitting in the queue for an environment.
+
+    Only jobs Calkit submitted for this project are counted; the limit is
+    about not monopolizing a shared queue with one project's iterated stage,
+    not about the user's jobs at large.
+    """
+    by_kind: dict[str, list[str]] = {}
+    for job_name, info in _load_jobs().items():
+        if job_name == exclude:
+            continue
+        # Records written before environments were tracked have no
+        # environment key. Count them: over-counting delays a submission,
+        # while under-counting is what floods the queue.
+        job_env = info.get("environment")
+        if job_env is not None and job_env != environment:
+            continue
+        if info.get("exit_code") is not None:
+            # Already observed to have finished, so it cannot be in the queue.
+            continue
+        by_kind.setdefault(info.get("kind", "slurm"), []).append(
+            info["job_id"]
+        )
+    return sum(
+        len(_active_job_ids(kind, job_ids))
+        for kind, job_ids in by_kind.items()
+    )
+
+
+@contextlib.contextmanager
+def _queue_lock():
+    """Cross-process mutex around checking for a slot and then submitting.
+
+    Built on SQLite's own locking---an IMMEDIATE transaction takes a write
+    lock that other connections block on---rather than a lock library, since
+    scheduler state already depends on SQLite locking correctly here, and it
+    keeps this to the standard library.
+    """
+    calkit.ensure_local_dir()
+    conn = sqlite3.connect(
+        QUEUE_LOCK_PATH, timeout=QUEUE_LOCK_TIMEOUT, isolation_level=None
+    )
+    try:
+        while True:
+            try:
+                conn.execute("BEGIN IMMEDIATE")
+                break
+            except sqlite3.OperationalError as e:
+                # Another process is mid-submit. It holds the lock only for as
+                # long as one submit command takes, so retrying is better than
+                # failing the stage. Only contention is worth retrying though:
+                # a read-only file or a bad disk would never clear, and
+                # retrying those forever would hang the stage with no
+                # explanation instead of surfacing the real problem.
+                msg = str(e).lower()
+                if "locked" not in msg and "busy" not in msg:
+                    raise
+                time.sleep(1)
+        try:
+            yield
+        finally:
+            conn.rollback()
+    finally:
+        conn.close()
+
+
+def _parse_max_concurrent_jobs(environment: str, value: object) -> int | None:
+    """Interpret an environment's ``max_concurrent_jobs``, or ``None``.
+
+    The value is read straight out of ``calkit.yaml`` rather than through
+    ``SlurmEnvironment``, so its ``ge=1`` constraint never runs and a
+    hand-edited file can put anything here. Reject what we cannot honor
+    instead of guessing: silently reading a bad value as "no limit" would
+    flood the queue this setting exists to protect, and a negative one would
+    leave the wait below with a condition that can never come true.
+    """
+    if value is None:
+        return None
+    # bool is an int subclass, and `max_concurrent_jobs: true` is a mistake,
+    # not a limit of one.
+    if isinstance(value, bool) or not isinstance(value, int):
+        raise_error(
+            f"Environment '{environment}' has an invalid "
+            f"max_concurrent_jobs ({value!r}); expected a positive integer"
+        )
+    # 0 means no limit, matching `calkit update slurm-env
+    # --max-concurrent-jobs 0`.
+    if value == 0:
+        return None
+    if value < 0:
+        raise_error(
+            f"Environment '{environment}' has a negative "
+            f"max_concurrent_jobs ({value}); it must be at least 1"
+        )
+    return value
+
+
+@contextlib.contextmanager
+def _queue_slot(environment: str, name: str, max_jobs: int | None):
+    """Hold a slot in the environment's queue for the duration of the block.
+
+    Blocks until this project has fewer than ``max_jobs`` jobs queued or
+    running, then keeps the lock while the caller submits and records the new
+    job, so concurrent waiters see it and do not overshoot the limit.
+    """
+    max_jobs = _parse_max_concurrent_jobs(environment, max_jobs)
+    if max_jobs is None:
+        yield
+        return
+    announced = False
+    while True:
+        # Check without the lock first. A capped sweep leaves one waiting
+        # process per unsubmitted case, and each poll would otherwise hold the
+        # exclusive lock across a scheduler query---with enough waiters the
+        # lock stays busy and slots go unclaimed long after they free up.
+        # A full queue is the common answer while waiting and needs no lock to
+        # act on; the authoritative check is still made under it below.
+        n_queued = _count_queued_jobs(environment, exclude=name)
+        if n_queued < max_jobs:
+            with _queue_lock():
+                n_queued = _count_queued_jobs(environment, exclude=name)
+                if n_queued < max_jobs:
+                    yield
+                    return
+        if not announced:
+            typer.echo(
+                f"Waiting for a free slot in environment '{environment}' "
+                f"({n_queued} of {max_jobs} jobs queued)"
+            )
+            announced = True
+        # Mocked jobs run locally and finish in seconds, so a full interval
+        # would dominate the runtime of anything using the mock scheduler.
+        interval = (
+            MOCK_QUEUE_SLOT_POLL_INTERVAL
+            if _mock_enabled()
+            else QUEUE_SLOT_POLL_INTERVAL
+        )
+        time.sleep(
+            interval
+            * (
+                1
+                + random.uniform(
+                    -QUEUE_SLOT_POLL_JITTER, QUEUE_SLOT_POLL_JITTER
+                )
+            )
+        )
+
+
 def _cancel(kind: str, job_id: str) -> tuple[bool, str]:
     if _mock_enabled():
         return _mock_cancel(job_id)
@@ -655,6 +859,11 @@ def run_batch(
     the same name, we'll wait for it to finish. The only exception is if the
     dependencies have changed, in which case any queued or running jobs will
     be canceled and a new one submitted.
+
+    If the environment sets ``max_concurrent_jobs``, submission waits until
+    this project has fewer than that many jobs queued or running, so an
+    iterated stage does not put all of its jobs into a shared cluster's queue
+    at once.
     """
     if args is None:
         args = []
@@ -872,51 +1081,63 @@ def run_batch(
                     shutil.rmtree(out)
             except Exception as e:
                 raise_error(f"Error deleting '{out}': {e}")
-    # Submit and record atomically with respect to Ctrl+C: a SIGINT here would
-    # otherwise kill the submit command mid-flight or leave a job submitted but
-    # never written to the database (so a re-run would resubmit it). Defer the
-    # signal until the job is safely recorded, then act on it.
+    # Wait for room in the queue before submitting, if the environment caps
+    # how many of this project's jobs may be queued at once. This is done
+    # after deleting outputs (which can be slow) so the slot is held only for
+    # as long as submitting takes.
+    max_jobs = env.get("max_concurrent_jobs")
     pid = None
     interrupted: list[bool] = []
-    prev_sigint = signal.signal(
-        signal.SIGINT, lambda *_: interrupted.append(True)
-    )
-    try:
-        if _mock_enabled():
-            job_id = uuid.uuid4().hex[:12]
-            pid = _mock_submit(
-                job_id=job_id, job_command=job_command, log_path=log_path
-            )
-        else:
-            # start_new_session shields the submit command from the terminal's
-            # Ctrl+C so the submission always completes.
-            p = subprocess.run(
-                submit_cmd,
-                input=submit_input,
-                capture_output=True,
-                check=False,
-                text=True,
-                start_new_session=True,
-            )
-            if p.returncode != 0:
-                raise_error(f"Failed to submit new job: {p.stderr}")
-            job_id = p.stdout.strip()
-        typer.echo(f"Submitted job with ID: {job_id}")
-        new_job = {
-            "kind": kind,
-            "job_id": job_id,
-            "deps": deps,
-            "target": target,
-            "args": args,
-            "setup": setup_cmds,
-            "dep_md5s": current_dep_md5s,
-            "submitted_at": calkit.utcnow().isoformat(),
-        }
-        if pid is not None:
-            new_job["pid"] = pid
-        _record_job(name, new_job)
-    finally:
-        signal.signal(signal.SIGINT, prev_sigint)
+    # The wait for a slot sits outside the SIGINT deferral below so Ctrl+C
+    # still works while waiting---that wait can be long, and deferring the
+    # signal through it would make it uninterruptible.
+    with _queue_slot(environment, name, max_jobs):
+        # Submit and record atomically with respect to Ctrl+C: a SIGINT here
+        # would otherwise kill the submit command mid-flight or leave a job
+        # submitted but never written to the database (so a re-run would
+        # resubmit it). Defer the signal until the job is safely recorded,
+        # then act on it. The queue slot is held across both, so a concurrent
+        # waiter counts this job before deciding it has room to submit.
+        prev_sigint = signal.signal(
+            signal.SIGINT, lambda *_: interrupted.append(True)
+        )
+        try:
+            if _mock_enabled():
+                job_id = uuid.uuid4().hex[:12]
+                pid = _mock_submit(
+                    job_id=job_id, job_command=job_command, log_path=log_path
+                )
+            else:
+                # start_new_session shields the submit command from the
+                # terminal's Ctrl+C so the submission always completes.
+                p = subprocess.run(
+                    submit_cmd,
+                    input=submit_input,
+                    capture_output=True,
+                    check=False,
+                    text=True,
+                    start_new_session=True,
+                )
+                if p.returncode != 0:
+                    raise_error(f"Failed to submit new job: {p.stderr}")
+                job_id = p.stdout.strip()
+            typer.echo(f"Submitted job with ID: {job_id}")
+            new_job = {
+                "kind": kind,
+                "environment": environment,
+                "job_id": job_id,
+                "deps": deps,
+                "target": target,
+                "args": args,
+                "setup": setup_cmds,
+                "dep_md5s": current_dep_md5s,
+                "submitted_at": calkit.utcnow().isoformat(),
+            }
+            if pid is not None:
+                new_job["pid"] = pid
+            _record_job(name, new_job)
+        finally:
+            signal.signal(signal.SIGINT, prev_sigint)
     if interrupted:
         typer.echo(
             f"Interrupted after submitting job '{name}' ({job_id}); it will "

{calkit_python-0.42.0 → calkit_python-0.42.1}/calkit/cli/update.py

@@ -928,6 +928,16 @@ def update_slurm_env(
             "--set-default-setup", help="Replace default setup list."
         ),
     ] = [],
+    max_concurrent_jobs: Annotated[
+        int | None,
+        typer.Option(
+            "--max-concurrent-jobs",
+            help=(
+                "Maximum number of this project's jobs allowed in the queue "
+                "at once, or 0 to remove the limit."
+            ),
+        ),
+    ] = None,
 ) -> None:
     """Update a SLURM environment."""
     ck_info, env = _load_env(env_name)
@@ -959,6 +969,15 @@ def update_slurm_env(
         env["default_setup"] = cmds
     elif "default_setup" in env:
         del env["default_setup"]
+    if max_concurrent_jobs is not None:
+        if max_concurrent_jobs < 0:
+            raise_error("--max-concurrent-jobs cannot be negative")
+        # 0 is the way to clear the limit, since omitting the option means
+        # "leave it alone" rather than "unlimited".
+        if max_concurrent_jobs == 0:
+            env.pop("max_concurrent_jobs", None)
+        else:
+            env["max_concurrent_jobs"] = max_concurrent_jobs
     calkit.save_calkit_info(ck_info)
     typer.echo(f"Updated slurm environment '{env_name}'")