blaxel 0.2.26rc119__py3-none-any.whl → 0.2.27__py3-none-any.whl

blaxel/core/common/settings.py

@@ -1,11 +1,33 @@
 import os
 import platform
+from pathlib import Path
 from typing import Dict
 
+import tomli
+
 from ..authentication import BlaxelAuth, auth
 from .logger import init_logger
 
 
+def _get_package_version() -> str:
+    """Get the package version from pyproject.toml."""
+    try:
+        # Navigate up from this file to the project root
+        current_file = Path(__file__)
+        project_root = current_file.parent.parent.parent.parent.parent
+        pyproject_path = project_root / "pyproject.toml"
+
+        if pyproject_path.exists():
+            with open(pyproject_path, "rb") as f:
+                pyproject_data = tomli.load(f)
+                return pyproject_data.get("project", {}).get("version", "unknown")
+        else:
+            return "unknown"
+    except Exception as e:
+        print(f"Warning: Failed to read package version: {e}")
+        return "unknown"
+
+
 def _get_os_arch() -> str:
     """Get OS and architecture information."""
     try:
@@ -34,6 +56,33 @@ def _get_os_arch() -> str:
         return "unknown/unknown"
 
 
+def _get_commit_hash() -> str:
+    """Get commit hash from pyproject.toml."""
+    try:
+        # Try to read from pyproject.toml (build-time injection)
+        current_file = Path(__file__)
+        project_root = current_file.parent.parent.parent.parent.parent
+        pyproject_path = project_root / "pyproject.toml"
+
+        if pyproject_path.exists():
+            with open(pyproject_path, "rb") as f:
+                pyproject_data = tomli.load(f)
+                # Check multiple possible locations for commit
+                commit = None
+                if "tool" in pyproject_data and "blaxel" in pyproject_data["tool"]:
+                    commit = pyproject_data["tool"]["blaxel"].get("commit")
+                if not commit and "project" in pyproject_data:
+                    commit = pyproject_data["project"].get("commit")
+                if not commit and "build" in pyproject_data:
+                    commit = pyproject_data["build"].get("commit")
+
+                if commit:
+                    return commit[:7] if len(commit) > 7 else commit
+    except Exception:
+        pass
+
+    return "unknown"
+
 class Settings:
     auth: BlaxelAuth
 
@@ -41,6 +90,7 @@ class Settings:
         init_logger(self.log_level)
         self.auth = auth(self.env, self.base_url)
         self._headers = None
+        self._version = None
 
     @property
     def env(self) -> str:
@@ -67,30 +117,20 @@ class Settings:
         return "https://run.blaxel.dev"
 
 
-    @property
-    def sentry_dsn(self) -> str:
-        """Get the Sentry DSN (injected at build time)."""
-        import blaxel
-        return blaxel.__sentry_dsn__
-
     @property
     def version(self) -> str:
-        """Get the package version (injected at build time)."""
-        import blaxel
-        return blaxel.__version__ or "unknown"
-
-    @property
-    def commit(self) -> str:
-        """Get the commit hash (injected at build time)."""
-        import blaxel
-        return blaxel.__commit__ or "unknown"
+        """Get the package version."""
+        if self._version is None:
+            self._version = _get_package_version()
+        return self._version
 
     @property
     def headers(self) -> Dict[str, str]:
         """Get the headers for API requests."""
         headers = self.auth.get_headers()
         os_arch = _get_os_arch()
-        headers["User-Agent"] = f"blaxel/sdk/python/{self.version} ({os_arch}) blaxel/{self.commit}"
+        commit_hash = _get_commit_hash()
+        headers["User-Agent"] = f"blaxel/sdk/python/{self.version} ({os_arch}) blaxel/{commit_hash}"
         return headers
 
 

blaxel/core/common/webhook.py

