icinga2-checks/check_vllm.py

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": 3,
    }

    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:
            if not json_data.get('text') or not len(json_data['text']):
                text_result = f"did not receive valid JSON\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"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=200, help='Request timeout in seconds (default: 200)')
    parser.add_argument('-v', '--no-verify', action='store_false', help='Do not verify SSL')
    parser.add_argument('-n', '--does-not-contain', type=str, help='Response must not contain this string')
    args = parser.parse_args()

    try:
        main(args.critical, args.warning, args.timeout, args.target_url, args.does_not_contain, args.no_verify)
    except Exception as e:
        print(f'UNKNOWN: exception "{e}"')
        print(traceback.format_exc())
        sys.exit(nagios.STATE_UNKNOWN)