Change in API from s3 to https

valence-labs · Jul 16, 2024 · a89a672 · a89a672
1 parent 6acf7fc
commit a89a672
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 16 deletions.
diff --git a/openqdc/utils/download_api.py b/openqdc/utils/download_api.py
@@ -35,7 +35,7 @@ class FileSystem:
     public_endpoint: Optional[AbstractFileSystem] = None
     private_endpoint: Optional[AbstractFileSystem] = None
     local_endpoint: AbstractFileSystem = LocalFileSystem()
-    endpoint_url = "https://874f02b9d981bd6c279e979c0d91c4b4.r2.cloudflarestorage.com"
+    endpoint_url = "https://storage.openqdc.org"  # "https://874f02b9d981bd6c279e979c0d91c4b4.r2.cloudflarestorage.com"
 
     def __init__(self):
         load_dotenv()  # load environment variables from .env
@@ -89,19 +89,29 @@ def get_default_endpoint(self, endpoint: str) -> AbstractFileSystem:
         """
         if endpoint == "private":
             # return fsspec.filesystem("gs")
-            return fsspec.filesystem("s3", key=self.KEY, secret=self.SECRET, endpoint_url=self.endpoint_url)
-        elif endpoint == "public":
-            return (
-                fsspec.filesystem(
-                    "s3",
-                    key="a046308c078b0134c9e261aa91f63ab2",
-                    secret="d5b32f241ad8ee8d0a3173cd51b4f36d6869f168b21acef75f244a81dc10e1fb",
-                    endpoint_url=self.endpoint_url,
-                )
-                if os.environ.get("OPENQDC_DOWNLOAD_API") is None
-                else fsspec.filesystem("https")
+            return fsspec.filesystem(
+                "s3",
+                key=self.KEY,
+                secret=self.SECRET,
+                endpoint_url="https://874f02b9d981bd6c279e979c0d91c4b4.r2.cloudflarestorage.com",
             )
-            # return fsspec.filesystem("https")
+        elif endpoint == "public":
+            return fsspec.filesystem("https")
+            # return fsspec.filesystem(
+            #    "s3",
+            #    **ioqdc.request_s3fs_config()
+            # )
+            # return (
+            #    fsspec.filesystem(
+            #        "s3",
+            #        key="a046308c078b0134c9e261aa91f63ab2",
+            #        secret="d5b32f241ad8ee8d0a3173cd51b4f36d6869f168b21acef75f244a81dc10e1fb",
+            #        endpoint_url=self.endpoint_url,
+            #    )
+            #    if os.environ.get("OPENQDC_DOWNLOAD_API") is None
+            #    else
+            #    fsspec.filesystem("https")
+            # )
         else:
             return self.local_endpoint
 

diff --git a/openqdc/utils/io.py b/openqdc/utils/io.py
@@ -3,6 +3,7 @@
 import json
 import os
 import pickle as pkl
+from os.path import join as p_join
 from typing import Dict, List, Optional
 
 import fsspec
@@ -23,14 +24,17 @@
     "~/.cache/openqdc" if "OPENQDC_CACHE_DIR" not in os.environ else os.path.normpath(os.environ["OPENQDC_CACHE_DIR"])
 )
 
-_OPENQDC_DOWNLOAD_API = {"s3": "/openqdc/v1", "gs": "https://storage.googleapis.com/qmdata-public/openqdc"}
+_OPENQDC_DOWNLOAD_API = {
+    "s3": "https://storage.openqdc.org/v1",
+    "gs": "https://storage.googleapis.com/qmdata-public/openqdc",
+}
 
 
 def set_cache_dir(d):
     r"""
     Optionally set the _OPENQDC_CACHE_DIR directory.
 
-    Args:
+    Args:ßß
         d (str): path to a local folder.
     """
     if d is None:
@@ -56,7 +60,7 @@ def get_remote_cache(write_access=False) -> str:
     Returns the entry point based on the write access.
     """
     if write_access:
-        remote_cache = "/openqdc/v1"  # "gs://qmdata-public/openqdc"
+        remote_cache = "openqdc/v1"  # "gs://qmdata-public/openqdc"
         # remote_cache = "gs://qmdata-public/openqdc"
     else:
         remote_cache = _OPENQDC_DOWNLOAD_API.get(os.environ.get("OPENQDC_DOWNLOAD_API", "s3"))
@@ -89,6 +93,15 @@ def pull_locally(local_path, overwrite=False):
     return local_path
 
 
+def request_s3fs_config():
+    import httpx
+
+    response = httpx.get(p_join(_OPENQDC_DOWNLOAD_API["s3"], "config.json"))
+    response.raise_for_status()
+    config = response.json()
+    return config
+
+
 def copy_exists(local_path):
     remote_path = local_path.replace(get_local_cache(), get_remote_cache())
     return os.path.exists(local_path) or API.exists(remote_path)