oopsie-data-tools 0.2.0__py3-none-any.whl

oopsie_data_tools/utils/hf_upload.py

@@ -0,0 +1,235 @@
+"""Validation + HuggingFace upload helpers.
+
+Backs ``oopsie-data validate``, ``oopsie-data upload`` and ``oopsie-data submissions``
+(:mod:`oopsie_data_tools.cli`), which is the only entry point — the parallel
+``scripts/validate_and_upload/`` copies of this pipeline are gone.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from collections import Counter
+
+from oopsie_data_tools.utils.contributor_config import read_contributor_config
+from oopsie_data_tools.utils.hf_limits import FILE_LIMIT
+from oopsie_data_tools.utils.log import setup_logger
+from oopsie_data_tools.utils.validation.errors import EpisodeValidationError
+from oopsie_data_tools.utils.validation.validation_utils import (
+    validate_h5_file,
+    validate_session_dir,
+)
+
+logger = logging.getLogger(__name__)
+
+
+# ── HuggingFace authentication ────────────────────────────────────────────────
+
+
+def resolve_hf_target() -> tuple[str, str]:
+    """Resolve ``(hf_token, repo_id)`` lazily.
+
+    Read only when actually uploading so validation and ``--help`` work on a fresh
+    checkout without a filled-in contributor config. ``HF_TOKEN`` in the environment
+    overrides the config token.
+    """
+    lab_id, config_token = read_contributor_config()
+    token = os.environ.get("HF_TOKEN", "").strip() or config_token
+    return token, f"OopsieData-Submissions/{lab_id}"
+
+
+def hf_login(token: str) -> str:
+    from huggingface_hub import login, whoami
+
+    login(token=token, add_to_git_credential=False)
+    info = whoami(token=token)
+    logger.info("[auth]  Logged in as: %s", info["name"])
+    return info["name"]
+
+
+# ── Validation ────────────────────────────────────────────────────────────────
+
+
+def run_validation(base_path: str, episode_id: str | None = None, log_path: str | None = None) -> int:
+    """Validate a single episode or a session dir.
+
+    Args:
+        base_path: Session directory, or the directory holding ``<episode_id>.h5``.
+        episode_id: Optional zero-padded episode id; validates just that episode.
+        log_path: Optional file to mirror log output into.
+
+    Returns:
+        A shell-style exit code: 0 if everything passed, 1 otherwise.
+    """
+    # Route this module's pass/fail lines to the log file too, so --log-path captures
+    # the single-file path (validate_session_dir already logs through its own logger).
+    if log_path is not None:
+        setup_logger(__name__, log_path)
+
+    target = os.path.join(base_path, f"{episode_id}.h5") if episode_id else base_path
+    if os.path.isfile(target):
+        try:
+            validate_h5_file(target, strict_annotation_check=True, log_path=log_path)
+            logger.info("%s passed", os.path.basename(target))
+            return 0
+        except EpisodeValidationError as e:
+            logger.error("Validation failed: %s", e)
+            return 1
+        except Exception as e:
+            logger.error("Unexpected error: %s", e)
+            return 1
+
+    if os.path.isdir(target):
+        return validate_session_dir(target, strict_annotation_check=True, log_path=log_path)
+
+    logger.error("Path does not exist: %s", target)
+    return 1
+
+
+# ── Repo creation ─────────────────────────────────────────────────────────────
+
+
+def ensure_repo(api, repo: str) -> None:
+    try:
+        api.repo_info(repo_id=repo, repo_type="dataset")
+        logger.info("[hf]    Repo already exists: https://huggingface.co/datasets/%s", repo)
+    except Exception:
+        logger.info("[hf]    Creating repo: %s", repo)
+        api.create_repo(repo_id=repo, repo_type="dataset", private=False)
+        logger.info("[hf]    Created: https://huggingface.co/datasets/%s", repo)
+
+
+# ── Upload ────────────────────────────────────────────────────────────────────
+
+
+def check_folder_size(samples_dir: str, suggest_fix: bool = True) -> list[tuple[str, int]]:
+    """Find directories under ``samples_dir`` that exceed the HF per-directory file limit.
+
+    Args:
+        samples_dir: Session directory to check, recursively.
+        suggest_fix: Log the command that repairs the layout. Set False when the caller is
+            about to run it anyway (``upload --with-restructure``), so the output does not
+            tell the user to do by hand what is happening on the next line.
+
+    Returns:
+        A list of ``(directory, file_count)`` pairs; empty if the layout is fine.
+    """
+    oversized = [
+        (dirpath, len(filenames))
+        for dirpath, _, filenames in os.walk(samples_dir)
+        if len(filenames) > FILE_LIMIT
+    ]
+    if not oversized:
+        return []
+
+    logger.error("[precheck] The following directories exceed %d files:", FILE_LIMIT)
+    for d, n in oversized:
+        logger.error("             %s  (%d files)", d, n)
+    if suggest_fix:
+        logger.error(
+            "[precheck] HuggingFace Hub enforces a per-directory file limit.\n"
+            "           Restructure the folder first, then re-run the upload,\n"
+            "           or pass --with-restructure to do both in one step:\n\n"
+            "             oopsie-data restructure --source %s\n",
+            samples_dir,
+        )
+    return oversized
+
+
+def upload_dataset(api, repo: str, samples_dir: str) -> None:
+    logger.info("[upload] Uploading %s → %s", samples_dir, repo)
+    logger.info("[upload] Files to upload:")
+    total_bytes = 0
+    for root, _, files in os.walk(samples_dir):
+        for f in files:
+            fpath = os.path.join(root, f)
+            size = os.path.getsize(fpath)
+            rel = os.path.relpath(fpath, samples_dir)
+            total_bytes += size
+            logger.info("           %s  (%.1f MB)", rel, size / 1e6)
+
+    logger.info("[upload] Total size: %.2f GB", total_bytes / 1e9)
+    logger.info("[upload] Uploading (this may take several minutes)...")
+
+    api.upload_large_folder(
+        folder_path=samples_dir,
+        repo_id=repo,
+        repo_type="dataset",
+    )
+
+    logger.info("[upload] Done!")
+    logger.info("[upload] Dataset URL: https://huggingface.co/datasets/%s", repo)
+
+    # Post-upload confirmation: read the repo back so the user sees their data landed.
+    try:
+        remote_h5 = [
+            f for f in api.list_repo_files(repo_id=repo, repo_type="dataset")
+            if f.endswith(".h5")
+        ]
+        local_h5 = sum(
+            1 for _r, _d, files in os.walk(samples_dir) for fn in files if fn.endswith(".h5")
+        )
+        logger.info(
+            "[upload] Confirmed: %d episode(s) now in the repo (from %d local .h5).",
+            len(remote_h5), local_h5,
+        )
+        if local_h5 and len(remote_h5) < local_h5:
+            logger.warning(
+                "[upload] Remote episode count (%d) is below local (%d) — "
+                "re-run the upload if this is unexpected.",
+                len(remote_h5), local_h5,
+            )
+    except Exception as e:
+        logger.warning("[upload] Could not confirm upload via repo listing: %s", e)
+
+
+# ── Submissions query ─────────────────────────────────────────────────────────
+
+
+def query_submissions(lab_id: str | None = None) -> int:
+    """Report what has landed in ``OopsieData-Submissions/<lab_id>``, downloading nothing.
+
+    Args:
+        lab_id: Lab to query. Defaults to the ``lab_id`` in the contributor config.
+
+    Returns:
+        A shell-style exit code. A repo that does not exist yet is not an error — it is
+        created on the first successful upload — so that case still returns 0.
+    """
+    from huggingface_hub import HfApi
+
+    if lab_id:
+        # An explicit --lab-id must not require a filled-in contributor config, but the
+        # token still comes from it unless $HF_TOKEN overrides.
+        try:
+            _, config_token = read_contributor_config()
+        except RuntimeError:
+            config_token = ""
+    else:
+        lab_id, config_token = read_contributor_config()
+
+    hf_token = os.environ.get("HF_TOKEN", "").strip() or config_token
+    repo = f"OopsieData-Submissions/{lab_id}"
+    api = HfApi(token=hf_token or None)
+
+    try:
+        api.repo_info(repo_id=repo, repo_type="dataset")
+    except Exception:
+        logger.info("No submissions repo found yet at https://huggingface.co/datasets/%s", repo)
+        logger.info("(It is created automatically on your first successful upload.)")
+        return 0
+
+    files = api.list_repo_files(repo_id=repo, repo_type="dataset")
+    h5 = [f for f in files if f.endswith(".h5") or f.endswith(".hdf5")]
+    mp4 = [f for f in files if f.endswith(".mp4")]
+    by_dir = Counter(f.split("/")[0] if "/" in f else "(root)" for f in h5)
+
+    logger.info("Repo:           https://huggingface.co/datasets/%s", repo)
+    logger.info("Episodes (.h5): %d", len(h5))
+    logger.info("Videos (.mp4):  %d", len(mp4))
+    logger.info("Total files:    %d", len(files))
+    if by_dir:
+        logger.info("Episodes by top-level folder:")
+        for name, count in sorted(by_dir.items()):
+            logger.info("  %-32s %d", name, count)
+    return 0

oopsie_data_tools/utils/log.py

@@ -0,0 +1,35 @@
+"""Attach a log file to a named logger, without stacking duplicate handlers."""
+
+from __future__ import annotations
+
+import logging
+import os
+
+
+def setup_logger(name: str, log_file: str | os.PathLike) -> logging.Logger:
+    """Add a file handler for ``log_file`` to the named logger, once.
+
+    Pass ``__name__`` of the calling module so the logger stays within the
+    ``oopsie_data_tools`` hierarchy and does not affect other packages.
+
+    Calling this repeatedly for the same file is safe — validation calls it once per
+    episode. Calling it for a *different* file adds a second handler rather than being
+    ignored: the guard used to be "does this logger have any handler at all", so the second
+    ``--log-path`` in a process silently wrote nothing.
+    """
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.INFO)
+
+    target = os.path.abspath(os.fspath(log_file))
+    for handler in logger.handlers:
+        if isinstance(handler, logging.FileHandler) and (
+            os.path.abspath(handler.baseFilename) == target
+        ):
+            return logger
+
+    handler = logging.FileHandler(target)
+    handler.setFormatter(
+        logging.Formatter("%(asctime)s %(levelname)s %(name)s — %(message)s")
+    )
+    logger.addHandler(handler)
+    return logger

oopsie_data_tools/utils/migrate_taxonomy_v2.py

@@ -0,0 +1,215 @@
+"""Offline migration of annotation attrs from taxonomy v1 to v2.
+
+Deliberately **not** wired into the ``oopsie-data`` CLI. Readers upcast v1 files on the fly
+(see :func:`~oopsie_data_tools.annotation_tool.annotation_schema.read_annotation_attrs`), so
+nobody *needs* to run this — it exists for the one-off case of normalizing a directory in
+place, and making it a supported command would imply an ongoing guarantee we do not want.
+
+Run it ad hoc::
+
+    python -m oopsie_data_tools.utils.migrate_taxonomy_v2 <path>            # dry run
+    python -m oopsie_data_tools.utils.migrate_taxonomy_v2 <path> --apply
+
+The mapping tables live in ``annotation_schema`` rather than here, so this and the read-time
+upcast cannot disagree about what a v1 value means.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any, Mapping
+
+import h5py
+
+from oopsie_data_tools.annotation_tool.annotation_schema import (
+    ANNOTATION_SCHEMA_V2,
+    OUTCOME_SUCCESS,
+    SEVERITIES,
+    SIDE_EFFECT_CATEGORIES,
+    TAXONOMY_SCHEMA_V2,
+    V1_SUCCESS_CATEGORY_TO_OUTCOME,
+    as_value_list,
+    decode_attr,
+    normalize_category,
+    normalize_severity,
+    parse_taxonomy,
+    success_to_outcome,
+)
+from oopsie_data_tools.utils.h5 import find_episode_files
+
+
+def unmapped_values(attrs: Mapping[str, Any]) -> list:
+    """v1 category/severity values with no slug, reported so an operator can look at them.
+
+    These are preserved verbatim rather than folded into ``other``: they are real human
+    labels, and relabelling one on a guess is not something a later run could undo.
+    """
+    taxonomy = parse_taxonomy(attrs.get("taxonomy", ""))
+    out = []
+    raw_categories = taxonomy.get("side_effect_category", taxonomy.get("failure_category"))
+    for value in as_value_list(raw_categories):
+        if normalize_category(value) not in SIDE_EFFECT_CATEGORIES:
+            out.append(decode_attr(value))
+    severity = taxonomy.get("severity")
+    if severity is not None and decode_attr(severity):
+        if normalize_severity(severity) not in SEVERITIES:
+            out.append(decode_attr(severity))
+    return out
+
+
+def is_already_v2(attrs: Mapping[str, Any]) -> bool:
+    """Whether this subgroup is already v2, by schema string or by a usable outcome."""
+    if decode_attr(attrs.get("schema")) == ANNOTATION_SCHEMA_V2:
+        return True
+    outcome = decode_attr(parse_taxonomy(attrs.get("taxonomy", "")).get("outcome")).lower()
+    return outcome in OUTCOME_SUCCESS
+
+
+def migrate_annotation_attrs(attrs: Mapping[str, Any]) -> "dict | None":
+    """v1 attr set -> the v2 attr set, or ``None`` when there is nothing to do.
+
+    ``None`` covers both "already v2" and "not actually an annotation", which keeps the
+    migration idempotent: a second run rewrites nothing.
+
+    Pure — takes and returns plain mappings — so the mapping can be tested without HDF5.
+    """
+    if is_already_v2(attrs):
+        return None
+    # A subgroup with neither a success score nor a taxonomy is not an annotation. Inventing
+    # a success for it would turn a file that fails strict validation into one that passes.
+    if attrs.get("success") is None and not decode_attr(attrs.get("taxonomy", "")):
+        return None
+
+    taxonomy = parse_taxonomy(attrs.get("taxonomy", ""))
+
+    outcome = success_to_outcome(attrs.get("success"))
+    if outcome is None:
+        return None
+    legacy_category = decode_attr(taxonomy.get("success_category")).lower()
+    if outcome == "success" and legacy_category:
+        outcome = V1_SUCCESS_CATEGORY_TO_OUTCOME.get(legacy_category, outcome)
+
+    raw_categories = taxonomy.get("side_effect_category", taxonomy.get("failure_category"))
+
+    out = dict(attrs)
+    out["schema"] = ANNOTATION_SCHEMA_V2
+    out["taxonomy_schema"] = TAXONOMY_SCHEMA_V2
+    # The v1 spelling is dropped rather than left alongside the new one, so no file ever
+    # carries both and readers never have to pick.
+    out.pop("failure_description", None)
+    out["episode_description"] = decode_attr(attrs.get("failure_description", ""))
+    out["taxonomy"] = json.dumps(
+        {
+            "outcome": outcome,
+            "side_effect_category": [
+                normalize_category(c) for c in as_value_list(raw_categories)
+            ],
+            "severity": normalize_severity(taxonomy.get("severity", "")),
+        },
+        ensure_ascii=False,
+    )
+    return out
+
+
+def migrate_h5_file(h5_path: "str | Path", *, apply: bool = True) -> dict:
+    """Migrate every annotator subgroup in one episode file.
+
+    Returns a report: ``{"path", "migrated": [...], "skipped": {annotator: reason},
+    "unmapped": {annotator: [...]}, "error": str | None}``.
+    """
+    path = Path(h5_path)
+    report: dict = {
+        "path": str(path),
+        "migrated": [],
+        "skipped": {},
+        "unmapped": {},
+        "error": None,
+    }
+
+    try:
+        # Read-only first, so a dry run — and a file with nothing to migrate — never opens
+        # for write and never touches mtime.
+        with h5py.File(path, "r") as f:
+            ea = f.get("episode_annotations")
+            if not isinstance(ea, h5py.Group):
+                report["skipped"]["*"] = "no episode_annotations group"
+                return report
+            planned = {}
+            for name in ea.keys():
+                sub = ea[name]
+                if not isinstance(sub, h5py.Group):
+                    report["skipped"][name] = "not a group"
+                    continue
+                attrs = dict(sub.attrs)
+                new_attrs = migrate_annotation_attrs(attrs)
+                if new_attrs is None:
+                    report["skipped"][name] = (
+                        "already v2" if is_already_v2(attrs) else "not an annotation"
+                    )
+                    continue
+                leftover = unmapped_values(attrs)
+                if leftover:
+                    report["unmapped"][name] = leftover
+                planned[name] = (attrs, new_attrs)
+
+        if not planned or not apply:
+            report["migrated"] = sorted(planned)
+            return report
+
+        with h5py.File(path, "r+") as f:
+            ea = f["episode_annotations"]
+            for name, (old_attrs, new_attrs) in planned.items():
+                sub = ea[name]
+                for key in set(old_attrs) - set(new_attrs):
+                    del sub.attrs[key]
+                for key, value in new_attrs.items():
+                    sub.attrs[key] = value
+        report["migrated"] = sorted(planned)
+    except Exception as e:  # noqa: BLE001 - one bad file must not stop the walk
+        report["error"] = f"{type(e).__name__}: {e}"
+
+    return report
+
+
+def migrate_tree(root: "str | Path", *, apply: bool = True) -> list:
+    """Migrate every episode file under *root*; one report per file."""
+    return [migrate_h5_file(p, apply=apply) for p in find_episode_files(Path(root))]
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("path", type=Path, help="episode .h5 file or a directory of them")
+    parser.add_argument(
+        "--apply",
+        action="store_true",
+        help="actually write the changes (without this it is a dry run)",
+    )
+    args = parser.parse_args()
+
+    if args.path.is_dir():
+        reports = migrate_tree(args.path, apply=args.apply)
+    else:
+        reports = [migrate_h5_file(args.path, apply=args.apply)]
+
+    migrated = 0
+    for report in reports:
+        if report["error"]:
+            print(f"ERROR {report['path']}: {report['error']}")
+            continue
+        if report["migrated"]:
+            migrated += 1
+            print(f"{'migrated' if args.apply else 'would migrate'} {report['path']}: "
+                  f"{', '.join(report['migrated'])}")
+        for name, values in report["unmapped"].items():
+            print(f"  unmapped values kept verbatim in {name}: {values}")
+
+    verb = "Migrated" if args.apply else "Would migrate"
+    print(f"\n{verb} {migrated} of {len(reports)} file(s).")
+    if not args.apply:
+        print("Dry run — re-run with --apply to write.")
+
+
+if __name__ == "__main__":
+    main()

