Bedrock fix

sagemaker-serve/example_notebooks/bedrock_nova_deployment.ipynb

@@ -0,0 +1,269 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Deploy a Nova Model to Amazon Bedrock\n",
+    "\n",
+    "This notebook demonstrates how to fine-tune an Amazon Nova model using the SageMaker SDK\n",
+    "and deploy it to Amazon Bedrock using `BedrockModelBuilder`.\n",
+    "\n",
+    "The workflow:\n",
+    "1. Fine-tune Nova Micro using `SFTTrainer`\n",
+    "2. Create a `BedrockModelBuilder` from the completed training job\n",
+    "3. Deploy to Bedrock — the builder automatically:\n",
+    "   - Detects the model as Nova\n",
+    "   - Reads the checkpoint URI from the training job manifest\n",
+    "   - Calls `CreateCustomModel`\n",
+    "   - Polls until the model is Active\n",
+    "   - Calls `CreateCustomModelDeployment`\n",
+    "   - Polls until the deployment is Active\n",
+    "4. Clean up resources\n",
+    "\n",
+    "**Prerequisites:**\n",
+    "- AWS credentials with SageMaker and Bedrock access in `us-east-1`\n",
+    "- `sagemaker-serve` and `sagemaker-train` packages installed\n",
+    "- An IAM role with Bedrock and SageMaker permissions"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import json\n",
+    "import time\n",
+    "import random\n",
+    "import boto3\n",
+    "\n",
+    "REGION = \"us-east-1\"\n",
+    "os.environ[\"AWS_DEFAULT_REGION\"] = REGION\n",
+    "\n",
+    "from sagemaker.core.helper.session_helper import get_execution_role\n",
+    "\n",
+    "role_arn = get_execution_role()\n",
+    "account_id = boto3.client(\"sts\").get_caller_identity()[\"Account\"]\n",
+    "bucket = f\"sagemaker-{REGION}-{account_id}\"\n",
+    "\n",
+    "print(f\"Region:  {REGION}\")\n",
+    "print(f\"Account: {account_id}\")\n",
+    "print(f\"Role:    {role_arn}\")\n",
+    "print(f\"Bucket:  {bucket}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 1: Prepare training data\n",
+    "\n",
+    "Upload a small JSONL dataset in the chat-messages format that Nova expects."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "s3 = boto3.client(\"s3\", region_name=REGION)\n",
+    "\n",
+    "train_key = \"nova-example/train.jsonl\"\n",
+    "train_uri = f\"s3://{bucket}/{train_key}\"\n",
+    "\n",
+    "rows = []\n",
+    "for i in range(50):\n",
+    "    rows.append(json.dumps({\n",
+    "        \"messages\": [\n",
+    "            {\"role\": \"user\", \"content\": f\"What is {i+1} + {i+1}?\"},\n",
+    "            {\"role\": \"assistant\", \"content\": f\"The answer is {(i+1)*2}.\"}\n",
+    "        ]\n",
+    "    }))\n",
+    "\n",
+    "s3.put_object(Bucket=bucket, Key=train_key, Body=\"\\n\".join(rows).encode())\n",
+    "print(f\"Uploaded {len(rows)} examples to {train_uri}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 2: Fine-tune Nova Micro with SFTTrainer\n",
+    "\n",
+    "This launches a SageMaker training job. It typically takes 15-30 minutes to complete."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sagemaker.train.sft_trainer import SFTTrainer\n",
+    "\n",
+    "trainer = SFTTrainer(\n",
+    "    model=\"nova-textgeneration-micro\",\n",
+    "    training_dataset=train_uri,\n",
+    "    accept_eula=True,\n",
+    "    model_package_group=\"nova-example-models\",\n",
+    ")\n",
+    "\n",
+    "# Set wait=True to block until training completes\n",
+    "trainer.train(wait=True)\n",
+    "\n",
+    "training_job = trainer._latest_training_job\n",
+    "print(f\"Training job: {training_job.training_job_name}\")\n",
+    "print(f\"Status:       {training_job.training_job_status}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 3: Deploy to Bedrock with BedrockModelBuilder\n",
+    "\n",
+    "The builder handles the full deployment flow:\n",
+    "- Fetches the model package from the training job\n",
+    "- Detects it as a Nova model\n",
+    "- Reads the checkpoint URI from the training output manifest\n",
+    "- Creates a Bedrock custom model and polls until Active\n",
+    "- Creates a deployment and polls until Active"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sagemaker.serve.bedrock_model_builder import BedrockModelBuilder\n",
+    "\n",
+    "builder = BedrockModelBuilder(model=training_job)\n",
+    "\n",
+    "print(f\"Model package:    {builder.model_package}\")\n",
+    "print(f\"S3 artifacts:     {builder.s3_model_artifacts}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "rand = random.randint(1000, 9999)\n",
+    "custom_model_name = f\"nova-example-{rand}-{int(time.time())}\"\n",
+    "deployment_name = f\"{custom_model_name}-dep\"\n",
+    "\n",
+    "print(f\"Deploying as: {custom_model_name}\")\n",
+    "print(f\"This will poll for model creation and deployment — may take several minutes...\")\n",
+    "\n",
+    "response = builder.deploy(\n",
+    "    custom_model_name=custom_model_name,\n",
+    "    role_arn=role_arn,\n",
+    "    deployment_name=deployment_name,\n",
+    ")\n",
+    "\n",
+    "deployment_arn = response.get(\"customModelDeploymentArn\")\n",
+    "print(f\"\\nDeployment ARN: {deployment_arn}\")\n",
+    "print(\"Deployment is Active and ready for inference.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 4: Test inference (optional)\n",
+    "\n",
+    "Once the deployment is Active, you can invoke it via the Bedrock Runtime API."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "bedrock_runtime = boto3.client(\"bedrock-runtime\", region_name=REGION)\n",
+    "\n",
+    "# Get the model ARN from the deployment\n",
+    "bedrock = boto3.client(\"bedrock\", region_name=REGION)\n",
+    "dep_info = bedrock.get_custom_model_deployment(\n",
+    "    customModelDeploymentIdentifier=deployment_arn\n",
+    ")\n",
+    "model_arn = dep_info.get(\"modelArn\")\n",
+    "print(f\"Model ARN: {model_arn}\")\n",
+    "\n",
+    "# Invoke\n",
+    "invoke_response = bedrock_runtime.invoke_model(\n",
+    "    modelId=deployment_arn,\n",
+    "    contentType=\"application/json\",\n",
+    "    body=json.dumps({\n",
+    "        \"messages\": [{\"role\": \"user\", \"content\": \"What is 7 + 7?\"}]\n",
+    "    }),\n",
+    ")\n",
+    "\n",
+    "result = json.loads(invoke_response[\"body\"].read())\n",
+    "print(f\"Response: {result}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 5: Cleanup\n",
+    "\n",
+    "Delete the deployment and custom model to avoid ongoing charges."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "bedrock = boto3.client(\"bedrock\", region_name=REGION)\n",
+    "\n",
+    "# Delete deployment first\n",
+    "if deployment_arn:\n",
+    "    try:\n",
+    "        bedrock.delete_custom_model_deployment(\n",
+    "            customModelDeploymentIdentifier=deployment_arn\n",
+    "        )\n",
+    "        print(f\"Deleted deployment: {deployment_arn}\")\n",
+    "    except Exception as e:\n",
+    "        print(f\"Failed to delete deployment: {e}\")\n",
+    "\n",
+    "# Then delete the custom model\n",
+    "if model_arn:\n",
+    "    try:\n",
+    "        bedrock.delete_custom_model(modelIdentifier=model_arn)\n",
+    "        print(f\"Deleted custom model: {model_arn}\")\n",
+    "    except Exception as e:\n",
+    "        print(f\"Failed to delete custom model: {e}\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}