Skip to content

Commit

Permalink
SDK DataLoaders 2: barebones Python support (#5328)
Browse files Browse the repository at this point in the history
Exposes raw `DataLoader`s to the Python SDK through 2 new methods:
`RecordingStream.log_file_from_path` &
`RecordingStream.log_file_from_contents`.

Everything streams asynchronously, as expected.

There's no way to configure the behavior of the `DataLoader` at all,
yet. That comes next.

```python
rr.log_file_from_path(filepath)
```

![image](https://github.com/rerun-io/rerun/assets/2910679/0fe2d39c-f069-44a6-b836-e31001b3adaa)


---

Part of series of PR to expose configurable `DataLoader`s to our SDKs:
- #5327 
- #5328 
- #5330
- #5337
- #5351
- #5355
  • Loading branch information
teh-cmc committed Feb 29, 2024
1 parent 588eb3c commit 0f90a9c
Show file tree
Hide file tree
Showing 8 changed files with 189 additions and 2 deletions.
11 changes: 11 additions & 0 deletions examples/python/log_file/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
<!--[metadata]
title = "Log file example"
-->

Demonstrates how to log any file from the SDK using the [`DataLoader`](https://www.rerun.io/docs/howto/open-any-file) machinery.

Usage:
```bash
python examples/python/log_file/main.py examples/assets
```

46 changes: 46 additions & 0 deletions examples/python/log_file/main.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
#!/usr/bin/env python3
"""
Demonstrates how to log any file from the SDK using the `DataLoader` machinery.
See <https://www.rerun.io/docs/howto/open-any-file> for more information.
Usage:
```
python examples/python/log_file/main.py -- examples/assets
```
"""
from __future__ import annotations

import argparse
from pathlib import Path

import rerun as rr # pip install rerun-sdk

parser = argparse.ArgumentParser(
description="Demonstrates how to log any file from the SDK using the `DataLoader` machinery."
)
rr.script_add_args(parser)
parser.add_argument(
"--from-contents",
action="store_true",
default=False,
help="Log the contents of the file directly (files only -- not supported by external loaders).",
)
parser.add_argument("filepaths", nargs="+", type=Path, help="The filepaths to be loaded and logged.")
args = parser.parse_args()

rr.script_setup(args, "rerun_example_log_file")

for filepath in args.filepaths:
if not args.from_contents:
# Either log the file using its path…
rr.log_file_from_path(filepath)
else:
# …or using its contents if you already have them loaded for some reason.
try:
with open(filepath, "rb") as file:
rr.log_file_from_contents(filepath, file.read())
except Exception:
pass

rr.script_teardown(args)
1 change: 1 addition & 0 deletions examples/python/log_file/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
rerun-sdk
8 changes: 7 additions & 1 deletion rerun_py/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,13 @@ re_error.workspace = true
re_log.workspace = true
re_log_types.workspace = true
re_memory.workspace = true
rerun = { workspace = true, features = ["analytics", "run", "server", "sdk"] }
rerun = { workspace = true, features = [
"analytics",
"data_loaders",
"run",
"server",
"sdk",
] }
re_web_viewer_server = { workspace = true, optional = true }
re_ws_comms = { workspace = true, optional = true }

Expand Down
2 changes: 2 additions & 0 deletions rerun_py/docs/gen_common_index.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,8 @@ class Section:
title="Logging functions",
func_list=[
"log",
"log_file_from_path",
"log_file_from_contents",
],
),
Section(
Expand Down
6 changes: 5 additions & 1 deletion rerun_py/rerun_sdk/rerun/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,8 +77,10 @@
"get_recording_id",
"get_thread_local_data_recording",
"is_enabled",
"log_components",
"log",
"log_components",
"log_file_from_contents",
"log_file_from_path",
"memory_recording",
"new_entity_path",
"reset_time",
Expand All @@ -105,6 +107,8 @@
escape_entity_path_part,
log,
log_components,
log_file_from_contents,
log_file_from_path,
new_entity_path,
)
from .any_value import AnyValues
Expand Down
67 changes: 67 additions & 0 deletions rerun_py/rerun_sdk/rerun/_log.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations

from pathlib import Path
from typing import Any, Iterable

import pyarrow as pa
Expand Down Expand Up @@ -282,6 +283,72 @@ def log_components(
)


@catch_and_log_exceptions()
def log_file_from_path(
file_path: str | Path,
*,
recording: RecordingStream | None = None,
) -> None:
r"""
Logs the file at the given `path` using all `DataLoader`s available.
A single `path` might be handled by more than one loader.
This method blocks until either at least one `DataLoader` starts
streaming data in or all of them fail.
See <https://www.rerun.io/docs/howto/open-any-file> for more information.
Parameters
----------
file_path:
Path to the file to be logged.
recording:
Specifies the [`rerun.RecordingStream`][] to use. If left unspecified,
defaults to the current active data recording, if there is one. See
also: [`rerun.init`][], [`rerun.set_global_data_recording`][].
"""

bindings.log_file_from_path(Path(file_path), recording=recording)


@catch_and_log_exceptions()
def log_file_from_contents(
file_path: str | Path,
file_contents: bytes,
*,
recording: RecordingStream | None = None,
) -> None:
r"""
Logs the given `file_contents` using all `DataLoader`s available.
A single `path` might be handled by more than one loader.
This method blocks until either at least one `DataLoader` starts
streaming data in or all of them fail.
See <https://www.rerun.io/docs/howto/open-any-file> for more information.
Parameters
----------
file_path:
Path to the file that the `file_contents` belong to.
file_contents:
Contents to be logged.
recording:
Specifies the [`rerun.RecordingStream`][] to use. If left unspecified,
defaults to the current active data recording, if there is one. See
also: [`rerun.init`][], [`rerun.set_global_data_recording`][].
"""

bindings.log_file_from_contents(Path(file_path), file_contents, recording=recording)


def escape_entity_path_part(part: str) -> str:
r"""
Escape an individual part of an entity path.
Expand Down
50 changes: 50 additions & 0 deletions rerun_py/src/python_bridge.rs
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,8 @@ fn rerun_bindings(_py: Python<'_>, m: &PyModule) -> PyResult<()> {

// log any
m.add_function(wrap_pyfunction!(log_arrow_msg, m)?)?;
m.add_function(wrap_pyfunction!(log_file_from_path, m)?)?;
m.add_function(wrap_pyfunction!(log_file_from_contents, m)?)?;

// misc
m.add_function(wrap_pyfunction!(version, m)?)?;
Expand Down Expand Up @@ -962,6 +964,54 @@ fn log_arrow_msg(
Ok(())
}

#[pyfunction]
#[pyo3(signature = (
file_path,
recording=None,
))]
fn log_file_from_path(
py: Python<'_>,
file_path: std::path::PathBuf,
recording: Option<&PyRecordingStream>,
) -> PyResult<()> {
let Some(recording) = get_data_recording(recording) else {
return Ok(());
};

recording
.log_file_from_path(file_path)
.map_err(|err| PyRuntimeError::new_err(err.to_string()))?;

py.allow_threads(flush_garbage_queue);

Ok(())
}

#[pyfunction]
#[pyo3(signature = (
file_path,
file_contents,
recording=None,
))]
fn log_file_from_contents(
py: Python<'_>,
file_path: std::path::PathBuf,
file_contents: &[u8],
recording: Option<&PyRecordingStream>,
) -> PyResult<()> {
let Some(recording) = get_data_recording(recording) else {
return Ok(());
};

recording
.log_file_from_contents(file_path, std::borrow::Cow::Borrowed(file_contents))
.map_err(|err| PyRuntimeError::new_err(err.to_string()))?;

py.allow_threads(flush_garbage_queue);

Ok(())
}

// --- Misc ---

/// Return a verbose version string
Expand Down

0 comments on commit 0f90a9c

Please sign in to comment.