feat: log project annotations as a zipped MLflow artifact

docs/user-guide/cli-tools.md

@@ -552,7 +552,8 @@ jabs-cli cross-validation DIRECTORY --behavior BEHAVIOR \
     [--grouping-pattern REGEX] \
     [--classifier {catboost|random_forest|xgboost}] \
     [--report-file FILE] \
-    [--mlflow [ENV_FILE]] [--mlflow-experiment NAME] [--mlflow-tag KEY=VALUE] [--mlflow-no-report]
+    [--mlflow [ENV_FILE]] [--mlflow-experiment NAME] [--mlflow-tag KEY=VALUE] \
+    [--mlflow-no-report] [--mlflow-no-annotations]
 ```
 
 - `DIRECTORY`: Path to the JABS project directory.
@@ -562,7 +563,7 @@ jabs-cli cross-validation DIRECTORY --behavior BEHAVIOR \
 - `--grouping-pattern REGEX`: Regular expression applied to each video filename to derive a grouping key. Only used with `--grouping-strategy filename`. If omitted, the pattern saved in the project is used.
 - `--classifier {catboost|random_forest|xgboost}`: Classifier to evaluate. Defaults to `xgboost`. The available choices depend on which classifier libraries are installed; see [Classifier Types](classifier-types.md).
 - `--report-file FILE`: Where to write the training report. The format is chosen by extension: `.md` (Markdown) or `.json` (JSON). If omitted, a timestamped Markdown file is written to the current directory (`<behavior>_<timestamp>_training_report.md`).
-- `--mlflow`, `--mlflow-experiment`, `--mlflow-tag`, `--mlflow-no-report`: Optional MLflow logging (see [MLflow logging](#mlflow-logging)).
+- `--mlflow`, `--mlflow-experiment`, `--mlflow-tag`, `--mlflow-no-report`, `--mlflow-no-annotations`: Optional MLflow logging (see [MLflow logging](#mlflow-logging)).
 
 ### Grouping strategies
 
@@ -602,7 +603,7 @@ jabs-cli cross-validation /path/to/project --behavior grooming \
 
 ### MLflow logging
 
-The cross-validation command can optionally log each run to an [MLflow](https://mlflow.org/) tracking server, recording aggregate metrics, run parameters, descriptive tags, and the training report as an artifact. This is opt-in and off by default.
+The cross-validation command can optionally log each run to an [MLflow](https://mlflow.org/) tracking server, recording aggregate metrics, run parameters, descriptive tags, and — as artifacts — the training report and a zip of the project's annotations. This is opt-in and off by default.
 
 #### Installing the MLflow extra
 
@@ -688,7 +689,10 @@ Each invocation creates one MLflow run named `<behavior>-cv-<timestamp>`.
 
 **Tags:** auto-derived `behavior`, `classifier`, `cv_grouping_strategy`, and `jabs_git` (the short git SHA of the JABS checkout, when available). Any `--mlflow-tag` entries are merged on top, so a user tag wins over an auto tag with the same key.
 
-**Artifact:** the generated training report file, unless `--mlflow-no-report` is passed.
+**Artifacts:**
+
+- The generated training report file, unless `--mlflow-no-report` is passed.
+- `annotations.zip` — a zip of the project's `jabs/annotations` directory, unless `--mlflow-no-annotations` is passed. This captures the label set the run was computed from, so a run's metrics can be traced back to (and reproduced from) the exact annotations. Unpacking it recreates an `annotations/` directory. If the project has no annotations directory, or it holds no archivable files, no artifact is uploaded and a warning is logged. Symlinks are not archived (zipping one would store the target's content, publishing files from outside the project); any that are skipped are logged.
 
 #### Free-form tags
 
@@ -701,14 +705,19 @@ jabs-cli cross-validation /path/to/project --behavior grooming --mlflow settings
 
 Each entry is `KEY=VALUE`; only the first `=` splits the entry, so values may contain `=`.
 
-#### Skipping the report artifact
+#### Skipping artifacts
 
-To log metrics and parameters only (no report upload):
+Each artifact has its own opt-out flag. To log metrics and parameters only:
 
 ```bash
