From c1cc8297c363e5c9ddf6f38b64d8a39dc0d3e29e Mon Sep 17 00:00:00 2001
From: Cyberes <cyberes@evulid.cc>
Date: Wed, 20 Sep 2023 15:11:50 -0600
Subject: [PATCH] check_vllm: create new check

---
 check_vllm.py    | 99 ++++++++++++++++++++++++++++++++++++++++++++++++
 requirements.txt |  7 ++--
 2 files changed, 103 insertions(+), 3 deletions(-)
 create mode 100644 check_vllm.py

diff --git a/check_vllm.py b/check_vllm.py
new file mode 100644
index 0000000..3ba3b04
--- /dev/null
+++ b/check_vllm.py
@@ -0,0 +1,99 @@
+import argparse
+import json
+import sys
+import time
+import traceback
+
+import requests
+import tiktoken
+import urllib3
+
+from checker import nagios
+from checker import print_icinga2_check_status
+
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+
+def main(critical, warning, timeout, target_url, does_not_contain, verify):
+    headers = {'Content-Type': 'application/json'}
+    data = {
+        "prompt": "Test prompt",
+        "stream": False,
+        "temperature": 0,
+        "max_tokens": 50,
+    }
+
+    start_time = time.time()
+    try:
+        response = requests.post(target_url, headers=headers, data=json.dumps(data), verify=verify, timeout=timeout)
+        json_data = response.json()
+        error = None
+    except Exception as e:
+        response = None
+        json_data = None
+        error = e
+    end_time = time.time()
+    elapsed_time = round(end_time - start_time, 2)
+    exit_code = nagios.STATE_OK
+
+    if response:
+        perfdata = {
+            'generation_time': {
+                'value': elapsed_time,
+                'min': 0,
+                'unit': 's'
+            }
+        }
+        if response.status_code != 200:
+            text_result = f"failed with status code {response.status_code}.\n{response.text}"
+            exit_code = nagios.STATE_CRIT
+        else:
+            tokenizer = tiktoken.get_encoding("cl100k_base")
+            num_tokens = len(tokenizer.encode(json_data['text'][0]))
+            perfdata.update({
+                'tokens': {
+                    'value': num_tokens,
+                    'min': 0,
+                }
+            })
+
+            if not len(json_data['text'][0]):
+                text_result = f"response was empty.\n{response.text}"
+                exit_code = nagios.STATE_CRIT
+            elif elapsed_time >= warning:
+                text_result = f"response was {elapsed_time} seconds"
+                exit_code = nagios.STATE_WARN
+            elif elapsed_time >= critical:
+                text_result = f"response was {elapsed_time} seconds"
+                exit_code = nagios.STATE_CRIT
+            elif does_not_contain and does_not_contain in json_data['text'][0]:
+                text_result = f"\"{does_not_contain}\" was in the response:\n{json_data['text'][0]}"
+                exit_code = nagios.STATE_CRIT
+            else:
+                text_result = f"generation time was {elapsed_time} seconds and generated {num_tokens} tokens"
+                exit_code = nagios.STATE_OK
+    else:
+        perfdata = {}
+        text_result = f"CRITICAL - request failed - {error.__class__.__name__}: {error}"
+        exit_code = nagios.STATE_CRIT
+
+    print_icinga2_check_status(text_result, exit_code, perfdata)
+    sys.exit(exit_code)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Check DNS-over-TLS connectivity and response time')
+    parser.add_argument('-c', '--critical', type=float, default=60, help='Critical threshold for response time in seconds (default: 60)')
+    parser.add_argument('-w', '--warning', type=float, default=30, help='Warning threshold for response time in seconds (default: 30)')
+    parser.add_argument('-t', '--target-url', type=str, required=True, help='Your completion API endpoint')
+    parser.add_argument('-m', '--timeout', type=float, default=300, help='Request timeout in seconds (default: 300)')
+    parser.add_argument('-v', '--verify', action='store_false', help='Do not verify SSL')
+    parser.add_argument('-n', '--does-not-contain', type=str, help='Response must not contain')
+    args = parser.parse_args()
+
+    try:
+        main(args.critical, args.warning, args.timeout, args.target_url, args.does_not_contain, args.verify)
+    except Exception as e:
+        print(f'UNKNOWN: exception "{e}"')
+        print(traceback.format_exc())
+        sys.exit(nagios.STATE_UNKNOWN)
diff --git a/requirements.txt b/requirements.txt
index 51ec0e3..9d00332 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -15,7 +15,8 @@ hurry.filesize
 certifi~=2022.12.7
 cloudflarepycli~=1.7.0
 zfslib~=0.11.0
-dnspython
+dnspython~=2.3.0
 cf-speedtest
-pytz
-proxmoxer==2.0.1
\ No newline at end of file
+pytz~=2023.3
+proxmoxer==2.0.1
+tiktoken~=0.5.1
\ No newline at end of file