diff --git a/README.md b/README.md index 633a530..e339782 100644 --- a/README.md +++ b/README.md @@ -85,6 +85,8 @@ Once you've installed `airoa-lineage` in your project, you can start tracking da ### Basic Usage ```python +from openlineage.client.run import Dataset + from airoa_lineage.facets import CommonRunFacet from airoa_lineage.teleop import TeleopSession @@ -99,6 +101,9 @@ common_facet = CommonRunFacet( ) session = TeleopSession(namespace="production", common_facet=common_facet) +# Define output dataset +output_ds = Dataset(namespace="production", name="teleop_rosbag") + # Track teleoperation session with nominal time run_id = session.start( nominal_start_time="2025-10-22T00:00:00+00:00", @@ -108,7 +113,7 @@ print(f"Session started: {run_id}") # Collect robot data (sensor readings, camera feeds, motor commands)... -session.complete() +session.complete(output_datasets=[output_ds]) # Or if interrupted: session.cancel() # See examples/teleop_session.py for complete example with device facets and error handling diff --git a/docs/architecture.md b/docs/architecture.md index 82a4415..10022fc 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -57,11 +57,11 @@ airoa-lineage is built on the following architectural principles: **Lifecycle**: ``` -TeleopSession.start() # Emits START event +TeleopSession.start() # Emits START event └─> Data collection... -TeleopSession.running() # Optional: Emits RUNNING events for progress updates +TeleopSession.running() # Optional: Emits RUNNING events for progress updates └─> More collection... -TeleopSession.complete() # Emits COMPLETE event +TeleopSession.complete(output_datasets) # Emits COMPLETE event ``` **API Documentation**: See docstrings in [session.py](../src/airoa_lineage/teleop/session.py) diff --git a/examples/data_conversion.py b/examples/data_conversion.py index de511f0..0b1f215 100644 --- a/examples/data_conversion.py +++ b/examples/data_conversion.py @@ -135,17 +135,17 @@ def main(): print("[2/4] Simulating data conversion...") # Phase 1: Reading rosbag files - session.running(message="Phase 1: Reading rosbag files") + session.running() time.sleep(2) print(" ✓ Phase 1 completed: Rosbag files read") # Phase 2: Extracting topics - session.running(message="Phase 2: Extracting topics") + session.running() time.sleep(2) print(" ✓ Phase 2 completed: Topics extracted") # Phase 3: Converting to LeRobot format - session.running(message="Phase 3: Converting to LeRobot format") + session.running() time.sleep(1) print(" ✓ Phase 3 completed: Conversion done") diff --git a/examples/teleop_session.py b/examples/teleop_session.py index 3f0053d..78b0210 100644 --- a/examples/teleop_session.py +++ b/examples/teleop_session.py @@ -21,6 +21,8 @@ import time from datetime import datetime, timedelta, timezone +from openlineage.client.run import Dataset + from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet from airoa_lineage.marquez_client import MarquezClient from airoa_lineage.teleop import TeleopSession @@ -52,6 +54,10 @@ def main(): nominal_start = now - timedelta(days=1) nominal_end = nominal_start + timedelta(hours=5) + # Output dataset definition + output_dataset_name = "teleop_rosbag_2025_12_18" + output_ds = Dataset(namespace=namespace, name=output_dataset_name) + # Create common facet with robot and repository information common_facet = CommonRunFacet( robotId=robot_id, @@ -101,6 +107,9 @@ def main(): print(f" Start: {nominal_start.isoformat()}") print(f" End: {nominal_end.isoformat()}") print() + print("Datasets:") + print(f" Output: {namespace}/{output_dataset_name}") + print() # ========== START Event ========== print("[1/4] Sending START event...") @@ -134,11 +143,11 @@ def main(): print() # ========== COMPLETE Event ========== - print("[3/4] Sending COMPLETE event...") + print("[3/4] Sending COMPLETE event with output datasets...") - # Send COMPLETE event using the session - # This sends an OpenLineage COMPLETE event to Marquez - session.complete() + # Send COMPLETE event with output datasets + # This sends an OpenLineage COMPLETE event to Marquez with output datasets + session.complete(output_datasets=[output_ds]) print("✓ COMPLETE event sent successfully") print() @@ -173,6 +182,7 @@ def main(): print(f"✓ Repository: {repository_uri}") print(f"✓ Commit: {repository_hash} ({repository_branch})") print(f"✓ Tag: {repository_tag}") + print(f"✓ Output: {output_dataset_name}") print(f"✓ Run ID: {session.run_id}") print() print("Next steps:") @@ -189,6 +199,11 @@ def main(): f" curl {marquez_url}/api/v1/namespaces/{namespace}/jobs/{job_name}/runs/{session.run_id}" ) print() + print("4. View output dataset lineage:") + print( + f" {marquez_url.replace(':9000', ':3000')}/lineage/dataset/{namespace}/{output_dataset_name}" + ) + print() print("=" * 60) diff --git a/src/airoa_lineage/teleop/session.py b/src/airoa_lineage/teleop/session.py index b185988..b8469c1 100644 --- a/src/airoa_lineage/teleop/session.py +++ b/src/airoa_lineage/teleop/session.py @@ -1,8 +1,9 @@ """Teleoperation session management with OpenLineage tracking.""" -from typing import Dict, Optional +from typing import Dict, List, Optional from openlineage.client.facet import BaseFacet +from openlineage.client.run import Dataset from airoa_lineage.core import BaseSession from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet @@ -110,6 +111,7 @@ def __init__( super().__init__(namespace, job_name, marquez_url, run_id, facet_prefix) self.common_facet = common_facet self.device_facet = device_facet + self._output_datasets: List[Dataset] = [] def _get_session_facets(self) -> Dict[str, BaseFacet]: """Return session-specific facets.""" @@ -121,3 +123,49 @@ def _get_session_facets(self) -> Dict[str, BaseFacet]: def _get_producer(self) -> str: """Return producer identifier.""" return "airoa-teleop-system" + + def _get_outputs(self) -> List[Dataset]: + """Return output datasets.""" + return self._output_datasets + + def complete( # type: ignore[override] + self, + output_datasets: List[Dataset], + ) -> None: + """ + Send COMPLETE event to Marquez. + + This method sends an OpenLineage COMPLETE event to track the successful + completion of the teleoperation data collection session. + + Args: + output_datasets: List of output datasets (required) + + Raises: + RuntimeError: If session was not started or already completed + + Examples: + >>> from airoa_lineage.teleop import TeleopSession + >>> from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet + >>> from openlineage.client.run import Dataset + >>> common_facet = CommonRunFacet( + ... robotId="hsr001", + ... location="weblab", + ... repositoryHash="df110d5", + ... repositoryUri="https://github.com/user/repo.git", + ... repositoryTag="v1.0.0", + ... repositoryBranch="main" + ... ) + >>> device_facet = DeviceRunFacet(hostname="operator-pc-001") + >>> output_ds = Dataset(namespace="airoa_teleop", name="teleop_rosbag") + >>> session = TeleopSession( + ... namespace="airoa_teleop", + ... common_facet=common_facet, + ... device_facet=device_facet + ... ) + >>> session.start() + >>> # Collect teleoperation data... + >>> session.complete(output_datasets=[output_ds]) + """ + self._output_datasets = output_datasets + return super().complete() diff --git a/tests/unit/teleop/test_session.py b/tests/unit/teleop/test_session.py index db25c7c..6cebc67 100644 --- a/tests/unit/teleop/test_session.py +++ b/tests/unit/teleop/test_session.py @@ -4,6 +4,7 @@ import uuid import pytest +from openlineage.client.run import Dataset from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet from airoa_lineage.teleop import TeleopSession @@ -40,6 +41,8 @@ def test_initialization_with_defaults(self, common_facet, device_facet): # Session should not be started or completed assert session._started is False assert session._completed is False + # Output datasets should be empty list + assert session._output_datasets == [] def test_initialization_with_custom_values(self): """Test initialization with custom values.""" @@ -121,6 +124,7 @@ class TestTeleopSessionComplete: def test_complete_without_start_raises_error(self, common_facet, device_facet): """Test that calling complete() without start() raises RuntimeError.""" with BaseSessionTestHelper(TeleopSession) as helper: + output_ds = Dataset(namespace="test", name="output_dataset") session, mock_client = helper.create_session( namespace="test_namespace", common_facet=common_facet, @@ -129,7 +133,7 @@ def test_complete_without_start_raises_error(self, common_facet, device_facet): # complete() without start() should raise error with pytest.raises(RuntimeError, match="not started"): - session.complete() + session.complete(output_datasets=[output_ds]) # Verify emit was never called mock_client.emit.assert_not_called() @@ -137,17 +141,18 @@ def test_complete_without_start_raises_error(self, common_facet, device_facet): def test_complete_twice_raises_error(self, common_facet, device_facet): """Test that calling complete() twice raises RuntimeError.""" with BaseSessionTestHelper(TeleopSession) as helper: + output_ds = Dataset(namespace="test", name="output_dataset") session, mock_client = helper.create_session( namespace="test_namespace", common_facet=common_facet, device_facet=device_facet, ) session.start() - session.complete() + session.complete(output_datasets=[output_ds]) # Second complete() should raise error with pytest.raises(RuntimeError, match="already completed"): - session.complete() + session.complete(output_datasets=[output_ds]) # Verify emit was called twice (start + complete, not 3 times) assert mock_client.emit.call_count == 2 @@ -155,17 +160,43 @@ def test_complete_twice_raises_error(self, common_facet, device_facet): def test_complete_success(self, common_facet, device_facet): """Test successful complete() after start().""" with BaseSessionTestHelper(TeleopSession) as helper: + output_ds = Dataset(namespace="test", name="output_dataset") session, mock_client = helper.create_session( namespace="test_namespace", common_facet=common_facet, device_facet=device_facet, ) session.start() - session.complete() + session.complete(output_datasets=[output_ds]) # Verify session state helper.assert_complete_success(session, mock_client) + def test_complete_with_output_datasets(self, common_facet, device_facet): + """Test that output datasets are included in COMPLETE event.""" + with BaseSessionTestHelper(TeleopSession) as helper: + output_ds = Dataset(namespace="test", name="teleop_rosbag") + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) + session.start() + session.complete(output_datasets=[output_ds]) + + # Verify emit was called twice + assert mock_client.emit.call_count == 2 + + # Check COMPLETE event (second call) + complete_event = helper.get_emitted_event(mock_client, call_index=1) + + # Verify output datasets are in the event + assert len(complete_event.outputs) == 1 + assert complete_event.outputs[0].namespace == "test" + assert complete_event.outputs[0].name == "teleop_rosbag" + # Verify inputs are empty + assert len(complete_event.inputs) == 0 + class TestTeleopSessionNominalTime: """Test TeleopSession with nominal time facets.""" @@ -289,13 +320,14 @@ def test_start_includes_common_facet(self, common_facet, device_facet): def test_complete_includes_common_facet(self, common_facet, device_facet): """Test that COMPLETE event includes common facet.""" with BaseSessionTestHelper(TeleopSession) as helper: + output_ds = Dataset(namespace="test", name="output_dataset") session, mock_client = helper.create_session( namespace="test_namespace", common_facet=common_facet, device_facet=device_facet, ) session.start() - session.complete() + session.complete(output_datasets=[output_ds]) # Verify emit was called twice assert mock_client.emit.call_count == 2 @@ -398,13 +430,14 @@ def test_start_includes_device_facet(self, common_facet, device_facet): def test_complete_includes_device_facet(self, common_facet, device_facet): """Test that COMPLETE event includes device facet.""" with BaseSessionTestHelper(TeleopSession) as helper: + output_ds = Dataset(namespace="test", name="output_dataset") session, mock_client = helper.create_session( namespace="test_namespace", common_facet=common_facet, device_facet=device_facet, ) session.start() - session.complete() + session.complete(output_datasets=[output_ds]) # Verify emit was called twice assert mock_client.emit.call_count == 2 @@ -504,13 +537,14 @@ def test_complete_without_prefix_uses_default_keys( ): """Test that COMPLETE event uses 'common' and 'device' keys without prefix.""" with BaseSessionTestHelper(TeleopSession) as helper: + output_ds = Dataset(namespace="test", name="output_dataset") session, mock_client = helper.create_session( namespace="test_namespace", common_facet=common_facet, device_facet=device_facet, ) session.start() - session.complete() + session.complete(output_datasets=[output_ds]) # Verify emit was called twice assert mock_client.emit.call_count == 2 @@ -527,6 +561,7 @@ def test_complete_without_prefix_uses_default_keys( def test_complete_with_prefix_uses_prefixed_keys(self, common_facet, device_facet): """Test that COMPLETE event uses prefixed keys when facet_prefix is set.""" with BaseSessionTestHelper(TeleopSession) as helper: + output_ds = Dataset(namespace="test", name="output_dataset") session, mock_client = helper.create_session( namespace="test_namespace", common_facet=common_facet, @@ -534,7 +569,7 @@ def test_complete_with_prefix_uses_prefixed_keys(self, common_facet, device_face facet_prefix="airoa", ) session.start() - session.complete() + session.complete(output_datasets=[output_ds]) # Verify emit was called twice assert mock_client.emit.call_count == 2