-
Notifications
You must be signed in to change notification settings - Fork 252
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
2 changed files
with
167 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,164 @@ | ||
#!/usr/bin/env python3 | ||
|
||
""" | ||
Upload an .rrd to Google Cloud. | ||
Installation | ||
------------ | ||
Requires the following packages: | ||
pip install google-cloud-storage | ||
Before running, you have to authenticate via the Google Cloud CLI: | ||
- Install it (https://cloud.google.com/storage/docs/gsutil_install) | ||
- Set up credentials (https://cloud.google.com/storage/docs/gsutil_install#authenticate) | ||
If you get this error: | ||
File "…/site-packages/cryptography/hazmat/primitives/asymmetric/utils.py", line 6, in <module> | ||
from cryptography.hazmat.bindings._rust import asn1 | ||
pyo3_runtime.PanicException: Python API call failed | ||
Then run `python3 -m pip install cryptography==38.0.4` | ||
(https://levelup.gitconnected.com/fix-attributeerror-module-lib-has-no-attribute-openssl-521a35d83769) | ||
Usage | ||
----- | ||
Use the script: | ||
python3 scripts/upload_rrd.py --help | ||
or the just command: | ||
pixi run upload-image --help | ||
""" | ||
|
||
from __future__ import annotations | ||
|
||
import argparse | ||
import hashlib | ||
import logging | ||
import re | ||
import sys | ||
from io import BytesIO | ||
from pathlib import Path | ||
|
||
import requests | ||
import tqdm | ||
from google.cloud import storage | ||
|
||
|
||
class Uploader: | ||
def __init__(self): | ||
gcs = storage.Client("rerun-open") | ||
self.bucket = gcs.bucket("rerun-rrd") | ||
|
||
def upload_data( | ||
self, data: bytes, gcs_path: str, content_type: str | None = None, content_encoding: str | None = None | ||
) -> None: | ||
""" | ||
Low-level upload of data. | ||
Parameters | ||
---------- | ||
data: | ||
The data to upload. | ||
gcs_path: | ||
The path of the object. | ||
content_type: | ||
The content type of the object. | ||
content_encoding: | ||
The content encoding of the object. | ||
""" | ||
|
||
logging.info(f"Uploading {gcs_path} (size: {len(data)}, type: {content_type}, encoding: {content_encoding})") | ||
destination = self.bucket.blob(gcs_path) | ||
destination.content_type = content_type | ||
destination.content_encoding = content_encoding | ||
|
||
if destination.exists(): | ||
logging.warn(f"blob {gcs_path} already exists in GCS, skipping upload") | ||
return | ||
|
||
stream = BytesIO(data) | ||
destination.upload_from_file(stream) | ||
|
||
|
||
def data_hash(data: bytes) -> str: | ||
"""Compute a sha1 hash digest of some data.""" | ||
return hashlib.sha1(data).hexdigest() | ||
|
||
|
||
def download_file(url: str, path: Path) -> None: | ||
path.parent.mkdir(parents=True, exist_ok=True) | ||
logging.info("Downloading %s to %s", url, path) | ||
response = requests.get(url, stream=True) | ||
with tqdm.tqdm.wrapattr( | ||
open(path, "wb"), | ||
"write", | ||
miniters=1, | ||
total=int(response.headers.get("content-length", 0)), | ||
desc=f"Downloading {path.name}", | ||
) as f: | ||
for chunk in response.iter_content(chunk_size=4096): | ||
f.write(chunk) | ||
|
||
|
||
DESCRIPTION = """Upload an .rrd to static.rerun.io. | ||
pixi run upload-rrd --version 0.15.0 path/to/recording.rrd | ||
The version is used for two things: | ||
A) used as a folder name in the GCS bucket. | ||
B) used to generate a link to the correct version of the Rerun web viewer. | ||
""" | ||
|
||
|
||
def main() -> None: | ||
parser = argparse.ArgumentParser(description=DESCRIPTION, formatter_class=argparse.RawTextHelpFormatter) | ||
parser.add_argument("path", type=str, help="Recording .rrd to upload") | ||
parser.add_argument( | ||
"--name", type=str, required=False, help="Name of the recording. If not supplied, the file name is used." | ||
) | ||
parser.add_argument("--version", type=str, required=True, help="The Rerun version, e.g. '0.15.0'.") | ||
parser.add_argument("--debug", action="store_true", help="Enable debug logging.") | ||
args = parser.parse_args() | ||
|
||
if args.debug: | ||
logging.basicConfig(level=logging.DEBUG) | ||
else: | ||
logging.basicConfig(level=logging.INFO) | ||
|
||
try: | ||
if not args.path.endswith(".rrd"): | ||
raise RuntimeError("File path expected to have .rrd extension") | ||
|
||
file_path = Path(args.path) | ||
name = args.name or file_path.stem | ||
version = args.version | ||
|
||
# Check if user put `v0.15.0` instead of `0.15.0`: | ||
if m := re.match(r"v(\d+\.v\d+\..*)", version): | ||
version = m.group(1) | ||
raise RuntimeError("Version should be in the format '{version}', without a leading 'v'") | ||
|
||
file_data = file_path.read_bytes() | ||
digest = data_hash(file_data) | ||
gcp_path = f"{version}/{name}_{digest}.rrd" | ||
|
||
uploader = Uploader() | ||
uploader.upload_data(file_data, gcp_path, content_type="application/octet-stream") | ||
|
||
recording_url = f"https://static.rerun.io/{gcp_path}" | ||
print(f"Recording at: {recording_url}") | ||
print(f"View it at: https://rerun.io/viewer/version/{version}/?recording={recording_url}") | ||
|
||
except RuntimeError as e: | ||
print(f"Error: {e.args[0]}", file=sys.stderr) | ||
return | ||
|
||
|
||
if __name__ == "__main__": | ||
main() |