scorebook 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

scorebook/inference/bedrock.py

@@ -0,0 +1,305 @@
+"""
+AWS Bedrock batch inference implementation for Scorebook.
+
+This module provides utilities for running batch inference using AWS Bedrock's
+Model Invocation Jobs, supporting large-scale asynchronous processing. It handles
+API communication, request formatting, response processing, and S3 operations.
+"""
+
+import asyncio
+import json
+import os
+import tempfile
+import uuid
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Tuple
+
+import boto3
+from botocore.config import Config
+from botocore.exceptions import ClientError
+from tqdm.asyncio import tqdm
+
+
+async def batch(
+    items: List[Any],
+    model: Optional[str] = None,
+    aws_region: Optional[str] = None,
+    aws_profile: Optional[str] = None,
+    bucket: Optional[str] = None,
+    input_prefix: Optional[str] = None,
+    output_prefix: Optional[str] = None,
+    role_arn: Optional[str] = None,
+    **hyperparameters: Any,
+) -> List[Any]:
+    """Process multiple inference requests in batch using AWS Bedrock.
+
+    This asynchronous function handles batch processing of inference requests,
+    optimizing for cost and throughput using AWS Bedrock's Model Invocation Jobs.
+
+    Args:
+        items: List of preprocessed items to process.
+        model: Bedrock model ID (e.g., 'us.anthropic.claude-3-5-sonnet-20241022-v2:0').
+        aws_region: AWS region for Bedrock and S3.
+        aws_profile: AWS profile name for authentication.
+        bucket: S3 bucket name for input/output data.
+        input_prefix: S3 prefix for input data.
+        output_prefix: S3 prefix for output data.
+        role_arn: IAM role ARN for Bedrock execution.
+        hyperparameters: Additional parameters for the batch requests.
+
+    Returns:
+        A list of raw model responses.
+    """
+    # Set up AWS session and clients
+    session_kwargs = {}
+    if aws_profile:
+        session_kwargs["profile_name"] = aws_profile
+    if aws_region:
+        session_kwargs["region_name"] = aws_region
+
+    session = boto3.Session(**session_kwargs)
+
+    boto_config = Config(region_name=aws_region, retries={"max_attempts": 10, "mode": "adaptive"})
+
+    s3_client = session.client("s3", config=boto_config)
+    bedrock_client = session.client("bedrock", config=boto_config)
+
+    # Upload batch data
+    input_uri = await _upload_batch(
+        items, s3_client, bucket, input_prefix, model, **hyperparameters
+    )
+
+    # Start batch job
+    job_arn = await _start_batch_job(
+        bedrock_client, model, input_uri, bucket, output_prefix, role_arn
+    )
+
+    # Wait for completion with progress tracking
+    await _wait_for_completion(bedrock_client, job_arn, len(items))
+
+    # Retrieve results
+    results = await _get_batch_results(s3_client, bedrock_client, job_arn)
+
+    return results
+
+
+async def _upload_batch(
+    items: List[Any],
+    s3_client: Any,
+    bucket: Optional[str],
+    input_prefix: Optional[str],
+    model: Optional[str],
+    **hyperparameters: Any,
+) -> str:
+    """Create a JSONL file from preprocessed items and upload to S3 for batch processing."""
+
+    # Generate unique run ID and key
+    run_id = datetime.utcnow().strftime("%Y%m%dT%H%M%S") + "-" + uuid.uuid4().hex[:8]
+
+    if input_prefix:
+        input_key = f"{input_prefix.rstrip('/')}/inputs-{run_id}.jsonl"
+    else:
+        input_key = f"inputs-{run_id}.jsonl"
+
+    # Create temp JSONL file
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f:
+        for i, item in enumerate(items):
+            # Construct batch request in Bedrock format
+            record = {
+                "recordId": f"rec-{i:04d}",
+                "modelInput": _build_claude_messages_payload(item, **hyperparameters),
+            }
+            f.write(json.dumps(record, separators=(",", ":")) + "\n")
+        file_path = f.name
+
+    # Upload to S3
+    try:
+        body = open(file_path, "rb").read()
+        s3_client.put_object(
+            Bucket=bucket,
+            Key=input_key,
+            Body=body,
+            StorageClass="INTELLIGENT_TIERING",
+            ContentType="application/json",
+        )
+        input_uri = f"s3://{bucket}/{input_key}"
+    except Exception as e:
+        raise Exception(f"Failed to upload file to S3: {e}")
+    finally:
+        # Clean up temp file
+        os.unlink(file_path)
+
+    return input_uri
+
+
+def _build_claude_messages_payload(item: Any, **hyperparameters: Any) -> Dict[str, Any]:
+    """Build Claude messages payload for Bedrock batch processing."""
+
+    # item is a list of messages from our preprocessor
+    messages = item
+
+    # Convert to Bedrock format and extract system message
+    bedrock_messages = []
+    system_content = None
+
+    for msg in messages:
+        if msg["role"] == "system":
+            system_content = msg["content"]
+        else:
+            bedrock_messages.append(
+                {"role": msg["role"], "content": [{"type": "text", "text": msg["content"]}]}
+            )
+
+    payload = {
+        "anthropic_version": "bedrock-2023-05-31",
+        "max_tokens": 256,
+        "messages": bedrock_messages,
+    }
+
+    if system_content:
+        payload["system"] = system_content
+
+    payload.update(hyperparameters)
+    return payload
+
+
+async def _start_batch_job(
+    bedrock_client: Any,
+    model: Optional[str],
+    input_uri: str,
+    bucket: Optional[str],
+    output_prefix: Optional[str],
+    role_arn: Optional[str],
+) -> str:
+    """Start a Bedrock Model Invocation Job."""
+
+    # Generate unique job name and output URI
+    run_id = datetime.utcnow().strftime("%Y%m%dT%H%M%S") + "-" + uuid.uuid4().hex[:8]
+    job_name = f"bedrock-batch-{run_id}"
+
+    if output_prefix:
+        output_uri = f"s3://{bucket}/{output_prefix.rstrip('/')}/job-{run_id}/"
+    else:
+        output_uri = f"s3://{bucket}/job-{run_id}/"
+
+    try:
+        response = bedrock_client.create_model_invocation_job(
+            jobName=job_name,
+            modelId=model,
+            roleArn=role_arn,
+            inputDataConfig={"s3InputDataConfig": {"s3Uri": input_uri}},
+            outputDataConfig={"s3OutputDataConfig": {"s3Uri": output_uri}},
+            tags=[{"key": "project", "value": "scorebook-batch"}],
+        )
+        job_arn: str = response["jobArn"]
+        return job_arn
+    except ClientError as e:
+        error_info = e.response.get("Error", {})
+        raise Exception(f"Failed to create batch job: {error_info}")
+
+
+async def _wait_for_completion(bedrock_client: Any, job_arn: str, total_items: int) -> None:
+    """Wait for batch job completion with progress tracking."""
+
+    # Initialize progress bar
+    pbar = tqdm(total=total_items, desc="Batch processing", unit="requests")
+
+    terminal_states = {"Completed", "Failed", "Stopped"}
+    sleep_time = 15
+
+    while True:
+        try:
+            desc = bedrock_client.get_model_invocation_job(jobIdentifier=job_arn)
+            status = desc["status"]
+
+            # Get progress if available
+            job_state = desc.get("jobState", {})
+            progress = job_state.get("percentComplete")
+
+            # Update progress bar
+            if progress is not None:
+                completed = int((progress / 100) * total_items)
+                pbar.n = completed
+                pbar.set_postfix(status=status, progress=f"{progress}%")
+            else:
+                pbar.set_postfix(status=status)
+
+            pbar.refresh()
+
+            if status in terminal_states:
+                if status == "Completed":
+                    pbar.n = pbar.total
+                    pbar.set_postfix(status="COMPLETED")
+                else:
+                    pbar.close()
+                    error_msg = desc.get("failureMessage", f"Job ended with status {status}")
+                    raise Exception(f"Batch job failed: {error_msg}")
+                break
+
+            # Wait before checking again
+            await asyncio.sleep(sleep_time)
+
+        except Exception as e:
+            pbar.close()
+            raise e
+
+    pbar.close()
+
+
+async def _get_batch_results(s3_client: Any, bedrock_client: Any, job_arn: str) -> List[str]:
+    """Download and parse batch results from S3."""
+
+    # Get job details to find output location
+    desc = bedrock_client.get_model_invocation_job(jobIdentifier=job_arn)
+    output_uri = desc["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"]
+
+    bucket_name, prefix = s3_uri_to_bucket_and_prefix(output_uri)
+
+    # Find the output JSONL file
+    try:
+        response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=prefix)
+        contents = response.get("Contents", [])
+
+        # Look for the output JSONL file
+        output_key = None
+        for obj in contents:
+            if obj["Key"].endswith(".jsonl.out"):
+                output_key = obj["Key"]
+                break
+
+        if not output_key:
+            raise Exception("No output JSONL file found in S3")
+
+        # Download and parse results
+        obj_response = s3_client.get_object(Bucket=bucket_name, Key=output_key)
+        content = obj_response["Body"].read().decode("utf-8")
+
+        results = []
+        for line in content.strip().split("\n"):
+            if line.strip():
+                result_obj = json.loads(line)
+                # Extract text from Claude response format
+                model_output = result_obj.get("modelOutput", {})
+                content_list = model_output.get("content", [])
+                if content_list and len(content_list) > 0:
+                    text = content_list[0].get("text", "")
+                    results.append(text)
+                else:
+                    results.append("")
+
+        return results
+
+    except Exception as e:
+        raise Exception(f"Failed to retrieve batch results: {e}")
+
+
+def s3_uri_to_bucket_and_prefix(s3_uri: str) -> Tuple[str, str]:
+    """Parse S3 URI to bucket and prefix."""
+    # Parse S3 URI
+    if s3_uri.startswith("s3://"):
+        uri_parts = s3_uri[5:].split("/", 1)
+        bucket_name = uri_parts[0]
+        prefix = uri_parts[1] if len(uri_parts) > 1 else ""
+    else:
+        raise ValueError(f"Invalid S3 URI: {s3_uri}")
+    return bucket_name, prefix

