calkit-python 0.26.0__tar.gz → 0.26.2__tar.gz

{calkit_python-0.26.0 → calkit_python-0.26.2}/PKG-INFO

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

{calkit_python-0.26.0 → calkit_python-0.26.2}/calkit/__init__.py

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

{calkit_python-0.26.0 → calkit_python-0.26.2}/calkit/cli/check.py

@@ -79,6 +79,11 @@ def check_environment(
             ),
             platform=env.get("platform"),
             deps=env.get("deps", []),
+            env_vars=env.get("env_vars", []),
+            ports=env.get("ports", []),
+            gpus=env.get("gpus"),
+            user=env.get("user"),
+            wdir=env.get("wdir"),
             quiet=not verbose,
         )
     elif env["kind"] == "conda":
@@ -164,6 +169,14 @@ def check_docker_env(
         str | None,
         typer.Option("--platform", help="Which platform(s) to build for."),
     ] = None,
+    user: Annotated[
+        str | None,
+        typer.Option("--user", help="Which user to run the container as."),
+    ] = None,
+    wdir: Annotated[
+        str | None,
+        typer.Option("--wdir", help="Working directory inside the container."),
+    ] = None,
     deps: Annotated[
         list[str],
         typer.Option(
@@ -172,6 +185,30 @@ def check_docker_env(
             help="Declare an explicit dependency for this Docker image.",
         ),
     ] = [],
+    env_vars: Annotated[
+        list[str],
+        typer.Option(
+            "--env-var",
+            "-e",
+            help="Declare an explicit environment variable for the container.",
+        ),
+    ] = [],
+    ports: Annotated[
+        list[str],
+        typer.Option(
+            "--port",
+            "-p",
+            help="Declare an explicit port for the container.",
+        ),
+    ] = [],
+    gpus: Annotated[
+        str | None,
+        typer.Option(
+            "--gpus",
+            "-g",
+            help="Declare an explicit GPU requirement for the container.",
+        ),
+    ] = None,
     quiet: Annotated[
         bool, typer.Option("--quiet", "-q", help="Be quiet.")
     ] = False,
@@ -277,6 +314,18 @@ def check_docker_env(
     inspect = get_docker_inspect()
     inspect["DockerfileMD5"] = dockerfile_md5
     inspect["DepsMD5s"] = deps_md5s
+    if platform is not None:
+        inspect["Platform"] = platform
+    if wdir is not None:
+        inspect["WorkDir"] = wdir
+    if user is not None:
+        inspect["User"] = user
+    if env_vars:
+        inspect["EnvVars"] = env_vars
+    if ports:
+        inspect["Ports"] = ports
+    if gpus:
+        inspect["GPUs"] = gpus
     lock_dir = os.path.dirname(lock_fpath)
     if lock_dir:
         os.makedirs(lock_dir, exist_ok=True)

{calkit_python-0.26.0 → calkit_python-0.26.2}/calkit/cli/main.py

@@ -1198,6 +1198,17 @@ def run_in_env(
         ]
         if platform:
             docker_cmd += ["--platform", platform]
+        env_vars = env.get("env-vars", {})
+        if env_vars:
+            for key, value in env_vars.items():
+                if isinstance(value, str):
+                    value = os.path.expandvars(value)
+                docker_cmd += ["-e", f"{key}={value}"]
+        if (gpus := env.get("gpus")) is not None:
+            docker_cmd += ["--gpus", gpus]
+        if ports := env.get("ports"):
+            for port in ports:
+                docker_cmd += ["-p", port]
         docker_user = env.get("user")
         if docker_user is None:
             try:

{calkit_python-0.26.0 → calkit_python-0.26.2}/calkit/cli/new.py

@@ -506,7 +506,14 @@ def new_question(
 def new_notebook(
     path: Annotated[str, typer.Argument(help="Notebook path (relative)")],
     title: Annotated[str, typer.Option("--title")],
-    description: Annotated[str, typer.Option("--description")] = None,
+    description: Annotated[str | None, typer.Option("--description")] = None,
+    stage_name: Annotated[
+        str | None,
+        typer.Option(
+            "--stage",
+            help="Name of the pipeline stage that runs this notebook.",
+        ),
+    ] = None,
     commit: Annotated[bool, typer.Option("--commit")] = False,
 ):
     """Add a new notebook."""
@@ -516,8 +523,7 @@ def new_notebook(
         raise ValueError("Path is not a file")
     if not path.endswith(".ipynb"):
         raise ValueError("Path does not have .ipynb extension")
-    # TODO: Add option to create stages that run `calkit nb clean` and
-    # `calkit nb execute`
+    # TODO: Add option to create stage if one doesn't exist
     ck_info = calkit.load_calkit_info()
     notebooks = ck_info.get("notebooks", [])
     paths = [f.get("path") for f in notebooks]
@@ -526,6 +532,12 @@ def new_notebook(
     obj = dict(path=path, title=title)
     if description is not None:
         obj["description"] = description
+    if stage_name is not None:
+        if stage_name not in ck_info.get("pipeline", {}).get("stages", {}):
+            raise ValueError(
+                f"Stage '{stage_name}' does not exist in the pipeline"
+            )
+        obj["stage"] = stage_name
     notebooks.append(obj)
     ck_info["notebooks"] = notebooks
     with open("calkit.yaml", "w") as f:
@@ -533,7 +545,7 @@ def new_notebook(
     if commit:
         repo = git.Repo()
         repo.git.add("calkit.yaml")
-        repo.git.commit(["-m", f"Add notebook {path}"])
+        repo.git.commit(["calkit.yaml", "-m", f"Add notebook {path}"])
 
 
 @new_app.command("docker-env")
@@ -574,6 +586,30 @@ def new_docker_env(
             "--add-layer", help="Add a layer (options: miniforge, foampy)."
         ),
     ] = [],
+    env_vars: Annotated[
+        list[str],
+        typer.Option(
+            "--env-var", help="Environment variables to set in the container."
+        ),
+    ] = [],
+    gpus: Annotated[str | None, typer.Option("--gpus")] = None,
+    args: Annotated[
+        list[str],
+        typer.Option(
+            "--arg",
+            help="Arguments to use when running container.",
+        ),
+    ] = [],
+    deps: Annotated[
+        list[str],
+        typer.Option(
+            "--dep",
+            help=(
+                "Path to add as a dependency, i.e., "
+                "a file that gets added to the container."
+            ),
+        ),
+    ] = [],
     wdir: Annotated[
         str, typer.Option("--wdir", help="Working directory.")
     ] = "/work",
@@ -587,6 +623,16 @@ def new_docker_env(
         str | None,
         typer.Option("--platform", help="Which platform(s) to build for."),
     ] = None,
+    ports: Annotated[
+        list[str],
+        typer.Option(
+            "--port",
+            help=(
+                "Ports to expose in the container, e.g., '8080:80'. "
+                "Can be specified multiple times."
+            ),
+        ),
+    ] = [],
     description: Annotated[
         str | None, typer.Option("--description", help="Description.")
     ] = None,
@@ -615,8 +661,11 @@ def new_docker_env(
     if path is not None and base and os.path.isfile(path) and not overwrite:
         raise_error("Output path already exists (use -f to overwrite)")
     if image_name is None:
-        typer.echo("No image name specified; using environment name")
-        image_name = name
+        typer.echo(
+            "No image name specified; using project and environment name"
+        )
+        _, project_name = calkit.detect_project_name().split("/")
+        image_name = f"{project_name}-{name}"
     repo = git.Repo()
     if base and path is not None:
         txt = "FROM " + base + "\n\n"
@@ -658,6 +707,21 @@ def new_docker_env(
         env["platform"] = platform
     if user:
         env["user"] = user
+    if gpus:
+        env["gpus"] = gpus
+    if env_vars:
+        env["env_vars"] = {}
+        for var in env_vars:
+            if "=" not in var:
+                raise_error(f"Invalid environment variable format: {var}")
+            key, value = var.split("=", 1)
+            env["env_vars"][key] = value
+    if args:
+        env["args"] = args
+    if deps:
+        env["deps"] = deps
+    if ports:
+        env["ports"] = ports
     envs[name] = env
     ck_info["environments"] = envs
     with open("calkit.yaml", "w") as f:
@@ -1788,6 +1852,88 @@ def new_latex_stage(
     )
 
 
+class NotebookStorage(str, Enum):
+    git = "git"
+    dvc = "dvc"
+    none = None
+
+
+@new_app.command(name="jupyter-notebook-stage")
+def new_jupyter_notebook_stage(
+    name: StageArgs.name,
+    environment: StageArgs.environment,
+    notebook_path: Annotated[
+        str, typer.Option("--notebook-path", help="Path to notebook.")
+    ],
+    inputs: StageArgs.inputs = [],
+    outputs: StageArgs.outputs = [],
+    outs_git: StageArgs.outs_git = [],
+    outs_git_no_delete: StageArgs.outs_git_no_delete = [],
+    outs_no_delete: StageArgs.outs_no_delete = [],
+    outs_no_store: StageArgs.outs_no_store = [],
+    outs_no_store_no_delete: StageArgs.outs_no_store_no_delete = [],
+    html_storage: Annotated[
+        NotebookStorage,
+        typer.Option(
+            "--html-storage",
+            help=("In what system to store the HTML output of the notebook."),
+        ),
+    ] = NotebookStorage.dvc,
+    cleaned_ipynb_storage: Annotated[
+        NotebookStorage,
+        typer.Option(
+            "--cleaned-ipynb-storage",
+            help=(
+                "In what system to store the cleaned ipynb output of "
+                "the notebook."
+            ),
+        ),
+    ] = NotebookStorage.git,
+    executed_ipynb_storage: Annotated[
+        NotebookStorage,
+        typer.Option(
+            "--executed-ipynb-storage",
+            help=(
+                "In what system to store the executed ipynb output of "
+                "the notebook."
+            ),
+        ),
+    ] = NotebookStorage.dvc,
+    overwrite: StageArgs.overwrite = False,
+    no_check: StageArgs.no_check = False,
+    no_commit: StageArgs.no_commit = False,
+):
+    """Add a stage to the pipeline that runs a Jupyter notebook."""
+    ck_outs = _to_ck_outs(
+        outputs=outputs,
+        outs_git=outs_git,
+        outs_git_no_delete=outs_git_no_delete,
+        outs_no_delete=outs_no_delete,
+        outs_no_store=outs_no_store,
+        outs_no_store_no_delete=outs_no_store_no_delete,
+    )
+    try:
+        stage = calkit.models.pipeline.JupyterNotebookStage(
+            kind="jupyter-notebook",
+            environment=environment,
+            inputs=inputs,  # type: ignore
+            outputs=ck_outs,
+            notebook_path=notebook_path,
+            html_storage=html_storage.value,
+            cleaned_ipynb_storage=cleaned_ipynb_storage.value,
+            executed_ipynb_storage=executed_ipynb_storage.value,
+        )
+    except Exception as e:
+        raise_error(f"Invalid stage specification: {e}")
+    _save_stage(
+        stage=stage,
+        name=name,
+        overwrite=overwrite,
+        no_check=no_check,
+        no_commit=no_commit,
+    )
+
+
 @new_app.command(name="release")
 def new_release(
     name: Annotated[

{calkit_python-0.26.0 → calkit_python-0.26.2}/calkit/cli/notebooks.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import os
 import subprocess
+from pathlib import PurePosixPath
 
 import typer
 from typing_extensions import Annotated
@@ -78,7 +79,8 @@ def execute_notebook(
     )
     folder = os.path.dirname(fpath_out_exec)
     os.makedirs(folder, exist_ok=True)
-    fpath_out_exec = os.path.abspath(fpath_out_exec)
+    fname = os.path.basename(fpath_out_exec)
+    fpath_out_exec = PurePosixPath(fpath_out_exec).as_posix()
     cmd = [
         "jupyter",
         "nbconvert",
@@ -86,8 +88,10 @@ def execute_notebook(
         "--execute",
         "--to",
         "notebook",
+        "--output-dir",
+        PurePosixPath(folder).as_posix(),
         "--output",
-        fpath_out_exec,
+        fname,
     ]
     run_in_env(cmd=cmd, env_name=env_name, no_check=no_check)
     for to_fmt in to:
@@ -101,7 +105,7 @@ def execute_notebook(
                 raise_error(f"Invalid output format: '{to}'")
             folder = os.path.dirname(fpath_out)
             os.makedirs(folder, exist_ok=True)
-            fpath_out = os.path.abspath(fpath_out)
+            fname_out = os.path.basename(fpath_out)
             # Now convert without executing or checking the environment
             cmd = [
                 "jupyter",
@@ -109,8 +113,10 @@ def execute_notebook(
                 fpath_out_exec,
                 "--to",
                 to_fmt,
+                "--output-dir",
+                PurePosixPath(folder).as_posix(),
                 "--output",
-                fpath_out,
+                fname_out,
             ]
             typer.echo(f"Exporting {to_fmt}")
             run_in_env(cmd=cmd, env_name=env_name, no_check=True)

{calkit_python-0.26.0 → calkit_python-0.26.2}/calkit/magics.py

@@ -13,6 +13,20 @@ from IPython.core.magic import Magics, cell_magic, magics_class
 
 import calkit
 
+# This code used to ensure we can import from the project root in stage
+# scripts
+PROJECT_ROOT_PATH_PY_CODE = """
+import os
+import sys
+
+sys.path.insert(
+    0,
+    os.path.abspath(
+        os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../../")
+    ),
+)
+"""
+
 
 def _parse_string_arg(val: str):
     try:
@@ -191,6 +205,7 @@ class Calkit(Magics):
         # First, let's write this cell out to a script, ensuring that we
         # load the important state at the top
         script_txt = "# This script was automatically generated by Calkit\n\n"
+        script_txt += PROJECT_ROOT_PATH_PY_CODE + "\n\n"
         script_txt += "import calkit\n\n"
         if args.dep:
             for d in args.dep:

{calkit_python-0.26.0 → calkit_python-0.26.2}/calkit/pipeline.py

@@ -65,6 +65,9 @@ def to_dvc(
         outs = []
         if env_fpath is not None:
             deps.append(env_fpath)
+        # Docker envs sometimes have deps, so add those too
+        if env.get("deps", []):
+            deps += env["deps"]
         outs.append({lock_fpath: dict(cache=False, persist=True)})
         stage = dict(cmd=cmd, deps=deps, outs=outs, always_changed=True)
         stage["desc"] = (

{calkit_python-0.26.0 → calkit_python-0.26.2}/calkit/tests/cli/test_main.py

@@ -28,6 +28,7 @@ def test_run_in_env(tmp_dir):
     # First create a new Docker environment for this bare project
     subprocess.check_call(
         "calkit new docker-env "
+        "--image my-image "
         "--name my-image "
         "--from ubuntu "
         "--add-layer uv "

{calkit_python-0.26.0 → calkit_python-0.26.2}/docs/notebooks.md

@@ -80,8 +80,10 @@ of environments a project can use, and they can be of any type.
 
 ## Adding a notebook to the pipeline
 
-A notebook can be added to the pipeline by editing the project's `calkit.yaml`
-file directly, using a `jupyter-notebook` stage.
+A notebook can be added to the pipeline either with
+`calkit new jupyter-notebook-stage` or
+by editing the project's `calkit.yaml`
+file directly.
 For example:
 
 ```yaml