-jabs-cli cross-validation /path/to/project --behavior grooming --mlflow --mlflow-no-report
+jabs-cli cross-validation /path/to/project --behavior grooming --mlflow \
+    --mlflow-no-report --mlflow-no-annotations
 ```
 
+Pass only `--mlflow-no-annotations` to keep the report but skip the annotations upload — worth doing for a project with a very large annotations directory, since the archive is uploaded on every run.
+
+Archiving the annotations is best-effort: if the zip cannot be written, a warning is logged and the run still gets its metrics, parameters, and report artifact.
+
 #### Exit codes and failure handling
 
 MLflow logging happens **after** the cross-validation results are printed and the report is saved, so a logging failure never costs you the results:

src/jabs/classifier/mlflow_logging.py

@@ -25,6 +25,8 @@
 import math
 import os
 import subprocess
+import tempfile
+import zipfile
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -37,6 +39,10 @@
 
 logger = logging.getLogger(__name__)
 
+#: File name used for the zipped annotations artifact attached to each run. Kept
+#: constant across runs so the artifact is easy to find and compare in the UI.
+ANNOTATIONS_ARCHIVE_NAME = "annotations.zip"
+
 
 def mlflow_available() -> bool:
     """Return True if an import spec for the optional ``mlflow`` package is found.
@@ -130,6 +136,83 @@ def load_env_file(env_file: Path | None, *, override: bool = True) -> dict[str,
     return values
 
 
+def archive_annotations(annotations_dir: Path, output_path: Path) -> Path | None:
+    """Zip a JABS project's annotations directory for upload as an MLflow artifact.
+
+    Captures the label set the cross-validation run was computed from, so a run's
+    metrics can be traced back to (and reproduced from) the exact annotations.
+
+    Archive members are stored under a single top-level directory named after
+    ``annotations_dir`` (normally ``annotations/``), so unpacking the zip recreates
+    the directory rather than scattering JSON files into the current directory.
+    Files are added in sorted order to keep archives reproducible.
+
+    Symlinks are **not** archived, and neither are files reached through one. The
+    archive is uploaded to a tracking server, so it must contain only what really
+    lives in the project: zipping a symlink stores the *target's* content, which
+    would publish files from outside the project. Skipped symlinks are logged.
+
+    Args:
+        annotations_dir: The project's ``jabs/annotations`` directory.
+        output_path: Destination path for the ``.zip`` file.
+
+    Returns:
+        ``output_path`` if an archive was written, or None if ``annotations_dir``
+        does not exist or holds no archivable files -- there is then nothing worth
+        uploading.
+
+    Raises:
+        OSError: If reading the annotations or writing the archive fails.
+    """
+    if not annotations_dir.is_dir():
+        logger.warning(
+            "Annotations directory not found, no annotations artifact: %s", annotations_dir
+        )
+        return None
+
+    real_root = annotations_dir.resolve()
+    files: list[Path] = []
+    skipped_links = 0
+    for path in sorted(annotations_dir.rglob("*")):
+        # Skip symlinks themselves, and any entry whose real location falls
+        # outside the annotations directory -- the latter guards against a
+        # directory symlink being traversed (which pathlib does not currently do,
+        # but is not contractually guaranteed across Python versions).
+        if path.is_symlink() or not path.resolve().is_relative_to(real_root):
+            skipped_links += 1
+            continue
+        if path.is_file():
+            files.append(path)
+
+    if skipped_links:
+        logger.warning(
+            "Skipped %d symlinked entr%s under %s: symlinks are not archived",
+            skipped_links,
+            "y" if skipped_links == 1 else "ies",
+            annotations_dir,
+        )
+
+    if not files:
+        logger.warning(
+            "No archivable annotation files in %s, no annotations artifact", annotations_dir
+        )
+        return None
+
+    root = Path(annotations_dir.name)
+    with zipfile.ZipFile(output_path, "w", compression=zipfile.ZIP_DEFLATED) as archive:
+        for path in files:
+            archive.write(path, arcname=str(root / path.relative_to(annotations_dir)))
+
+    logger.info(
+        "Archived %d annotation file(s) from %s into %s (%d bytes)",
+        len(files),
+        annotations_dir,
+        output_path.name,
+        output_path.stat().st_size,
+    )
+    return output_path
+
+
 def _git_sha() -> str | None:
     """Short git SHA of the jabs checkout, or None if unavailable."""
     try:
@@ -266,23 +349,29 @@ def log_cross_validation_to_mlflow(
     *,
     report_data: TrainingReportData,
     report_file: Path | None = None,
+    annotations_dir: Path | None = None,
     env_file: Path | None = None,
     experiment_name: str | None = None,
     run_name: str | None = None,
     tags: dict[str, str] | None = None,
     log_report_artifact: bool = True,
+    log_annotations_artifact: bool = True,
 ) -> tuple[str, str]:
     """Create one MLflow run for a cross-validation run and return its ids.
 
     Logs aggregate CV metrics, curated params, auto-derived plus caller tags,
-    and (optionally) the training report as an artifact. Connection config comes
-    from the environment; ``env_file``, if given, is loaded into it first.
+    and (optionally) the training report and a zip of the project's annotations
+    as artifacts. Connection config comes from the environment; ``env_file``, if
+    given, is loaded into it first.
 
     Args:
         report_data: Completed training report data.
         report_file: Path to the saved training report to upload as an artifact.
             Ignored if None or missing on disk, or if ``log_report_artifact`` is
             False.
+        annotations_dir: The project's ``jabs/annotations`` directory, zipped and
+            uploaded as ``annotations.zip``. Ignored if None, if the directory is
+            missing or empty, or if ``log_annotations_artifact`` is False.
         env_file: Optional ``.env`` file with ``MLFLOW_*`` connection settings.
             If None, connection config comes from the ambient environment.
         experiment_name: Explicit MLflow experiment name. If None, the experiment is
@@ -294,12 +383,15 @@ def log_cross_validation_to_mlflow(
         tags: Caller-supplied run tags; merged over the auto-derived tags (so a
             user tag with the same key wins).
         log_report_artifact: Whether to upload the training report artifact.
+        log_annotations_artifact: Whether to upload the zipped annotations artifact.
 
     Returns:
         A ``(run_id, tracking_uri)`` tuple for the created MLflow run.
 
     Raises:
-        MlflowLoggingError: If the ``mlflow`` package is not installed.
+        MlflowLoggingError: If the ``mlflow`` package is not installed. A failure
+            to archive the annotations is *not* fatal: it is logged as a warning
+            and the run keeps its metrics, params, and report artifact.
     """
     try:
         import mlflow
@@ -340,6 +432,24 @@ def log_cross_validation_to_mlflow(
         if log_report_artifact and report_file is not None and Path(report_file).is_file():
             mlflow.log_artifact(str(report_file))
 
+        if log_annotations_artifact and annotations_dir is not None:
+            # The annotations are supplementary provenance, so a failure to
+            # archive them must not cost the run its metrics, params, or report.
+            try:
+                with tempfile.TemporaryDirectory() as staging_dir:
+                    archive = archive_annotations(
+                        Path(annotations_dir), Path(staging_dir) / ANNOTATIONS_ARCHIVE_NAME
+                    )
+                    if archive is not None:
+                        mlflow.log_artifact(str(archive))
+            except OSError:
+                logger.warning(
+                    "Failed to archive annotations from %s; run %s has no annotations artifact",
+                    annotations_dir,
+                    run.info.run_id,
+                    exc_info=True,
+                )
+
         run_id = run.info.run_id
 
     tracking_uri = mlflow.get_tracking_uri()

src/jabs/resources/docs/user_guide/cli-tools.md

@@ -552,7 +552,8 @@ jabs-cli cross-validation DIRECTORY --behavior BEHAVIOR \
     [--grouping-pattern REGEX] \
     [--classifier {catboost|random_forest|xgboost}] \
     [--report-file FILE] \
-    [--mlflow [ENV_FILE]] [--mlflow-experiment NAME] [--mlflow-tag KEY=VALUE] [--mlflow-no-report]
+    [--mlflow [ENV_FILE]] [--mlflow-experiment NAME] [--mlflow-tag KEY=VALUE] \
+    [--mlflow-no-report] [--mlflow-no-annotations]
 ```
 
 - `DIRECTORY`: Path to the JABS project directory.
