From de54e6b9a60d6338ac837b89e6e32a68eb5e50bd Mon Sep 17 00:00:00 2001
From: Xin Yang <xyang19@gmail.com>
Date: Tue, 23 Jul 2024 22:51:45 -0700
Subject: [PATCH] Update

---
 tests/integration/llm/client.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/integration/llm/client.py b/tests/integration/llm/client.py
index 34958c1ec9..9724d4de5d 100644
--- a/tests/integration/llm/client.py
+++ b/tests/integration/llm/client.py
@@ -907,7 +907,9 @@ def awscurl_run(data,
     if tokenizer:
         command = f"TOKENIZER={tokenizer} {command}"
     if output:
-        output_path = os.path.join(os.path.curdir, "outputs", "output")
+        output_dir = os.path.join(os.path.curdir, "outputs")
+        os.mkdir(output_dir)
+        output_path = os.path.join(output_dir, "output")
         command = f"{command} -o {output_path}"
     LOGGER.info(f"Running command {command}")
     res = sp.run(command, shell=True, stdout=sp.PIPE, stderr=sp.PIPE)
@@ -915,6 +917,8 @@ def awscurl_run(data,
         raise ValueError("Found error result in awscurl_run")
     if dataset:
         shutil.rmtree(dataset_dir)
+    if output:
+        shutil.rmtree(output_dir)
 
 
 def send_image_json(img_url, data):