feat: OSMO backend integrations

.github/workflows/ci.yml

@@ -171,12 +171,65 @@ jobs:
       - name: Run tests (parallel)
         run: >
           pytest --tb=short -x -m "not glaas and not live_glaas"
+          --ignore=tests/backends/osmo
           --ignore=tests/execution/runtime/test_sitecustomize_perf.py
 
       - name: Run runtime perf guards (serial)
         if: matrix.python-version == '3.12'
         run: pytest --tb=short -x tests/execution/runtime/test_sitecustomize_perf.py -n 0
 
+  test-osmo-e2e:
+    runs-on: ubuntu-latest
+    timeout-minutes: 90
+    needs: [build-rust-binaries]
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python 3.11
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Set up Rust
+        uses: dtolnay/rust-toolchain@stable
+
+      - name: Download Rust binaries
+        uses: actions/download-artifact@v4
+        with:
+          name: roar-rust-binaries-linux-x86_64
+          path: rust/target/release/
+
+      - name: Make tracer and proxy executables runnable
+        run: |
+          chmod +x \
+            rust/target/release/roar-tracer \
+            rust/target/release/roar-proxy \
+            rust/target/release/roar-tracer-ebpf \
+            rust/target/release/roard \
+            rust/target/release/roar-tracer-preload
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install maturin
+          pip install -e .[dev]
+          pip install "ray[default]==2.44.1"
+
+      - name: Build native hash extension
+        run: ./scripts/build_hash_native.sh
+
+      - name: Verify native hash extension import
+        run: python -c "import roar._hash_native as m; print(m.__name__)"
+
+      - name: Run OSMO e2e tests (serial)
+        env:
+          OSMO_DOCKERHUB_USERNAME: ${{ secrets.DOCKER_HUB_USERNAME }}
+          OSMO_DOCKERHUB_PASSWORD: ${{ secrets.DOCKER_HUB_PASSWORD }}
+          OSMO_PRELOAD_PULL_RETRIES: "5"
+          OSMO_QUICK_START_CHART_VERSION: "1.0.1"
+          OSMO_TEST_PYTHON_IMAGE: public.ecr.aws/docker/library/python:3.11-slim
+        run: pytest --tb=short -x -m "osmo_e2e" tests/backends/osmo -n 0
+
   test-macos:
     needs: [build-rust-binaries-macos]
     strategy:
@@ -252,6 +305,7 @@ jobs:
       - name: Run tests (parallel)
         run: >
           pytest --tb=short -x -m "not glaas and not live_glaas and not ebpf and not large_pipeline"
+          --ignore=tests/backends/osmo
           --ignore=tests/execution/runtime/test_sitecustomize_perf.py
 
   test-tracer-privileged:

.gitignore