@@ -0,0 +1,187 @@
+"""Webhook signature verification for async-sidecar callbacks."""
+
+import hashlib
+import hmac
+import time
+from typing import Optional, Protocol
+
+
+class RequestLike(Protocol):
+    """Protocol for request-like objects with body and headers."""
+
+    @property
+    def body(self) -> bytes:
+        """Raw request body as bytes."""
+        ...
+
+    @property
+    def headers(self) -> dict[str, str]:
+        """Request headers as dictionary."""
+        ...
+
+
+class AsyncSidecarCallback:
+    """Callback payload from async-sidecar."""
+
+    def __init__(
+        self,
+        status_code: int,
+        response_body: str,
+        response_length: int,
+        timestamp: int,
+    ):
+        self.status_code = status_code
+        self.response_body = response_body
+        self.response_length = response_length
+        self.timestamp = timestamp
+
+
+def verify_webhook_signature(
+    body: bytes | str,
+    signature: str,
+    secret: str,
+    timestamp: Optional[str] = None,
+    max_age: int = 300,
+) -> bool:
+    """
+    Verify the HMAC-SHA256 signature of a webhook callback from async-sidecar.
+
+    Args:
+        body: The raw request body (bytes or string)
+        signature: The X-Blaxel-Signature header value (format: "sha256=<hex_digest>")
+        secret: The secret key used to sign the webhook (same as CALLBACK_SECRET in async-sidecar)
+        timestamp: Optional X-Blaxel-Timestamp header value for replay attack prevention
+        max_age: Maximum age of the webhook in seconds (default: 300 = 5 minutes)
+
+    Returns:
+        True if the signature is valid, False otherwise
+
+    Example:
+        ```python
+        from blaxel.core import verify_webhook_signature
+        from flask import Flask, request
+
+        app = Flask(__name__)
+
+        @app.route('/webhook', methods=['POST'])
+        def webhook():
+            is_valid = verify_webhook_signature(
+                body=request.get_data(),
+                signature=request.headers.get('X-Blaxel-Signature', ''),
+                secret='your-callback-secret'
+            )
+
+            if not is_valid:
+                return {'error': 'Invalid signature'}, 401
+
+            data = request.json
+            # Process callback...
+            return {'received': True}
+        ```
+    """
+    if not body or not signature or not secret:
+        return False
+
+    try:
+        # Verify timestamp if provided (prevents replay attacks)
+        if timestamp:
+            request_time = int(timestamp)
+            current_time = int(time.time())
+            age = abs(current_time - request_time)
+
+            if age > max_age:
+                return False
+
+        # Convert body to bytes if string
+        body_bytes = body.encode("utf-8") if isinstance(body, str) else body
+
+        # Extract hex signature from "sha256=<hex>" format
+        expected_signature = signature.replace("sha256=", "")
+
+        # Compute HMAC-SHA256 signature
+        computed_signature = hmac.new(
+            secret.encode("utf-8"), body_bytes, hashlib.sha256
+        ).hexdigest()
+
+        # Timing-safe comparison to prevent timing attacks
+        return hmac.compare_digest(expected_signature, computed_signature)
+
+    except (ValueError, TypeError):
+        # Invalid signature format or other error
+        return False
+
+
+def verify_webhook_from_request(
+    request: RequestLike,
+    secret: str,
+    max_age: int = 300,
+) -> bool:
+    """
+    Helper to verify webhook from a request object (Flask, FastAPI, etc.).
+
+    Args:
+        request: Request object with `body` and `headers` attributes
+        secret: The callback secret
+        max_age: Optional maximum age in seconds (default: 300)
+
+    Returns:
+        True if the signature is valid, False otherwise
+
+    Example with Flask:
+        ```python
+        from blaxel.core import verify_webhook_from_request
+        from flask import Flask, request
+
+        app = Flask(__name__)
+
+        @app.route('/webhook', methods=['POST'])
+        def webhook():
+            if not verify_webhook_from_request(request, 'your-callback-secret'):
+                return {'error': 'Invalid signature'}, 401
+
+            data = request.json
+            print(f"Received callback: {data}")
+            return {'received': True}
+        ```
+
+    Example with FastAPI:
+        ```python
+        from blaxel.core import verify_webhook_signature
+        from fastapi import FastAPI, Request, HTTPException
+
+        app = FastAPI()
+
+        @app.post('/webhook')
+        async def webhook(request: Request):
+            body = await request.body()
+            signature = request.headers.get('x-blaxel-signature', '')
+
+            if not verify_webhook_signature(body, signature, 'your-callback-secret'):
+                raise HTTPException(status_code=401, detail='Invalid signature')
+
+            data = await request.json()
+            return {'received': True}
+        ```
+    """
+    signature = request.headers.get("x-blaxel-signature", "")
+    timestamp = request.headers.get("x-blaxel-timestamp")
+
+    if not signature:
+        return False
+
+    return verify_webhook_signature(
+        body=request.body,
+        signature=signature,
+        secret=secret,
+        timestamp=timestamp,
+        max_age=max_age,
+    )
+
+
+__all__ = [
+    "verify_webhook_signature",
+    "verify_webhook_from_request",
+    "AsyncSidecarCallback",
+    "RequestLike",
+]
+

