action deploy

Author: gabriben

.gitignore

@@ -7,4 +7,5 @@ src/wraval.egg-info/
 **__pycache__/
 prompts/*
 .idea
-src/wraval/custom_prompts/*
+src/wraval/custom_prompts/*
+src/wraval/testing.py

config/settings.toml

@@ -1,7 +1,10 @@
 [default]
 region = 'us-east-1'
-data_dir = "./data"
-# 's3://llm-finetune-us-east-1-{aws_account}/eval/tones/'
+data_dir = 's3://llm-finetune-us-east-1-{aws_account}/eval/tones/'
+# "./data"
+deploy_bucket_name = 's3://llm-finetune-us-east-1-{aws_account}'
+deploy_bucket_prefix = 'models'
+sagemaker_execution_role_arn = 'arn:aws:iam::{aws_account}:role/sagemaker-execution-role-us-east-1'
 
 [haiku-3]
 model = 'anthropic.claude-3-haiku-20240307-v1:0'

pyproject.toml

@@ -26,7 +26,9 @@ dependencies = [
     "botocore~=1.34.162",
     "sagemaker",
     "numpy",
-    "requests"
+    "requests",
+    "bitsandbytes",
+    "accelerate"
 ]
 
 [project.scripts]

src/wraval/actions/deploy.py src/wraval/actions/action_deploy.pysrc/wraval/actions/deploy.py renamed to src/wraval/actions/action_deploy.py

@@ -6,21 +6,22 @@
 import boto3
 import json
 
-MODEL_DIRECTORY = "model_artifacts"
+def cleanup_endpoints(endpoint_name):
 
-def parse_args():
-    arg_parser = ArgumentParser()
-    arg_parser.add_argument("--model_name", type=str, required=True, choices=(
-        "Qwen/Qwen2.5-1.5B-Instruct", 
-        "microsoft/Phi-3.5-mini-instruct"
-        )
-    )
-    arg_parser.add_argument("--bucket_name", type=str, required=True)
-    arg_parser.add_argument("--bucket_prefix", type=str, required=True)
-    arg_parser.add_argument("--sagemaker_execution_role_arn", type=str, required=True)
-    return arg_parser.parse_args()
+    sagemaker_client = boto3.client("sagemaker", region_name='us-east-1')
+
+    endpoints = sagemaker_client.list_endpoints()['Endpoints']
+    endpoints_configs = sagemaker_client.list_endpoint_configs()['EndpointConfigs']
 
-def load_artifacts(args):
+    endpoints_names = [e['EndpointName'] for e in endpoints]
+    endpoints_configs_names = [e['EndpointConfigName'] for e in endpoints_configs]
+
+    if endpoint_name in endpoints_names:
+        sagemaker_client.delete_endpoint(EndpointConfigName=endpoint_name)
+    if endpoint_name in endpoints_configs_names:
+        sagemaker_client.delete_endpoint_config(EndpointConfigName=endpoint_name)
+
+def load_artifacts(settings):
     bnb_config = BitsAndBytesConfig(
         load_in_4bit=True,
         bnb_4bit_quant_type="nf4",
@@ -29,27 +30,27 @@ def load_artifacts(args):
     )
 
     model = AutoModelForCausalLM.from_pretrained(
-        args.model_name,
+        settings.hf_name,
         device_map="auto",
         quantization_config=bnb_config
     )
 
-    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    tokenizer = AutoTokenizer.from_pretrained(settings.hf_name)
 
     model.save_pretrained(MODEL_DIRECTORY)
     tokenizer.save_pretrained(MODEL_DIRECTORY)
 
-def write_model_to_s3(args, model_name):
+def write_model_to_s3(settings, model_name):
     tar_file_name = f"{model_name}.tar.gz"
 
     with tarfile.open(tar_file_name, "w:gz") as tar:
         tar.add(MODEL_DIRECTORY, arcname=".")
 
-    s3_model_path = f"{args.bucket_prefix}/{tar_file_name}"
+    s3_model_path = f"{settings.deploy_bucket_prefix}/{tar_file_name}"
     s3_client = boto3.client("s3")
-    s3_client.upload_file(tar_file_name, args.bucket_name, s3_model_path)
+    s3_client.upload_file(tar_file_name, settings.deploy_bucket_name, s3_model_path)
 
-    s3_uri = f"s3://{args.bucket_name}/{s3_model_path}"
+    s3_uri = f"s3://{settings.deploy_bucket_name}/{s3_model_path}"
     print(f"Model uploaded to: {s3_uri}")
     return s3_uri
 
@@ -89,17 +90,15 @@ def validate_deployment(predictor):
         print(f"Validation failed: {e}")
         raise e
 
-def deploy():
-    args = parse_args()
-    load_artifacts(args)
-    sanitized_model_name = args.model_name.split('/')[1].replace('.', '-')
-    s3_uri = write_model_to_s3(args, sanitized_model_name)
+def deploy(settings, cleanup_endpoints=False):
+    sanitized_model_name = settings.hf_name.split('/')[1].replace('.', '-')
+    if cleanup_endpoints:
+        cleanup_endpoints(sanitized_model_name)
+    load_artifacts(settings)
+    s3_uri = write_model_to_s3(settings, sanitized_model_name)
     predictor = deploy_endpoint(
-        s3_uri, 
-        args.sagemaker_execution_role_arn, 
+        s3_uri,
+        settings.sagemaker_execution_role_arn, 
         sanitized_model_name
     )
-    validate_deployment(predictor)
-
-if __name__ == "__main__":
-    deploy()
+    validate_deployment(predictor)

src/wraval/main.py

@@ -11,6 +11,7 @@
 from wraval.actions.action_llm_judge import judge
 from wraval.actions.aws_utils import get_current_aws_account_id
 from wraval.actions.action_results import show_results
+from wraval.actions.action_deploy import deploy
 import os
 
 
@@ -32,6 +33,8 @@ def get_settings(args):
     ## add the AWS account you are logged into, if necessary.
     settings.model = settings.model.format(aws_account=settings.aws_account)
     settings.data_dir = settings.data_dir.format(aws_account=settings.aws_account)
+    settings.deploy_bucket_name = settings.deploy_bucket_name.format(aws_account=settings.aws_account)
+    settings.sagemaker_execution_role_arn = settings.sagemaker_execution_role_arn.format(aws_account=settings.aws_account)
 
     if args.custom_prompts:
         settings.custom_prompts = True
@@ -54,6 +57,7 @@ def parse_args() -> argparse.Namespace:
             "human_judge_upload",
             "human_judge_parsing",
             "show_results",
+            "deploy"
         ],
         help="Action to perform (generate data or run inference)",
     )
@@ -84,6 +88,10 @@ def parse_args() -> argparse.Namespace:
         "--custom-prompts", default=False, help="Load custom prompts from a prompt folder"
     )
 
+    parser.add_argument(
+        "--cleanup_endpoints", action='store_true'
+    )
+
     return parser.parse_args()
 
 
@@ -117,6 +125,9 @@ def handle_judge(args, settings):
 def handle_show_results(args, settings):
     show_results(settings, args.type)
 
+def handle_deploy(args, settings):
+    deploy(settings, args.cleanup_endpoints)
+
 
 def main():
     args = parse_args()
@@ -131,6 +142,8 @@ def main():
             handle_judge(args, settings)
         case "show_results":
             handle_show_results(args, settings)
+        case "deploy":
+            handle_deploy(args, settings)            
         case _:
             raise ValueError(f"Unknown action: {args.action}")