src/fpgm/data/ids.py

"""DROID raw bucket URL derivation from episode uuids.

The public GCS bucket keys objects by lab, calendar date, outcome
(``success``/``failure``) and a strftime-derived directory name. Everything
here is pure stdlib -- the network probe in :func:`resolve_episode_url` does a
lazy ``import requests`` -- so this module (and its uuid parsing logic) is
testable without numpy/h5py/requests installed, per :mod:`tests.test_ids`.
"""

from __future__ import annotations

from urllib.parse import quote

from fpgm.types import EpisodeId, EpisodeNotFoundError

#: DROID raw bucket root for the public, anonymous-access 1.0.1 release.
DROID_RAW_BUCKET_ROOT = "https://storage.googleapis.com/gresearch/robotics/droid_raw/1.0.1"

#: Whether an episode is a "success" or "failure" trajectory is not encoded
#: anywhere in the uuid, so callers who don't already know it must probe both.
#: "success" is tried first since it is the overwhelmingly common outcome.
OUTCOME_PROBE_ORDER = ("success", "failure")


def episode_dirname(episode_id: EpisodeId) -> str:
    """Derive the per-episode directory name used inside the DROID raw bucket.

    This reproduces the directory name Python's own ``datetime.strftime``
    produced when the episode was recorded:
    ``strftime('%a_%b_%e_%H:%M:%S_%Y').replace(' ', '_')``.

    The trap is glibc's ``%e``: it pads single-digit days with a *space*, not
    a zero, and that space then gets replaced by the trailing ``.replace(' ',
    '_')`` just like the literal underscores in the format string do -- so
    single-digit days produce a **double** underscore
    (e.g. ``"Fri_Oct__6_19:50:14_2023"``), not a single one. Verified against
    PointWorld's own ``real/droid_paths.txt`` for both single- and
    double-digit days.
    """
    raw = episode_id.timestamp.strftime("%a_%b_%e_%H:%M:%S_%Y")
    return raw.replace(" ", "_")


def episode_date(episode_id: EpisodeId) -> str:
    """The ``YYYY-MM-DD`` bucket path component."""
    return episode_id.timestamp.strftime("%Y-%m-%d")


def episode_base_url(episode_id: EpisodeId, outcome: str) -> str:
    """Build the (unverified) bucket URL for one success/failure hypothesis.

    The colon-heavy ``episode_dirname`` and the plus-heavy uuid both need
    percent-encoding for use in a URL; ``urllib.parse.quote``'s defaults
    (which treat ``:`` and ``+`` as reserved) handle both correctly.
    """
    if outcome not in OUTCOME_PROBE_ORDER:
        raise ValueError(f"outcome must be one of {OUTCOME_PROBE_ORDER}, got {outcome!r}")
    dirname = quote(episode_dirname(episode_id))
    return (
        f"{DROID_RAW_BUCKET_ROOT}/{episode_id.lab}/{outcome}/"
        f"{episode_date(episode_id)}/{dirname}/"
    )


def metadata_filename(episode_id: EpisodeId) -> str:
    """``metadata_<uuid>.json``, with the uuid's ``+`` percent-encoded for URLs."""
    return f"metadata_{quote(episode_id.uuid)}.json"


#: PointWorld's ``cameras.json`` records the episode's real bucket directory as a
#: ``gs://`` URL. It is worth preferring over the derived one: across all 42935
#: episodes the uuid-derived path is wrong for 1.52% of them (465 use underscores
#: where the dirname normally has colons, 189 sit under a calendar date one day
#: off their own timestamp -- both concentrated in IRIS, where 14% of episodes are
#: unreachable without it). It also states the outcome, saving the HEAD probe that
#: :func:`resolve_episode_url` needs for the 16.4% of episodes under ``failure/``.
_GS_PREFIX = "gs://gresearch/robotics/droid_raw/1.0.1/"


def bucket_url_from_scene_path(scene_path: str) -> str:
    """Convert a ``cameras.json`` ``scene_path`` into an https base URL (trailing slash).

    Raises:
        ValueError: if ``scene_path`` is not a DROID raw 1.0.1 ``gs://`` URL --
            silently falling back would turn a catalogue that disagrees with this
            code's assumptions into a mysterious 404 much later.
    """
    if not scene_path.startswith(_GS_PREFIX):
        raise ValueError(f"not a DROID raw 1.0.1 scene_path: {scene_path!r}")
    rel = scene_path[len(_GS_PREFIX):].strip("/")
    return f"{DROID_RAW_BUCKET_ROOT}/" + "/".join(quote(p) for p in rel.split("/")) + "/"


def resolve_episode_url(uuid: str, timeout: float = 10.0) -> str:
    """Probe the bucket for ``uuid`` under ``success/`` then ``failure/``.

    Returns the base URL (trailing slash) of whichever outcome directory
    actually contains the episode. Network access is isolated in this one
    function via a lazy ``requests`` import so the rest of :mod:`fpgm.data.ids`
    stays free of third-party dependencies.

    Raises:
        EpisodeNotFoundError: neither outcome directory has this uuid.
    """
    import requests

    episode_id = EpisodeId(uuid)
    for outcome in OUTCOME_PROBE_ORDER:
        base_url = episode_base_url(episode_id, outcome)
        probe_url = base_url + metadata_filename(episode_id)
        try:
            resp = requests.head(probe_url, timeout=timeout, allow_redirects=True)
        except requests.RequestException:
            continue
        if resp.status_code == 200:
            return base_url
    raise EpisodeNotFoundError(
        f"episode {uuid!r} not found under success/ or failure/ in the DROID raw bucket"
    )