scorebook/inference/openai.py

@@ -8,17 +8,19 @@ API communication, request formatting, and response processing.
 
 import asyncio
 import json
+import logging
 import tempfile
 from typing import Any, List
 
-from openai import OpenAI
-from tqdm.asyncio import tqdm
+from openai import AsyncOpenAI
+
+logger = logging.getLogger(__name__)
 
 
 async def responses(
     items: List[Any], model: str = "gpt-4.1-nano", client: Any = None, **hyperparameters: Any
 ) -> List[Any]:
-    """Process multiple inference requests using OpenAI's API.
+    """Process multiple inference requests using OpenAI's Async API.
 
     This asynchronous function handles multiple inference requests,
     manages the API communication, and processes the responses.
@@ -35,13 +37,67 @@ async def responses(
     Raises:
         NotImplementedError: Currently not implemented.
     """
-    if client is None:
-        client = OpenAI()
+    logger.debug("OpenAI responses function called with %d items", len(items))
+    logger.debug("Using model: %s", model)
+    logger.debug("Hyperparameters: %s", hyperparameters)
 
-    results = []
-    for item in items:
-        response = client.responses.create(model=model, input=item)
-        results.append(response)
+    if client is None:
+        logger.debug("Creating new AsyncOpenAI client")
+        client = AsyncOpenAI()
+
+    # Create all tasks concurrently for true parallelism
+    tasks = []
+    for i, item in enumerate(items):
+        logger.debug(
+            "Processing item %d: %s",
+            i,
+            str(item)[:100] + "..." if len(str(item)) > 100 else str(item),
+        )
+
+        # Handle string input from preprocessor - convert to proper messages format
+        if isinstance(item, str):
+            # Convert the string format to proper OpenAI messages array
+            messages = [{"role": "user", "content": item}]
+            logger.debug(
+                "Converted string to messages format: %s",
+                (
+                    messages[0]["content"][:100] + "..."
+                    if len(messages[0]["content"]) > 100
+                    else messages[0]["content"]
+                ),
+            )
+        elif isinstance(item, list):
+            # Already in proper messages format
+            messages = item
+            logger.debug("Item %d already in messages format", i)
+        else:
+            # Fallback: treat as user message
+            messages = [{"role": "user", "content": str(item)}]
+            logger.debug("Item %d converted to fallback format", i)
+
+        logger.debug("Creating OpenAI task %d with messages: %s", i, messages)
+        task = client.chat.completions.create(model=model, messages=messages, **hyperparameters)
+        tasks.append(task)
+
+    logger.debug("Created %d tasks, waiting for OpenAI responses...", len(tasks))
+    # Wait for all requests to complete in parallel
+    results = await asyncio.gather(*tasks)
+    logger.debug("Received %d responses from OpenAI", len(results))
+
+    for i, result in enumerate(results):
+        logger.debug("Response %d type: %s", i, type(result))
+        try:
+            if hasattr(result, "choices") and result.choices:
+                content = result.choices[0].message.content
+                logger.debug(
+                    "Response %d content: %s",
+                    i,
+                    content[:100] + "..." if content and len(content) > 100 else content,
+                )
+            else:
+                logger.debug("Response %d has no choices or unexpected format", i)
+        except Exception as e:
+            logger.error("Error logging response %d: %s", i, e)
 
     return results
 
@@ -70,40 +126,23 @@ async def batch(
         NotImplementedError: Currently not implemented.
     """
     if client is None:
-        client = OpenAI()
+        client = AsyncOpenAI()
 
-    file_id = _upload_batch(items, client)
-    batch_id = _start_batch(file_id, client)
-
-    # Initialize progress bar
-    pbar = tqdm(total=len(items), desc="Batch processing", unit="requests")
+    file_id = await _upload_batch(items, client)
+    batch_id = await _start_batch(file_id, client)
 
     awaiting_batch = True
     while awaiting_batch:
         batch_object = await _get_batch(batch_id, client)
         batch_status = batch_object.status
 
-        if hasattr(batch_object, "request_counts") and batch_object.request_counts:
-            completed = batch_object.request_counts.completed
-            total = batch_object.request_counts.total
-            pbar.n = completed
-            pbar.set_postfix(status=batch_status, completed=f"{completed}/{total}")
-        else:
-            pbar.set_postfix(status=batch_status)
-
-        pbar.refresh()
-
         if batch_status == "completed":
             awaiting_batch = False
-            pbar.n = pbar.total
-            pbar.set_postfix(status="completed")
         elif batch_status == "failed":
             raise Exception("Batch processing failed")
         else:
             await asyncio.sleep(60)
 
-    pbar.close()
-
     # Get the final batch object to access output_file_id
     final_batch_object = await _get_batch(batch_id, client)
     output_file_id = final_batch_object.output_file_id
@@ -112,7 +151,7 @@ async def batch(
     return batch_result
 
 
-def _upload_batch(items: List[Any], client: Any) -> str:
+async def _upload_batch(items: List[Any], client: Any) -> str:
     """Create a .jsonl file from preprocessed items and upload to OpenAI for batch processing.
 
     Args:
@@ -121,10 +160,9 @@ def _upload_batch(items: List[Any], client: Any) -> str:
     Returns:
         The file ID returned by OpenAI after uploading.
     """
-    print("Uploading batch...")
     # Instantiate OpenAI client
     if client is None:
-        client = OpenAI()
+        client = AsyncOpenAI()
 
     # Create temp .jsonl file
     with tempfile.NamedTemporaryFile(mode="w+", suffix=".jsonl", delete=False) as f:
@@ -141,13 +179,13 @@ def _upload_batch(items: List[Any], client: Any) -> str:
 
     # Upload file to OpenAI
     with open(file_path, "rb") as upload_file:
-        response = client.files.create(file=upload_file, purpose="batch")
+        response = await client.files.create(file=upload_file, purpose="batch")
 
     return str(response.id)
 
 
-def _start_batch(file_id: str, client: Any) -> str:
-    batch_response = client.batches.create(
+async def _start_batch(file_id: str, client: Any) -> str:
+    batch_response = await client.batches.create(
         input_file_id=file_id,
         endpoint="/v1/chat/completions",
         completion_window="24h",
@@ -156,13 +194,13 @@ def _start_batch(file_id: str, client: Any) -> str:
 
 
 async def _get_batch(batch_id: str, client: Any) -> Any:
-    batch_object = client.batches.retrieve(batch_id)
+    batch_object = await client.batches.retrieve(batch_id)
     return batch_object
 
 
 async def _get_results_file(output_file_id: str, client: Any) -> List[str]:
     """Download and parse the batch results file from OpenAI."""
-    response = client.files.content(output_file_id)
+    response = await client.files.content(output_file_id)
 
     # Parse the JSONL content
     content = response.content.decode("utf-8")