blaxel/core/jobs/__init__.py

@@ -2,12 +2,21 @@ import argparse
 import asyncio
 import os
 import sys
+import time
 from logging import getLogger
-from typing import Any, Awaitable, Callable, Dict
+from typing import Any, Awaitable, Callable, Dict, List
 
 import requests
 
 from ..client import client
+from ..client.api.jobs import (
+    create_job_execution,
+    delete_job_execution,
+    get_job_execution,
+    list_job_executions,
+)
+from ..client.models.create_job_execution_request import CreateJobExecutionRequest
+from ..client.models.job_execution import JobExecution
 from ..common.internal import get_forced_url, get_global_unique_hash
 from ..common.settings import settings
 
@@ -96,7 +105,7 @@ class BlJob:
 
     def call(self, url, input_data, headers: dict = {}, params: dict = {}):
         body = {"tasks": input_data}
-        
+
         # Merge settings headers with provided headers
         merged_headers = {**settings.headers, "Content-Type": "application/json", **headers}
 
@@ -110,7 +119,7 @@ class BlJob:
     async def acall(self, url, input_data, headers: dict = {}, params: dict = {}):
         logger.debug(f"Job Calling: {self.name}")
         body = {"tasks": input_data}
-        
+
         # Merge settings headers with provided headers
         merged_headers = {**settings.headers, "Content-Type": "application/json", **headers}
 
@@ -151,6 +160,346 @@ class BlJob:
                 )
         return response.text
 
