Use tokio::time::Duration and tokio::time::Instant for timekeeping

Summary: Previously we had to use u64 for serialization reasons but those reasons no longer exist

Differential Revision: D80556690

Author: thomasywang

hyperactor/src/channel/sim.rs

@@ -22,7 +22,6 @@ use super::*;
 use crate::channel;
 use crate::clock::Clock;
 use crate::clock::RealClock;
-use crate::clock::SimClock;
 use crate::data::Serialized;
 use crate::mailbox::MessageEnvelope;
 use crate::simnet;
@@ -129,7 +128,7 @@ pub(crate) struct MessageDeliveryEvent {
     src_addr: Option<ChannelAddr>,
     dest_addr: ChannelAddr,
     data: Serialized,
-    duration_ms: u64,
+    duration: tokio::time::Duration,
 }
 
 impl MessageDeliveryEvent {
@@ -139,7 +138,7 @@ impl MessageDeliveryEvent {
             src_addr,
             dest_addr,
             data,
-            duration_ms: 100,
+            duration: tokio::time::Duration::from_millis(100),
         }
     }
 }
@@ -158,8 +157,8 @@ impl Event for MessageDeliveryEvent {
         Ok(())
     }
 
-    fn duration_ms(&self) -> u64 {
-        self.duration_ms
+    fn duration(&self) -> tokio::time::Duration {
+        self.duration
     }
 
     fn summary(&self) -> String {
@@ -178,12 +177,12 @@ impl Event for MessageDeliveryEvent {
                 src: src_addr.clone(),
                 dst: self.dest_addr.clone(),
             };
-            self.duration_ms = topology
+            self.duration = topology
                 .lock()
                 .await
                 .topology
                 .get(&edge)
-                .map_or_else(|| 1, |v| v.latency.as_millis() as u64);
+                .map_or_else(|| tokio::time::Duration::from_millis(1), |v| v.latency);
         }
     }
 }
@@ -332,7 +331,7 @@ impl<M: RemoteMessage> Tx<M> for SimTx<M> {
                         self.dst_addr.clone(),
                         data,
                     )),
