hyenaDNA-training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f0dfbb7c-2280-4b01-949a-91238a892d0b",
   "metadata": {},
   "source": [
    "### 0. Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82308943-5300-4cc0-856c-879e1a6f24fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -qU pip\n",
    "%pip install -qU sagemaker boto3 awscli boto3 ipywidgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3e8d899-38ef-4e5b-b798-782d97c6144e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "from pathlib import Path\n",
    "import os\n",
    "from time import strftime\n",
    "from functools import partial\n",
    "import importlib\n",
    "\n",
    "import utilities as u\n",
    "\n",
    "import boto3\n",
    "import sagemaker\n",
    "from sagemaker.experiments.run import Run\n",
    "from sagemaker.inputs import TrainingInput\n",
    "from sagemaker.pytorch import PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc08ea8e-787f-42ae-bdba-4f3d8ba9c044",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sagemaker.__version__, boto3.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46afa5b6-7d8e-4d4e-9be7-bedaf774694b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "boto_session = boto3.session.Session()\n",
    "sagemaker_session = sagemaker.session.Session(boto_session)\n",
    "omics = boto3.client(\"omics\")\n",
    "\n",
    "REGION_NAME = sagemaker_session.boto_region_name\n",
    "S3_BUCKET = sagemaker_session.default_bucket()\n",
    "\n",
    "EXPERIMENT_NAME = \"hyenaDNA-pretraining-v2\"\n",
    "\n",
    "SAGEMAKER_EXECUTION_ROLE = sagemaker.session.get_execution_role(sagemaker_session)\n",
    "print(f\"Assumed SageMaker role is {SAGEMAKER_EXECUTION_ROLE}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87cd25e8-4256-49b5-97ef-b98ed7f0657e",
   "metadata": {},
   "source": [
    "### 1. Read the data from AWS HealthOmics\n",
    "\n",
    "To train the HyenaDNA model, we will directly utilize the data from the Healthomics Sequence Store. Let's retrieve the read set URI from the Healthomics Sequence Store. Make sure to replace your sequence store id in the following cell. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0bebb2c-f925-4532-8c03-60195f26d657",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "seq_store_id = \"4308389581\" # replace with your sequence store id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4de9f99-49cd-425c-b724-b27ed95ad750",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "seq_store_info = omics.get_sequence_store(id=seq_store_id)\n",
    "s3_uri = seq_store_info[\"s3Access\"][\"s3Uri\"]\n",
    "s3_arn = seq_store_info[\"s3Access\"][\"s3AccessPointArn\"]\n",
    "key_arn = seq_store_info[\"sseConfig\"][\"keyArn\"]\n",
    "s3_uri, s3_arn, key_arn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "207ee04b-a852-4533-beff-544e38fe5952",
   "metadata": {},
   "outputs": [],
   "source": [
    "S3_DATA_URI = f\"{s3_uri}readSet/\"\n",
    "S3_DATA_URI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6120bb71-a43e-4e15-bc36-800ae4c320ed",
   "metadata": {},
   "source": [
    "For the training job to access the objects in the above S3 access point, `s3uri`, you must add a policy\n",
    "to this execution role (`SAGEMAKER_EXECUTION_ROLE`). The output of the following cell is the policy that\n",
    "you should attach to this role:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63c953ee-7029-4562-b008-76b8f7594e81",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(json.dumps({\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        { \n",
    "            \"Sid\": \"S3DirectAccess\",\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [ \n",
    "                \"s3:GetObject\",\n",
    "                \"s3:ListBucket\"\n",
    "            ], \n",
    "            \"Resource\": \"*\",\n",
    "            \"Condition\": {\n",
    "                \"StringEquals\": {\n",
    "                    \"s3:DataAccessPointArn\": s3_arn\n",
    "                } \n",
    "            }\n",
    "        },\n",
    "        { \n",
    "            \"Sid\": \"DefaultSequenceStoreKMSDecrypt\",\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": \"kms:Decrypt\",\n",
    "            \"Resource\": key_arn\n",
    "        }\n",
    "    ] \n",
    "}, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ab4d0b7-0eb9-4757-aa89-aec2ca107115",
   "metadata": {},
   "source": [
    "### 2. Training\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7a68d1b-edb1-49ee-acae-076bdecff096",
   "metadata": {},
   "source": [
    "### 2.1 Define the training container "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7e3deae-0de2-499e-a60a-c76aaef138ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "pytorch_image_uri = f\"763104351884.dkr.ecr.{REGION_NAME}.amazonaws.com/pytorch-training:2.2.0-gpu-py310-cu121-ubuntu20.04-sagemaker\"\n",
    "pytorch_image_uri"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c52a3e14-2cab-472a-96d9-808c39c4f8c8",
   "metadata": {},
   "source": [
    "#### 2.2 Define the training job parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfe0c7f6-0a53-483a-80ae-793ad58d7ee5",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_ID = 'LongSafari/hyenadna-small-32k-seqlen-hf'\n",
    "TRAINING_JOB_NAME = 'hyenaDNA-pretraining'\n",
    "\n",
    "# Additional training parameters\n",
    "hyperparameters = {\n",
    "    \"species\" : \"mouse\",\n",
    "    \"epochs\": 150,\n",
    "    \"model_checkpoint\": MODEL_ID,\n",
    "    \"max_length\": 32_000,\n",
    "    \"batch_size\": 4, \n",
    "    \"learning_rate\": 6e-4,\n",
    "    \"weight_decay\" : 0.1,\n",
    "    \"log_level\" : \"INFO\",\n",
    "    \"log_interval\" : 100\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3178423a-828b-4101-b540-5efe22b5e223",
   "metadata": {},
   "source": [
    "#### 2.3 Define Metrics to track\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ad98da5-1d11-4eb4-8718-2e4cb4834872",
   "metadata": {},
   "outputs": [],
   "source": [
    "metric_definitions = [\n",
    "    {\"Name\": \"epoch\", \"Regex\": \"Epoch: ([0-9.]*)\"},\n",
    "    {\"Name\": \"step\", \"Regex\": \"Step: ([0-9.]*)\"},\n",
    "    {\"Name\": \"train_loss\", \"Regex\": \"Train Loss: ([0-9.e-]*)\"},\n",
    "    {\"Name\": \"train_perplexity\", \"Regex\": \"Train Perplexity: ([0-9.e-]*)\"},\n",
    "    {\"Name\": \"eval_loss\", \"Regex\": \"Eval Average Loss: ([0-9.e-]*)\"},\n",
    "    {\"Name\": \"eval_perplexity\", \"Regex\": \"Eval Perplexity: ([0-9.e-]*)\"}\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82c9636d-1707-42ee-aa77-3f63c095285c",
   "metadata": {},
   "source": [
    "#### 2.4 Define the tensorboard configurations to track the training results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d22067bf-c6b5-478d-b8ca-65a0bd487768",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.debugger import TensorBoardOutputConfig\n",
    "\n",
    "LOG_DIR=\"/opt/ml/output/tensorboard\"\n",
    "\n",
    "output_path = os.path.join(\n",
    "    \"s3://\", S3_BUCKET, \"sagemaker-output\", \"training\", TRAINING_JOB_NAME\n",
    ")\n",
    "\n",
    "tensorboard_output_config = TensorBoardOutputConfig(\n",
    "    s3_output_path=os.path.join(output_path, 'tensorboard'),\n",
    "    container_local_output_path=LOG_DIR\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af18fb5e-e1cf-447b-84fa-eb80b347c9e1",
   "metadata": {},
   "source": [
    "#### 2.4 Define Estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcd8126d-605c-40c6-bce1-a68602f80e82",
   "metadata": {},
   "outputs": [],
   "source": [
    "hyenaDNA_estimator = PyTorch(\n",
    "    base_job_name=TRAINING_JOB_NAME,\n",
    "    entry_point=\"train_hf_accelerate.py\",\n",
    "    source_dir=\"scripts/\",\n",
    "    instance_type=\"ml.g5.12xlarge\",\n",
    "    instance_count=1,\n",
    "    image_uri=pytorch_image_uri,\n",
    "    role=SAGEMAKER_EXECUTION_ROLE,\n",
    "    hyperparameters=hyperparameters,\n",
    "    metric_definitions=metric_definitions,\n",
    "    sagemaker_session=sagemaker_session,\n",
    "    distribution={\"torch_distributed\": {\"enabled\": True}},\n",
    "    tags=[{\"Key\": \"project\", \"Value\": \"genomics-model-pretraining\"}],\n",
    "    keep_alive_period_in_seconds=1800,\n",
    "    tensorboard_output_config=tensorboard_output_config,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea04af1c-82ce-44ca-bce7-76d9776c6d95",
   "metadata": {},
   "source": [
    "#### 2.5 Start Training with Distributed Data Parallel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2b5544d-cffb-4a1a-bbe5-5d260c31bf2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "with Run(\n",
    "    experiment_name=EXPERIMENT_NAME,\n",
    "    sagemaker_session=sagemaker_session,\n",
    ") as run:\n",
    "    hyenaDNA_estimator.fit(\n",
    "        {\n",
    "            \"data\": TrainingInput(\n",
    "                s3_data=S3_DATA_URI, input_mode=\"File\"\n",
    "            ),\n",
    "        },\n",
    "        wait=True,\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5efe1b4b-f7d5-402c-93ed-0844a64fd581",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_job_name = hyenaDNA_estimator.latest_training_job.name\n",
    "training_job_name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59941ec0-9bb4-44f6-994d-eaa77b2adb10",
   "metadata": {},
   "source": [
    "### 5. Training Results "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c6e66b6-92ce-4f8c-bc75-e884ee71b98a",
   "metadata": {},
   "source": [
    "* In our training process we had pushed the training resulsts to Tensorboard. You can see them using SageMaker tensorboad application. To obtain the link to the tensorboard, please replace the sagemaker studio profile name below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a2164bf-8dfe-4a47-912e-d29f7e11ce07",
   "metadata": {},
   "outputs": [],
   "source": [
    "user_profile = \"shamika\" # replace with your sagemaker studio profile name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c71a2a93-8574-443b-a5ff-ced57c00c9ff",
   "metadata": {},
   "source": [
    "Execuate following cell to get link to the the tensorboard application"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e743bb7b-6791-45fe-b678-4c9bee380e14",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.interactive_apps.tensorboard import TensorBoardApp\n",
    "\n",
    "with open(\"/opt/ml/metadata/resource-metadata.json\", \"r\") as f:\n",
    "    app_metadata = json.loads(f.read())\n",
    "    sm_user_profile_name = app_metadata[\"SpaceName\"]\n",
    "    sm_domain_id = app_metadata[\"DomainId\"]\n",
    "\n",
    "tb_app = TensorBoardApp(REGION_NAME)\n",
    "tb_app.get_app_url(\n",
    "    training_job_name=training_job_name,\n",
    "    create_presigned_domain_url=True,           \n",
    "    domain_id=sm_domain_id,                 \n",
    "    user_profile_name=user_profile, \n",
    "    open_in_default_web_browser=False,\n",
    "    optional_create_presigned_url_kwargs={} \n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c23d90b-f961-4f82-aa0c-46f7a6124c3a",
   "metadata": {},
   "source": [
    "### 6. Deploy trained model to an realtime endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93e65ac9-4ebf-4866-8578-8d8706fe1036",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.pytorch.model import PyTorchModel\n",
    "from sagemaker.serializers import JSONSerializer\n",
    "from sagemaker.estimator import Estimator\n",
    "\n",
    "training_job_name = \"hyenaDNA-pretraining-2024-04-06-06-23-26-412\"\n",
    "#attached_estimator = Estimator.attach(training_job_name)\n",
    "\n",
    "model_data = attached_estimator.model_data\n",
    "model_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f72c1f3-f658-4009-8fdd-c7e4d8bf548d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Deploy the model to create a real-time endpoint\n",
    "endpoint_name = 'hyenaDNA-pretrained-mouse-ep'  \n",
    "pytorch_deployment_uri = f\"763104351884.dkr.ecr.{REGION_NAME}.amazonaws.com/pytorch-inference:2.2.0-gpu-py310-cu118-ubuntu20.04-sagemaker\"\n",
    "\n",
    "hyenaDNAModel = PyTorchModel(\n",
    "    model_data=model_data,\n",
    "    role=SAGEMAKER_EXECUTION_ROLE,\n",
    "    image_uri=pytorch_deployment_uri,\n",
    "    entry_point=\"inference.py\",\n",
    "    source_dir=\"scripts/\",\n",
    "    sagemaker_session=sagemaker_session,\n",
    "    name=endpoint_name,\n",
    "    env = {\n",
    "        'MMS_MAX_REQUEST_SIZE': '2000000000',\n",
    "        'MMS_MAX_RESPONSE_SIZE': '2000000000',\n",
    "        'MMS_DEFAULT_RESPONSE_TIMEOUT': '900',\n",
    "        'TS_MAX_RESPONSE_SIZE':'2000000000',\n",
    "        'TS_MAX_REQUEST_SIZE':'2000000000',\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4808b94a-3079-40c9-bbe1-3fa596250803",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#real_time_endpoint_name = \"hyenaDNA-mouse-pretrained-real-ep\"\n",
    "env = {\n",
    "    'SAGEMAKER_MODEL_SERVER_TIMEOUT':'7200', \n",
    "    'TS_MAX_RESPONSE_SIZE':'2000000000',\n",
    "    'TS_MAX_REQUEST_SIZE':'2000000000',\n",
    "    'MMS_MAX_RESPONSE_SIZE':'2000000000',\n",
    "    'MMS_MAX_REQUEST_SIZE':'2000000000'\n",
    "}\n",
    "\n",
    "# deploy the endpoint endpoint\n",
    "realtime_predictor = hyenaDNAModel.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type=\"ml.g5.8xlarge\",\n",
    "    endpoint_name=endpoint_name,\n",
    "    env=env,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "477dc76d-980c-4a21-8be7-0feba1ecf108",
   "metadata": {},
   "source": [
    "### 7. Test the realtime endpoint\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c26d0e45-8810-4b4c-adee-bf204933cd55",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "sample_genome_data = []\n",
    "with open(\"./sample_mouse_data.json\") as file:\n",
    "    for line in file:\n",
    "        sample_genome_data.append(json.loads(line))\n",
    "len(sample_genome_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c05b22d7-fbdf-49c3-86a7-5abc110ad21c",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(sample_genome_data[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecdd3233-f291-46e6-8818-138178857abd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.deserializers import JSONDeserializer\n",
    "from sagemaker.serializers import JSONSerializer\n",
    "\n",
    "\n",
    "data = [sample_genome_data[0]]\n",
    "realtime_predictor.serializer = JSONSerializer()\n",
    "realtime_predictor.deserializer = JSONDeserializer()\n",
    "embeddings = realtime_predictor.predict(data=data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b79e9f9-e097-4e74-b7fb-938ca55fb615",
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b53aedd-ce70-4e15-8ccd-7c0a08fecc24",
   "metadata": {},
   "source": [
    "#### 7. Cleanup\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44c9045b-1e15-41c9-85dc-775233388ac2",
   "metadata": {},
   "outputs": [],
   "source": [
    "realtime_predictor.delete_endpoint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7814bfb-1c45-4796-8f1d-6cb760927558",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}