@@ -562,7 +563,7 @@ jabs-cli cross-validation DIRECTORY --behavior BEHAVIOR \
 - `--grouping-pattern REGEX`: Regular expression applied to each video filename to derive a grouping key. Only used with `--grouping-strategy filename`. If omitted, the pattern saved in the project is used.
 - `--classifier {catboost|random_forest|xgboost}`: Classifier to evaluate. Defaults to `xgboost`. The available choices depend on which classifier libraries are installed; see [Classifier Types](classifier-types.md).
 - `--report-file FILE`: Where to write the training report. The format is chosen by extension: `.md` (Markdown) or `.json` (JSON). If omitted, a timestamped Markdown file is written to the current directory (`<behavior>_<timestamp>_training_report.md`).
-- `--mlflow`, `--mlflow-experiment`, `--mlflow-tag`, `--mlflow-no-report`: Optional MLflow logging (see [MLflow logging](#mlflow-logging)).
+- `--mlflow`, `--mlflow-experiment`, `--mlflow-tag`, `--mlflow-no-report`, `--mlflow-no-annotations`: Optional MLflow logging (see [MLflow logging](#mlflow-logging)).
 
 ### Grouping strategies
 
@@ -602,7 +603,7 @@ jabs-cli cross-validation /path/to/project --behavior grooming \
 
 ### MLflow logging
 
-The cross-validation command can optionally log each run to an [MLflow](https://mlflow.org/) tracking server, recording aggregate metrics, run parameters, descriptive tags, and the training report as an artifact. This is opt-in and off by default.
+The cross-validation command can optionally log each run to an [MLflow](https://mlflow.org/) tracking server, recording aggregate metrics, run parameters, descriptive tags, and — as artifacts — the training report and a zip of the project's annotations. This is opt-in and off by default.
 
 #### Installing the MLflow extra
 
@@ -688,7 +689,10 @@ Each invocation creates one MLflow run named `<behavior>-cv-<timestamp>`.
 
 **Tags:** auto-derived `behavior`, `classifier`, `cv_grouping_strategy`, and `jabs_git` (the short git SHA of the JABS checkout, when available). Any `--mlflow-tag` entries are merged on top, so a user tag wins over an auto tag with the same key.
 
-**Artifact:** the generated training report file, unless `--mlflow-no-report` is passed.
+**Artifacts:**
+
+- The generated training report file, unless `--mlflow-no-report` is passed.
+- `annotations.zip` — a zip of the project's `jabs/annotations` directory, unless `--mlflow-no-annotations` is passed. This captures the label set the run was computed from, so a run's metrics can be traced back to (and reproduced from) the exact annotations. Unpacking it recreates an `annotations/` directory. If the project has no annotations directory, or it holds no archivable files, no artifact is uploaded and a warning is logged. Symlinks are not archived (zipping one would store the target's content, publishing files from outside the project); any that are skipped are logged.
 
 #### Free-form tags
 
@@ -701,14 +705,19 @@ jabs-cli cross-validation /path/to/project --behavior grooming --mlflow settings
 
 Each entry is `KEY=VALUE`; only the first `=` splits the entry, so values may contain `=`.
 
-#### Skipping the report artifact
+#### Skipping artifacts
 
-To log metrics and parameters only (no report upload):
+Each artifact has its own opt-out flag. To log metrics and parameters only:
 
 ```bash
-jabs-cli cross-validation /path/to/project --behavior grooming --mlflow --mlflow-no-report
+jabs-cli cross-validation /path/to/project --behavior grooming --mlflow \
+    --mlflow-no-report --mlflow-no-annotations
 ```
 
+Pass only `--mlflow-no-annotations` to keep the report but skip the annotations upload — worth doing for a project with a very large annotations directory, since the archive is uploaded on every run.
+
+Archiving the annotations is best-effort: if the zip cannot be written, a warning is logged and the run still gets its metrics, parameters, and report artifact.
+
 #### Exit codes and failure handling
 
 MLflow logging happens **after** the cross-validation results are printed and the report is saved, so a logging failure never costs you the results:

src/jabs/scripts/cli/cli.py

@@ -364,7 +364,8 @@ def prune(ctx: click.Context, directory: Path, behavior: str | None):
     default=None,
     metavar="ENV_FILE",
     help="Enable opt-in MLflow logging of the cross-validation results (aggregate "
-    "metrics, params, and the training report artifact). Optionally takes a path to a "
+    "metrics, params, the training report artifact, and a zip of the project's "
+    "annotations). Optionally takes a path to a "
     ".env file holding MLFLOW_* settings (tracking URI, experiment, auth, TLS); with no "
     "path, those are read from the ambient environment. Absent leaves the command's "
     "behavior unchanged. Requires the optional 'mlflow' extra "
@@ -392,8 +393,14 @@ def prune(ctx: click.Context, directory: Path, behavior: str | None):
     "--mlflow-no-report",
     "mlflow_no_report",
     is_flag=True,
-    help="With --mlflow, skip uploading the training report artifact (metrics + params "
-    "only). No-op without --mlflow.",
+    help="With --mlflow, skip uploading the training report artifact. No-op without --mlflow.",
+)
+@click.option(
+    "--mlflow-no-annotations",
+    "mlflow_no_annotations",
+    is_flag=True,
+    help="With --mlflow, skip uploading the zipped jabs/annotations directory. "
+    "No-op without --mlflow.",
 )
 @click.pass_context
 def cross_validation(
@@ -409,6 +416,7 @@ def cross_validation(
     mlflow_experiment: str | None,
     mlflow_tags: tuple[str, ...],
     mlflow_no_report: bool,
+    mlflow_no_annotations: bool,
 ):
     """Run leave-one-group-out cross-validation for a JABS project."""
     if report_file is not None and report_file.suffix.lower() not in {".md", ".json"}:
@@ -471,6 +479,7 @@ def cross_validation(
             mlflow_experiment=mlflow_experiment,
             mlflow_tags=parsed_mlflow_tags,
             mlflow_log_report=not mlflow_no_report,
+            mlflow_log_annotations=not mlflow_no_annotations,
         )
     except MlflowLoggingError:
         # Cross-validation and the report succeeded; only the optional MLflow

src/jabs/scripts/cli/cross_validation.py

@@ -35,6 +35,7 @@ def run_cross_validation(
     mlflow_experiment: str | None = None,
     mlflow_tags: dict[str, str] | None = None,
     mlflow_log_report: bool = True,
+    mlflow_log_annotations: bool = True,
 ) -> None:
     """Run cross-validation for a JABS project from the command line.
 
@@ -64,6 +65,9 @@ def run_cross_validation(
           over the auto-derived tags.
         mlflow_log_report (bool): Whether to upload the training report as an MLflow
           artifact. Only used when ``mlflow_enabled`` is True.
+        mlflow_log_annotations (bool): Whether to upload a zip of the project's
+          ``jabs/annotations`` directory as an MLflow artifact, capturing the labels
+          the run was computed from. Only used when ``mlflow_enabled`` is True.
 
     Raises:
         MlflowLoggingError: If MLflow logging is requested but fails. The
@@ -259,10 +263,12 @@ def progress_callback():
             run_id, tracking_uri = log_cross_validation_to_mlflow(
                 report_data=training_data,
                 report_file=report_file,
+                annotations_dir=project.annotation_dir,
                 env_file=mlflow_env_file,
                 experiment_name=mlflow_experiment,
                 tags=mlflow_tags,
                 log_report_artifact=mlflow_log_report,
+                log_annotations_artifact=mlflow_log_annotations,
             )
         except Exception as e:
             console.print(f"\nWarning: MLflow logging failed: {e}", style="bold yellow")