@@ -221,6 +221,7 @@ CLAUDE.md
 # roar
 roar/bin/*
 .roar/
+.tmp-osmo-e2e/
 
 .ralph
 # Ray benchmark results

pyproject.toml

@@ -41,6 +41,7 @@ dependencies = [
     "cryptography>=42.0.0",
     "dependency-injector>=4.40.0",
     "msgpack>=1.0.0",
+    "PyYAML>=6.0",
     "sqlalchemy>=2.0.0",
     "pysqlite3-binary>=0.5.0; sys_platform == 'linux' and platform_machine == 'x86_64'",  # Fallback when stdlib sqlite3 unavailable
     "pydantic>=2.0.0",
@@ -74,6 +75,7 @@ roar-worker = "roar.execution.runtime.worker_bootstrap:main"
 
 [project.entry-points."roar.execution_backends"]
 ray = "roar.backends.ray.plugin:register"
+osmo = "roar.backends.osmo.plugin:register"
 
 [tool.maturin]
 manifest-path = "rust/crates/artifact-hash-py/Cargo.toml"
@@ -103,6 +105,7 @@ markers = [
     "happy_path: Happy path tests for core functionality",
     "diagnostic: Opt-in diagnostics or aspirational performance budgets outside the default gate",
     "large_pipeline: Stress-style pipeline coverage with larger DAG fixtures",
+    "osmo_e2e: OSMO end-to-end tests requiring a Docker Compose managed KIND harness",
     "ray_e2e: Ray end-to-end tests requiring a running Docker cluster",
     "ray_contract: User-facing Ray contract tests using `roar run ray job submit ...`",
     "ray_diagnostic: Diagnostic Ray tests that intentionally inspect internal runtime details",

roar/backends/osmo/__init__.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from .export import OsmoLineageBundleExport, export_osmo_lineage_bundle
+from .host_execution import OsmoAttachOptions, attach_osmo_workflow, execute_osmo_workflow_submit
+from .lineage import discover_downloaded_lineage_bundles, reconstitute_osmo_lineage_bundles
+from .plugin import OSMO_EXECUTION_BACKEND, register
+from .runtime_bundle import OsmoRuntimeBundle, build_osmo_runtime_bundle
+from .workflow import (
+    PreparedOsmoWorkflow,
+    prepare_osmo_workflow_for_lineage,
+    resolve_roar_install_requirement,
+)
+
+__all__ = [
+    "OSMO_EXECUTION_BACKEND",
+    "OsmoAttachOptions",
+    "OsmoLineageBundleExport",
+    "OsmoRuntimeBundle",
+    "PreparedOsmoWorkflow",
+    "attach_osmo_workflow",
+    "build_osmo_runtime_bundle",
+    "discover_downloaded_lineage_bundles",
+    "execute_osmo_workflow_submit",
+    "export_osmo_lineage_bundle",
+    "prepare_osmo_workflow_for_lineage",
+    "reconstitute_osmo_lineage_bundles",
+    "register",
+    "resolve_roar_install_requirement",
+]

roar/backends/osmo/config.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from roar.execution.framework.contract import BackendConfigAdapter, ConfigurableKeySpec
+
+
+class OsmoBackendConfig(BaseModel):
+    """OSMO backend configuration."""
+
+    model_config = ConfigDict(
+        strict=False,
+        validate_assignment=True,
+        extra="ignore",
+        revalidate_instances="never",
+    )
+
+    enabled: bool = True
+    auto_prepare_submissions: bool = True
+    force_json_output: bool = True
+    wait_for_completion: bool = False
+    download_declared_outputs: bool = False
+    download_directory: str = ".roar/osmo/downloads"
+    ingest_lineage_bundles: bool = False
+    lineage_bundle_dataset_name: str = "roar-lineage"
+    lineage_bundle_filename: str = "roar-fragments.json"
+    runtime_install_requirement: str = ""
+    runtime_install_local_path: str = ""
+    runtime_install_remote_path: str = "/tmp/roar-osmo-install.whl"
+    query_timeout_seconds: int = Field(default=12 * 60, ge=1)
+    poll_interval_seconds: float = Field(default=5.0, gt=0.0)
+
+
+OSMO_CONFIGURABLE_KEYS = {
+    "osmo.enabled": ConfigurableKeySpec(
+        value_type=bool,
+        default=True,
+        description="Enable automatic OSMO workflow submit handling in roar run",
+    ),
+    "osmo.auto_prepare_submissions": ConfigurableKeySpec(
+        value_type=bool,
+        default=True,
+        description="Rewrite local OSMO workflow submits through a temporary Roar-instrumented workflow",
+    ),
+    "osmo.force_json_output": ConfigurableKeySpec(
+        value_type=bool,
+        default=True,
+        description="Append --format-type json to OSMO workflow submit commands when missing",
+    ),
+    "osmo.wait_for_completion": ConfigurableKeySpec(
+        value_type=bool,
+        default=False,
+        description="Poll submitted OSMO workflows to a terminal state before completing roar run",
+    ),
+    "osmo.download_declared_outputs": ConfigurableKeySpec(
+        value_type=bool,
+        default=False,
+        description="Download declared dataset outputs after successful waited OSMO workflow completion",
+    ),
+    "osmo.download_directory": ConfigurableKeySpec(
+        value_type=str,
+        default=".roar/osmo/downloads",
+        description="Directory for downloaded OSMO dataset outputs, relative to the repo root when not absolute",
+    ),
+    "osmo.ingest_lineage_bundles": ConfigurableKeySpec(
+        value_type=bool,
+        default=False,
+        description="Ingest downloaded OSMO output bundles named by osmo.lineage_bundle_filename into the local Roar DB",
+    ),
+    "osmo.lineage_bundle_dataset_name": ConfigurableKeySpec(
+        value_type=str,
+        default="roar-lineage",
+        description="Dataset name Roar uses by default when downloading returned OSMO lineage bundles",
+    ),
+    "osmo.lineage_bundle_filename": ConfigurableKeySpec(
+        value_type=str,
+        default="roar-fragments.json",
+        description="Filename Roar treats as a downloaded OSMO lineage bundle when ingest_lineage_bundles is enabled",
+    ),
+    "osmo.runtime_install_requirement": ConfigurableKeySpec(
+        value_type=str,
+        default="",
+        description="Pinned requirement, wheel URL, or install target used by OSMO runtime wrappers when bootstrapping Roar remotely; packaged roar-cli wheels are expected to include tracer binaries",
+    ),
+    "osmo.runtime_install_local_path": ConfigurableKeySpec(
+        value_type=str,
+        default="",
+        description="Optional local wheel or artifact path injected into prepared OSMO workflows and installed by the wrapper; roar-cli wheels should be built with packaged binaries",
+    ),
+    "osmo.runtime_install_remote_path": ConfigurableKeySpec(
+        value_type=str,
+        default="/tmp/roar-osmo-install.whl",
+        description="Remote path inside OSMO tasks used for an injected runtime install artifact",
+    ),
+    "osmo.query_timeout_seconds": ConfigurableKeySpec(
+        value_type=int,
+        default=12 * 60,
+        description="Maximum time to wait for a submitted OSMO workflow to reach a terminal state",
+    ),
+    "osmo.poll_interval_seconds": ConfigurableKeySpec(
+        value_type=float,
+        default=5.0,
+        description="Polling interval in seconds when waiting for OSMO workflow completion",
+    ),
+}
+
+OSMO_INIT_TEMPLATE = """\
+[osmo]
+# Enable OSMO workflow-submit recognition in roar run
+enabled = true
+# Rewrite local workflow submits through a temporary Roar-instrumented workflow
+auto_prepare_submissions = true
+# Append --format-type json when missing so Roar can capture workflow metadata
+force_json_output = true
+# Optionally wait for submitted workflows to reach a terminal state
+wait_for_completion = false
+# Optionally download declared dataset outputs after a successful waited run
+download_declared_outputs = false
+# Optionally ingest downloaded lineage bundles back into the local Roar DB
+ingest_lineage_bundles = false
+# Standard dataset name Roar expects for returned lineage bundles
+lineage_bundle_dataset_name = "roar-lineage"
+# Optional pinned requirement or wheel URL installed by the injected OSMO wrapper
+# For roar-cli, use a packaged wheel or index source that includes bundled tracer binaries
+runtime_install_requirement = ""
+# Optional local wheel path injected into the workflow and installed by the wrapper
+# For roar-cli, prefer a wheel built through scripts/build_wheel_with_bins.sh
+runtime_install_local_path = ""
+# Remote install-artifact path used inside OSMO tasks when runtime_install_local_path is set
+runtime_install_remote_path = "/tmp/roar-osmo-install.whl"
+"""
+
+
+def normalize_osmo_backend_config(section: Mapping[str, Any] | None) -> dict[str, Any]:
+    return OsmoBackendConfig.model_validate(dict(section or {})).model_dump()
+
+
+def load_osmo_backend_config(start_dir: str | None = None) -> dict[str, Any]:
+    try:
+        from roar.integrations.config import load_config
+
+        config = load_config(start_dir=start_dir)
+    except Exception:
+        return dict(OSMO_BACKEND_CONFIG.default_values)
+
+    section = config.get("osmo", {})
+    if not isinstance(section, Mapping):
+        return dict(OSMO_BACKEND_CONFIG.default_values)
+    return normalize_osmo_backend_config(section)
+
+
+OSMO_BACKEND_CONFIG = BackendConfigAdapter(
+    section_name="osmo",
+    default_values=OsmoBackendConfig().model_dump(),
+    configurable_keys=OSMO_CONFIGURABLE_KEYS,
+    init_template=OSMO_INIT_TEMPLATE,
+    normalize_section=normalize_osmo_backend_config,
+)
+
+
+__all__ = [
+    "OSMO_BACKEND_CONFIG",
+    "OSMO_CONFIGURABLE_KEYS",
+    "OSMO_INIT_TEMPLATE",
+    "OsmoBackendConfig",
+    "load_osmo_backend_config",
+    "normalize_osmo_backend_config",
+]