-                    time: SimClock.millis_since_start(RealClock.now()),
+                    time: RealClock.now(),
                 }),
                 _ => handle.send_event(Box::new(MessageDeliveryEvent::new(
                     self.src_addr.clone(),
@@ -551,7 +550,10 @@ mod tests {
             .await
             .unwrap();
 
-        assert_eq!(SimClock.millis_since_start(RealClock.now()), 0);
+        assert_eq!(
+            SimClock.duration_since_start(RealClock.now()),
+            tokio::time::Duration::ZERO
+        );
         // Fast forward real time to 5 seconds
         tokio::time::advance(tokio::time::Duration::from_secs(5)).await;
         {

hyperactor/src/clock.rs

@@ -197,11 +197,7 @@ impl Clock for SimClock {
 
         simnet_handle()
             .unwrap()
-            .send_event(SleepEvent::new(
-                tx.bind(),
-                mailbox,
-                duration.as_millis() as u64,
-            ))
+            .send_event(SleepEvent::new(tx.bind(), mailbox, duration))
             .unwrap();
         rx.recv().await.unwrap();
     }
@@ -212,11 +208,7 @@ impl Clock for SimClock {
 
         simnet_handle()
             .unwrap()
-            .send_nonadvanceable_event(SleepEvent::new(
-                tx.bind(),
-                mailbox,
-                duration.as_millis() as u64,
-            ))
+            .send_nonadvanceable_event(SleepEvent::new(tx.bind(), mailbox, duration))
             .unwrap();
         rx.recv().await.unwrap();
     }
@@ -247,11 +239,7 @@ impl Clock for SimClock {
 
         simnet_handle()
             .unwrap()
-            .send_event(SleepEvent::new(
-                tx.bind(),
-                mailbox,
-                duration.as_millis() as u64,
-            ))
+            .send_event(SleepEvent::new(tx.bind(), mailbox, duration))
             .unwrap();
 
         let fut = f;
@@ -290,14 +278,19 @@ impl SimClock {
     }
 
     /// Advance the sumulator's time to the specified instant
-    pub fn advance_to(&self, millis: u64) {
+    pub fn advance_to(&self, time: tokio::time::Instant) {
         let mut guard = SIM_TIME.now.lock().unwrap();
-        *guard = SIM_TIME.start + tokio::time::Duration::from_millis(millis);
+        *guard = time;
     }
 
     /// Get the number of milliseconds elapsed since the start of the simulation
-    pub fn millis_since_start(&self, instant: tokio::time::Instant) -> u64 {
-        instant.duration_since(SIM_TIME.start).as_millis() as u64
+    pub fn duration_since_start(&self, instant: tokio::time::Instant) -> tokio::time::Duration {
+        instant.duration_since(SIM_TIME.start)
+    }
+
+    /// Instant marking the start of the simulation
+    pub fn start(&self) -> tokio::time::Instant {
+        SIM_TIME.start.clone()
     }
 }
 
@@ -347,10 +340,16 @@ mod tests {
     #[tokio::test]
     async fn test_sim_clock_simple() {
         let start = SimClock.now();
-        assert_eq!(SimClock.millis_since_start(start), 0);
-        SimClock.advance_to(10000);
+        assert_eq!(
+            SimClock.duration_since_start(start),
+            tokio::time::Duration::ZERO
+        );
+        SimClock.advance_to(SimClock.start() + tokio::time::Duration::from_millis(10000));
         let end = SimClock.now();
-        assert_eq!(SimClock.millis_since_start(end), 10000);
+        assert_eq!(
+            SimClock.duration_since_start(end),
+            tokio::time::Duration::from_millis(10000)
+        );
         assert_eq!(
             end.duration_since(start),
             tokio::time::Duration::from_secs(10)
@@ -360,7 +359,7 @@ mod tests {
     #[tokio::test]
     async fn test_sim_clock_system_time() {
         let start = SimClock.system_time_now();
-        SimClock.advance_to(10000);
+        SimClock.advance_to(SimClock.start() + tokio::time::Duration::from_millis(10000));
         let end = SimClock.system_time_now();
         assert_eq!(
             end.duration_since(start).unwrap(),

hyperactor/src/simnet.rs

@@ -70,7 +70,7 @@ const OPERATIONAL_MESSAGE_BUFFER_SIZE: usize = 8;
 pub trait Address: Hash + Debug + Eq + PartialEq + Ord + PartialOrd + Clone {}
 impl<A: Hash + Debug + Eq + PartialEq + Ord + PartialOrd + Clone> Address for A {}
 
-type SimulatorTimeInstant = u64;
+type SimulatorTimeInstant = tokio::time::Instant;
 
 /// The unit of execution for the simulator.
 /// Using handle(), simnet can schedule executions in the network.
@@ -99,7 +99,7 @@ pub trait Event: Send + Sync + Debug {
 
     /// The latency of the event. This could be network latency, induced latency (sleep), or
     /// GPU work latency.
-    fn duration_ms(&self) -> u64;
+    fn duration(&self) -> tokio::time::Duration;
 
     /// Read the simnet config and update self accordingly.
     async fn read_simnet_config(&mut self, _topology: &Arc<Mutex<SimNetConfig>>) {}
@@ -126,8 +126,8 @@ impl Event for NodeJoinEvent {
         self.handle().await
     }
 
-    fn duration_ms(&self) -> u64 {
-        0
+    fn duration(&self) -> tokio::time::Duration {
+        tokio::time::Duration::ZERO
     }
 
     fn summary(&self) -> String {
@@ -139,15 +139,19 @@ impl Event for NodeJoinEvent {
 pub(crate) struct SleepEvent {
     done_tx: OncePortRef<()>,
     mailbox: Mailbox,
-    duration_ms: u64,
+    duration: tokio::time::Duration,
 }
 
 impl SleepEvent {
-    pub(crate) fn new(done_tx: OncePortRef<()>, mailbox: Mailbox, duration_ms: u64) -> Box<Self> {
+    pub(crate) fn new(
+        done_tx: OncePortRef<()>,
+        mailbox: Mailbox,
+        duration: tokio::time::Duration,
+    ) -> Box<Self> {
         Box::new(Self {
             done_tx,
             mailbox,
-            duration_ms,
+            duration,
         })
     }
 }
@@ -166,12 +170,12 @@ impl Event for SleepEvent {
         Ok(())
     }
 
-    fn duration_ms(&self) -> u64 {
-        self.duration_ms
+    fn duration(&self) -> tokio::time::Duration {
+        self.duration
     }
 
     fn summary(&self) -> String {
-        format!("Sleeping for {} ms", self.duration_ms)
+        format!("Sleeping for {} ms", self.duration.as_millis())
     }
 }
 
@@ -200,8 +204,8 @@ impl Event for TorchOpEvent {
         Ok(())
     }
 
-    fn duration_ms(&self) -> u64 {
-        100
+    fn duration(&self) -> tokio::time::Duration {
+        tokio::time::Duration::from_millis(100)
     }
 
     fn summary(&self) -> String {
@@ -561,22 +565,20 @@ impl SimNet {
         // Get latency
         event.read_simnet_config(&self.config).await;
         ScheduledEvent {
-            time: SimClock.millis_since_start(
-                SimClock.now() + tokio::time::Duration::from_millis(event.duration_ms()),
-            ),
+            time: SimClock.now() + event.duration(),
             event,
         }
     }
 
     /// Schedule the event into the network.
     fn schedule_event(&mut self, scheduled_event: ScheduledEvent, advanceable: bool) {
-        let start_at = SimClock.millis_since_start(SimClock.now());
+        let start_at = SimClock.now();
         let end_at = scheduled_event.time;
 
         self.records.push(SimulatorEventRecord {
             summary: scheduled_event.event.summary(),
-            start_at,
-            end_at,
+            start_at: SimClock.duration_since_start(start_at).as_millis() as u64,
+            end_at: SimClock.duration_since_start(end_at).as_millis() as u64,
         });
 
         if advanceable {
@@ -604,7 +606,7 @@ impl SimNet {
     ) -> Vec<SimulatorEventRecord> {
         // The simulated number of milliseconds the training script
         // has spent waiting for the backend to resolve a future
-        let mut training_script_waiting_time: u64 = 0;
+        let mut training_script_waiting_time = tokio::time::Duration::from_millis(0);
         // Duration elapsed while only non_advanceable_events has events
         let mut debounce_timer: Option<tokio::time::Instant> = None;
         'outer: loop {
@@ -638,9 +640,7 @@ impl SimNet {
                         .scheduled_events
                         .first_key_value()
                         .is_some_and(|(time, _)| {
-                            *time
-                                > SimClock.millis_since_start(RealClock.now())
-                                    + training_script_waiting_time
+                            *time > RealClock.now() + training_script_waiting_time
                         })
                 {
                     tokio::task::yield_now().await;
@@ -705,8 +705,7 @@ impl SimNet {
                     continue;
                 };
                 if training_script_state_rx.borrow().is_waiting() {
-                    let advanced_time =
-                        scheduled_time - SimClock.millis_since_start(SimClock.now());
+                    let advanced_time = scheduled_time - SimClock.now();
                     training_script_waiting_time += advanced_time;
                 }
                 SimClock.advance_to(scheduled_time);
@@ -749,9 +748,9 @@ pub struct SimulatorEventRecord {
     /// Event dependent summary for user
     pub summary: String,
     /// The time at which the message delivery was started.
-    pub start_at: SimulatorTimeInstant,
+    pub start_at: u64,
     /// The time at which the message was delivered to the receiver.
-    pub end_at: SimulatorTimeInstant,
+    pub end_at: u64,
 }
 
 /// A configuration for the network topology.
@@ -805,7 +804,7 @@ mod tests {
         src_addr: SimAddr,
         dest_addr: SimAddr,
         data: Serialized,
-        duration_ms: u64,
+        duration: tokio::time::Duration,
         dispatcher: Option<TestDispatcher>,
     }
 
@@ -823,8 +822,8 @@ mod tests {
             }
             Ok(())
         }
-        fn duration_ms(&self) -> u64 {
-            self.duration_ms
+        fn duration(&self) -> tokio::time::Duration {
+            self.duration
         }
 
         fn summary(&self) -> String {
@@ -840,12 +839,12 @@ mod tests {
                 src: self.src_addr.addr().clone(),
                 dst: self.dest_addr.addr().clone(),
             };
-            self.duration_ms = config
+            self.duration = config
                 .lock()
                 .await
                 .topology
                 .get(&edge)
-                .map_or_else(|| 1, |v| v.latency.as_millis() as u64);
+                .map_or_else(|| tokio::time::Duration::from_millis(1), |v| v.latency);
         }
     }
 
@@ -860,7 +859,7 @@ mod tests {
                 src_addr,
                 dest_addr,
                 data,
-                duration_ms: 1,
+                duration: tokio::time::Duration::from_millis(1),
                 dispatcher,
             }
         }
@@ -1132,12 +1131,18 @@ mod tests {
         start();
 
         let start = SimClock.now();
-        assert_eq!(SimClock.millis_since_start(start), 0);
+        assert_eq!(
+            SimClock.duration_since_start(start),
+            tokio::time::Duration::ZERO
+        );
 
         SimClock.sleep(tokio::time::Duration::from_secs(10)).await;
 
         let end = SimClock.now();
-        assert_eq!(SimClock.millis_since_start(end), 10000);
+        assert_eq!(
+            SimClock.duration_since_start(end),
+            tokio::time::Duration::from_secs(10)
+        );
     }
 
     #[tokio::test]