diff --git a/samples/tools/autogenbench/autogenbench/run_cmd.py b/samples/tools/autogenbench/autogenbench/run_cmd.py index a8674ea710b1..d8f9faeacfd6 100644 --- a/samples/tools/autogenbench/autogenbench/run_cmd.py +++ b/samples/tools/autogenbench/autogenbench/run_cmd.py @@ -485,7 +485,14 @@ def run_scenario_in_docker(work_dir, env, timeout=TASK_TIMEOUT, docker_image=Non # Create and run the container container = client.containers.run( - image, command=["sh", "run.sh"], working_dir="/workspace", environment=env, detach=True, volumes=volumes + image, + command=["sh", "run.sh"], + working_dir="/workspace", + environment=env, + detach=True, + remove=True, + auto_remove=True, + volumes=volumes, ) # Read the logs in a streaming fashion. Keep an eye on the time to make sure we don't need to stop. @@ -494,23 +501,45 @@ def run_scenario_in_docker(work_dir, env, timeout=TASK_TIMEOUT, docker_image=Non logs = container.logs(stream=True) log_file = open(os.path.join(work_dir, "console_log.txt"), "wt", encoding="utf-8") stopping = False + exiting = False - for chunk in logs: # When streaming it should return a generator - # Stream the data to the log file and the console - chunk = chunk.decode("utf-8") - log_file.write(chunk) - log_file.flush() - sys.stdout.reconfigure(encoding="utf-8") - sys.stdout.write(chunk) - sys.stdout.flush() - - # Check if we need to terminate - if not stopping and time.time() - start_time >= docker_timeout: + while True: + try: + chunk = next(logs) # Manually step the iterator so it is captures with the try-catch + + # Stream the data to the log file and the console + chunk = chunk.decode("utf-8") + log_file.write(chunk) + log_file.flush() + sys.stdout.reconfigure(encoding="utf-8") + sys.stdout.write(chunk) + sys.stdout.flush() + + # Check if we need to terminate + if not stopping and time.time() - start_time >= docker_timeout: + container.stop() + + # Don't exit the loop right away, as there are things we may still want to read from the logs + # but remember how we got here. + stopping = True + except KeyboardInterrupt: + log_file.write("\nKeyboard interrupt (Ctrl-C). Attempting to exit gracefully.\n") + log_file.flush() + sys.stdout.write("\nKeyboard interrupt (Ctrl-C). Attempting to exit gracefully.\n") + sys.stdout.flush() + + # Start the exit process, and give it a minute, but keep iterating container.stop() + exiting = True + docker_timeout = time.time() - start_time + 60 + except StopIteration: + break - # Don't exit the loop right away, as there are things we may still want to read from the logs - # but remember how we got here. - stopping = True + # Clean up the container + try: + container.remove() + except docker.errors.APIError: + pass if stopping: # By this line we've exited the loop, and the container has actually stopped. log_file.write("\nDocker timed out.\n") @@ -518,6 +547,9 @@ def run_scenario_in_docker(work_dir, env, timeout=TASK_TIMEOUT, docker_image=Non sys.stdout.write("\nDocker timed out.\n") sys.stdout.flush() + if exiting: # User hit ctrl-C + sys.exit(1) + def build_default_docker_image(docker_client, image_tag): for segment in docker_client.api.build( diff --git a/samples/tools/autogenbench/autogenbench/version.py b/samples/tools/autogenbench/autogenbench/version.py index 3b93d0be0c9e..27fdca497c37 100644 --- a/samples/tools/autogenbench/autogenbench/version.py +++ b/samples/tools/autogenbench/autogenbench/version.py @@ -1 +1 @@ -__version__ = "0.0.2" +__version__ = "0.0.3"