wehub-resource-sync
105 lines
2.8 KiB
Python
105 lines
2.8 KiB
Python
"""Getting Started workflow: Catalog SDK regions (Python only).
|
|
|
|
The Log step lives in `tutorials/getting_started_log` because it uses only the
|
|
Logging SDK and is therefore available in Python, Rust, and C++.
|
|
"""
|
|
|
|
import math
|
|
import os
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import torch.multiprocessing
|
|
|
|
import rerun as rr
|
|
|
|
# Rerun's tokio runtime is not fork-safe; DataLoader workers must use `spawn`.
|
|
torch.multiprocessing.set_start_method("spawn", force=True)
|
|
|
|
# Run from a fresh temp dir so the .rrd files this snippet writes don't
|
|
# collide with other snippets executing in parallel from the same cwd.
|
|
os.chdir(tempfile.mkdtemp())
|
|
|
|
# Materialize the .rrd that the catalog regions below register against.
|
|
# Same code as the Log step's three-language snippet, repeated here so this
|
|
# file runs end-to-end.
|
|
with rr.RecordingStream(
|
|
"rerun_example_getting_started", recording_id="run-1"
|
|
) as _rec:
|
|
_rec.save("run-1.rrd")
|
|
for _t in range(10):
|
|
_rec.set_time("step", sequence=_t)
|
|
_rec.log("/arm/shoulder", rr.Scalars(math.sin(_t * 0.5)))
|
|
_rec.log("/arm/elbow", rr.Scalars(math.cos(_t * 0.5)))
|
|
|
|
# Start an in-process catalog server on a random port so this snippet runs
|
|
# end-to-end. In a real workflow you'd run `rerun server` in a separate
|
|
# terminal, which is what the docs show.
|
|
_server = rr.server.Server()
|
|
server_url = _server.url()
|
|
|
|
|
|
# region: setup
|
|
# `server_url` is the catalog URL — defaults to "rerun+http://127.0.0.1:51234"
|
|
# when running `rerun server` locally.
|
|
client = rr.catalog.CatalogClient(server_url)
|
|
# endregion: setup
|
|
|
|
|
|
# region: ingest
|
|
dataset = client.create_dataset("demo", exist_ok=True)
|
|
dataset.register([Path("run-1.rrd").absolute().as_uri()]).wait()
|
|
# endregion: ingest
|
|
|
|
|
|
# region: annotate
|
|
with rr.RecordingStream(
|
|
"rerun_example_getting_started", recording_id="run-1"
|
|
) as ann:
|
|
ann.save("run-1-properties.rrd")
|
|
ann.send_property(
|
|
"episode", rr.AnyValues(success=True, task="pick_and_place")
|
|
)
|
|
|
|
dataset.register(
|
|
[Path("run-1-properties.rrd").absolute().as_uri()], layer_name="properties"
|
|
).wait()
|
|
# endregion: annotate
|
|
|
|
|
|
# region: query
|
|
df = dataset.filter_contents(["/arm/**"]).reader(index="step")
|
|
print(
|
|
df.select(
|
|
"rerun_segment_id",
|
|
"/arm/shoulder:Scalars:scalars",
|
|
"/arm/elbow:Scalars:scalars",
|
|
)
|
|
)
|
|
# endregion: query
|
|
|
|
# region: train
|
|
from torch.utils.data import DataLoader
|
|
|
|
from rerun.experimental.dataloader import (
|
|
DataSource,
|
|
Field,
|
|
NumericDecoder,
|
|
RerunIterableDataset,
|
|
)
|
|
|
|
ds = RerunIterableDataset(
|
|
source=DataSource(dataset=dataset),
|
|
index="step",
|
|
fields={
|
|
"shoulder": Field(
|
|
"/arm/shoulder:Scalars:scalars", decode=NumericDecoder()
|
|
),
|
|
"elbow": Field("/arm/elbow:Scalars:scalars", decode=NumericDecoder()),
|
|
},
|
|
)
|
|
|
|
for batch in DataLoader(ds, batch_size=4):
|
|
print(batch)
|
|
# endregion: train
|