foodforthought-cli 0.2.8__py3-none-any.whl → 0.3.0__py3-none-any.whl

ate/urdf/cloud.py

@@ -0,0 +1,491 @@
+"""
+Cloud client for URDF scan processing.
+
+This module provides the client-side logic for uploading videos to the
+cloud GPU pipeline and retrieving the generated URDF artifacts.
+"""
+
+import json
+import time
+import sys
+from pathlib import Path
+from typing import Optional, Dict, Any
+from dataclasses import dataclass
+
+import requests
+
+# Import the shared client for authentication
+from ..client import ATEClient, BASE_URL
+
+
+@dataclass
+class CloudJobResult:
+    """Result from a cloud scan job."""
+    success: bool
+    job_id: str
+    artifact_url: Optional[str] = None
+    error: Optional[str] = None
+    link_count: int = 0
+    joint_count: int = 0
+    processing_time_seconds: float = 0.0
+
+
+class CloudScanClient:
+    """Client for cloud-based URDF scan processing."""
+
+    def __init__(self, base_url: str = BASE_URL):
+        """Initialize the cloud scan client."""
+        self.client = ATEClient(base_url=base_url)
+        self.base_url = base_url
+
+    def check_credits(self) -> Dict[str, Any]:
+        """
+        Check the user's credit balance.
+
+        Returns:
+            Dict with 'balance' key, or 'error' if failed
+        """
+        try:
+            response = self.client.get("/scan/jobs")
+            return response.get("credits", {"balance": 0})
+        except Exception as e:
+            return {"error": str(e), "balance": 0}
+
+    def list_jobs(self) -> Dict[str, Any]:
+        """
+        List all scan jobs for the authenticated user.
+
+        Returns:
+            Dict with 'jobs' list and 'credits' info
+        """
+        try:
+            response = self.client.get("/scan/jobs")
+            return response
+        except Exception as e:
+            return {"error": str(e), "jobs": []}
+
+    def get_job(self, job_id: str) -> Dict[str, Any]:
+        """
+        Get details for a specific scan job.
+
+        Args:
+            job_id: The job ID to retrieve
+
+        Returns:
+            Dict with job details or error
+        """
+        try:
+            response = self.client.get(f"/scan/jobs/{job_id}")
+            return response
+        except Exception as e:
+            return {"error": str(e)}
+
+    def create_job(
+        self,
+        robot_name: Optional[str] = None,
+        scale_ref: Optional[str] = None,
+    ) -> tuple[str, str, str, Dict[str, Any]]:
+        """
+        Create a new cloud scan job.
+
+        Args:
+            robot_name: Optional robot name for the URDF
+            scale_ref: Optional scale reference (e.g., "gripper:85mm")
+
+        Returns:
+            Tuple of (job_id, upload_url, blob_token, credits_info)
+
+        Raises:
+            RuntimeError: If job creation fails
+        """
+        data = {}
+        if robot_name:
+            data["robotName"] = robot_name
+        if scale_ref:
+            data["scaleRef"] = scale_ref
+
+        try:
+            response = self.client.post("/scan/jobs", data=data)
+        except SystemExit:
+            raise RuntimeError("Failed to create job. Please check your authentication.")
+
+        if "error" in response:
+            raise RuntimeError(response["error"])
+
+        job_id = response.get("jobId")
+        upload_url = response.get("uploadUrl")  # Server returns full upload URL
+        blob_token = response.get("blobToken")  # Server returns blob token for auth
+        credits_info = response.get("credits", {})
+
+        if not job_id or not upload_url or not blob_token:
+            raise RuntimeError(f"Invalid response from server: {response}")
+
+        return job_id, upload_url, blob_token, credits_info
+
+    def upload_video(
+        self,
+        video_path: Path,
+        upload_url: str,
+        blob_token: str,
+        progress_callback: Optional[callable] = None,
+    ) -> bool:
+        """
+        Upload a video file to Vercel Blob storage.
+
+        Args:
+            video_path: Path to the video file
+            upload_url: Full URL for upload (from create_job)
+            blob_token: Authentication token for blob upload (from create_job)
+            progress_callback: Optional callback(bytes_sent, total_bytes)
+
+        Returns:
+            True if upload succeeded
+
+        Raises:
+            RuntimeError: If upload fails
+        """
+        if not video_path.exists():
+            raise RuntimeError(f"Video file not found: {video_path}")
+
+        file_size = video_path.stat().st_size
+
+        with open(video_path, "rb") as f:
+            try:
+                response = requests.put(
+                    upload_url,
+                    data=f.read(),
+                    headers={
+                        "Authorization": f"Bearer {blob_token}",
+                        "Content-Type": "video/mp4",
+                        "x-api-version": "7",
+                    },
+                    timeout=600,  # 10 minute timeout for large files
+                )
+                response.raise_for_status()
+                
+                # Extract the blob URL from response
+                result = response.json()
+                blob_url = result.get("url")
+                return blob_url  # Return the full public URL
+            except requests.exceptions.RequestException as e:
+                raise RuntimeError(f"Upload failed: {e}")
+
+    def update_video_url(self, job_id: str, video_url: str) -> bool:
+        """
+        Update the job's video URL after upload completes.
+        
+        The Vercel Blob upload returns the full public URL, which we need
+        to store in the job so the trigger route can use it for download.
+        
+        Args:
+            job_id: The job ID
+            video_url: The full Vercel Blob public URL
+            
+        Returns:
+            True if update succeeded
+            
+        Raises:
+            RuntimeError: If update fails
+        """
+        try:
+            response = self.client.put(f"/scan/jobs/{job_id}", data={"videoUrl": video_url})
+        except SystemExit:
+            raise RuntimeError("Failed to update video URL. Please check your authentication.")
+        
+        if "error" in response:
+            raise RuntimeError(f"Failed to update video URL: {response['error']}")
+        
+        return True
+
+    def start_processing(self, job_id: str) -> bool:
+        """
+        Signal the server to start GPU processing after upload is complete.
+
+        This calls the trigger endpoint which initiates the Modal GPU worker.
+
+        Args:
+            job_id: The job ID
+
+        Returns:
+            True if processing started
+
+        Raises:
+            RuntimeError: If trigger fails
+        """
+        try:
+            response = self.client.post(f"/scan/jobs/{job_id}/trigger")
+        except SystemExit:
+            raise RuntimeError("Failed to trigger processing. Please check your authentication.")
+
+        if "error" in response:
+            raise RuntimeError(f"Failed to start processing: {response['error']}")
+
+        return True
+
+    def poll_job(
+        self,
+        job_id: str,
+        timeout_seconds: int = 600,
+        poll_interval: float = 5.0,
+        status_callback: Optional[callable] = None,
+    ) -> CloudJobResult:
+        """
+        Poll a job until completion or timeout.
+
+        Args:
+            job_id: The job ID to poll
+            timeout_seconds: Maximum time to wait
+            poll_interval: Seconds between polls
+            status_callback: Optional callback(status_string)
+
+        Returns:
+            CloudJobResult with the final status
+        """
+        start_time = time.time()
+
+        while time.time() - start_time < timeout_seconds:
+            try:
+                response = self.client.get(f"/scan/jobs/{job_id}")
+            except SystemExit:
+                return CloudJobResult(
+                    success=False,
+                    job_id=job_id,
+                    error="Failed to poll job status",
+                )
+
+            status = response.get("status", "unknown")
+
+            if status_callback:
+                status_callback(status)
+
+            if status == "completed":
+                return CloudJobResult(
+                    success=True,
+                    job_id=job_id,
+                    artifact_url=response.get("artifactUrl"),
+                    link_count=response.get("linkCount", 0),
+                    joint_count=response.get("jointCount", 0),
+                    processing_time_seconds=response.get("processingTimeSeconds", 0),
+                )
+
+            if status == "failed":
+                return CloudJobResult(
+                    success=False,
+                    job_id=job_id,
+                    error=response.get("error", "Processing failed"),
+                )
+
+            time.sleep(poll_interval)
+
+        return CloudJobResult(
+            success=False,
+            job_id=job_id,
+            error=f"Timeout after {timeout_seconds} seconds",
+        )
+
+    def download_artifact(
+        self,
+        artifact_url: str,
+        output_dir: Path,
+        progress_callback: Optional[callable] = None,
+    ) -> Path:
+        """
+        Download the generated URDF artifact.
+
+        Args:
+            artifact_url: URL to the artifact zip
+            output_dir: Directory to extract to
+            progress_callback: Optional callback(bytes_downloaded, total_bytes)
+
+        Returns:
+            Path to the extracted directory
+
+        Raises:
+            RuntimeError: If download or extraction fails
+        """
+        import zipfile
+        import tempfile
+
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        # Download the artifact
+        try:
+            response = requests.get(artifact_url, stream=True, timeout=300)
+            response.raise_for_status()
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"Download failed: {e}")
+
+        total_size = int(response.headers.get("content-length", 0))
+
+        # Save to temp file
+        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
+            tmp_path = Path(tmp.name)
+            downloaded = 0
+
+            for chunk in response.iter_content(chunk_size=8192):
+                tmp.write(chunk)
+                downloaded += len(chunk)
+                if progress_callback and total_size:
+                    progress_callback(downloaded, total_size)
+
+        # Extract the zip
+        try:
+            with zipfile.ZipFile(tmp_path, "r") as zf:
+                zf.extractall(output_dir)
+        except zipfile.BadZipFile as e:
+            tmp_path.unlink()
+            raise RuntimeError(f"Failed to extract artifact: {e}")
+
+        # Cleanup
+        tmp_path.unlink()
+
+        return output_dir
+
+
+def run_cloud_pipeline(
+    video_path: Optional[str] = None,
+    output_dir: str = "./urdf_output",
+    robot_name: Optional[str] = None,
+    scale_ref: Optional[str] = None,
+    camera_id: int = 0,
+) -> Path:
+    """
+    Run the full cloud pipeline for URDF generation.
+
+    This is the main entry point for cloud-based scanning.
+
+    Args:
+        video_path: Path to existing video file (if None, will capture)
+        output_dir: Output directory for the URDF
+        robot_name: Optional robot name
+        scale_ref: Optional scale reference
+        camera_id: Camera ID for capture (if no video provided)
+
+    Returns:
+        Path to the generated URDF
+
+    Raises:
+        RuntimeError: If any step fails
+    """
+    print("\n" + "#" * 60)
+    print("# CLOUD URDF GENERATION PIPELINE")
+    print("#" * 60)
+
+    client = CloudScanClient()
+
+    # Check credits first
+    credits = client.check_credits()
+    if "error" in credits:
+        raise RuntimeError(f"Failed to check credits: {credits['error']}")
+
+    balance = credits.get("balance", 0)
+    print(f"\nCredit balance: {balance}")
+
+    if balance < 5:
+        raise RuntimeError(
+            f"Insufficient credits ({balance}). Cloud scan requires 5 credits.\n"
+            "Visit https://kindly.fyi/dashboard/billing to purchase credits."
+        )
+
+    # Handle video capture if no video provided
+    if not video_path:
+        print("\nCapturing video for cloud processing...")
+        from .capture import run_capture_only
+
+        video_path = run_capture_only(
+            output_dir=output_dir,
+            camera_id=camera_id,
+            robot_name=robot_name,
+        )
+        print(f"Captured: {video_path}")
+
+    video_path = Path(video_path)
+    if not video_path.exists():
+        raise RuntimeError(f"Video file not found: {video_path}")
+
+    print(f"\nVideo: {video_path} ({video_path.stat().st_size / 1024 / 1024:.1f} MB)")
+
+    # Create job
+    print("\nCreating cloud job...")
+    job_id, upload_url, blob_token, credits_info = client.create_job(robot_name, scale_ref)
+    print(f"Job ID: {job_id}")
+    print(f"Cost: {credits_info.get('cost', 5)} credits")
+
+    # Upload video
+    print("\nUploading video...")
+
+    def upload_progress(sent, total):
+        pct = int(100 * sent / total)
+        bar_len = 30
+        filled = int(bar_len * sent / total)
+        bar = "=" * filled + "-" * (bar_len - filled)
+        print(f"\r  [{bar}] {pct}% ({sent / 1024 / 1024:.1f} / {total / 1024 / 1024:.1f} MB)", end="", flush=True)
+
+    blob_url = client.upload_video(video_path, upload_url, blob_token, upload_progress)
+    print("\n  Upload complete!")
+    
+    # Update job with full blob URL for download
+    if blob_url:
+        client.update_video_url(job_id, blob_url)
+        print(f"  Stored blob URL: {blob_url[:60]}...")
+
+    # Trigger GPU processing
+    print("\nStarting GPU processing...")
+    client.start_processing(job_id)
+    print("  Processing triggered on cloud GPU (A10G)")
+
+    # Poll for completion
+    print("\nWaiting for processing to complete...")
+
+    last_status = None
+    def status_callback(status):
+        nonlocal last_status
+        if status != last_status:
+            print(f"  Status: {status}")
+            last_status = status
+
+    result = client.poll_job(job_id, status_callback=status_callback)
+
+    if not result.success:
+        raise RuntimeError(f"Cloud processing failed: {result.error}")
+
+    print(f"\nProcessing complete in {result.processing_time_seconds:.1f}s")
+    print(f"  Links: {result.link_count}")
+    print(f"  Joints: {result.joint_count}")
+
+    # Download artifact
+    print("\nDownloading artifact...")
+    output_path = Path(output_dir)
+
+    def download_progress(downloaded, total):
+        pct = int(100 * downloaded / total)
+        bar_len = 30
+        filled = int(bar_len * downloaded / total)
+        bar = "=" * filled + "-" * (bar_len - filled)
+        print(f"\r  [{bar}] {pct}%", end="", flush=True)
+
+    client.download_artifact(result.artifact_url, output_path, download_progress)
+    print("\n  Download complete!")
+
+    # Find the URDF file
+    urdf_files = list(output_path.glob("*.urdf"))
+    if urdf_files:
+        urdf_path = urdf_files[0]
+    else:
+        urdf_path = output_path / f"{robot_name or 'robot'}.urdf"
+
+    print("\n" + "#" * 60)
+    print("# CLOUD PIPELINE COMPLETE")
+    print("#" * 60)
+    print(f"\nOutput: {output_path}")
+    print(f"URDF: {urdf_path}")
+
+    return urdf_path
+
+
+__all__ = [
+    "CloudScanClient",
+    "CloudJobResult",
+    "run_cloud_pipeline",
+]