wehub-resource-sync
90 lines
2.4 KiB
Python
Executable File
90 lines
2.4 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
"""
|
|
Build and upload rerun_notebook wheels to GCS.
|
|
|
|
IMPORTANT: rerun_js must be built beforehand, otherwise this script will fail. Use `pixi run js-build-base`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import subprocess
|
|
import zipfile
|
|
|
|
from google.cloud.storage import Bucket
|
|
from google.cloud.storage import Client as Gcs
|
|
|
|
|
|
def run(
|
|
cmd: str,
|
|
*,
|
|
cwd: str | None = None,
|
|
env: dict[str, str] | None = None,
|
|
) -> None:
|
|
print(f"{cwd or ''}> {cmd}")
|
|
subprocess.check_output(cmd.split(), cwd=cwd, env=env)
|
|
|
|
|
|
def build_and_upload(bucket: Bucket, gcs_dir: str) -> str:
|
|
dist = "dist/"
|
|
|
|
# Build into `dist`
|
|
run(
|
|
f"hatch build -t wheel ../{dist}",
|
|
cwd="rerun_notebook",
|
|
)
|
|
|
|
pkg = os.listdir(dist)[0]
|
|
wheel = f"{dist}/{pkg}"
|
|
|
|
# Upload to GCS
|
|
print("Uploading to GCS…")
|
|
bucket.blob(f"{gcs_dir}/{pkg}").upload_from_filename(wheel)
|
|
|
|
return wheel
|
|
|
|
|
|
def publish_notebook_asset(bucket: Bucket, gcs_dir: str, wheel: str) -> None:
|
|
"""Extract widget.js and re_viewer_bg.wasm from the notebook wheel and upload to the web viewer bucket."""
|
|
|
|
with zipfile.ZipFile(wheel, "r") as archive:
|
|
archive.extract("rerun_notebook/static/widget.js", "extracted")
|
|
archive.extract("rerun_notebook/static/re_viewer_bg.wasm", "extracted")
|
|
|
|
for filename in ["widget.js", "re_viewer_bg.wasm"]:
|
|
local_path = f"extracted/rerun_notebook/static/{filename}"
|
|
blob = bucket.blob(f"{gcs_dir}/{filename}")
|
|
print(f"Uploading {local_path} to gs://{bucket.name}/{blob.name}")
|
|
blob.upload_from_filename(local_path)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Build and upload rerun_notebook wheels to GCS")
|
|
parser.add_argument("--dir", required=True, help="Upload the wheel to the given directory in GCS")
|
|
parser.add_argument(
|
|
"--notebook-dir",
|
|
required=False,
|
|
help="Upload notebook assets (widget.js, re_viewer_bg.wasm) to the given directory in the web viewer bucket",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
gcs = Gcs("rerun-open")
|
|
|
|
wheel = build_and_upload(
|
|
gcs.bucket("rerun-builds"),
|
|
args.dir,
|
|
)
|
|
|
|
if args.notebook_dir:
|
|
publish_notebook_asset(
|
|
gcs.bucket("rerun-web-viewer"),
|
|
args.notebook_dir,
|
|
wheel,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|