From 4bbe5d3567bfd94770b2310020308d91f7443060 Mon Sep 17 00:00:00 2001
From: Dushyant Behl <dushyantbehl@users.noreply.github.com>
Date: Fri, 23 Feb 2024 14:08:57 +0530
Subject: [PATCH] change default output path for aim run export

Signed-off-by: Dushyant Behl <dushyantbehl@users.noreply.github.com>
---
 tuning/tracker/aimstack_tracker.py | 19 ++++++++++++++-----
 tuning/tracker/tracker.py          |  3 ++-
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/tuning/tracker/aimstack_tracker.py b/tuning/tracker/aimstack_tracker.py
index 7640b1db7..caf8fc4a1 100644
--- a/tuning/tracker/aimstack_tracker.py
+++ b/tuning/tracker/aimstack_tracker.py
@@ -11,7 +11,7 @@ class CustomAimCallback(AimCallback):
 
     # A path to export run hash generated by Aim
     # This is used to link back to the expriments from outside aimstack
-    aim_run_hash_export_path = None
+    run_hash_export_path = None
 
     def on_init_end(self, args, state, control, **kwargs):
 
@@ -20,9 +20,18 @@ def on_init_end(self, args, state, control, **kwargs):
 
         self.setup() # initializes the run_hash
 
-        # store the run hash
-        if self.aim_run_hash_export_path:
-            with open(self.aim_run_hash_export_path, 'w') as f:
+        # Store the run hash
+        # Change default run hash path to output directory
+        if self.run_hash_export_path is None:
+            if args and args.output_dir:
+                # args.output_dir/.aim_run_hash
+                self.run_hash_export_path = os.path.join(
+                                                args.output_dir,
+                                                '.aim_run_hash'
+                                            )
+
+        if self.run_hash_export_path:
+            with open(self.run_hash_export_path, 'w') as f:
                 f.write('{\"run_hash\":\"'+str(self._run.hash)+'\"}\n')
 
     def on_train_begin(self, args, state, control, model=None, **kwargs):
@@ -60,7 +69,7 @@ def get_hf_callback(self):
         else:
             aim_callback = CustomAimCallback(experiment=exp)
 
-        aim_callback.aim_run_hash_export_path = hash_export_path
+        aim_callback.run_hash_export_path = hash_export_path
         self.hf_callback = aim_callback
         return self.hf_callback
 
diff --git a/tuning/tracker/tracker.py b/tuning/tracker/tracker.py
index a1d2d2ba0..3a2a18ed2 100644
--- a/tuning/tracker/tracker.py
+++ b/tuning/tracker/tracker.py
@@ -9,7 +9,8 @@ def __init__(self, name=None, tracker_config=None) -> None:
         else:
             self._name = name
 
-    def get_hf_callback():
+    # we use args here to denote any argument.
+    def get_hf_callback(self):
         return None
 
     def track(self, metric, name, stage):