diff --git a/python/README.md b/python/README.md index 41c6be681..70b2f9194 100644 --- a/python/README.md +++ b/python/README.md @@ -26,6 +26,7 @@ These are some scripts on the side that might be useful for debugging or experim And: * `play.py` - A working very primitive GTP engine that just uses the raw neural net to make moves. Also implements a ton of GTP extensions to draw colored heatmaps in GoGUI. This or `test.py` are good places to look if you're interested in how to use the raw pytorch checkpoints for direct inference in python, rather than using the C++ KataGo engine. +* `flask_wrapper_human_sl_server.py`, `flask_wrapper_analysis_engine.py` - Flask web servers that forward POSTed JSON queries to `humanslnet_server.py` or to the C++ analysis engine (command in `KATAGO_COMMAND`). ### Distributed backend diff --git a/python/flask_wrapper_analysis_engine.py b/python/flask_wrapper_analysis_engine.py new file mode 100644 index 000000000..9d62ef789 --- /dev/null +++ b/python/flask_wrapper_analysis_engine.py @@ -0,0 +1,86 @@ +from flask import Flask, request, jsonify +from subprocess import Popen, PIPE, STDOUT +import threading +import json +import os + +app = Flask(__name__) + +# Store the process globally +cli_process = None +lock = threading.Lock() + +def start_cli(): + global cli_process + command = os.environ.get('KATAGO_COMMAND') + + if not command: + raise ValueError('Environment variable KATAGO_COMMAND not set') + + with lock: + if cli_process is None: + cli_process = Popen(command, shell=True, stdin=PIPE, stdout=PIPE, stderr=STDOUT, text=True) + wait_for_ready_state(cli_process) + +def wait_for_ready_state(process): + """Reads lines from the CLI program until it outputs the ready state message.""" + try: + while True: + line = process.stdout.readline() + if not line: + break + print(f"Initial CLI Output: {line.strip()}") # Print the initial output for debugging + + if "Started, ready to begin handling requests" in line: + print("CLI program is ready.") + break # Stop reading once the ready state is detected + except Exception as e: + print(f"Error reading initial output: {str(e)}") + +@app.route('/run', methods=['POST']) +def run_cli(): + global cli_process + input_data = request.json + + if not input_data: + return jsonify({'error': 'No input provided'}), 400 + + analyze_turns = input_data.get("analyzeTurns", []) + num_responses_expected = len(analyze_turns) + + with lock: + if cli_process is None: + return jsonify({'error': 'CLI program is not running'}), 400 + + try: + # Send input to the process + input_json = json.dumps(input_data) + cli_process.stdin.write(input_json + '\n') + cli_process.stdin.flush() + + # Read the output from the process until all responses are received + responses = [] + while len(responses) < num_responses_expected: + output_line = cli_process.stdout.readline() + if not output_line: + break + output_line = output_line.strip() + + try: + # Attempt to parse the output as JSON + response_data = json.loads(output_line) + responses.append(response_data) + except json.JSONDecodeError: + # Ignore lines that are not valid JSON + print(f"Non-JSON line ignored: {output_line}") + + if len(responses) != num_responses_expected: + return jsonify({'error': 'Did not receive the expected number of responses from CLI program'}), 500 + + return jsonify(responses) + except Exception as e: + return jsonify({'error': f'An error occurred: {str(e)}'}), 500 + +if __name__ == '__main__': + start_cli() # Start the CLI program when the server starts + app.run(host='0.0.0.0', port=5000) diff --git a/python/flask_wrapper_human_sl_server.py b/python/flask_wrapper_human_sl_server.py new file mode 100644 index 000000000..f8a247c25 --- /dev/null +++ b/python/flask_wrapper_human_sl_server.py @@ -0,0 +1,119 @@ +from flask import Flask, request, jsonify +from subprocess import Popen, PIPE +import time +import threading +import json +from flask_cors import CORS, cross_origin +import os + +app = Flask(__name__) +cors = CORS(app) +app.config['CORS_HEADERS'] = 'Content-Type' + +app = Flask(__name__) + +# Store the process globally +cli_process = None +lock = threading.RLock() + +def start_cli(): + global cli_process + global monitor_output_thread + command = "python ./humanslnet_server.py -checkpoint ./b18c384nbt-humanv0.ckpt -device cpu -webserver True" + + with lock: + if cli_process is None: + cli_process = Popen(command, shell=True, stdin=PIPE, stdout=PIPE, stderr=PIPE, text=True) + + wait_for_ready_state(cli_process) + stderr_thread = threading.Thread(target=monitor_stderr, args=(cli_process,)) + + stderr_thread.daemon = True + + stderr_thread.start() + + +def wait_for_ready_state(process): + """Reads lines from the CLI program until it outputs the ready state message.""" + try: + while True: + line = process.stdout.readline() + if not line: + break + print(f"Initial CLI Output: {line.strip()}") # Print the initial output for debugging + + if "Ready to receive input" in line: + print("CLI program is ready.") + break # Stop reading once the ready state is detected + except Exception as e: + print(f"Error reading initial output: {str(e)}") + +def monitor_stderr(process): + global cli_process + while True: + error_output = process.stderr.readline() + if error_output == '' and process.poll() is not None: + break + if error_output: + print(f"ERROR: {error_output.strip()}") + if "ERROR" in error_output or "Exception" in error_output or "Error" in error_output: + print("Error detected, quitting the process...") + os._exit(1) + time.sleep(0.1) + + +def restart_process(): + global cli_process + with lock: + if cli_process: + cli_process.kill() + cli_process = None + start_cli() + +@app.route('/', methods=['POST']) +@cross_origin() +def run_cli(): + global cli_process + input_data = request.json + + if not input_data: + return jsonify({'error': 'No input provided'}), 400 + + with lock: + if cli_process is None: + start_cli() + return jsonify({'error': 'CLI program was not running, restarted'}), 400 + + try: + # Send input to the process + input_json = json.dumps(input_data) + cli_process.stdin.write(input_json + '\n') + cli_process.stdin.flush() + + # Read the output from the process until all responses are received + num_responses_expected = 1 + responses = [] + while len(responses) < num_responses_expected: + output_line = cli_process.stdout.readline() + if not output_line: + break + output_line = output_line.strip() + + try: + # Attempt to parse the output as JSON + response_data = json.loads(output_line) + responses.append(response_data) + except json.JSONDecodeError: + # Ignore lines that are not valid JSON + print(f"Non-JSON line ignored: {output_line}") + + if len(responses) != num_responses_expected: + return jsonify({'error': 'Did not receive the expected number of responses from CLI program'}), 500 + + return jsonify(responses) + except Exception as e: + return jsonify({'error': f'An error occurred: {str(e)}'}), 500 + +if __name__ == '__main__': + start_cli() # Start the CLI program when the server starts + app.run(host='0.0.0.0', port=5000) diff --git a/python/forever_run_flask_wrapper_sl_server.sh b/python/forever_run_flask_wrapper_sl_server.sh new file mode 100755 index 000000000..36794d53c --- /dev/null +++ b/python/forever_run_flask_wrapper_sl_server.sh @@ -0,0 +1,7 @@ +#!/bin/bash + +while true; do + python flask_wrapper_human_sl_server.py + echo "The server crashed with exit code $?. Respawning.." >&2 + sleep 1 +done diff --git a/python/humanslnet_server.py b/python/humanslnet_server.py index ca7b3e33f..3509568a3 100644 --- a/python/humanslnet_server.py +++ b/python/humanslnet_server.py @@ -14,11 +14,16 @@ def numpy_array_encoder(obj): return float(obj) raise TypeError(f'Object of type {obj.__class__.__name__} is not JSON serializable') +def write(output): + sys.stdout.write(json.dumps(output,default=numpy_array_encoder) + "\n") + sys.stdout.flush() + def main(): parser = argparse.ArgumentParser() parser.add_argument('-checkpoint', help='Checkpoint to test', required=True) parser.add_argument('-use-swa', help='Use SWA model', action="store_true", required=False) parser.add_argument('-device', help='Device to use, such as cpu or cuda:0', required=True) + parser.add_argument('-webserver', help='set if is used for the flask wrapper', required=False) args = parser.parse_args() model, swa_model, _ = load_model(args.checkpoint, use_swa=args.use_swa, device=args.device, pos_len=19, verbose=False) @@ -26,10 +31,10 @@ def main(): model = swa_model game_state = None - def write(output): - sys.stdout.write(json.dumps(output,default=numpy_array_encoder) + "\n") - sys.stdout.flush() + if args.webserver: + write("Ready to receive input") + # DEBUGGING # game_state = GameState(board_size=19, rules=GameState.RULES_JAPANESE) # sgfmeta = SGFMetadata() @@ -66,8 +71,6 @@ def write(output): elif data["command"] == "get_model_outputs": sgfmeta = SGFMetadata.of_dict(data["sgfmeta"]) - # features = Features(model.config, model.pos_len) - # foo = game_state.get_input_features(features) outputs = game_state.get_model_outputs(model, sgfmeta=sgfmeta) filtered_outputs = {} for key in outputs: @@ -78,5 +81,6 @@ def write(output): else: raise ValueError(f"Unknown command: {data['command']}") + if __name__ == "__main__": main() diff --git a/python/requirements.txt b/python/requirements.txt new file mode 100644 index 000000000..cc1f11109 --- /dev/null +++ b/python/requirements.txt @@ -0,0 +1,5 @@ +matplotlib==3.9.0 +numpy==2.0.1 +torch==2.4.0 +wxPython==4.2.1 +flask_cors==4.0.1