Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ This library is designed to be integrated into robot data collection systems, ET

- 🤖 **Teleoperation Session Management**: Track robot data collection sessions with `TeleopSession`
- 🔄 **Data Conversion Session Management**: Track data conversion pipelines with `ConversionSession`
- ☁️ **Wasabi Upload Management**: Track Wasabi cloud storage uploads with `WasabiUploadSession`
- 🔌 **USB Data Copy Management**: Track USB data copy operations with `USBCopySession`
- 📊 **OpenLineage Integration**: Full OpenLineage specification support (START/COMPLETE/RUNNING/FAIL)
- 🏷️ **Custom Run Facets**: Robot metadata (robotId, location, repository info) via `CommonRunFacet`
Expand Down Expand Up @@ -185,6 +186,14 @@ uv run python examples/teleop_session.py

Demonstrates a complete teleoperation session with nominal time tracking.

### Wasabi Upload Session

```bash
uv run python examples/wasabi_upload_session.py
```

Demonstrates a complete Wasabi cloud storage upload session with nominal time tracking.

### Batch ETL Job

```bash
Expand Down
24 changes: 24 additions & 0 deletions docs/architecture.md
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,28 @@ ConversionSession.complete(output_datasets) # Emits COMPLETE event
**API Documentation**: See docstrings in [session.py](../src/airoa_lineage/conversion/session.py)
**Usage Example**: [examples/data_conversion.py](../examples/data_conversion.py)

### WasabiUploadSession

@yuya-haruna yuya-haruna Nov 20, 2025

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wasabiという固有サービス名に限定しないほうが良いのかなとちょっと思いましたがNITSです。

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

おっしゃるとおりで、USB のところもどうしようかなと思いました。
現状では USB や Wasabi の方が直感的でわかりやすいですし、ストレージが変わるタイミングがあれば、Rename するだけでコストは掛からないので、一旦、わかりやすさをとった感じです。


**Location**: [src/airoa_lineage/wasabi_upload/session.py](../src/airoa_lineage/wasabi_upload/session.py)

**Purpose**: Manages the lifecycle of Wasabi cloud storage upload operations with OpenLineage tracking.

**Key Responsibilities**:
- Track Wasabi data upload sessions (START → COMPLETE)
- Automatic run_id generation and management
- Nominal time support for historical data processing
- State validation (prevents duplicate start/complete calls)

**Lifecycle**:
```
WasabiUploadSession.start() # Emits START event
└─> Data upload to Wasabi...
WasabiUploadSession.complete() # Emits COMPLETE event
```

**API Documentation**: See docstrings in [session.py](../src/airoa_lineage/wasabi_upload/session.py)
**Usage Example**: [examples/wasabi_upload_session.py](../examples/wasabi_upload_session.py)

### USBCopySession

**Location**: [src/airoa_lineage/usb_copy/session.py](../src/airoa_lineage/usb_copy/session.py)
Expand Down Expand Up @@ -247,6 +269,7 @@ For detailed event schemas, see the [OpenLineage specification](https://openline
- [MarquezClient](../src/airoa_lineage/marquez_client/client.py) - Event emission and REST API
- [TeleopSession](../src/airoa_lineage/teleop/session.py) - Robot data collection lifecycle
- [ConversionSession](../src/airoa_lineage/conversion/session.py) - Data conversion pipeline lifecycle
- [WasabiUploadSession](../src/airoa_lineage/wasabi_upload/session.py) - Wasabi cloud storage upload lifecycle
- [USBCopySession](../src/airoa_lineage/usb_copy/session.py) - USB data copy lifecycle
- [Timestamp Utilities](../src/airoa_lineage/utils/timestamps.py) - UTC timestamp generation
- [CommonRunFacet](../src/airoa_lineage/facets/common.py) - Robot metadata facet
Expand All @@ -257,6 +280,7 @@ For detailed event schemas, see the [OpenLineage specification](https://openline

- [examples/teleop_session.py](../examples/teleop_session.py) - Complete teleoperation session with nominal time
- [examples/data_conversion.py](../examples/data_conversion.py) - Data conversion pipeline with RUNNING events
- [examples/wasabi_upload_session.py](../examples/wasabi_upload_session.py) - Wasabi upload session with nominal time
- [examples/usb_copy_session.py](../examples/usb_copy_session.py) - USB data copy session with nominal time
- [examples/simple_batch_etl.py](../examples/simple_batch_etl.py) - Batch ETL job with input/output datasets

Expand Down
102 changes: 72 additions & 30 deletions docs/development.md
Original file line number Diff line number Diff line change
Expand Up @@ -138,15 +138,42 @@ Tests are organized in the `tests/unit/` directory:

```
tests/
├── conftest.py # Common fixtures (common_facet, device_facet, aws_job_facet)
└── unit/
├── marquez_client/
│ └── test_client.py
├── teleop/
│ └── test_session.py
├── conversion/
│ └── test_session.py
├── wasabi_upload/
│ └── test_session.py
├── usb_copy/
│ └── test_session.py
├── core/
│ └── test_base_session.py
├── helpers/ # Test utilities
│ ├── test_helper.py # BaseSessionTestHelper class
│ ├── base_session_test.py # BaseSessionTest base class
│ └── test_test_helper.py # Tests for the helper itself
└── utils/
└── test_timestamps.py
```

**Test Helpers:**

- `BaseSessionTestHelper`: Utility class for creating mocked sessions and asserting common patterns
- `create_session()`: Create session with mocked MarquezClient
- `get_emitted_event()`: Extract emitted event at specific call index
- `assert_start_success()`, `assert_complete_success()`: Verify session state
- `assert_common_facet()`, `assert_device_facet()`, `assert_aws_job_facet()`: Verify facets
- `assert_nominal_time()`, `assert_event_structure()`: Verify event structure

- `BaseSessionTest`: Base class for session tests with common test methods
- Subclasses define `SESSION_CLASS`, `DEFAULT_JOB_NAME`, `PRODUCER_NAME`
- Provides 19 common test methods for initialization, start, complete, nominal time, facets
- Reduces test duplication and improves maintainability

### Run All Tests

```bash
Expand Down Expand Up @@ -217,40 +244,44 @@ addopts = [
def test_1(self): # Avoid
```

2. **Follow AAA pattern (Arrange, Act, Assert):**
2. **Use BaseSessionTest for session tests:**
```python
def test_complete_success(self):
# Arrange
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")
session = TeleopSession(
namespace="test",
common_facet=common_facet,
device_facet=device_facet
)
session.start()

# Act
session.complete()

# Assert
assert session._completed is True
from tests.unit.helpers.base_session_test import BaseSessionTest
from airoa_lineage.teleop import TeleopSession

class TestTeleopSessionInitialization(BaseSessionTest):
"""Test TeleopSession initialization."""

SESSION_CLASS = TeleopSession
DEFAULT_JOB_NAME = "teleop-session"
PRODUCER_NAME = "airoa-teleop-system"

# Common tests (19 methods) are inherited from BaseSessionTest
# Only add session-specific tests here
def test_custom_device_facet(self, common_facet, device_facet):
session = TeleopSession(
namespace="test",
common_facet=common_facet,
device_facet=device_facet
)
assert session.device_facet == device_facet
```

3. **Use mocking for external dependencies:**
3. **Use BaseSessionTestHelper for custom tests:**
```python
@patch("airoa_lineage.teleop.session.MarquezClient")
def test_start_returns_run_id(self, mock_client_class):
mock_client = MagicMock()
mock_client_class.return_value = mock_client
# ... test code
from tests.unit.helpers.test_helper import BaseSessionTestHelper

def test_complete_success(self, common_facet):
with BaseSessionTestHelper(TeleopSession) as helper:
session, mock_client = helper.create_session(
namespace="test_namespace",
common_facet=common_facet,
)
session.start()
session.complete()

# Use helper assertions
helper.assert_complete_success(session, mock_client)
```

4. **Test both success and failure cases:**
Expand Down Expand Up @@ -477,13 +508,24 @@ airoa-lineage/
│ ├── __init__.py
│ └── timestamps.py # Timestamp utilities
├── tests/
│ ├── conftest.py # Common fixtures
│ └── unit/ # Unit tests (mirrors src/ structure)
│ ├── marquez_client/
│ │ └── test_client.py
│ ├── teleop/
│ │ └── test_session.py
│ ├── conversion/
│ │ └── test_session.py
│ ├── wasabi_upload/
│ │ └── test_session.py
│ ├── usb_copy/
│ │ └── test_session.py
│ ├── core/
│ │ └── test_base_session.py
│ ├── helpers/ # Test utilities
│ │ ├── test_helper.py
│ │ ├── base_session_test.py
│ │ └── test_test_helper.py
│ └── utils/
│ └── test_timestamps.py
├── .gitignore
Expand Down
187 changes: 187 additions & 0 deletions examples/wasabi_upload_session.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,187 @@
"""
Wasabi Data Upload Session Example

This example demonstrates a complete Wasabi data upload session:
1. Send START event when upload begins with nominal time period
2. Simulate data upload to Wasabi cloud storage
3. Send COMPLETE event when upload ends

It shows how to:
- Use WasabiUploadSession class to manage upload session lifecycle
- Create CommonRunFacet with robot and repository information
- Automatically track run_id without manual management
- Specify robot_id and location to identify the robot and its location
- Specify repository information (hash, URI, tag, branch) for traceability
- Specify nominal time period (the time range of data being processed)
- Use facet_prefix to namespace custom facets (e.g., "airoa_common")
- Track a complete Wasabi data upload session with OpenLineage
"""

import os
import time
from datetime import datetime, timedelta, timezone

from airoa_lineage.facets import CommonRunFacet
from airoa_lineage.marquez_client import MarquezClient
from airoa_lineage.wasabi_upload import WasabiUploadSession


def main():
"""Run a complete Wasabi data upload session simulation."""

# Session configuration
marquez_url = os.getenv("MARQUEZ_URL", "http://localhost:9000")
namespace = "airoa_examples"
job_name = "wasabi-data-upload"

# Robot and repository information
robot_id = "hsr001" # Robot identifier
location = "weblab" # Location identifier
repository_hash = "df110d5" # Git commit hash
repository_uri = "https://github.com/AIRoA/airoa-lineage.git" # Repository URL
repository_tag = "v1.0.0" # Git tag
repository_branch = "main" # Git branch

# Calculate nominal time period (data time range)
# nominal_start: 1 day ago from now
# nominal_end: 5 hours after nominal_start
now = datetime.now(timezone.utc)
nominal_start = now - timedelta(days=1)
nominal_end = nominal_start + timedelta(hours=5)

# Create common facet with robot and repository information
common_facet = CommonRunFacet(
robotId=robot_id,
location=location,
repositoryHash=repository_hash,
repositoryUri=repository_uri,
repositoryTag=repository_tag,
repositoryBranch=repository_branch,
)

# Create session instance
# This automatically generates a run_id and initializes the Marquez client
# Using facet_prefix="airoa" to namespace our custom facets
session = WasabiUploadSession(
namespace=namespace,
common_facet=common_facet,
job_name=job_name,
marquez_url=marquez_url,
facet_prefix="airoa",
)

print("=" * 60)
print("Wasabi Data Upload Session Example")
print("=" * 60)
print(f"Namespace: {namespace}")
print(f"Job Name: {job_name}")
print()
print("Robot & Repository Information:")
print(f" Robot ID: {robot_id}")
print(f" Location: {location}")
print(f" Repo Hash: {repository_hash}")
print(f" Repo URI: {repository_uri}")
print(f" Repo Tag: {repository_tag}")
print(f" Repo Branch: {repository_branch}")
print()
print(f"Run ID: {session.run_id}")
print(f"Marquez: {marquez_url}")
print()
print("Nominal Time Period (data time range):")
print(f" Start: {nominal_start.isoformat()}")
print(f" End: {nominal_end.isoformat()}")
print()

# ========== START Event ==========
print("[1/4] Sending START event...")

# Send START event with nominal time period
# This specifies the time range of the data being processed
session.start(
nominal_start_time=nominal_start.isoformat(),
nominal_end_time=nominal_end.isoformat(),
)

print("✓ START event sent successfully")
print()

# ========== Simulate Data Upload ==========
print("[2/4] Simulating data upload to Wasabi...")

# In a real scenario, this is where you would:
# - Prepare data files for upload
# - Connect to Wasabi S3-compatible API
# - Upload files to Wasabi bucket
# - Verify upload integrity
upload_duration = 3 # seconds
time.sleep(upload_duration)

print(f"✓ Data upload completed ({upload_duration} seconds)")
print(" - In a real scenario, this would:")
print(" • Prepare data files for upload")
print(" • Connect to Wasabi S3-compatible API")
print(" • Upload files to Wasabi bucket")
print(" • Verify upload integrity")
print()

# ========== COMPLETE Event ==========
print("[3/4] Sending COMPLETE event...")

# Send COMPLETE event using the session
# This sends an OpenLineage COMPLETE event to Marquez
session.complete()

print("✓ COMPLETE event sent successfully")
print()

# ========== Query Lineage ==========
print("[4/4] Querying job information...")
time.sleep(1) # Give Marquez time to process

try:
# Initialize client to query job information
client = MarquezClient(marquez_url, namespace=namespace)

# Query job information
job = client.get_job(job_name)
print("✓ Job information retrieved:")
print(f" - Latest run state: {job.get('latestRun', {}).get('state', 'N/A')}")
print(f" - Job type: {job.get('type', 'N/A')}")
print()

except Exception as e:
print(f"⚠ Failed to query job: {e}")
print(" (This is expected if Marquez is still processing the events)")
print()

# ========== Summary ==========
print("=" * 60)
print("Session Summary")
print("=" * 60)
print("✓ Wasabi data upload session completed successfully")
print(f"✓ Robot ID: {robot_id}")
print(f"✓ Location: {location}")
print(f"✓ Repository: {repository_uri}")
print(f"✓ Commit: {repository_hash} ({repository_branch})")
print(f"✓ Tag: {repository_tag}")
print(f"✓ Run ID: {session.run_id}")
print()
print("Next steps:")
print("1. View lineage in Marquez Web UI:")
print(
f" {marquez_url.replace(':9000', ':3000')}/lineage/job/{namespace}/{job_name}"
)
print()
print("2. Query job via API:")
print(f" curl {marquez_url}/api/v1/namespaces/{namespace}/jobs/{job_name}")
print()
print("3. Query this specific run:")
print(
f" curl {marquez_url}/api/v1/namespaces/{namespace}/jobs/{job_name}/runs/{session.run_id}"
)
print()
print("=" * 60)


if __name__ == "__main__":
main()
Loading