From 0019cdc8186e5bbf493025112017a2e0002a0e60 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Mon, 4 Nov 2024 15:34:27 -0500 Subject: [PATCH] Add registration support to python remote API --- Cargo.lock | 305 +++++++++++++++++++++++++++++++++++++++-- Cargo.toml | 1 + rerun_py/Cargo.toml | 4 + rerun_py/src/remote.rs | 123 ++++++++++++++++- 4 files changed, 419 insertions(+), 14 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 80faaaa751f6..0ab08efe6afc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1257,6 +1257,7 @@ dependencies = [ "android-tzdata", "iana-time-zone", "num-traits", + "serde", "windows-targets 0.52.6", ] @@ -1867,6 +1868,12 @@ dependencies = [ "rerun", ] +[[package]] +name = "doc-comment" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fea41bba32d969b513997752735605054bc0dfa92b4c56bf1189f2e174be7a10" + [[package]] name = "document-features" version = "0.2.10" @@ -2910,6 +2917,12 @@ dependencies = [ "unicode-segmentation", ] +[[package]] +name = "heck" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" + [[package]] name = "heck" version = "0.5.0" @@ -3042,7 +3055,7 @@ dependencies = [ "http 0.2.12", "http-cache", "http-cache-semantics", - "reqwest", + "reqwest 0.11.27", "reqwest-middleware", "serde", "task-local-extensions", @@ -3145,7 +3158,25 @@ dependencies = [ "hyper 0.14.31", "rustls 0.21.12", "tokio", - "tokio-rustls", + "tokio-rustls 0.24.1", +] + +[[package]] +name = "hyper-rustls" +version = "0.27.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08afdbb5c31130e3034af566421053ab03787c640246a446327f550d11bcb333" +dependencies = [ + "futures-util", + "http 1.1.0", + "hyper 1.5.0", + "hyper-util", + "rustls 0.23.16", + "rustls-native-certs", + "rustls-pki-types", + "tokio", + "tokio-rustls 0.26.0", + "tower-service", ] [[package]] @@ -3723,6 +3754,16 @@ dependencies = [ "rawpointer", ] +[[package]] +name = "md-5" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf" +dependencies = [ + "cfg-if", + "digest", +] + [[package]] name = "memchr" version = "2.7.4" @@ -3923,7 +3964,7 @@ dependencies = [ "hexf-parse", "indexmap 2.6.0", "log", - "rustc-hash", + "rustc-hash 1.1.0", "spirv", "termcolor", "thiserror", @@ -4203,7 +4244,7 @@ dependencies = [ "num-integer", "num-traits", "pyo3", - "rustc-hash", + "rustc-hash 1.1.0", ] [[package]] @@ -4427,6 +4468,36 @@ dependencies = [ "memchr", ] +[[package]] +name = "object_store" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6da452820c715ce78221e8202ccc599b4a52f3e1eb3eedb487b680c81a8e3f3" +dependencies = [ + "async-trait", + "base64 0.22.1", + "bytes", + "chrono", + "futures", + "humantime", + "hyper 1.5.0", + "itertools 0.13.0", + "md-5", + "parking_lot", + "percent-encoding", + "quick-xml 0.36.2", + "rand", + "reqwest 0.12.9", + "ring", + "serde", + "serde_json", + "snafu", + "tokio", + "tracing", + "url", + "walkdir", +] + [[package]] name = "objectron" version = "0.20.0-alpha.1+dev" @@ -4453,6 +4524,12 @@ version = "11.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b410bbe7e14ab526a0e86877eb47c6996a2bd7746f027ba551028c925390e4e9" +[[package]] +name = "openssl-probe" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff011a302c396a5197692431fc1948019154afc178baf7d8e37367442a4601cf" + [[package]] name = "option-ext" version = "0.2.0" @@ -5024,6 +5101,56 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f7649a7b4df05aed9ea7ec6f628c67c9953a43869b8bc50929569b2999d443fe" dependencies = [ "memchr", + "serde", +] + +[[package]] +name = "quinn" +version = "0.11.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8c7c5fdde3cdae7203427dc4f0a68fe0ed09833edc525a03456b153b79828684" +dependencies = [ + "bytes", + "pin-project-lite", + "quinn-proto", + "quinn-udp", + "rustc-hash 2.0.0", + "rustls 0.23.16", + "socket2", + "thiserror", + "tokio", + "tracing", +] + +[[package]] +name = "quinn-proto" +version = "0.11.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fadfaed2cd7f389d0161bb73eeb07b7b78f8691047a6f3e73caaeae55310a4a6" +dependencies = [ + "bytes", + "rand", + "ring", + "rustc-hash 2.0.0", + "rustls 0.23.16", + "slab", + "thiserror", + "tinyvec", + "tracing", +] + +[[package]] +name = "quinn-udp" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e346e016eacfff12233c243718197ca12f148c84e1e84268a896699b41c71780" +dependencies = [ + "cfg_aliases 0.2.1", + "libc", + "once_cell", + "socket2", + "tracing", + "windows-sys 0.59.0", ] [[package]] @@ -6649,7 +6776,7 @@ dependencies = [ "http 0.2.12", "http-body 0.4.6", "hyper 0.14.31", - "hyper-rustls", + "hyper-rustls 0.24.2", "ipnet", "js-sys", "log", @@ -6659,14 +6786,14 @@ dependencies = [ "percent-encoding", "pin-project-lite", "rustls 0.21.12", - "rustls-pemfile", + "rustls-pemfile 1.0.4", "serde", "serde_json", "serde_urlencoded", "sync_wrapper 0.1.2", "system-configuration", "tokio", - "tokio-rustls", + "tokio-rustls 0.24.1", "tower-service", "url", "wasm-bindgen", @@ -6676,6 +6803,51 @@ dependencies = [ "winreg", ] +[[package]] +name = "reqwest" +version = "0.12.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a77c62af46e79de0a562e1a9849205ffcb7fc1238876e9bd743357570e04046f" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-core", + "futures-util", + "h2 0.4.6", + "http 1.1.0", + "http-body 1.0.1", + "http-body-util", + "hyper 1.5.0", + "hyper-rustls 0.27.3", + "hyper-util", + "ipnet", + "js-sys", + "log", + "mime", + "once_cell", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls 0.23.16", + "rustls-native-certs", + "rustls-pemfile 2.2.0", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper 1.0.1", + "tokio", + "tokio-rustls 0.26.0", + "tokio-util", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams", + "web-sys", + "windows-registry", +] + [[package]] name = "reqwest-middleware" version = "0.2.5" @@ -6685,7 +6857,7 @@ dependencies = [ "anyhow", "async-trait", "http 0.2.12", - "reqwest", + "reqwest 0.11.27", "serde", "task-local-extensions", "thiserror", @@ -6780,6 +6952,7 @@ dependencies = [ "itertools 0.13.0", "mimalloc", "numpy", + "object_store", "once_cell", "parking_lot", "pyo3", @@ -6801,6 +6974,7 @@ dependencies = [ "re_ws_comms", "tokio", "tonic", + "url", "uuid", ] @@ -7137,6 +7311,12 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" +[[package]] +name = "rustc-hash" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "583034fd73374156e66797ed8e5b0d5690409c9226b22d87cb7f19821c05d152" + [[package]] name = "rustc_version" version = "0.4.1" @@ -7209,6 +7389,19 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-native-certs" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fcaf18a4f2be7326cd874a5fa579fae794320a0f388d365dca7e480e55f83f8a" +dependencies = [ + "openssl-probe", + "rustls-pemfile 2.2.0", + "rustls-pki-types", + "schannel", + "security-framework", +] + [[package]] name = "rustls-pemfile" version = "1.0.4" @@ -7218,6 +7411,15 @@ dependencies = [ "base64 0.21.7", ] +[[package]] +name = "rustls-pemfile" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dce314e5fee3f39953d46bb63bb8a46d40c2f8fb7cc5a3b6cab2bde9721d6e50" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "rustls-pki-types" version = "1.10.0" @@ -7266,6 +7468,15 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "schannel" +version = "0.1.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01227be5826fa0690321a2ba6c5cd57a19cf3f6a09e76973b58e61de6ab9d1c1" +dependencies = [ + "windows-sys 0.59.0", +] + [[package]] name = "scoped-tls" version = "1.0.1" @@ -7288,6 +7499,29 @@ dependencies = [ "untrusted", ] +[[package]] +name = "security-framework" +version = "2.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "897b2245f0b511c87893af39b033e5ca9cce68824c4d7e7630b5a1d339658d02" +dependencies = [ + "bitflags 2.6.0", + "core-foundation 0.9.4", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea4a292869320c0272d7bc55a5a6aafaff59b4f63404a003887b679a2e05b4b6" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "semver" version = "1.0.23" @@ -7599,6 +7833,28 @@ dependencies = [ "serde", ] +[[package]] +name = "snafu" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4de37ad025c587a29e8f3f5605c00f70b98715ef90b9061a815b9e59e9042d6" +dependencies = [ + "doc-comment", + "snafu-derive", +] + +[[package]] +name = "snafu-derive" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "990079665f075b699031e9c08fd3ab99be5029b96f3b78dc0709e8f77e4efebf" +dependencies = [ + "heck 0.4.1", + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "snippets" version = "0.20.0-alpha.1+dev" @@ -7772,6 +8028,9 @@ name = "sync_wrapper" version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7065abeca94b6a8a577f9bd45aa0867a2238b74e8eb67cf10d492bc39351394" +dependencies = [ + "futures-core", +] [[package]] name = "sysinfo" @@ -8123,6 +8382,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-rustls" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c7bc40d0e5a97695bb96e27995cd3a08538541b0a846f65bba7a359f36700d4" +dependencies = [ + "rustls 0.23.16", + "rustls-pki-types", + "tokio", +] + [[package]] name = "tokio-stream" version = "0.1.16" @@ -8378,7 +8648,7 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "deb68604048ff8fa93347f02441e4487594adc20bb8a084f9e564d2b827a0a9f" dependencies = [ - "rustc-hash", + "rustc-hash 1.1.0", ] [[package]] @@ -8602,7 +8872,7 @@ dependencies = [ "http-cache-reqwest", "image", "log", - "reqwest", + "reqwest 0.11.27", "reqwest-middleware", "thiserror", "tokio", @@ -9046,7 +9316,7 @@ dependencies = [ "parking_lot", "profiling", "raw-window-handle", - "rustc-hash", + "rustc-hash 1.1.0", "smallvec", "thiserror", "wgpu-hal", @@ -9086,7 +9356,7 @@ dependencies = [ "profiling", "raw-window-handle", "renderdoc-sys", - "rustc-hash", + "rustc-hash 1.1.0", "smallvec", "thiserror", "wasm-bindgen", @@ -9207,6 +9477,17 @@ dependencies = [ "syn 2.0.85", ] +[[package]] +name = "windows-registry" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e400001bb720a623c1c69032f8e3e4cf09984deec740f007dd2b03ec864804b0" +dependencies = [ + "windows-result", + "windows-strings", + "windows-targets 0.52.6", +] + [[package]] name = "windows-result" version = "0.2.0" diff --git a/Cargo.toml b/Cargo.toml index a34ff2217acd..48b826a22b25 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -216,6 +216,7 @@ notify = { version = "6.1.1", features = ["macos_kqueue"] } num-derive = "0.4" num-traits = "0.2" numpy = "0.22" +object_store = { version = "0.10.2" } once_cell = "1.17" # No lazy_static - use `std::sync::OnceLock` or `once_cell` instead ordered-float = "4.3.0" parking_lot = "0.12" diff --git a/rerun_py/Cargo.toml b/rerun_py/Cargo.toml index 02c54653f1e9..f5f1dd1d1798 100644 --- a/rerun_py/Cargo.toml +++ b/rerun_py/Cargo.toml @@ -33,10 +33,12 @@ extension-module = ["pyo3/extension-module"] nasm = ["re_video/nasm"] remote = [ + "dep:object_store", "dep:re_remote_store_types", "dep:re_ws_comms", "dep:tokio", "dep:tonic", + "dep:url", ] ## Support serving a web viewer over HTTP with `serve()`. @@ -79,6 +81,7 @@ rand = { workspace = true, features = ["std", "std_rng"] } uuid.workspace = true # Deps for remote feature +object_store = { workspace = true, optional = true, features = ["aws"] } re_remote_store_types = { workspace = true, optional = true } tokio = { workspace = true, optional = true } # Not used yet, but we will need it when we start streaming data @@ -86,6 +89,7 @@ tokio = { workspace = true, optional = true } tonic = { workspace = true, default-features = false, features = [ "transport", ], optional = true } +url = { workspace = true, optional = true } [build-dependencies] diff --git a/rerun_py/src/remote.rs b/rerun_py/src/remote.rs index 0e8ae7655632..7ce74085f446 100644 --- a/rerun_py/src/remote.rs +++ b/rerun_py/src/remote.rs @@ -1,7 +1,12 @@ #![allow(unsafe_op_in_unsafe_fn)] +use arrow::{array::ArrayData, pyarrow::PyArrowType}; // False positive due to #[pyfunction] macro -use pyo3::{exceptions::PyRuntimeError, prelude::*, Bound, PyResult}; -use re_remote_store_types::v0::{storage_node_client::StorageNodeClient, ListRecordingsRequest}; +use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyDict, Bound, PyResult}; +use re_chunk::TransportChunk; +use re_remote_store_types::v0::{ + storage_node_client::StorageNodeClient, EncoderVersion, ListRecordingsRequest, + RecordingMetadata, RecordingType, RegisterRecordingRequest, +}; /// Register the `rerun.remote` module. pub(crate) fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { @@ -65,6 +70,120 @@ impl PyConnection { .collect()) }) } + + /// Register a recording along with some metadata + #[pyo3(signature = ( + storage_url, + metadata = None + ))] + fn register( + &mut self, + storage_url: &str, + metadata: Option<&Bound<'_, PyDict>>, + ) -> PyResult { + self.runtime.block_on(async { + let storage_url = url::Url::parse(storage_url) + .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; + + let _obj = object_store::ObjectStoreScheme::parse(&storage_url) + .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; + + let metadata = metadata + .map(|metadata| { + let (schema, data): ( + Vec, + Vec>, + ) = metadata + .iter() + .map(|(key, value)| { + let key = key.to_string(); + let value = value.extract::()?; + let value_array = value.to_arrow2()?; + let field = arrow2::datatypes::Field::new( + key, + value_array.data_type().clone(), + true, + ); + Ok((field, value_array)) + }) + .collect::>>()? + .into_iter() + .unzip(); + + let schema = arrow2::datatypes::Schema::from(schema); + + let data = arrow2::chunk::Chunk::new(data); + + let metadata_tc = TransportChunk { + schema: schema.clone(), + data, + }; + + RecordingMetadata::try_from(EncoderVersion::V0, &metadata_tc) + .map_err(|err| PyRuntimeError::new_err(err.to_string())) + }) + .transpose()?; + + let request = RegisterRecordingRequest { + // TODO(jleibs): Description should really just be in the metadata + description: Default::default(), + url: storage_url.to_string(), + metadata, + typ: RecordingType::Rrd.into(), + }; + + let resp = self + .client + .register_recording(request) + .await + .map_err(|err| PyRuntimeError::new_err(err.to_string()))? + .into_inner(); + + let recording_id: String = resp.id.map_or("Unknown".to_owned(), |id| id.id); + + Ok(recording_id) + }) + } +} + +/// A type alias for metadata. +#[derive(FromPyObject)] +enum MetadataLike { + PyArrow(PyArrowType), + // TODO(jleibs): Support converting other primitives +} + +impl MetadataLike { + fn to_arrow2(&self) -> PyResult> { + match self { + Self::PyArrow(array) => { + let array = arrow2::array::from_data(&array.0); + if array.len() == 1 { + Ok(array) + } else { + Err(PyRuntimeError::new_err( + "Metadata must be a single array, not a list", + )) + } + } + } + } + + #[allow(dead_code)] + fn to_arrow(&self) -> PyResult> { + match self { + Self::PyArrow(array) => { + let array = arrow::array::make_array(array.0.clone()); + if array.len() == 1 { + Ok(array) + } else { + Err(PyRuntimeError::new_err( + "Metadata must be a single array, not a list", + )) + } + } + } + } } /// The info for a recording stored in the archive.