fleetmesh-node 0.2.0__py3-none-any.whl

node/planes/services/compute/executors/bacalhau.py

@@ -0,0 +1,605 @@
+"""Bacalhau REST client connector and Blossom CAS compute-over-data mounting bridge.
+
+Implements Phase 2 of the compute roadmap:
+  - BacalhauClientProtocol: abstract interface for Bacalhau daemon orchestrators.
+  - BacalhauRestClient: native HTTP/REST connector for Bacalhau v1 orchestrator API.
+  - MockBacalhauClient: test double for deterministic testing and dry-run fixtures.
+  - BlossomMountBridge: Blossom CAS input resolution and output packaging/uploading.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from typing import Any, Callable, Mapping, NamedTuple, Optional, Protocol, Sequence, TypedDict
+
+from node.transport import netfetch
+from ..common import ComputeError, resolve_blob
+from ..sandbox import ComputeLimits
+from ..types import BacalhauJobSummary
+from .base import ComputeExecutionResult, ComputeExecutor, ComputeQuote
+
+
+ProgressCallback = Callable[[str, float, Mapping[str, Any]], None]
+
+
+class BlossomStagingResult(NamedTuple):
+    """Result of staging Blossom CAS input blobs into container mount points."""
+    staged_paths: list[str]
+    total_bytes: int
+    files: dict[str, bytes]
+    source_cids: list[str]
+
+
+class BlossomUploadResult(NamedTuple):
+    """Result of packaging and committing container outputs to Blossom CAS."""
+    result_cid: str
+    result_hash: str
+    output_bytes: bytes
+    mounted_outputs: list[str]
+
+
+class BacalhauTaskSpec(TypedDict, total=False):
+    """Specification of a single Bacalhau container task."""
+    name: str
+    engine: str
+    image: str
+    entrypoint: list[str]
+    parameters: list[str]
+    env: dict[str, str]
+    input_sources: list[Mapping[str, Any]]
+    result_paths: list[Mapping[str, Any]]
+    resources: Mapping[str, Any]
+    timeouts: dict[str, int]
+
+
+class BacalhauClientProtocol(Protocol):
+    """Interface required for Bacalhau orchestrator connectors."""
+
+    def submit_job(self, job_payload: Mapping[str, Any], timeout_seconds: int = 300) -> str:
+        """Submit a batch job specification and return assigned JobID."""
+        ...
+
+    def get_job_status(self, job_id: str) -> Mapping[str, Any]:
+        """Query execution lifecycle state for the given job_id."""
+        ...
+
+    def get_job_results(self, job_id: str) -> Mapping[str, Any]:
+        """Fetch outputs, logs, and artifacts for a completed job."""
+        ...
+
+    def cancel_job(self, job_id: str, reason: str = '') -> bool:
+        """Abort or terminate a running Bacalhau job."""
+        ...
+
+    def is_alive(self) -> bool:
+        """Check if the Bacalhau daemon/orchestrator endpoint is reachable."""
+        ...
+
+
+class BacalhauRestClient:
+    """Live Bacalhau REST API v1 connector using Fleetmesh netfetch."""
+
+    def __init__(
+        self,
+        endpoint: str = 'http://127.0.0.1:1234',
+        timeout_seconds: float = 30.0,
+    ) -> None:
+        self.endpoint: str = endpoint.rstrip('/')
+        self.timeout_seconds: float = timeout_seconds
+
+    def is_alive(self) -> bool:
+        """Check if the Bacalhau daemon endpoint responds to health probes."""
+        for api_route in ('/api/v1/orchestrator/nodes', '/api/v1/version', '/api/v1/agent/alive'):
+            url = f"{self.endpoint}{api_route}"
+            try:
+                body, _ = netfetch.fetch(
+                    url,
+                    max_bytes=65536,
+                    timeout=2.0,
+                    method='GET',
+                    allow_private=True,
+                )
+                if body:
+                    return True
+            except Exception:
+                continue
+        return False
+
+    def build_bacalhau_job_spec(self, job_payload: Mapping[str, Any]) -> Mapping[str, Any]:
+        """Convert Fleetmesh job descriptor into standard Bacalhau v1 format."""
+        engine = str(job_payload.get('engine', 'docker')).lower()
+        image = str(job_payload.get('image', ''))
+        wasm_mod = job_payload.get('wasm_module', {})
+        entrypoint = job_payload.get('entrypoint') or ['/bin/sh', '-c']
+        parameters = job_payload.get('parameters') or []
+        env_vars = job_payload.get('env') or {}
+        resources = job_payload.get('resources') or {'cpu': 1, 'gpu': 0, 'memory_gib': 2.0, 'timeout_seconds': 300}
+        storage_mounts = job_payload.get('storage_mounts') or []
+        output_mounts = job_payload.get('output_mounts') or []
+
+        input_sources: list[Mapping[str, Any]] = []
+        for mount in storage_mounts:
+            cid = mount.get('source_cid', '')
+            mount_path = mount.get('mount_path', '/inputs')
+            input_sources.append({
+                'Source': {
+                    'Type': 'urlDownload' if cid.startswith('http') else 's3' if cid.startswith('s3://') else 'localDirectory',
+                    'Params': {
+                        'URL': cid if cid.startswith('http') else '',
+                        'Target': mount_path,
+                    }
+                }
+            })
+
+        result_paths = [
+            {'Name': m.get('name', 'output'), 'Path': m.get('container_path', '/outputs')}
+            for m in output_mounts
+        ] or [{'Name': 'outputs', 'Path': '/outputs'}]
+
+        task = {
+            'Name': 'main',
+            'Engine': {
+                'Type': engine,
+                'Params': {
+                    'Image': image or wasm_mod.get('sha256', ''),
+                    'Entrypoint': entrypoint,
+                    'Parameters': parameters,
+                    'EnvironmentVariables': [f"{k}={v}" for k, v in env_vars.items()],
+                }
+            },
+            'InputSources': input_sources,
+            'ResultPaths': result_paths,
+            'Resources': {
+                'CPU': str(resources.get('cpu', 1)),
+                'Memory': f"{int(resources.get('memory_gib', 2))}Gb",
+                'GPU': str(resources.get('gpu', 0)),
+            },
+            'Timeouts': {
+                'ExecutionTimeout': int(resources.get('timeout_seconds', 300)),
+            },
+            'Network': {
+                'Type': 'None',
+            },
+        }
+
+        job_id = str(job_payload.get('id') or os.urandom(8).hex())
+        return {
+            'Job': {
+                'Name': f"fleetmesh-{job_id[:16]}",
+                'Type': 'batch',
+                'Count': 1,
+                'Tasks': [task],
+            }
+        }
+
+    def submit_job(self, job_payload: Mapping[str, Any], timeout_seconds: int = 300) -> str:
+        """Submit a job specification to the Bacalhau orchestrator."""
+        spec = self.build_bacalhau_job_spec(job_payload)
+        data_bytes = json.dumps(spec).encode('utf-8')
+
+        # Try orchestrator API first, fall back to /api/v1/run
+        last_err: Optional[Exception] = None
+        for api_route in ('/api/v1/orchestrator/jobs', '/api/v1/run'):
+            url = f"{self.endpoint}{api_route}"
+            try:
+                body, _ = netfetch.fetch(
+                    url,
+                    max_bytes=1048576,
+                    timeout=self.timeout_seconds,
+                    method='POST',
+                    data=data_bytes,
+                    headers={'Content-Type': 'application/json', 'User-Agent': 'fleetmesh-bacalhau/1.0'},
+                    allow_private=True,
+                )
+                resp = json.loads(body.decode('utf-8'))
+                job_id = resp.get('JobID') or resp.get('ID') or (resp.get('Job') or {}).get('ID')
+                if job_id:
+                    return str(job_id)
+            except netfetch.FetchError as e:
+                last_err = e
+                continue
+            except Exception as e:
+                last_err = e
+                continue
+
+        raise ComputeError(f"Bacalhau job submission failed: {last_err}")
+
+    def get_job_status(self, job_id: str) -> Mapping[str, Any]:
+        """Fetch current state of a Bacalhau job."""
+        url = f"{self.endpoint}/api/v1/orchestrator/jobs/{job_id}"
+        try:
+            body, _ = netfetch.fetch(
+                url,
+                max_bytes=524288,
+                timeout=self.timeout_seconds,
+                method='GET',
+                allow_private=True,
+            )
+            data = json.loads(body.decode('utf-8'))
+            state = data.get('State') or (data.get('Job') or {}).get('State') or {}
+            status_str = state.get('StateType') or 'Running'
+            return {
+                'job_id': job_id,
+                'state': status_str,
+                'raw': data,
+            }
+        except Exception as e:
+            raise ComputeError(f"Failed to fetch Bacalhau job status for {job_id}: {e}")
+
+    def get_job_results(self, job_id: str) -> Mapping[str, Any]:
+        """Fetch results and output references for a completed job."""
+        url = f"{self.endpoint}/api/v1/orchestrator/jobs/{job_id}/results"
+        try:
+            body, _ = netfetch.fetch(
+                url,
+                max_bytes=1048576,
+                timeout=self.timeout_seconds,
+                method='GET',
+                allow_private=True,
+            )
+            return json.loads(body.decode('utf-8'))
+        except Exception as e:
+            raise ComputeError(f"Failed to fetch Bacalhau job results for {job_id}: {e}")
+
+    def cancel_job(self, job_id: str, reason: str = '') -> bool:
+        """Cancel an in-flight Bacalhau job."""
+        url = f"{self.endpoint}/api/v1/orchestrator/jobs/{job_id}/stop"
+        try:
+            data_bytes = json.dumps({'Reason': reason or 'cancelled by fleetmesh'}).encode('utf-8')
+            netfetch.fetch(
+                url,
+                max_bytes=65536,
+                timeout=self.timeout_seconds,
+                method='POST',
+                data=data_bytes,
+                headers={'Content-Type': 'application/json'},
+                allow_private=True,
+            )
+            return True
+        except Exception:
+            return False
+
+
+class MockBacalhauClient:
+    """Test double implementing BacalhauClientProtocol for isolated and mock tests."""
+
+    def __init__(self, simulate_failure: bool = False) -> None:
+        self.submitted_jobs: list[Mapping[str, Any]] = []
+        self.cancelled_jobs: list[str] = []
+        self.simulate_failure: bool = simulate_failure
+        self.call_history: list[str] = []
+
+    def is_alive(self) -> bool:
+        return True
+
+    def submit_job(self, job_payload: Mapping[str, Any], timeout_seconds: int = 300) -> str:
+        self.call_history.append('submit_job')
+        if self.simulate_failure:
+            raise ComputeError("Simulated Bacalhau daemon failure")
+        job_id = f"mock-bacalhau-{os.urandom(8).hex()}"
+        self.submitted_jobs.append({'job_id': job_id, 'payload': job_payload})
+        return job_id
+
+    def get_job_status(self, job_id: str) -> Mapping[str, Any]:
+        self.call_history.append(f'get_job_status:{job_id}')
+        if job_id in self.cancelled_jobs:
+            return {'job_id': job_id, 'state': 'Cancelled'}
+        return {'job_id': job_id, 'state': 'Completed'}
+
+    def get_job_results(self, job_id: str) -> Mapping[str, Any]:
+        self.call_history.append(f'get_job_results:{job_id}')
+        return {
+            'job_id': job_id,
+            'status': 'success',
+            'output_cid': f"bafkrei{hashlib.sha256(job_id.encode()).hexdigest()[:48]}",
+            'exit_code': 0,
+        }
+
+    def cancel_job(self, job_id: str, reason: str = '') -> bool:
+        self.call_history.append(f'cancel_job:{job_id}')
+        self.cancelled_jobs.append(job_id)
+        return True
+
+
+class BlossomMountBridge:
+    """Manages input staging from Blossom CAS and output artifact publishing."""
+
+    def resolve_and_stage_inputs(
+        self,
+        storage_mounts: Sequence[Mapping[str, Any]],
+        blob_store: Optional[Any] = None,
+        blossom_servers: Optional[Sequence[str]] = None,
+    ) -> BlossomStagingResult:
+        """Resolve inputs from Blossom CAS or local storage, staging them for execution."""
+        staged_paths: list[str] = []
+        source_cids: list[str] = []
+        files: dict[str, bytes] = {}
+        total_bytes = 0
+
+        for mount in storage_mounts:
+            cid = str(mount.get('source_cid', ''))
+            mount_path = str(mount.get('mount_path', '/inputs'))
+            if not cid:
+                continue
+
+            source_cids.append(cid)
+            spec = {'cid': cid} if cid.startswith('baf') else {'sha256': cid}
+            data = resolve_blob(spec, blob_store=blob_store, blossom_servers=blossom_servers)
+            if data is not None:
+                files[mount_path] = data
+                staged_paths.append(mount_path)
+                total_bytes += len(data)
+
+        return BlossomStagingResult(
+            staged_paths=staged_paths,
+            total_bytes=total_bytes,
+            files=files,
+            source_cids=source_cids,
+        )
+
+    def package_and_upload_outputs(
+        self,
+        raw_output_or_files: bytes | Mapping[str, bytes],
+        output_mounts: Sequence[Mapping[str, Any]],
+        blob_store: Optional[Any] = None,
+        blossom_servers: Optional[Sequence[str]] = None,
+    ) -> BlossomUploadResult:
+        """Hash, persist, and register job output artifacts into Blossom CAS."""
+        if isinstance(raw_output_or_files, bytes):
+            out_bytes = raw_output_or_files
+        else:
+            out_bytes = json.dumps(
+                {k: v.decode('utf-8', errors='replace') for k, v in raw_output_or_files.items()},
+                sort_keys=True
+            ).encode('utf-8')
+
+        result_hash = hashlib.sha256(out_bytes).hexdigest()
+
+        # Persist to local blob_store if present
+        result_cid = None
+        if blob_store is not None and hasattr(blob_store, 'put'):
+            try:
+                desc = blob_store.put(out_bytes)
+                result_cid = desc.get('cid') if isinstance(desc, dict) else str(desc)
+            except Exception:
+                pass
+
+        if not result_cid:
+            # Format canonical Blossom SHA-256 CID
+            result_cid = result_hash
+
+        # Push to remote Blossom servers if configured
+        if blossom_servers:
+            for srv in blossom_servers:
+                try:
+                    url = f"{srv.rstrip('/')}/upload"
+                    netfetch.fetch(
+                        url,
+                        max_bytes=1048576,
+                        timeout=15.0,
+                        method='PUT',
+                        data=out_bytes,
+                        headers={'Content-Type': 'application/octet-stream', 'X-SHA-256': result_hash},
+                        allow_private=True,
+                    )
+                except Exception:
+                    pass
+
+        mounted_outputs = [str(m.get('name', 'outputs')) for m in output_mounts] or ['default']
+        return BlossomUploadResult(
+            result_cid=result_cid,
+            result_hash=result_hash,
+            output_bytes=out_bytes,
+            mounted_outputs=mounted_outputs,
+        )
+
+
+class BacalhauExecutor(ComputeExecutor):
+    """Compute-over-data batch driver backed by Bacalhau / Docker containers.
+
+    Integrates Phase 2 capabilities:
+      - Mounts Blossom CAS CIDs directly into container volume mounts.
+      - Dispatches jobs to Bacalhau orchestrator or local sandbox.
+      - Emits 5-state lifecycle progress heartbeats (compute/1.1/progress).
+      - Honors cooperative cancellation tokens (compute/1.1/cancel).
+    """
+
+    engine_id: str = 'bacalhau'
+    wasi_profile_digest: str = ''
+
+    def __init__(
+        self,
+        endpoint: str = 'http://127.0.0.1:1234',
+        mock_mode: bool = False,
+        client: Optional[BacalhauClientProtocol] = None,
+        mount_bridge: Optional[BlossomMountBridge] = None,
+    ) -> None:
+        self.endpoint = endpoint.rstrip('/')
+        self.mock_mode = mock_mode
+        self.client = client
+        self.mount_bridge = mount_bridge or BlossomMountBridge()
+
+    def is_live_runnable(self) -> bool:
+        """Check whether Bacalhau execution can succeed on this node."""
+        if self.mock_mode:
+            return True
+        if self.client is not None:
+            try:
+                return self.client.is_alive()
+            except Exception:
+                return False
+        return False
+
+    def quote(
+        self,
+        module_spec: Mapping[str, Any],
+        input_spec: Mapping[str, Any],
+        limits: ComputeLimits,
+    ) -> ComputeQuote:
+        cpu_minutes = int(input_spec.get('cpu_minutes', 1))
+        gpu_minutes = int(input_spec.get('gpu_minutes', 0))
+        ram_gib = float(input_spec.get('ram_gib', 1.0))
+        fuel_est = int(cpu_minutes * 150 + gpu_minutes * 1200 + ram_gib * 50)
+        return ComputeQuote(
+            fuel_estimate=min(fuel_est, limits.max_fuel),
+            rate_unit='msat_per_job',
+            rate_amount=fuel_est,
+            engine_id=self.engine_id,
+        )
+
+    def execute(
+        self,
+        job_payload: Mapping[str, Any],
+        fuel_cap: int,
+        progress_callback: Optional[ProgressCallback] = None,
+        cancellation_token: Optional[Any] = None,
+    ) -> ComputeExecutionResult:
+        engine = job_payload.get('engine', 'docker')
+        image = job_payload.get('image', '')
+        wasm_module = job_payload.get('wasm_module', {})
+        entrypoint = job_payload.get('entrypoint') or ['/bin/sh', '-c']
+        parameters = job_payload.get('parameters') or []
+        storage_mounts = job_payload.get('storage_mounts') or []
+        output_mounts = job_payload.get('output_mounts') or []
+        resources = job_payload.get('resources') or {'cpu': 1, 'gpu': 0, 'memory_gib': 2.0, 'timeout_seconds': 60}
+
+        verif_mode = self.verification_mode(job_payload)
+
+        # Honesty gate. If not in mock_mode, an active and verified live backend must be present.
+        if not self.mock_mode and self.client is None and not self.is_live_runnable():
+            raise ComputeError('not implemented: bacalhau executor is a stub (mock_mode only)')
+
+        def _check_cancelled() -> None:
+            if cancellation_token is not None:
+                is_set = False
+                if hasattr(cancellation_token, 'is_set') and cancellation_token.is_set():
+                    is_set = True
+                elif isinstance(cancellation_token, bool) and cancellation_token:
+                    is_set = True
+                if is_set:
+                    raise ComputeError('job cancelled by requester')
+
+        # 1. State: queued
+        if progress_callback:
+            progress_callback('queued', 0.0, {'status': 'job queued for batch execution'})
+
+        # 2. State: downloading_inputs (stage Blossom CAS volumes)
+        _check_cancelled()
+        if progress_callback:
+            progress_callback('downloading_inputs', 0.2, {'status': 'staging CAS input volumes'})
+
+        blob_store = job_payload.get('blob_store')
+        blossom_servers = job_payload.get('blossom_servers')
+
+        staging_res = self.mount_bridge.resolve_and_stage_inputs(
+            storage_mounts=storage_mounts,
+            blob_store=blob_store,
+            blossom_servers=blossom_servers,
+        )
+        mounted_inputs = staging_res.files
+        total_in_bytes = staging_res.total_bytes
+
+        # 3. State: running (execute via Bacalhau orchestrator or deterministic pricing model)
+        _check_cancelled()
+        if progress_callback:
+            progress_callback('running', 0.5, {'status': 'executing container workload'})
+
+        client_results: Optional[Mapping[str, Any]] = None
+        if self.client is not None:
+            timeout_sec = int(resources.get('timeout_seconds') or 60)
+            job_id = self.client.submit_job(job_payload, timeout_seconds=timeout_sec)
+            try:
+                _check_cancelled()
+                status_data = self.client.get_job_status(job_id)
+                state_str = str(status_data.get('state', '')).lower()
+                if 'cancel' in state_str:
+                    raise ComputeError('job cancelled by requester')
+                if 'fail' in state_str or 'error' in state_str:
+                    raise ComputeError(f"Bacalhau job failed: {status_data}")
+
+                _check_cancelled()
+                client_results = self.client.get_job_results(job_id)
+            except ComputeError:
+                raise
+            except Exception as e:
+                if cancellation_token is not None and getattr(cancellation_token, 'is_set', lambda: False)():
+                    self.client.cancel_job(job_id, 'cancelled by fleetmesh session')
+                    raise ComputeError('job cancelled by requester')
+                raise ComputeError(f"Bacalhau job execution failed: {e}")
+
+        cpu = int(resources.get('cpu') or 1)
+        gpu = int(resources.get('gpu') or 0)
+        timeout_sec = int(resources.get('timeout_seconds') or 60)
+        cpu_sec = min(timeout_sec, 10)
+        gpu_sec = min(timeout_sec, 10) if gpu > 0 else 0
+        mem_gib = float(resources.get('memory_gib') or 1.0)
+        fuel_used = cpu * 150 + gpu * 1200 + int(mem_gib * 50)
+
+        if fuel_used > fuel_cap:
+            raise ComputeError(f"Bacalhau job exceeded fuel cap ({fuel_used} > {fuel_cap})")
+
+        # 4. State: uploading_outputs (package artifacts and commit to Blossom CAS)
+        _check_cancelled()
+        if progress_callback:
+            progress_callback('uploading_outputs', 0.8, {'status': 'packaging and uploading output artifacts'})
+
+        out_summary: BacalhauJobSummary = {
+            'status': 'completed',
+            'engine': engine,
+            'image': image or wasm_module.get('sha256', ''),
+            'entrypoint': entrypoint,
+            'parameters': parameters,
+            'mounted_paths': list(mounted_inputs.keys()),
+            'outputs': [m.get('name', 'output') for m in output_mounts] or ['default'],
+            'bytes_processed': total_in_bytes,
+        }
+        if client_results:
+            out_summary['results'] = client_results
+
+        upload_res = self.mount_bridge.package_and_upload_outputs(
+            raw_output_or_files=json.dumps(out_summary, sort_keys=True).encode('utf-8'),
+            output_mounts=output_mounts,
+            blob_store=blob_store,
+            blossom_servers=blossom_servers,
+        )
+        raw_output = upload_res.output_bytes
+        result_hash = upload_res.result_hash
+
+        # 5. State: completed
+        if progress_callback:
+            progress_callback('completed', 1.0, {
+                'status': 'job completed successfully',
+                'result_hash': result_hash,
+                'result_cid': upload_res.result_cid,
+            })
+
+        counters = {
+            'cpu_seconds': cpu_sec,
+            'gpu_seconds': gpu_sec,
+            'bytes_in': total_in_bytes,
+            'bytes_out': len(raw_output),
+        }
+
+        return ComputeExecutionResult(
+            result_hash=result_hash,
+            raw_output=raw_output,
+            fuel_used=fuel_used,
+            engine_id=self.engine_id,
+            wasi_profile_digest=f"bacalhau:{engine}:{image or 'wasm'}",
+            counters=counters,
+            verification_mode=verif_mode,
+        )
+
+
+__all__ = [
+    'ProgressCallback',
+    'BlossomStagingResult',
+    'BlossomUploadResult',
+    'BacalhauTaskSpec',
+    'BacalhauClientProtocol',
+    'BacalhauRestClient',
+    'MockBacalhauClient',
+    'BlossomMountBridge',
+    'BacalhauExecutor',
+]