check_vllm: create new check

This commit is contained in:
Cyberes 2023-09-20 15:11:50 -06:00
parent b69d83900c
commit c1cc8297c3
2 changed files with 103 additions and 3 deletions

99
check_vllm.py Normal file
View File

@ -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)

View File

@ -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
pytz~=2023.3
proxmoxer==2.0.1
tiktoken~=0.5.1