diff --git a/check_vllm.py b/check_vllm.py index 3ba3b04..55654e2 100644 --- a/check_vllm.py +++ b/check_vllm.py @@ -45,33 +45,40 @@ def main(critical, warning, timeout, target_url, does_not_contain, verify): } } if response.status_code != 200: - text_result = f"failed with status code {response.status_code}.\n{response.text}" + 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]}" + if not json_data.get('text'): + text_result = f"did not receive valid JSON\n{response.text}" 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 + if 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"CRITICAL - request failed - {error.__class__.__name__}: {error}" @@ -87,12 +94,12 @@ if __name__ == "__main__": 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') + 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.verify) + 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())