Merge pull request #12 from EetuaLaine/main

Restructure setup.py to ensure executability from any directory

Author: gabriben

.gitignore

@@ -8,4 +8,6 @@ src/wraval.egg-info/
 prompts/*
 .idea
 src/wraval/custom_prompts/*
-src/wraval/testing.py
+src/wraval/testing.py
+src/wraval/model_artifacts/*
+!src/wraval/model_artifacts/code/

README.md

@@ -4,6 +4,15 @@ WRAVAL helps in evaluating LLMs for writing assistant tasks like summarization,
 
 ## Quick start
 
+Disclaimer: this project requires a machine that supports bitsandbytes and CUDA.
+
+Before installing, execute the following to ensure correct dependencies:
+
+```bash
+pip install uv
+uv pip compile pyproject.toml -o requirements.txt
+```
+
 ```bash
 pip install -e .
 wraval generate

config/settings.toml

@@ -1,8 +1,7 @@
 [default]
 region = 'us-east-1'
 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_name = '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'
 

pyproject.toml

@@ -13,22 +13,23 @@ authors = [{ name = "Gabriel Benedict", email = "[email protected]" }]
 
 # Dependencies – see note below for loading from requirements.txt
 dependencies = [
-    "tqdm==4.66.4",
-    "pandas==2.2.3",
-    "beautifulsoup4==4.12.3",
-    "boto3==1.34.143",
-    "plotly==5.24.1",
-    "transformers==4.48.1",
-    "datasets==3.2.0",
-    "evaluate==0.4.3",
-    "dynaconf==3.2.7",
-    "torch~=2.6.0",
-    "botocore~=1.34.162",
-    "sagemaker",
+    "tqdm~=4.66.4",
+    "pandas~=2.2.3",
+    "beautifulsoup4~=4.12.3",
+    "boto3",
+    "plotly~=5.24.1",
+    "transformers==4.51.0",
+    "datasets~=3.2.0",
+    "evaluate~=0.4.3",
+    "dynaconf~=3.2.7",
+    "torch",
+    "botocore",
+    "sagemaker==2.244.2",
     "numpy",
     "requests",
-    "bitsandbytes",
-    "accelerate"
+    "bitsandbytes==0.45.5",
+    "accelerate",
+    "torchvision"
 ]
 
 [project.scripts]

requirements.txt

@@ -1,5 +1,7 @@
 # This file was autogenerated by uv via the following command:
 #    uv pip compile pyproject.toml -o requirements.txt
+accelerate==1.1.1
+    # via wraval (pyproject.toml)
 aiohappyeyeballs==2.6.1
     # via aiohttp
 aiohttp==3.11.18
@@ -14,6 +16,8 @@ antlr4-python3-runtime==4.9.3
     # via omegaconf
 anyio==4.9.0
     # via starlette
+async-timeout==5.0.1
+    # via aiohttp
 attrs==23.2.0
     # via
     #   aiohttp
@@ -22,6 +26,8 @@ attrs==23.2.0
     #   sagemaker
 beautifulsoup4==4.12.3
     # via wraval (pyproject.toml)
+bitsandbytes==0.44.1
+    # via wraval (pyproject.toml)
 boto3==1.34.143
     # via
     #   wraval (pyproject.toml)
@@ -56,6 +62,8 @@ dynaconf==3.2.7
     # via wraval (pyproject.toml)
 evaluate==0.4.3
     # via wraval (pyproject.toml)
+exceptiongroup==1.3.0
+    # via anyio
 fastapi==0.115.12
     # via sagemaker
 filelock==3.18.0
@@ -82,6 +90,7 @@ hf-xet==1.1.0
     # via huggingface-hub
 huggingface-hub==0.31.1
     # via
+    #   accelerate
     #   datasets
     #   evaluate
     #   tokenizers
@@ -131,15 +140,53 @@ networkx==3.4.2
 numpy==1.26.4
     # via
     #   wraval (pyproject.toml)
+    #   accelerate
+    #   bitsandbytes
     #   datasets
     #   evaluate
     #   pandas
     #   sagemaker
+    #   torchvision
     #   transformers
+nvidia-cublas-cu12==12.4.5.8
+    # via
+    #   nvidia-cudnn-cu12
+    #   nvidia-cusolver-cu12
+    #   torch
+nvidia-cuda-cupti-cu12==12.4.127
+    # via torch
+nvidia-cuda-nvrtc-cu12==12.4.127
+    # via torch
+nvidia-cuda-runtime-cu12==12.4.127
+    # via torch
+nvidia-cudnn-cu12==9.1.0.70
+    # via torch
+nvidia-cufft-cu12==11.2.1.3
+    # via torch
+nvidia-curand-cu12==10.3.5.147
+    # via torch
+nvidia-cusolver-cu12==11.6.1.9
+    # via torch
+nvidia-cusparse-cu12==12.3.1.170
+    # via
+    #   nvidia-cusolver-cu12
+    #   torch
+nvidia-cusparselt-cu12==0.6.2
+    # via torch
+nvidia-nccl-cu12==2.21.5
+    # via torch
+nvidia-nvjitlink-cu12==12.4.127
+    # via
+    #   nvidia-cusolver-cu12
+    #   nvidia-cusparse-cu12
+    #   torch
+nvidia-nvtx-cu12==12.4.127
+    # via torch
 omegaconf==2.2.3
     # via sagemaker
 packaging==25.0
     # via
+    #   accelerate
     #   datasets
     #   evaluate
     #   huggingface-hub
@@ -154,6 +201,8 @@ pandas==2.2.3
     #   sagemaker
 pathos==0.3.2
     # via sagemaker
+pillow==11.2.1
+    # via torchvision
 platformdirs==4.3.8
     # via
     #   sagemaker
@@ -171,7 +220,9 @@ propcache==0.3.1
 protobuf==4.25.7
     # via sagemaker
 psutil==7.0.0
-    # via sagemaker
+    # via
+    #   accelerate
+    #   sagemaker
 pyarrow==20.0.0
     # via datasets
 pydantic==2.11.4
@@ -190,6 +241,7 @@ pytz==2025.2
     # via pandas
 pyyaml==6.0.2
     # via
+    #   accelerate
     #   datasets
     #   huggingface-hub
     #   omegaconf
@@ -220,15 +272,15 @@ rpds-py==0.24.0
 s3transfer==0.10.4
     # via boto3
 safetensors==0.5.3
-    # via transformers
+    # via
+    #   accelerate
+    #   transformers
 sagemaker==2.236.0
     # via wraval (pyproject.toml)
 sagemaker-core==1.0.16
     # via sagemaker
 schema==0.7.7
     # via sagemaker
-setuptools==80.3.1
-    # via torch
 six==1.17.0
     # via
     #   google-pasta
@@ -250,6 +302,12 @@ tenacity==9.1.2
 tokenizers==0.21.1
     # via transformers
 torch==2.6.0
+    # via
+    #   wraval (pyproject.toml)
+    #   accelerate
+    #   bitsandbytes
+    #   torchvision
+torchvision==0.21.0
     # via wraval (pyproject.toml)
 tqdm==4.66.4
     # via
@@ -261,16 +319,22 @@ tqdm==4.66.4
     #   transformers
 transformers==4.48.1
     # via wraval (pyproject.toml)
+triton==3.2.0
+    # via torch
 typing-extensions==4.13.2
     # via
     #   anyio
+    #   exceptiongroup
     #   fastapi
     #   huggingface-hub
+    #   multidict
     #   pydantic
     #   pydantic-core
     #   referencing
+    #   rich
     #   torch
     #   typing-inspection
+    #   uvicorn
 typing-inspection==0.4.0
     # via pydantic
 tzdata==2025.2

setup.py

@@ -10,7 +10,11 @@
     package_dir={"": "src"},
     install_requires=required,
     data_files=[
-        ('config', ['config/settings.toml'])
+        ('config', ['config/settings.toml']),
+        ('model_artifacts/code', [
+            'src/wraval/model_artifacts/code/inference.py',
+            'src/wraval/model_artifacts/code/requirements.txt'
+        ])
     ],
     include_package_data=True,
     entry_points={

src/wraval/actions/action_deploy.py

@@ -1,22 +1,39 @@
-from argparse import ArgumentParser
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-from sagemaker.huggingface import HuggingFaceModel
-import torch
+import json
+import os
 import tarfile
+from argparse import ArgumentParser
+
 import boto3
-import json
+import torch
+from sagemaker.huggingface import HuggingFaceModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+PACKAGE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+MODEL_DIRECTORY = os.path.join(PACKAGE_DIR, "model_artifacts")
+CODE_PATH = "code"
 
-MODEL_DIRECTORY = '../../../model_artifacts'
+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",
+        "microsoft/Phi-4-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()
 
 def cleanup_endpoints(endpoint_name):
 
-    sagemaker_client = boto3.client("sagemaker", region_name='us-east-1')
+    sagemaker_client = boto3.client("sagemaker", region_name="us-east-1")
 
-    endpoints = sagemaker_client.list_endpoints()['Endpoints']
-    endpoints_configs = sagemaker_client.list_endpoint_configs()['EndpointConfigs']
+    endpoints = sagemaker_client.list_endpoints()["Endpoints"]
+    endpoints_configs = sagemaker_client.list_endpoint_configs()["EndpointConfigs"]
 
-    endpoints_names = [e['EndpointName'] for e in endpoints]
-    endpoints_configs_names = [e['EndpointConfigName'] for e in endpoints_configs]
+    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)
@@ -44,14 +61,14 @@ def load_artifacts(settings):
 
 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"{settings.deploy_bucket_prefix}/{tar_file_name}"
     s3_client = boto3.client("s3")
     s3_client.upload_file(tar_file_name, settings.deploy_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
@@ -92,15 +109,35 @@ def validate_deployment(predictor):
         print(f"Validation failed: {e}")
         raise e
 
-def deploy(settings, cleanup_endpoints=False):
-    sanitized_model_name = settings.hf_name.split('/')[1].replace('.', '-')
-    if cleanup_endpoints:
-        cleanup_endpoints(sanitized_model_name)
+def validate_model_directory():
+    endpoint_code_path = os.path.join(MODEL_DIRECTORY, CODE_PATH)
+    inference_script_name = "inference.py"
+    requirements_name = "requirements.txt"
+    if not os.path.isdir(endpoint_code_path):
+        raise ValueError(f"{endpoint_code_path} is missing.")
+    if not os.path.isfile(os.path.join(endpoint_code_path, inference_script_name)):
+        raise ValueError(f"{inference_script_name} is missing from the code directory.")
+    if not os.path.isfile(os.path.join(endpoint_code_path, requirements_name)):
+        raise ValueError(f"{requirements_name} is missing from the code directory.")
+
+
+def cleanup_model_directory():
+    for item in os.listdir(MODEL_DIRECTORY):
+        item_path = os.path.join(MODEL_DIRECTORY, item)
+        if item == CODE_PATH:
+            continue
+        if os.path.isfile(item_path):
+            os.remove(item_path)
+
+def deploy(settings):
+    validate_model_directory()
+    cleanup_model_directory()
+    sanitized_model_name = settings.hf_name.split("/")[1].replace(".", "-")
     load_artifacts(settings)
     s3_uri = write_model_to_s3(settings, sanitized_model_name)
     predictor = deploy_endpoint(
         s3_uri,
-        settings.sagemaker_execution_role_arn, 
+        settings.sagemaker_execution_role_arn,
         sanitized_model_name
     )
     validate_deployment(predictor)