oopsie_data_tools/utils/paths.py

@@ -0,0 +1,162 @@
+"""Resolve where oopsie configs live, so the package works when pip-installed.
+
+The two kinds of config are looked up separately, because they belong to different things.
+
+**Credentials** (``contributor_config.yaml``) belong to the person, so they are stored once
+per user and shared by every project:
+
+1. ``$OOPSIE_CONFIG_DIR`` — explicit override
+2. ``$XDG_CONFIG_HOME/oopsie-data`` (default ``~/.config/oopsie-data``)
+3. the repo's ``configs/`` directory, when running from a source checkout
+
+**Robot profiles** belong to the robot code that uses them, so they are looked up next to
+that code and never in the user config directory:
+
+1. ``$OOPSIE_ROBOT_PROFILES_DIR`` — explicit override
+2. ``./robot_profiles`` or ``./configs/robot_profiles``, relative to the working directory
+3. the repo's ``configs/robot_profiles``, when running from a source checkout
+
+In each chain the first location that actually exists wins; when none do, the first
+*writable* location is returned, which is where new files should be created. Note that
+profiles are usually loaded by explicit path (``load_robot_profile(path)``), so this lookup
+mainly backs the bundled example profiles.
+
+``oopsie-data --config-dir <dir> <command>`` sets the credential override for one invocation.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+ENV_CONFIG_DIR = "OOPSIE_CONFIG_DIR"
+ENV_PROFILES_DIR = "OOPSIE_ROBOT_PROFILES_DIR"
+
+_PACKAGE_ROOT = Path(__file__).resolve().parent.parent
+# Deliberately only the checkout's configs/. Nothing in either lookup chain resolves into
+# an installed package: config that lives in site-packages is read-only, invisible to the
+# user, and would silently shadow their own. Installed users get None here and fall through
+# to a location they control. ``oopsie-data new-profile`` covers the "I have no checkout to
+# copy a template from" case.
+_REPO_CONFIG_DIR = _PACKAGE_ROOT.parent / "configs"
+
+PROFILES_DIR_NAME = "robot_profiles"
+
+
+def _env_dir(name: str) -> Path | None:
+    value = os.environ.get(name, "").strip()
+    return Path(value).expanduser().resolve() if value else None
+
+
+def _unique(candidates: list[Path | None]) -> list[Path]:
+    """Drop Nones and repeats, keeping order.
+
+    Two entries can name the same directory — running from the checkout root makes
+    ``./configs/robot_profiles`` and the repo's own profile dir identical — and a chain
+    that lists it twice is confusing when printed by ``oopsie-data config``.
+    """
+    seen = set()
+    unique = []
+    for candidate in candidates:
+        if candidate is None:
+            continue
+        key = candidate.resolve()
+        if key not in seen:
+            seen.add(key)
+            unique.append(candidate)
+    return unique
+
+
+# ── credentials ───────────────────────────────────────────────────────────────
+
+
+def env_config_dir() -> Path | None:
+    """``$OOPSIE_CONFIG_DIR`` if set, else None."""
+    return _env_dir(ENV_CONFIG_DIR)
+
+
+def user_config_dir() -> Path:
+    """Per-user config location, where ``init`` stores credentials."""
+    xdg = os.environ.get("XDG_CONFIG_HOME", "").strip()
+    base = Path(xdg).expanduser() if xdg else Path.home() / ".config"
+    return base / "oopsie-data"
+
+
+def repo_config_dir() -> Path | None:
+    """The checkout's ``configs/`` directory, or None when not running from source."""
+    return _REPO_CONFIG_DIR if _REPO_CONFIG_DIR.is_dir() else None
+
+
+def config_search_dirs() -> list[Path]:
+    """Candidate directories for the contributor config, highest precedence first."""
+    return _unique([env_config_dir(), user_config_dir(), repo_config_dir()])
+
+
+def write_config_dir() -> Path:
+    """Where a new contributor config should be written."""
+    return env_config_dir() or user_config_dir()
+
+
+def contributor_config_path() -> Path:
+    """The contributor config that will be read, or where to create one."""
+    for candidate in config_search_dirs():
+        target = candidate / "contributor_config.yaml"
+        if target.is_file():
+            return target
+    return write_config_dir() / "contributor_config.yaml"
+
+
+def config_dir() -> Path:
+    """Directory the contributor config resolves to."""
+    return contributor_config_path().parent
+
+
+def config_dir_source() -> str:
+    """Human-readable reason the resolved config dir was chosen (for messages)."""
+    resolved = contributor_config_path().parent
+    if resolved == env_config_dir():
+        return f"${ENV_CONFIG_DIR}"
+    if resolved == user_config_dir():
+        return "user config dir"
+    if resolved == repo_config_dir():
+        return "repo checkout"
+    return str(resolved)
+
+
+# ── robot profiles ────────────────────────────────────────────────────────────
+
+
+def env_profiles_dir() -> Path | None:
+    """``$OOPSIE_ROBOT_PROFILES_DIR`` if set, else None."""
+    return _env_dir(ENV_PROFILES_DIR)
+
+
+def project_profiles_dirs() -> list[Path]:
+    """Profile directories relative to the working directory, in precedence order."""
+    cwd = Path.cwd()
+    return [cwd / PROFILES_DIR_NAME, cwd / "configs" / PROFILES_DIR_NAME]
+
+
+def profiles_search_dirs() -> list[Path]:
+    """Candidate directories for robot profiles, highest precedence first."""
+    repo_dir = repo_config_dir()
+    return _unique(
+        [
+            env_profiles_dir(),
+            *project_profiles_dirs(),
+            repo_dir / PROFILES_DIR_NAME if repo_dir is not None else None,
+        ]
+    )
+
+
+def write_profiles_dir() -> Path:
+    """Where a new robot profile should be written: the override, else project-local."""
+    return env_profiles_dir() or project_profiles_dirs()[0]
+
+
+def robot_profiles_dir() -> Path:
+    """The robot profile directory that will be read, or where to create one."""
+    for candidate in profiles_search_dirs():
+        if candidate.is_dir():
+            return candidate
+    return write_profiles_dir()