From e18cecd42d348b4f6bc45b682b336a868bb044af Mon Sep 17 00:00:00 2001
From: Alex Filby <afilby@nvidia.com>
Date: Mon, 20 Jul 2026 15:49:17 -0500
Subject: [PATCH] fix: use nanosecond experiment timestamps

Reduce experiment directory collisions by using time_ns for generated and reset experiment IDs. Add focused coverage for both paths.

Signed-off-by: Alex Filby <afilby@nvidia.com>
---
 nemo_run/run/experiment.py  |  4 ++--
 test/run/test_experiment.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/nemo_run/run/experiment.py b/nemo_run/run/experiment.py
index ba430413..d963dd8e 100644
--- a/nemo_run/run/experiment.py
+++ b/nemo_run/run/experiment.py
@@ -333,7 +333,7 @@ def __init__(
             assert id, "Cannot reconstruct an experiment without id."
 
         self._title = title
-        self._id = id or f"{title}_{int(time.time())}"
+        self._id = id or f"{title}_{time.time_ns()}"
         self._enable_goodbye_message = enable_goodbye_message
         self._threadpool_workers = threadpool_workers
         self._skip_status_at_exit = skip_status_at_exit
@@ -1068,7 +1068,7 @@ def reset(self) -> "Experiment":
             return self
 
         old_id, old_exp_dir, old_launched = self._id, self._exp_dir, self._launched
-        self._id = f"{self._title}_{int(time.time())}"
+        self._id = f"{self._title}_{time.time_ns()}"
         self._exp_dir = os.path.join(get_nemorun_home(), "experiments", self._title, self._id)
         self._launched = False
         self._live_progress = None
diff --git a/test/run/test_experiment.py b/test/run/test_experiment.py
index 4f160237..15977db0 100644
--- a/test/run/test_experiment.py
+++ b/test/run/test_experiment.py
@@ -82,6 +82,18 @@ def test_experiment_creation(temp_dir):
     assert isinstance(exp.executor, LocalExecutor)
 
 
+def test_experiment_creation_uses_nanosecond_timestamp(temp_dir):
+    timestamp_ns = 1_753_041_987_123_456_789
+
+    with patch("nemo_run.run.experiment.time.time_ns", return_value=timestamp_ns):
+        exp = Experiment("test-exp")
+
+    assert exp._id == f"test-exp_{timestamp_ns}"
+    assert exp._exp_dir == os.path.join(
+        temp_dir, "experiments", "test-exp", f"test-exp_{timestamp_ns}"
+    )
+
+
 def test_experiment_with_custom_id(temp_dir):
     """Test creating an experiment with a custom id."""
     exp = Experiment("test-exp", id="custom-id")
@@ -345,6 +357,22 @@ def test_reset_not_run_experiment(temp_dir):
             assert reset_exp is exp  # The implementation returns self now
 
 
+def test_reset_uses_nanosecond_timestamp(temp_dir):
+    timestamp_ns = 1_753_041_987_987_654_321
+    exp = Experiment("test-exp", id="test-exp_original")
+    exp._prepare()
+    Path(os.path.join(exp._exp_dir, Experiment._DONE_FILE)).touch()
+
+    with patch("nemo_run.run.experiment.time.time_ns", return_value=timestamp_ns):
+        reset_exp = exp.reset()
+
+    assert reset_exp is exp
+    assert exp._id == f"test-exp_{timestamp_ns}"
+    assert exp._exp_dir == os.path.join(
+        temp_dir, "experiments", "test-exp", f"test-exp_{timestamp_ns}"
+    )
+
+
 @patch("nemo_run.run.experiment.get_runner")
 def test_experiment_from_id(mock_get_runner, temp_dir):
     """Test reconstructing an experiment from its ID."""