vertexAI_batch_prediction/batch_prediction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "085beea0-10f8-4fb5-8af7-51596147a767",
   "metadata": {},
   "source": [
    "## **Setting Up a Batch Prediction Job on VertexAI**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "de87d0e1-23bf-4032-8a43-75951a1f4419",
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.cloud import aiplatform\n",
    "\n",
    "\"\"\" Setting the location \"\"\"\n",
    "aiplatform.init(location='us-west1')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ce9eaf8-465c-4ff3-b9da-7de8e37d860a",
   "metadata": {},
   "source": [
    "### Running a Batch prediction job required registering the model in Model registry. Later find the registered mode id under version details of the registerted model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6b434fe9-4c25-4faf-8cd7-7910f5fb9ae1",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" Add model id \"\"\"\n",
    "registered_model_id = \"$MODEL ID$\" \n",
    "model_id = f'projects/ml-ops-segment-anything/locations/us-west1/models/{registered_model_id}'\n",
    "model = aiplatform.Model(model_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f4cf7bb7-98f1-474d-8a57-0657626ff45b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['jsonl', 'bigquery', 'csv', 'tf-record', 'tf-record-gzip', 'file-list']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.supported_input_storage_formats"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bad4fdbd-2966-48e5-8ef3-b98a8acca4b5",
   "metadata": {},
   "source": [
    "### Here we are going to preprocess the batch of images and prepare a JSON file with all the preprocessed data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ad6b1b05-9646-4884-aed9-8e393fd31894",
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.cloud import storage\n",
    "import base64\n",
    "import json\n",
    "\n",
    "# Initialize a client\n",
    "storage_client = storage.Client()\n",
    "\n",
    "\"\"\" Add your Bucket Name \"\"\"\n",
    "bucket_name = \"$BUCKET_NAME$\"\n",
    "\n",
    "# Create a bucket object\n",
    "bucket = storage_client.get_bucket(f\"{bucket_name}\")\n",
    "\n",
    "# Create blob objects from the filepath\n",
    "blobs = bucket.list_blobs(prefix=\"batch-prediction-images\")\n",
    "\n",
    "# Iterate over the blobs and filter based on file extension\n",
    "image_extensions = ('.png', '.jpg', '.jpeg')\n",
    "\n",
    "# collecting all image file names \n",
    "image_blobs = [blob for blob in blobs if blob.name.lower().endswith(image_extensions)]\n",
    "\n",
    "# Download all images and create a base64 String of each image\n",
    "image_base64_dict = {\"instances\": []}\n",
    "for image_blob in image_blobs:\n",
    "    # Creating base64 string of the image\n",
    "    base64_str = base64.b64encode(image_blob.download_as_bytes()).decode('utf-8')\n",
    "    \n",
    "    image_base64_dict[\"instances\"].append({\n",
    "        'image': base64_str,\n",
    "        'file_path': f\"gs://{bucket_name}/batch-prediction-images/{image_blob.name}\"\n",
    "    })\n",
    "\n",
    "# Save the preprocessed data into a JSON file\n",
    "with open(\"batch_predict.jsonl\", \"w\") as file:\n",
    "    json.dump(image_base64_dict, file)\n",
    "\n",
    "# Push the JSON file into cloud storage\n",
    "blob = bucket.blob(\"batch_predict.jsonl\")\n",
    "blob.upload_from_filename(\"batch-prediction-images/batch_predict.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "38404927-fc8f-4b43-b4d5-8cd53bb5c9fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" Batch Prediction Job Name \"\"\"\n",
    "job_display_name = 'sam-batch-prediction'\n",
    "\n",
    "\"\"\" Cloud Storage Preprocessed JSON data path\"\"\"\n",
    "gcs_source = f'gs://{bucket_name}/batch-prediction-images/batch_predict.jsonl'\n",
    "\n",
    "\"\"\" Batch prediction JSON results storage path\"\"\"\n",
    "gcs_destination_prefix = f'gs://{bucket}/batch-prediction-images'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "862f0c02-fd7e-41e6-9dff-2d40be17369b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating BatchPredictionJob\n",
      "BatchPredictionJob created. Resource name: projects/633534855904/locations/us-west1/batchPredictionJobs/3963534908969713664\n",
      "To use this BatchPredictionJob in another session:\n",
      "bpj = aiplatform.BatchPredictionJob('projects/633534855904/locations/us-west1/batchPredictionJobs/3963534908969713664')\n",
      "View Batch Prediction Job:\n",
      "https://console.cloud.google.com/ai/platform/locations/us-west1/batch-predictions/3963534908969713664?project=633534855904\n",
      "BatchPredictionJob projects/633534855904/locations/us-west1/batchPredictionJobs/3963534908969713664 current state:\n",
      "JobState.JOB_STATE_RUNNING\n",
      "BatchPredictionJob projects/633534855904/locations/us-west1/batchPredictionJobs/3963534908969713664 current state:\n",
      "JobState.JOB_STATE_RUNNING\n",
      "BatchPredictionJob projects/633534855904/locations/us-west1/batchPredictionJobs/3963534908969713664 current state:\n",
      "JobState.JOB_STATE_RUNNING\n",
      "BatchPredictionJob projects/633534855904/locations/us-west1/batchPredictionJobs/3963534908969713664 current state:\n",
      "JobState.JOB_STATE_RUNNING\n",
      "BatchPredictionJob projects/633534855904/locations/us-west1/batchPredictionJobs/3963534908969713664 current state:\n",
      "JobState.JOB_STATE_RUNNING\n",
      "BatchPredictionJob projects/633534855904/locations/us-west1/batchPredictionJobs/3963534908969713664 current state:\n",
      "JobState.JOB_STATE_RUNNING\n",
      "BatchPredictionJob projects/633534855904/locations/us-west1/batchPredictionJobs/3963534908969713664 current state:\n",
      "JobState.JOB_STATE_RUNNING\n",
      "BatchPredictionJob projects/633534855904/locations/us-west1/batchPredictionJobs/3963534908969713664 current state:\n",
      "JobState.JOB_STATE_RUNNING\n",
      "BatchPredictionJob projects/633534855904/locations/us-west1/batchPredictionJobs/3963534908969713664 current state:\n",
      "JobState.JOB_STATE_SUCCEEDED\n",
      "BatchPredictionJob run completed. Resource name: projects/633534855904/locations/us-west1/batchPredictionJobs/3963534908969713664\n"
     ]
    }
   ],
   "source": [
    "\"\"\" Setup the Batch Prediction Job \"\"\"\n",
    "\n",
    "#  Refer to GCP documentation for different machine types and accelerator types\n",
    "batch_prediction_job = model.batch_predict(\n",
    "    job_display_name=job_display_name,\n",
    "    machine_type= \"n1-standard-2\", \n",
    "    accelerator_type= \"NVIDIA_TESLA_T4\", \n",
    "    accelerator_count= 1,\n",
    "    gcs_source=gcs_source,\n",
    "    gcs_destination_prefix=gcs_destination_prefix,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": "pytorch-gpu.1-13.m108",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/pytorch-gpu.1-13:m108"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}