calkit-python 0.29.2__tar.gz → 0.30.0__tar.gz

{calkit_python-0.29.2 → calkit_python-0.30.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: calkit-python
-Version: 0.29.2
+Version: 0.30.0
 Summary: Reproducibility simplified.
 Project-URL: Homepage, https://calkit.org
 Project-URL: Issues, https://github.com/calkit/calkit/issues

{calkit_python-0.29.2 → calkit_python-0.30.0}/calkit/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.29.2"
+__version__ = "0.30.0"
 
 from .core import *  # noqa: F403, I001
 from . import git  # noqa: F401

{calkit_python-0.29.2 → calkit_python-0.30.0}/calkit/cli/check.py

@@ -3,21 +3,12 @@
 from __future__ import annotations
 
 import functools
-import hashlib
 import json
 import os
 import platform as _platform
 import subprocess
-import warnings
 from typing import Annotated
 
-from calkit.environments import get_env_lock_fpath
-
-# See https://github.com/calkit/calkit/issues/346
-with warnings.catch_warnings():
-    warnings.filterwarnings("ignore", category=UserWarning)
-    import checksumdir
-
 import dotenv
 import git
 import typer
@@ -27,6 +18,8 @@ import calkit.matlab
 import calkit.pipeline
 from calkit.check import check_reproducibility
 from calkit.cli import raise_error, warn
+from calkit.core import get_md5
+from calkit.environments import get_env_lock_fpath
 
 check_app = typer.Typer(no_args_is_help=True)
 
@@ -288,14 +281,6 @@ def check_docker_env(
             resp[key] = out[0].get(key)
         return resp
 
-    def get_md5(path: str, exclude_files: list[str] | None = None) -> str:
-        if os.path.isdir(path):
-            return checksumdir.dirhash(dep, excluded_files=exclude_files)
-        else:
-            with open(path) as f:
-                content = f.read()
-            return hashlib.md5(content.encode()).hexdigest()
-
     outfile = open(os.devnull, "w") if quiet else None
     typer.echo(f"Checking for existing image with tag {tag}", file=outfile)
     # First call Docker inspect

{calkit_python-0.29.2 → calkit_python-0.30.0}/calkit/cli/main.py

@@ -51,6 +51,7 @@ from calkit.cli.new import new_app
 from calkit.cli.notebooks import notebooks_app
 from calkit.cli.office import office_app
 from calkit.cli.overleaf import overleaf_app
+from calkit.cli.slurm import slurm_app
 from calkit.cli.update import update_app
 from calkit.environments import get_env_lock_fpath
 from calkit.models import Procedure
@@ -77,6 +78,7 @@ app.add_typer(update_app, name="update", help="Update objects.")
 app.add_typer(check_app, name="check", help="Check things.")
 app.add_typer(overleaf_app, name="overleaf", help="Interact with Overleaf.")
 app.add_typer(cloud_app, name="cloud", help="Interact with a Calkit Cloud.")
+app.add_typer(slurm_app, name="slurm", help="Work with SLURM.")
 
 # Constants for version control auto-ignore
 AUTO_IGNORE_SUFFIXES = [".DS_Store", ".env", ".pyc", ".synctex.gz"]
@@ -187,8 +189,11 @@ def clone(
     no_dvc_pull: Annotated[
         bool, typer.Option("--no-dvc-pull", help="Do not pull DVC objects.")
     ] = False,
-    recursive: Annotated[
-        bool, typer.Option("--recursive", help="Recursively clone submodules.")
+    non_recursive: Annotated[
+        bool,
+        typer.Option(
+            "--no-recursive", help="Do not recursively clone submodules."
+        ),
     ] = False,
 ):
     """Clone a Git repo and by default configure and pull from the DVC
@@ -219,7 +224,7 @@ def clone(
             url = url.replace("https://github.com/", "git@github.com:")
     # Git clone
     cmd = ["git", "clone", url]
-    if recursive:
+    if not non_recursive:
         cmd.append("--recursive")
     if location is not None:
         cmd.append(location)
@@ -552,6 +557,9 @@ def save(
             help="Additional DVC args to pass when pushing.",
         ),
     ] = [],
+    no_recursive: Annotated[
+        bool, typer.Option("--no-recursive", help="Do not push to submodules.")
+    ] = False,
     verbose: Annotated[
         bool, typer.Option("--verbose", "-v", help="Print verbose output.")
     ] = False,
@@ -606,7 +614,10 @@ def save(
         if verbose and not any_dvc:
             typer.echo("Not pushing to DVC since no DVC files were staged")
         push(
-            no_dvc=not any_dvc, git_args=git_push_args, dvc_args=dvc_push_args
+            no_dvc=not any_dvc,
+            git_args=git_push_args,
+            dvc_args=dvc_push_args,
+            no_recursive=no_recursive,
         )
 
 
@@ -629,6 +640,12 @@ def pull(
             help="Force pull, potentially overwriting local changes.",
         ),
     ] = False,
+    no_recursive: Annotated[
+        bool,
+        typer.Option(
+            "--no-recursive", help="Do not recursively pull from submodules."
+        ),
+    ] = False,
 ):
     """Pull with both Git and DVC."""
     typer.echo("Git pulling")
@@ -638,7 +655,10 @@ def pull(
         if "-f" not in dvc_args and "--force" not in dvc_args:
             dvc_args.append("-f")
     try:
-        subprocess.check_call(["git", "pull"] + git_args)
+        git_cmd = ["git", "pull"]
+        if not no_recursive and "--recurse-submodules" not in git_args:
+            git_cmd.append("--recurse-submodules")
+        subprocess.check_call(git_cmd + git_args)
     except subprocess.CalledProcessError:
         raise_error("Git pull failed")
     typer.echo("DVC pulling")
@@ -667,11 +687,17 @@ def push(
         list[str],
         typer.Option("--dvc-arg", help="Additional DVC args."),
     ] = [],
+    no_recursive: Annotated[
+        bool, typer.Option("--no-recursive", help="Do not push to submodules.")
+    ] = False,
 ):
     """Push with both Git and DVC."""
     typer.echo("Pushing to Git remote")
     try:
-        subprocess.check_call(["git", "push"] + git_args)
+        git_cmd = ["git", "push"]
+        if not no_recursive and "--recurse-submodules" not in git_args:
+            git_cmd.append("--recurse-submodules=on-demand")
+        subprocess.check_call(git_cmd + git_args)
     except subprocess.CalledProcessError:
         raise_error("Git push failed")
     if not no_dvc:

calkit_python-0.30.0/calkit/cli/slurm.py

@@ -0,0 +1,344 @@
+"""CLI for working with SLURM."""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import socket
+import subprocess
+import time
+
+import typer
+from typing_extensions import Annotated
+
+import calkit
+from calkit.cli import raise_error
+
+slurm_app = typer.Typer(no_args_is_help=True)
+
+
+@slurm_app.command(name="batch")
+def run_sbatch(
+    name: Annotated[
+        str,
+        typer.Option("--name", "-n", help="Job name."),
+    ],
+    script: Annotated[
+        str,
+        typer.Argument(help="Path to the SLURM script to run."),
+    ],
+    environment: Annotated[
+        str,
+        typer.Option(
+            "--environment",
+            "-e",
+            help="Calkit (slurm) environment to use for the job.",
+        ),
+    ],
+    args: Annotated[
+        list[str] | None,
+        typer.Argument(
+            help=(
+                "Arguments for sbatch, the first of which should be the "
+                "script."
+            )
+        ),
+    ] = None,
+    deps: Annotated[
+        list[str],
+        typer.Option(
+            "--dep",
+            "-d",
+            help=(
+                "Additional dependencies to track, which if changed signify"
+                " a job is invalid."
+            ),
+        ),
+    ] = [],
+    outs: Annotated[
+        list[str],
+        typer.Option(
+            "--out",
+            "-o",
+            help=(
+                "Non-persistent output files or directories produced by the "
+                "job, which will be deleted before submitting a new job."
+            ),
+        ),
+    ] = [],
+    sbatch_opts: Annotated[
+        list[str],
+        typer.Option(
+            "--sbatch-option",
+            "-s",
+            help="Additional options to pass to sbatch (no spaces allowed).",
+        ),
+    ] = [],
+) -> None:
+    """Submit a SLURM batch job for the project.
+
+    Duplicates are not allowed, so if one is already running or queued with
+    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 cancelled and a new one submitted.
+    """
+
+    def check_job_running_or_queued(job_id: str) -> bool:
+        p = subprocess.run(
+            ["squeue", "--job", job_id], capture_output=True, text=True
+        )
+        if p.returncode != 0:
+            return False
+        return len(p.stdout.strip().split("\n")) > 1
+
+    def cancel_job(job_id: str, reason: str) -> None:
+        typer.echo(f"{reason}; canceling existing job ID {job_id}")
+        p = subprocess.run(
+            ["scancel", job_id], capture_output=True, text=True, check=False
+        )
+        if p.returncode != 0:
+            raise_error(
+                f"Failed to cancel existing job ID {job_id}: {p.stderr}"
+            )
+
+    if args is None:
+        args = []
+    cmd = (
+        [
+            "sbatch",
+            "--parsable",
+            "--job-name",
+            name,
+            "-o",
+            ".calkit/slurm/logs/%j.out",
+        ]
+        + sbatch_opts
+        + [script]
+        + args
+    )
+    if not os.path.isfile(script):
+        raise_error(f"SLURM script '{script}' does not exist")
+    if environment != "_system":
+        ck_info = calkit.load_calkit_info()
+        env = ck_info.get("environments", {}).get(environment, {})
+        env_kind = env.get("kind")
+        if env_kind != "slurm":
+            raise_error(
+                f"Environment '{environment}' is not a slurm environment"
+            )
+        # Check host matches
+        env_host = env.get("host", "localhost")
+        if env_host != "localhost" and env_host != socket.gethostname():
+            raise_error(
+                f"Environment '{environment}' is for host '{env_host}', but "
+                f"this is '{socket.gethostname()}'"
+            )
+    deps = [script] + deps
+    slurm_dir = os.path.join(".calkit", "slurm")
+    logs_dir = os.path.join(slurm_dir, "logs")
+    os.makedirs(logs_dir, exist_ok=True)
+    jobs_path = os.path.join(slurm_dir, "jobs.json")
+    if os.path.isfile(jobs_path):
+        with open(jobs_path, "r") as f:
+            jobs = json.load(f)
+    else:
+        jobs = {}
+    typer.echo("Computing MD5s for dependencies")
+    current_dep_md5s = {}
+    for dep in deps:
+        if not os.path.exists(dep):
+            raise_error(f"Dependency path '{dep}' does not exist.")
+        current_dep_md5s[dep] = calkit.get_md5(dep)
+    # See if there is a job with this name
+    if name in jobs:
+        job_info = jobs[name]
+        job_id = job_info["job_id"]
+        job_deps = job_info["deps"]
+        job_args = job_info.get("args", [])
+        running_or_queued = check_job_running_or_queued(job_id)
+        should_wait = True
+        if running_or_queued:
+            typer.echo(
+                f"Job '{name}' with is already running or queued with ID "
+                f"{job_id}"
+            )
+            # Check if args have changed
+            if job_args != args:
+                should_wait = False
+                cancel_job(
+                    job_id=job_id,
+                    reason=f"Arguments for job '{name}' have changed",
+                )
+            # Check if dependency paths have changed
+            if set(job_deps) != set(deps):
+                should_wait = False
+                cancel_job(
+                    job_id=job_id,
+                    reason=f"Dependencies for job '{name}' have changed",
+                )
+            # Check dependency hashes
+            job_dep_md5s = job_info.get("dep_md5s", {})
+            for dep_path, md5 in current_dep_md5s.items():
+                job_md5 = job_dep_md5s.get(dep_path)
+                if md5 != job_md5:
+                    should_wait = False
+                    cancel_job(
+                        job_id=job_id,
+                        reason=(
+                            f"Dependency '{dep_path}' for job '{name}' has "
+                            "changed"
+                        ),
+                    )
+                    break
+            # Wait for the job to finish if it's running or queued and valid
+            if should_wait:
+                typer.echo("Waiting for job to finish")
+            while running_or_queued and should_wait:
+                running_or_queued = check_job_running_or_queued(job_id)
+                time.sleep(1)
+            if should_wait:
+                raise typer.Exit(0)
+    # Job is not running or queued, so we can submit
+    # But first, delete any non-persistent outputs
+    for out in outs:
+        if os.path.exists(out):
+            typer.echo(f"Deleting output path '{out}'")
+            try:
+                if os.path.isfile(out):
+                    os.remove(out)
+                else:
+                    shutil.rmtree(out)
+            except Exception as e:
+                raise_error(f"Error deleting '{out}': {e}")
+    p = subprocess.run(cmd, capture_output=True, check=False, text=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 = {
+        "job_id": job_id,
+        "deps": deps,
+        "args": args,
+        "dep_md5s": current_dep_md5s,
+    }
+    jobs[name] = new_job
+    with open(jobs_path, "w") as f:
+        json.dump(jobs, f, indent=4)
+    # Now wait for job to finish
+    typer.echo("Waiting for job to finish")
+    running_or_queued = True
+    while running_or_queued:
+        running_or_queued = check_job_running_or_queued(job_id)
+        time.sleep(1)
+
+
+@slurm_app.command(name="queue")
+def get_queue() -> None:
+    """List SLURM jobs submitted via Calkit."""
+    slurm_dir = os.path.join(".calkit", "slurm")
+    jobs_path = os.path.join(slurm_dir, "jobs.json")
+    if os.path.isfile(jobs_path):
+        with open(jobs_path, "r") as f:
+            jobs = json.load(f)
+    else:
+        jobs = {}
+    if len(jobs) == 0:
+        typer.echo("No jobs found for this project")
+        raise typer.Exit(0)
+    job_ids = [j["job_id"] for j in jobs.values()]
+    subprocess.run(
+        ["squeue", "-j", ",".join(job_ids)],
+        capture_output=False,
+        text=True,
+    )
+
+
+@slurm_app.command(name="cancel")
+def cancel_jobs(
+    names: Annotated[
+        list[str],
+        typer.Argument(help="Names of jobs to cancel."),
+    ],
+) -> None:
+    """Cancel SLURM jobs by their name in the project."""
+    slurm_dir = os.path.join(".calkit", "slurm")
+    jobs_path = os.path.join(slurm_dir, "jobs.json")
+    if os.path.isfile(jobs_path):
+        with open(jobs_path, "r") as f:
+            jobs = json.load(f)
+    else:
+        jobs = {}
+    if len(jobs) == 0:
+        typer.echo("No jobs found for this project")
+        raise typer.Exit(0)
+    # Get any job IDs that are actually running or queued
+    job_ids = [j["job_id"] for j in jobs.values()]
+    p = subprocess.run(
+        ["squeue", "-h", "-o", "%A", "-j", ",".join(job_ids)],
+        capture_output=True,
+        text=True,
+    )
+    running_or_queued_ids = p.stdout.strip().split("\n")
+    running_or_queued_ids = [j for j in running_or_queued_ids if j]
+    for name in names:
+        if name not in jobs:
+            typer.echo(f"No job named '{name}' found for this project")
+            continue
+        job_info = jobs[name]
+        job_id = job_info["job_id"]
+        if job_id not in running_or_queued_ids:
+            typer.echo(
+                f"Job '{name}' (last submitted ID: {job_id}) is not "
+                "running or queued"
+            )
+            continue
+        p = subprocess.run(
+            ["scancel", job_id], capture_output=True, text=True, check=False
+        )
+        if p.returncode != 0:
+            raise_error(f"Failed to cancel job ID {job_id}: {p.stderr}")
+        typer.echo(f"Cancelled job '{name}' with ID {job_id}")
+
+
+@slurm_app.command(name="logs")
+def get_logs(
+    name: Annotated[
+        str,
+        typer.Argument(help="Name of the job to get logs for."),
+    ],
+    follow: Annotated[
+        bool,
+        typer.Option(
+            "--follow", "-f", help="Follow the log output like tail -f."
+        ),
+    ] = False,
+) -> None:
+    """Get the logs for a SLURM job by its name in the project."""
+    slurm_dir = os.path.join(".calkit", "slurm")
+    jobs_path = os.path.join(slurm_dir, "jobs.json")
+    if os.path.isfile(jobs_path):
+        with open(jobs_path, "r") as f:
+            jobs = json.load(f)
+    else:
+        jobs = {}
+    if len(jobs) == 0:
+        typer.echo("No jobs found for this project")
+        raise typer.Exit(0)
+    if name not in jobs:
+        raise_error(f"No job named '{name}' found for this project")
+    job_info = jobs[name]
+    job_id = job_info["job_id"]
+    log_path = os.path.join(slurm_dir, "logs", f"{job_id}.out")
+    if not os.path.isfile(log_path):
+        raise_error(f"No log file found for job '{name}' with ID {job_id}")
+    if follow:
+        p = subprocess.Popen(["tail", "-f", log_path])
+        try:
+            p.wait()
+        except KeyboardInterrupt:
+            p.terminate()
+            raise typer.Exit(0)
+    else:
+        with open(log_path, "r") as f:
+            typer.echo(f.read())

{calkit_python-0.29.2 → calkit_python-0.30.0}/calkit/config.py

@@ -56,8 +56,11 @@ def supports_keyring() -> bool:
 KEYRING_SUPPORTED = supports_keyring()
 
 
-def get_env() -> Literal["local", "staging", "production"]:
-    return os.getenv("CALKIT_ENV", "production")
+def get_env() -> Literal["test", "local", "staging", "production"]:
+    env = os.getenv("CALKIT_ENV", "production")
+    if env not in ["test", "local", "staging", "production"]:
+        raise ValueError(f"{env} is not a valid environment name")
+    return env  # type: ignore
 
 
 def set_env(name: Literal["local", "staging", "production"]) -> None:
@@ -184,10 +187,10 @@ class Settings(BaseSettings):
             dotenv_settings,
             YamlConfigSettingsSource(settings_cls),
             KeyringSecretsSource(settings_cls),
-        )
+        )  # type: ignore
 
     def write(self) -> None:
-        base_dir = os.path.dirname(self.model_config["yaml_file"])
+        base_dir = os.path.dirname(self.model_config["yaml_file"])  # type: ignore
         os.makedirs(base_dir, exist_ok=True)
         cfg = self.model_dump()
         # Remove anything that should be in the keyring
@@ -205,8 +208,10 @@ class Settings(BaseSettings):
                         except keyring.errors.KeyringError:
                             # Ignore errors when deleting secrets
                             pass
-        with open(self.model_config["yaml_file"], "w") as f:
+        with open(self.model_config["yaml_file"], "w") as f:  # type: ignore
             yaml.safe_dump(cfg, f)
+        # Ensure permissions are user read/write only
+        os.chmod(self.model_config["yaml_file"], 0o600)  # type: ignore
 
 
 def read() -> Settings:

{calkit_python-0.29.2 → calkit_python-0.30.0}/calkit/core.py

@@ -15,7 +15,12 @@ import re
 import socket
 import subprocess
 import uuid
+import warnings
 
+# See https://github.com/calkit/calkit/issues/346
+with warnings.catch_warnings():
+    warnings.filterwarnings("ignore", category=UserWarning)
+    import checksumdir
 import psutil
 import requests
 
@@ -572,3 +577,12 @@ def get_system_info() -> dict:
     system_info_str = json.dumps(system_info, sort_keys=True).encode()
     system_info["id"] = hashlib.sha1(system_info_str).hexdigest()
     return system_info
+
+
+def get_md5(path: str, exclude_files: list[str] | None = None) -> str:
+    if os.path.isdir(path):
+        return checksumdir.dirhash(path, excluded_files=exclude_files)
+    else:
+        with open(path) as f:
+            content = f.read()
+        return hashlib.md5(content.encode()).hexdigest()

{calkit_python-0.29.2 → calkit_python-0.30.0}/calkit/models/pipeline.py

@@ -13,6 +13,7 @@ from pydantic import (
     Discriminator,
     ValidationError,
     field_validator,
+    model_validator,
 )
 from typing_extensions import Annotated
 
@@ -109,6 +110,7 @@ class StageIteration(BaseModel):
 class Stage(BaseModel):
     """A stage in the pipeline."""
 
+    name: str | None = None
     kind: Literal[
         "python-script",
         "latex",
@@ -366,6 +368,54 @@ class JuliaCommandStage(Stage):
         return cmd
 
 
+class SBatchStage(Stage):
+    kind: Literal["sbatch"] = "sbatch"
+    script_path: str
+    args: list[str] = []
+    sbatch_options: list[str] = []
+
+    @property
+    def dvc_deps(self) -> list[str]:
+        return [self.script_path] + super().dvc_deps
+
+    @property
+    def dvc_outs(self) -> list[str | dict]:
+        # All outputs must be persistent, since ``calkit slurm batch``
+        # handles deletion
+        outs = super().dvc_outs
+        final_outs = []
+        for out in outs:
+            if isinstance(out, str):
+                final_outs.append({out: {"persist": True}})
+            elif isinstance(out, dict):
+                k = list(out.keys())[0]
+                v = out[k]
+                v["persist"] = True
+                final_outs.append({k: v})
+        return final_outs
+
+    @property
+    def dvc_cmd(self) -> str:
+        cmd = f"calkit slurm batch --name {self.name}"
+        if self.environment != "_system":
+            cmd += f" --environment {self.environment}"
+        for dep in self.dvc_deps:
+            if dep != self.script_path:
+                cmd += f" --dep {dep}"
+        for out in self.outputs:
+            # Determine if this is a non-persistent output
+            if isinstance(out, str):
+                cmd += f" --out {out}"
+            elif isinstance(out, PathOutput) and out.delete_before_run:
+                cmd += f" --out {out.path}"
+        for opt in self.sbatch_options:
+            cmd += f" -s {opt}"
+        cmd += f" -- {self.script_path}"
+        for arg in self.args:
+            cmd += f" {arg}"
+        return cmd
+
+
 class JupyterNotebookStage(Stage):
     """A stage that runs a Jupyter notebook.
 
@@ -565,9 +615,22 @@ class Pipeline(BaseModel):
                 | JupyterNotebookStage
                 | JuliaScriptStage
                 | JuliaCommandStage
+                | SBatchStage
             ),
             Discriminator("kind"),
         ],
     ]
     # Do not allow extra keys
     model_config = ConfigDict(extra="forbid")
+
+    @model_validator(mode="after")
+    def set_stage_names(self):
+        """Set the name field of each stage to match its key in the dict."""
+        for stage_name, stage in self.stages.items():
+            if stage.name is not None and stage.name != stage_name:
+                raise ValueError(
+                    f"Stage name '{stage.name}' does not match key "
+                    f"'{stage_name}'"
+                )
+            stage.name = stage_name
+        return self