Google cloud storage

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/<install-hash>` |
-| `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/<install-hash>` |
+| `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

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(

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,

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,

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 = [

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,