diff --git a/.github/Dockerfile.base b/.github/Dockerfile.base
index c0a01e6d69..cdbdee7cc7 100644
--- a/.github/Dockerfile.base
+++ b/.github/Dockerfile.base
@@ -28,7 +28,10 @@ RUN apt-get update && apt-get install -y \
     graphviz \
     patchelf \
     libyaml-cpp-dev \
-    libboost-all-dev
+    libboost-all-dev \
+    curl \
+    jq \
+    sudo
 
 # Install clang 17
 RUN wget https://apt.llvm.org/llvm.sh && \
diff --git a/.github/Dockerfile.ird b/.github/Dockerfile.ird
index d7528b9987..4ffcb9e997 100644
--- a/.github/Dockerfile.ird
+++ b/.github/Dockerfile.ird
@@ -13,7 +13,8 @@ RUN apt-get update && apt-get install -y \
     ssh \
     sudo \
     wget \
-    vim
+    vim \
+    gh
 
 # Create a directory for the toolchain and set permissions
 RUN mkdir -p $TTMLIR_TOOLCHAIN_DIR && \
diff --git a/test/lit.cfg.py b/test/lit.cfg.py
index defdda42ef..d789add18e 100644
--- a/test/lit.cfg.py
+++ b/test/lit.cfg.py
@@ -11,27 +11,16 @@
 import lit.formats
 import lit.util
 from lit.llvm import llvm_config
-from lit.formats.base import TestFormat
 
+# Add the current folder to PYTHONPATH
+current_folder = os.path.dirname(__file__)
+sys.path.append(current_folder)
+os.environ["PYTHONPATH"] = current_folder + os.pathsep + os.environ.get("PYTHONPATH", "")
+from timestamped_test_format import TimestampedTestFormat
 
-# Configuration file for the 'lit' test runner.
-
-class TimestampedTestFormat(TestFormat):
-    def __init__(self):
-        super().__init__('Timestamped Test')
-
-    def execute(self, test, lit_config):
-        start_timestamp = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S%z")
-        test.setMetric('start_timestamp', start_timestamp)
 
-        # Run the test
-        result = lit.TestRunner.executeShTest(test, lit_config)
-
-        end_timestamp = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S%z")
-        test.setMetric('end_timestamp', end_timestamp)
+# Configuration file for the 'lit' test runner.
 
-        return result
-    
 def set_system_desc_features(system_desc):
     config.available_features.add(system_desc["chip_descs"][0]["arch"])
     if len(system_desc["chip_desc_indices"]) > 1:
@@ -40,8 +29,8 @@ def set_system_desc_features(system_desc):
 
 # name: The name of this test suite.
 config.name = "TTMLIR"
-
 config.test_format = TimestampedTestFormat()
+# config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)
 
 # Stablehlo tests can be optionally enabled.
 if config.enable_stablehlo:
diff --git a/test/timestamped_test_format.py b/test/timestamped_test_format.py
new file mode 100644
index 0000000000..75a29507dd
--- /dev/null
+++ b/test/timestamped_test_format.py
@@ -0,0 +1,22 @@
+import datetime
+import lit.Test
+from lit.formats.shtest import ShTest
+from lit.Test import toMetricValue
+import lit.TestRunner
+
+from unittest.mock import patch
+
+class TimestampedTestFormat(ShTest):
+    def __init__(self):
+        super().__init__()
+
+    def execute(self, test, lit_config):
+        start_timestamp = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S%z")
+        # Run the test
+        result = lit.TestRunner.executeShTest(test, lit_config, True)
+        end_timestamp = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S%z")
+        result.addMetric("start_timestamp", toMetricValue(start_timestamp))
+        result.addMetric("end_timestamp", toMetricValue(end_timestamp))
+
+        return result
+    
\ No newline at end of file