diff --git a/README.md b/README.md index 0ec66c01..cd77ab32 100644 --- a/README.md +++ b/README.md @@ -127,12 +127,13 @@ pip install -e ".[mujoco,curobo]" You may wish to specify some environment variables to configure behavior. Environment variables beginning with the `MLSPACES` prefix can be used to customize MolmoSpaces behavior. -| Environment Variable | Effect | Default | -|---|---|---| -| `MLSPACES_ASSETS_DIR` | Where to place downloaded assets | `~/.cache/molmospaces/assets/` | -| `MLSPACES_FORCE_INSTALL` | Override existing assets | `True` | -| `MLSPACES_PINNED_ASSETS_FILE` | A `.json` file containing pinned versions for each asset, used to override the versions specified in [molmo_spaces_constants.py](molmo_spaces/molmo_spaces_constants.py). | | -| `MUJOCO_EGL_DEVICE_ID` | The rendering device; indices do not always match `CUDA_VISIBLE_DEVICES`. See [here](https://github.com/allenai/molmospaces/issues/66) for details. | `0`| +| Environment Variable | Effect | Default | +|-------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------| +| `MLSPACES_ASSETS_DIR` | Where to place downloaded assets | `~/.cache/molmospaces/assets/` | +| `MLSPACES_FORCE_INSTALL` | Override existing assets | `True` | +| `MLSPACES_PINNED_ASSETS_FILE` | A `.json` file containing pinned versions for each asset, used to override the versions specified in [molmo_spaces_constants.py](molmo_spaces/molmo_spaces_constants.py). | | +| `MUJOCO_EGL_DEVICE_ID` | The rendering device; indices do not always match `CUDA_VISIBLE_DEVICES`. See [here](https://github.com/allenai/molmospaces/issues/66) for details. | `0` | +| `MLSPACES_RESOURCE_STORAGE` | Where to download assets from. Options: `{gc, hf, r2}` for Google Cloud, HuggingFace, or Cloudflare R2 | `gc` | ### Quick Test diff --git a/molmo_spaces/molmo_spaces_constants.py b/molmo_spaces/molmo_spaces_constants.py index 867dd56e..725d8b1e 100644 --- a/molmo_spaces/molmo_spaces_constants.py +++ b/molmo_spaces/molmo_spaces_constants.py @@ -13,11 +13,13 @@ from collections import defaultdict from pathlib import Path from copy import deepcopy +from typing import Literal import compress_json from molmospaces_resources import ( HFRemoteStorage, R2RemoteStorage, + GCRemoteStorage, ResourceManager, setup_resource_manager, str2bool, @@ -82,7 +84,20 @@ def resource_manager_log_level(log_level=logging.DEBUG): else None ) -USE_HUGGING_FACE = False # If True, HF_TOKEN needs to exist in the environment +# If using hf, HF_TOKEN may need to exist in the environment +_valid_storage_types = {"gc", "hf", "r2"} + +StorageTypes = Literal["gc", "hf", "r2"] + +DEFAULT_STORAGE: StorageTypes = "gc" + +_raw_storage_type = os.getenv("MLSPACES_RESOURCE_STORAGE", DEFAULT_STORAGE) +if _raw_storage_type not in _valid_storage_types: + raise ValueError( + f"Invalid MLSPACES_RESOURCE_STORAGE={_raw_storage_type!r}. Must be one of: {_valid_storage_types!r}" + ) + +STORAGE_IN_USE: StorageTypes = _raw_storage_type # type: ignore[assignment] DATA_TYPE_TO_SOURCE_TO_VERSION = dict( robots={ @@ -216,11 +231,16 @@ def register_user_grasp_library(root_name: str, path: Path, object_library: str) def _select_storage(): - return ( - HFRemoteStorage("allenai/molmospaces", repo_prefix="mujoco", token=os.getenv("HF_TOKEN")) - if USE_HUGGING_FACE - else R2RemoteStorage("mujoco-thor-resources") - ) + if STORAGE_IN_USE == "hf": + return HFRemoteStorage( + "allenai/molmospaces", repo_prefix="mujoco", token=os.getenv("HF_TOKEN") + ) + elif STORAGE_IN_USE == "r2": + return R2RemoteStorage("mujoco-thor-resources") + elif STORAGE_IN_USE == "gc": + return GCRemoteStorage("molmospaces-mujoco") + + raise NotImplementedError(f"Storage type {STORAGE_IN_USE} not supported") def get_resource_manager( diff --git a/molmo_spaces_isaac/src/molmo_spaces_isaac/downloader/main.py b/molmo_spaces_isaac/src/molmo_spaces_isaac/downloader/main.py index 7c9ba439..ce989d9e 100644 --- a/molmo_spaces_isaac/src/molmo_spaces_isaac/downloader/main.py +++ b/molmo_spaces_isaac/src/molmo_spaces_isaac/downloader/main.py @@ -5,7 +5,7 @@ from typing import Literal import tyro -from molmospaces_resources import HFRemoteStorage, R2RemoteStorage, ResourceManager +from molmospaces_resources import HFRemoteStorage, R2RemoteStorage, GCRemoteStorage, ResourceManager logger = logging.getLogger("molmospaces_resources") logger.setLevel(logging.DEBUG) @@ -84,8 +84,8 @@ class DownloadArgs: # If not provided, uses HF_TOKEN from environment hf_token: str | None = None - # Use R2 remote storage (HuggingFace by default) - use_r2: bool = False + # Storage to use (HuggingFace by default) + storage: list[Literal["hf", "r2", "gc"]] = field(default_factory=lambda: "hf") def main() -> int: @@ -112,14 +112,19 @@ def main() -> int: dataset_id ] - manager = ResourceManager( - remote_storage=R2RemoteStorage(f"{TYPE_TO_PREFIX[args.type]}-thor-resources") - if args.use_r2 - else HFRemoteStorage( + if args.storage == "hf": + remote_storage = HFRemoteStorage( repo_id="allenai/molmospaces", repo_prefix=TYPE_TO_PREFIX[args.type], token=args.hf_token or os.getenv("HF_TOKEN"), - ), + ) + elif args.storage == "r2": + remote_storage = R2RemoteStorage(f"{TYPE_TO_PREFIX[args.type]}-thor-resources") + else: + remote_storage = GCRemoteStorage(f"molmospaces-{TYPE_TO_PREFIX[args.type]}") + + manager = ResourceManager( + remote_storage=remote_storage, data_type_to_source_to_version=sources_to_version, symlink_dir=args.install_dir, cache_dir=args.cache_dir / args.type, diff --git a/molmo_spaces_maniskill/src/molmo_spaces_maniskill/downloader/main.py b/molmo_spaces_maniskill/src/molmo_spaces_maniskill/downloader/main.py index 7c9ba439..ce989d9e 100644 --- a/molmo_spaces_maniskill/src/molmo_spaces_maniskill/downloader/main.py +++ b/molmo_spaces_maniskill/src/molmo_spaces_maniskill/downloader/main.py @@ -5,7 +5,7 @@ from typing import Literal import tyro -from molmospaces_resources import HFRemoteStorage, R2RemoteStorage, ResourceManager +from molmospaces_resources import HFRemoteStorage, R2RemoteStorage, GCRemoteStorage, ResourceManager logger = logging.getLogger("molmospaces_resources") logger.setLevel(logging.DEBUG) @@ -84,8 +84,8 @@ class DownloadArgs: # If not provided, uses HF_TOKEN from environment hf_token: str | None = None - # Use R2 remote storage (HuggingFace by default) - use_r2: bool = False + # Storage to use (HuggingFace by default) + storage: list[Literal["hf", "r2", "gc"]] = field(default_factory=lambda: "hf") def main() -> int: @@ -112,14 +112,19 @@ def main() -> int: dataset_id ] - manager = ResourceManager( - remote_storage=R2RemoteStorage(f"{TYPE_TO_PREFIX[args.type]}-thor-resources") - if args.use_r2 - else HFRemoteStorage( + if args.storage == "hf": + remote_storage = HFRemoteStorage( repo_id="allenai/molmospaces", repo_prefix=TYPE_TO_PREFIX[args.type], token=args.hf_token or os.getenv("HF_TOKEN"), - ), + ) + elif args.storage == "r2": + remote_storage = R2RemoteStorage(f"{TYPE_TO_PREFIX[args.type]}-thor-resources") + else: + remote_storage = GCRemoteStorage(f"molmospaces-{TYPE_TO_PREFIX[args.type]}") + + manager = ResourceManager( + remote_storage=remote_storage, data_type_to_source_to_version=sources_to_version, symlink_dir=args.install_dir, cache_dir=args.cache_dir / args.type, diff --git a/pyproject.toml b/pyproject.toml index 68477b23..12e9cc71 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,7 +72,7 @@ dependencies = [ "wandb>=0.18.10", "websockets", "pytest", # note (yejin): moved from dev deps since needed in Docker builds - "molmospaces-resources==0.0.1b4", + "molmospaces-resources==0.0.1", ] [project.optional-dependencies] dev = [ diff --git a/scripts/assets/usda_downloader.py b/scripts/assets/usda_downloader.py index 8bdb7841..d83879c1 100644 --- a/scripts/assets/usda_downloader.py +++ b/scripts/assets/usda_downloader.py @@ -4,7 +4,7 @@ from pxr import Sdf, Usd, Ar -from molmospaces_resources import setup_resource_manager, R2RemoteStorage +from molmospaces_resources import setup_resource_manager, GCRemoteStorage from molmo_spaces.molmo_spaces_constants import ASSETS_DIR, DATA_CACHE_DIR from molmospaces_resources import ResourceManager, SourceInfo @@ -36,7 +36,7 @@ def get_resource_manager() -> ResourceManager: global _RESOURCE_MANAGER if _RESOURCE_MANAGER is None: _RESOURCE_MANAGER = setup_resource_manager( - R2RemoteStorage("isaac-thor-resources"), + GCRemoteStorage("molmospaces-isaac"), symlink_dir=ISAAC_ASSETS_DIR, versions=ISAAC_DATA_TYPE_TO_SOURCE_TO_VERSION, cache_dir=ISAAC_DATA_CACHE_DIR,