From 061a857b3d0581e53bbe8b72d1e6aedcc89680b4 Mon Sep 17 00:00:00 2001
From: afourney <adamfo@microsoft.com>
Date: Fri, 29 Mar 2024 18:16:41 -0700
Subject: [PATCH] AutoGenBench: Handle Ctrl-C more gracefully.  (#2174)

* Prints the version of AutoGenBench from the command line, closing i1458

* Added autogenbench version to timestamp.txt

* Attempting to fix formatting.

* Add a gitignore for autogenbench

* Generalize to read all template dirs from Templates

* AutoGenBench logs telemetry when available.

* Remove spaces if present from template names.

* Bump version.

* Fixed formatting.

* Allow native warning to be skipped. Mount autogen repo in Docker if it can be found (experimental).

* Native execution now occurs in a venv.

* Bump version.

* Fixed a prompt escaping bug evident in GAIA task '6f37996b-2ac7-44b0-8e68-6d28256631b4'

* Updated all scenarios to use template discovery.

* Update with main version of runtime_logging.

* Better handling of Ctrl-C and cleanup of unused containers.

* Even stronger hinting that containers should be removed.

---------

Co-authored-by: gagb <gagb@users.noreply.github.com>
---
 .../autogenbench/autogenbench/run_cmd.py      | 62 ++++++++++++++-----
 .../autogenbench/autogenbench/version.py      |  2 +-
 2 files changed, 48 insertions(+), 16 deletions(-)

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"