From 8080828a84e945297cfdd92b3591cdfc409044ba Mon Sep 17 00:00:00 2001
From: Jason Lee <lixiaochuan20@gmail.com>
Date: Mon, 18 Mar 2024 00:02:41 +0800
Subject: [PATCH] update wandb settings

---
 lib_run_single.py | 12 +++++++-----
 run.py            |  6 +++---
 settings.json     |  2 +-
 3 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/lib_run_single.py b/lib_run_single.py
index ff9972d..ab9816e 100644
--- a/lib_run_single.py
+++ b/lib_run_single.py
@@ -14,8 +14,8 @@ with open("./settings.json", "r") as file:
     data = json.load(file)
 time_limit = data["time_limit"]
 
-@timeout(time_limit, use_signals=False)
-def run_single_example(agent, env, example, max_steps, instruction, args, example_result_dir, scores):
+# @timeout(time_limit, use_signals=False)
+def run_single_example(agent, env, example, max_steps, instruction, args, example_result_dir, scores, run):
     agent.reset()
     obs = env.reset(task_config=example)
     done = False
@@ -46,7 +46,7 @@ def run_single_example(agent, env, example, max_steps, instruction, args, exampl
             str_table.add_data(wandb.Image(data_or_path=os.path.join(example_result_dir, f"step_{step_idx + 1}_{action_timestamp}.png"), caption=f"step_{step_idx + 1}_{action_timestamp}"),
                             thisrun_a11tree,
                             response, action, action_timestamp, done)
-            wandb.log({"Reward": reward}) 
+            run.log({"Reward": reward}) 
             with open(os.path.join(example_result_dir, "traj.jsonl"), "a") as f:
                 f.write(json.dumps({
                     "step_num": step_idx + 1,
@@ -62,11 +62,13 @@ def run_single_example(agent, env, example, max_steps, instruction, args, exampl
                 logger.info("The episode is done.")
                 break
         step_idx += 1
-    wandb.log({"str_trajectory": str_table})
+    # wandb.log({"str_trajectory": str_table})
+    run.log({"str_trajectory": str_table})
     result = env.evaluate()
     logger.info("Result: %.2f", result)
     scores.append(result)
     with open(os.path.join(example_result_dir, "result.txt"), "w", encoding="utf-8") as f:
         f.write(f"{result}\n")
     env.controller.end_recording(os.path.join(example_result_dir, "recording.mp4"))
-    wandb.log({"Result": result})
+    run.log({"Result": result})
+    # wandb.log({"Result": result})
diff --git a/run.py b/run.py
index 4284169..505ae54 100644
--- a/run.py
+++ b/run.py
@@ -148,7 +148,7 @@ def test(
 
     for domain in tqdm(test_all_meta, desc="Domain"):
         for example_id in tqdm(test_all_meta[domain], desc="Example", leave=False):
-            wandb.init(project=f"OSworld-{args.action_space}-{args.observation_type}-{args.model}", group=f"{domain}", 
+            run = wandb.init(project=f"OSworld-{args.action_space}-{args.observation_type}-{args.model}", group=f"{domain}", 
                     name=f"{example_id}")
             # example setting
             config_file = os.path.join(args.test_config_base_dir, f"examples/{domain}/{example_id}.json")
@@ -164,7 +164,7 @@ def test(
             # wandb each example config settings
             cfg_args["instruction"] = instruction
             cfg_args["start_time"] = datetime.datetime.now().strftime("%Y:%m:%d-%H:%M:%S")
-            wandb.config.update(cfg_args)   
+            run.config.update(cfg_args)   
 
             example_result_dir = os.path.join(
                 args.result_dir,
@@ -178,7 +178,7 @@ def test(
             # example start running
             try:
                 lib_run_single.run_single_example(agent, env, example, max_steps, instruction, args, example_result_dir,
-                                                  scores)
+                                                  scores, run)
             except Exception as e:
                 logger.error(f"Exception in {domain}/{example_id}: {e}")
                 wandb.log({"Exception": wandb.Table(data=[[f"Exception in {domain}/{example_id}: {e}"]], columns=["Error"])})
diff --git a/settings.json b/settings.json
index 23bab77..7ee7a21 100644
--- a/settings.json
+++ b/settings.json
@@ -1,3 +1,3 @@
 {
-    "time_limit": "10"
+    "time_limit": "600"
 }
\ No newline at end of file