+    def create_execution(self, request: CreateJobExecutionRequest) -> str:
+        """
+        Create a new execution for this job and return the execution ID.
+
+        Args:
+            request: The job execution request containing tasks and optional execution ID
+
+        Returns:
+            str: The execution ID
+
+        Raises:
+            Exception: If no execution ID is returned or the request fails
+        """
+        logger.debug(f"Creating execution for job: {self.name}")
+
+        response = create_job_execution.sync_detailed(
+            job_id=self.name,
+            client=client,
+            body=request,
+        )
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to create job execution: {response.status_code}")
+
+        # The API returns executionId at the root level
+        if response.parsed and hasattr(response.parsed, "to_dict"):
+            response_dict = response.parsed.to_dict()
+        else:
+            # Parse the raw response to get executionId
+            import json
+
+            response_dict = json.loads(response.content)
+
+        # Check for both camelCase (API) and snake_case (parsed) formats
+        execution_id = response_dict.get("execution_id") or response_dict.get("executionId")
+        if not execution_id:
+            raise Exception("No execution ID returned from create job execution")
+
+        logger.debug(f"Created execution: {execution_id}")
+        return execution_id
+
+    async def acreate_execution(self, request: CreateJobExecutionRequest) -> str:
+        """
+        Create a new execution for this job and return the execution ID (async).
+
+        Args:
+            request: The job execution request containing tasks and optional execution ID
+
+        Returns:
+            str: The execution ID
+
+        Raises:
+            Exception: If no execution ID is returned or the request fails
+        """
+        logger.debug(f"Creating execution for job: {self.name}")
+
+        response = await create_job_execution.asyncio_detailed(
+            job_id=self.name,
+            client=client,
+            body=request,
+        )
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to create job execution: {response.status_code}")
+
+        # The API returns executionId at the root level
+        if response.parsed and hasattr(response.parsed, "to_dict"):
+            response_dict = response.parsed.to_dict()
+        else:
+            # Parse the raw response to get executionId
+            import json
+
+            response_dict = json.loads(response.content)
+
+        # Check for both camelCase (API) and snake_case (parsed) formats
+        execution_id = response_dict.get("execution_id") or response_dict.get("executionId")
+        if not execution_id:
+            raise Exception("No execution ID returned from create job execution")
+
+        logger.debug(f"Created execution: {execution_id}")
+        return execution_id
+
+    def get_execution(self, execution_id: str) -> JobExecution:
+        """
+        Get a specific execution by ID.
+
+        Args:
+            execution_id: The execution ID
+
+        Returns:
+            JobExecution: The execution object
+
+        Raises:
+            Exception: If the execution is not found or the request fails
+        """
+        logger.debug(f"Getting execution {execution_id} for job: {self.name}")
+
+        response = get_job_execution.sync_detailed(
+            job_id=self.name,
+            execution_id=execution_id,
+            client=client,
+        )
+
+        if response.status_code == 404:
+            raise Exception(f"Execution '{execution_id}' not found for job '{self.name}'")
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to get job execution: {response.status_code}")
+
+        if not response.parsed:
+            raise Exception("No execution data returned")
+
+        return response.parsed
+
+    async def aget_execution(self, execution_id: str) -> JobExecution:
+        """
+        Get a specific execution by ID (async).
+
+        Args:
+            execution_id: The execution ID
+
+        Returns:
+            JobExecution: The execution object
+
+        Raises:
+            Exception: If the execution is not found or the request fails
+        """
+        logger.debug(f"Getting execution {execution_id} for job: {self.name}")
+
+        response = await get_job_execution.asyncio_detailed(
+            job_id=self.name,
+            execution_id=execution_id,
+            client=client,
+        )
+
+        if response.status_code == 404:
+            raise Exception(f"Execution '{execution_id}' not found for job '{self.name}'")
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to get job execution: {response.status_code}")
+
+        if not response.parsed:
+            raise Exception("No execution data returned")
+
+        return response.parsed
+
+    def list_executions(self, limit: int = 20, offset: int = 0) -> List[JobExecution]:
+        """
+        List all executions for this job.
+
+        Args:
+            limit: Maximum number of executions to return
+            offset: Offset for pagination
+
+        Returns:
+            List[JobExecution]: List of execution objects
+        """
+        logger.debug(f"Listing executions for job: {self.name}")
+
+        response = list_job_executions.sync_detailed(
+            job_id=self.name,
+            client=client,
+            limit=limit,
+            offset=offset,
+        )
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to list job executions: {response.status_code}")
+
+        return response.parsed or []
+
+    async def alist_executions(self, limit: int = 20, offset: int = 0) -> List[JobExecution]:
+        """
+        List all executions for this job (async).
+
+        Args:
+            limit: Maximum number of executions to return
+            offset: Offset for pagination
+
+        Returns:
+            List[JobExecution]: List of execution objects
+        """
+        logger.debug(f"Listing executions for job: {self.name}")
+
+        response = await list_job_executions.asyncio_detailed(
+            job_id=self.name,
+            client=client,
+            limit=limit,
+            offset=offset,
+        )
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to list job executions: {response.status_code}")
+
+        return response.parsed or []
+
+    def get_execution_status(self, execution_id: str) -> str:
+        """
+        Get the status of a specific execution.
+
+        Args:
+            execution_id: The execution ID
+
+        Returns:
+            str: The execution status
+        """
+        execution = self.get_execution(execution_id)
+        return execution.status if execution.status else "UNKNOWN"
+
+    async def aget_execution_status(self, execution_id: str) -> str:
+        """
+        Get the status of a specific execution (async).
+
+        Args:
+            execution_id: The execution ID
+
+        Returns:
+            str: The execution status
+        """
+        execution = await self.aget_execution(execution_id)
+        return execution.status if execution.status else "UNKNOWN"
+
+    def cancel_execution(self, execution_id: str) -> None:
+        """
+        Cancel a specific execution.
+
+        Args:
+            execution_id: The execution ID
+
+        Raises:
+            Exception: If the cancellation fails
+        """
+        logger.debug(f"Cancelling execution {execution_id} for job: {self.name}")
+
+        response = delete_job_execution.sync_detailed(
+            job_id=self.name,
+            execution_id=execution_id,
+            client=client,
+        )
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to cancel job execution: {response.status_code}")
+
+    async def acancel_execution(self, execution_id: str) -> None:
+        """
+        Cancel a specific execution (async).
+
+        Args:
+            execution_id: The execution ID
+
+        Raises:
+            Exception: If the cancellation fails
+        """
+        logger.debug(f"Cancelling execution {execution_id} for job: {self.name}")
+
+        response = await delete_job_execution.asyncio_detailed(
+            job_id=self.name,
+            execution_id=execution_id,
+            client=client,
+        )
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to cancel job execution: {response.status_code}")
+
+    def wait_for_execution(
+        self,
+        execution_id: str,
+        max_wait: int = 360,
+        interval: int = 3,
+    ) -> JobExecution:
+        """
+        Wait for an execution to complete.
+
+        Args:
+            execution_id: The execution ID to wait for
+            max_wait: Maximum time to wait in seconds (default: 360 = 6 minutes)
+            interval: Polling interval in seconds (default: 3 seconds)
+
+        Returns:
+            JobExecution: The completed execution
+
+        Raises:
+            Exception: If the execution doesn't complete within max_wait seconds
+        """
+        logger.debug(f"Waiting for execution {execution_id} to complete (max {max_wait}s)")
+
+        start_time = time.time()
+
+        while time.time() - start_time < max_wait:
+            execution = self.get_execution(execution_id)
+            status = execution.status
+
+            # Terminal states (Kubernetes-style: succeeded, failed, cancelled)
+            if status in ["succeeded", "failed", "cancelled"]:
+                logger.debug(f"Execution {execution_id} finished with status: {status}")
+                return execution
+
+            # Wait before polling again
+            time.sleep(interval)
+
+        raise Exception(f"Execution {execution_id} did not complete within {max_wait}s")
+
+    async def await_for_execution(
+        self,
+        execution_id: str,
+        max_wait: int = 360,
+        interval: int = 3,
+    ) -> JobExecution:
+        """
+        Wait for an execution to complete (async).
+
+        Args:
+            execution_id: The execution ID to wait for
+            max_wait: Maximum time to wait in seconds (default: 360 = 6 minutes)
+            interval: Polling interval in seconds (default: 3 seconds)
+
+        Returns:
+            JobExecution: The completed execution
+
+        Raises:
+            Exception: If the execution doesn't complete within max_wait seconds
+        """
+        logger.debug(f"Waiting for execution {execution_id} to complete (max {max_wait}s)")
+
+        start_time = time.time()
+
+        while time.time() - start_time < max_wait:
+            execution = await self.aget_execution(execution_id)
+            status = execution.status
+
+            # Terminal states (Kubernetes-style: succeeded, failed, cancelled)
+            if status in ["succeeded", "failed", "cancelled"]:
+                logger.debug(f"Execution {execution_id} finished with status: {status}")
+                return execution
+
+            # Wait before polling again
+            await asyncio.sleep(interval)
+
+        raise Exception(f"Execution {execution_id} did not complete within {max_wait}s")
+
     def __str__(self):
         return f"Job {self.name}"