Migrate to uv, lint fixes

Author: amrit110

.github/workflows/code_checks.yml

@@ -9,7 +9,7 @@ on:
       - .pre-commit-config.yaml
       - .github/workflows/code_checks.yml
       - '**.py'
-      - poetry.lock
+      - uv.lock
       - pyproject.toml
       - '**.ipynb'
   pull_request:
@@ -20,29 +20,30 @@ on:
       - .pre-commit-config.yaml
       - .github/workflows/code_checks.yml
       - '**.py'
-      - poetry.lock
+      - uv.lock
       - pyproject.toml
       - '**.ipynb'
 
 jobs:
   run-code-check:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4.1.1
-      - name: Install and configure Poetry
-        uses: snok/install-poetry@v1
+      - uses: actions/checkout@v4.2.2
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
         with:
-          virtualenvs-create: true
-          virtualenvs-in-project: true
-      - uses: actions/[email protected]
+          # Install a specific version of uv.
+          version: "0.5.21"
+          enable-cache: true
+      - name: "Set up Python"
+        uses: actions/setup-python@v5
         with:
-          python-version: '3.10'
-          cache: 'poetry'
+          python-version-file: ".python-version"
+      - name: Install the project
+        run: uv sync --dev
       - name: Install dependencies and check code
         run: |
-          poetry env use '3.10'
           source .venv/bin/activate
-          poetry install --with test --all-extras
           pre-commit run --all-files
       - name: pip-audit (gh-action-pip-audit)
         uses: pypa/[email protected]

.gitignore

@@ -81,9 +81,6 @@ target/
 profile_default/
 ipython_config.py
 
-# pyenv
-.python-version
-
 # pipenv
 #   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
 #   However, in case of collaboration, if having platform-specific dependencies or dependencies

.python-version

@@ -0,0 +1 @@
+3.10

examples/inference/llm/chat_completions.py

@@ -1,6 +1,10 @@
+"""Example of how to use the OpenAI API to generate chat completions."""
+
 from openai import OpenAI
 
-# The url is located in the .vLLM_model-variant_url file in the corresponding model directory.
+
+# The url is located in the .vLLM_model-variant_url file in the corresponding
+# model directory.
 client = OpenAI(base_url="http://gpuXXX:XXXX/v1", api_key="EMPTY")
 
 # Update the model path accordingly

examples/inference/llm/completions.py

@@ -1,6 +1,10 @@
+"""Example of how to use the OpenAI API to generate completions."""
+
 from openai import OpenAI
 
-# The url is located in the .vLLM_model-variant_url file in the corresponding model directory.
+
+# The url is located in the .vLLM_model-variant_url file in the corresponding
+# model directory.
 client = OpenAI(base_url="http://gpuXXX:XXXX/v1", api_key="EMPTY")
 
 # Update the model path accordingly

examples/inference/vlm/vision_completions.py

@@ -1,6 +1,10 @@
+"""Example of using the OpenAI API to generate completions for vision tasks."""
+
 from openai import OpenAI
 
-# The url is located in the .vLLM_model-variant_url file in the corresponding model directory.
+
+# The url is located in the .vLLM_model-variant_url file in the corresponding
+# model directory.
 client = OpenAI(base_url="http://gpuXXX:XXXX/v1", api_key="EMPTY")
 
 # Update the model path accordingly

examples/logits/logits.py

@@ -1,6 +1,10 @@
+"""Example of how to get logits from the model."""
+
 from openai import OpenAI
 
-# The url is located in the .vLLM_model-variant_url file in the corresponding model directory.
+
+# The url is located in the .vLLM_model-variant_url file in the corresponding
+# model directory.
 client = OpenAI(base_url="http://gpuXXX:XXXXX/v1", api_key="EMPTY")
 
 completion = client.completions.create(

poetry.lock

@@ -1,8 +1,11 @@
+"""Calculate the average prompt and generation throughput from a log file."""
+
 import argparse
 import re
 
 
-def filter_throughput(log_file_path):
+def filter_throughput(log_file_path: str) -> None:
+    """Filter log file for non-zero entries and calculate the avg throughput."""
     avg_prompt_throughput = []
     avg_generation_throughput = []
     # Define a regular expression pattern to extract throughput values
@@ -33,7 +36,8 @@ def filter_throughput(log_file_path):
     )
 
 
-def main():
+def main() -> None:
+    """Run the main function."""
     # Create the parser
     parser = argparse.ArgumentParser(
         description="Filter log file for non-zero throughput entries."

profile/avg_throughput.py

@@ -1,7 +1,11 @@
+"""Testing script."""
+
 import time
+from typing import List, Union
 
 import requests
 
+
 # Change the ENDPOINT and MODEL_PATH to match your setup
 ENDPOINT = "http://gpuXXX:XXXX/v1"
 MODEL_PATH = "Meta-Llama-3-70B"
@@ -71,19 +75,20 @@
 ]
 
 
-def send_request(prompt):
+def send_request(prompt: List[str]) -> Union[float, None]:
+    """Send a request to the API."""
     data = {"model": f"{MODEL_PATH}", "prompt": prompt, "max_tokens": 100}
     start_time = time.time()
     response = requests.post(f"{ENDPOINT}/completions", headers=HEADERS, json=data)
     duration = time.time() - start_time
     if response.status_code == 200:
         return duration
-    else:
-        return None
+    return None
 
 
-def main():
-    for i in range(10):
+def main() -> None:
+    """Run main function."""
+    for _ in range(10):
         print("Sending 20x requests 0-52...")
         send_request(PROMPTS * 20)
     print("Done!")