Assets lazy loading

Lazy download of large asset files from a public S3 bucket.

Most asset files (configs, poses, and the simplified default NeuroMechFly meshes) are small enough to ship inside the flygym package. The high-resolution fullsize meshes -- especially the FlyBody .obj meshes, which are an order of magnitude larger than everything else combined -- would bloat the package and the git repository, so they are hosted on an institution-managed S3 bucket and pulled in the first time they are needed, similar to how PyTorch downloads pretrained weights.

Downloaded files are cached on disk (see :func:get_cache_root) so the download happens only once per machine. The bucket is public and served over a standard S3-compatible HTTP endpoint, so plain urllib is enough -- no extra dependencies (boto3 etc.) are required.

The bucket stores each remotely hosted asset directory as a flat, versioned sub-prefix of :data:S3_ROOT_PREFIX, so future revisions can be uploaded under a new name without disturbing existing releases. Bump the version constants below to point a release at a new version. Example:

bucket:  flygym_assets/neuromechfly_fullsize_meshes_20260623a/<file>
cache:   ~/.cache/flygym_assets/neuromechfly_fullsize_meshes_20260623a/<file>

get_cache_root()

Return the directory under which downloaded assets are cached.

Resolution order:

1. $FLYGYM_ASSET_CACHE_DIR if set (useful for CI caching or shared, read-only installs);
2. $XDG_CACHE_HOME/flygym_assets if XDG_CACHE_HOME is set;
3. ~/.cache/flygym_assets otherwise.

The directory is named flygym_assets to match the bucket's top-level prefix (:data:S3_ROOT_PREFIX).

Source code in src/flygym/utils/assets_lazy_loading.py

def get_cache_root() -> Path:
    """Return the directory under which downloaded assets are cached.

    Resolution order:

    1. ``$FLYGYM_ASSET_CACHE_DIR`` if set (useful for CI caching or shared,
       read-only installs);
    2. ``$XDG_CACHE_HOME/flygym_assets`` if ``XDG_CACHE_HOME`` is set;
    3. ``~/.cache/flygym_assets`` otherwise.

    The directory is named ``flygym_assets`` to match the bucket's top-level
    prefix (:data:`S3_ROOT_PREFIX`).
    """
    env = os.environ.get("FLYGYM_ASSET_CACHE_DIR")
    if env:
        return Path(env).expanduser()
    # Per the XDG Base Directory spec, a relative XDG_CACHE_HOME is invalid and
    # must be ignored (as is an unset/empty value).
    xdg = os.environ.get("XDG_CACHE_HOME")
    if xdg and os.path.isabs(xdg):
        return Path(xdg) / S3_ROOT_PREFIX
    return Path.home() / ".cache" / S3_ROOT_PREFIX

lazy_load_asset_dir(rel_path)

Return the absolute local path to a bucket asset directory, downloading it from S3 on first use.

Parameters:

Name	Type	Description	Default
rel_path	PathLike | str	Path of the directory within the bucket, relative to :data:S3_ROOT_PREFIX (e.g. "neuromechfly_fullsize_meshes_20260623a", as defined by each fly model's *_FULLSIZE_MESH_DIR constant).	required

The directory is cached under :func:get_cache_root keyed by rel_path. If the cached copy already exists it is returned as-is (no network access); otherwise the whole directory is downloaded into a temporary location and moved into place atomically, so an interrupted or concurrent download never leaves a partial cache.

Raises:

Type	Description
FileNotFoundError	If rel_path does not exist in the bucket.

Source code in src/flygym/utils/assets_lazy_loading.py

def lazy_load_asset_dir(rel_path: os.PathLike | str) -> Path:
    """Return the absolute local path to a bucket asset directory, downloading it
    from S3 on first use.

    Args:
        rel_path: Path of the directory within the bucket, relative to
            :data:`S3_ROOT_PREFIX` (e.g. ``"neuromechfly_fullsize_meshes_20260623a"``,
            as defined by each fly model's ``*_FULLSIZE_MESH_DIR`` constant).

    The directory is cached under :func:`get_cache_root` keyed by ``rel_path``. If
    the cached copy already exists it is returned as-is (no network access);
    otherwise the whole directory is downloaded into a temporary location and moved
    into place atomically, so an interrupted or concurrent download never leaves a
    partial cache.

    Raises:
        FileNotFoundError: If ``rel_path`` does not exist in the bucket.
    """
    rel_path = Path(rel_path)
    cache_dir = get_cache_root() / rel_path
    if cache_dir.is_dir():
        return cache_dir

    cache_dir.parent.mkdir(parents=True, exist_ok=True)
    staging = Path(tempfile.mkdtemp(dir=cache_dir.parent, suffix=".partial"))
    try:
        _download_prefix(f"{S3_ROOT_PREFIX}/{rel_path.as_posix()}", staging)
        try:
            staging.replace(cache_dir)
        except OSError:
            # Another process finished downloading the same asset while we were
            # working: os.replace cannot move onto the now-populated directory.
            # Their copy is equivalent to ours, so use it instead of failing.
            if not cache_dir.is_dir():
                raise
    finally:
        shutil.rmtree(staging, ignore_errors=True)
    return cache_dir

prefetch_meshes()

Eagerly download all remotely hosted meshes into the cache.

Useful for warming a CI cache or preparing an offline environment. Returns the list of local directories that now hold the assets.

Source code in src/flygym/utils/assets_lazy_loading.py

def prefetch_meshes() -> list[Path]:
    """Eagerly download all remotely hosted meshes into the cache.

    Useful for warming a CI cache or preparing an offline environment. Returns the
    list of local directories that now hold the assets.
    """
    # Imported lazily: each model owns its mesh-version constant, and those modules
    # import from this one, so a top-level import here would be circular.
    from flygym.compose.fly.neuromechfly import NEUROMECHFLY_FULLSIZE_MESH_DIR
    from flygym.compose.fly.flybody import FLYBODY_FULLSIZE_MESH_DIR
    from flygym.compose.fly.musculoskeletal import MUSCULOSKELETAL_MESH_DIR

    return [
        lazy_load_asset_dir(NEUROMECHFLY_FULLSIZE_MESH_DIR),
        lazy_load_asset_dir(FLYBODY_FULLSIZE_MESH_DIR),
        lazy_load_asset_dir(MUSCULOSKELETAL_MESH_DIR),
    ]