Track timely master: un-shred arrange with multi-capability stamps

Cargo.toml

@@ -23,9 +23,9 @@ rust-version = "1.86"
 
 [workspace.dependencies]
 differential-dataflow = { path = "differential-dataflow", default-features = false, version = "0.25.1" }
-timely = { version = "0.31", default-features = false }
+#timely = { version = "0.31", default-features = false }
 columnar = { version = "0.13", default-features = false }
-#timely = { git = "https://github.com/TimelyDataflow/timely-dataflow", default-features = false }
+timely = { git = "https://github.com/TimelyDataflow/timely-dataflow", default-features = false }
 #timely = { path = "../timely-dataflow/timely/", default-features = false }
 
 [workspace.lints.clippy]

diagnostics/src/logging.rs

@@ -72,15 +72,15 @@ impl ClientInput {
     pub fn connect(&mut self, client_id: usize, elapsed: Duration) {
         let _ = self
             .sender
-            .send(Event::Messages(self.time, vec![(client_id, elapsed, 1)]));
+            .send(Event::Messages(timely::progress::Stamp::from_elem(self.time), vec![(client_id, elapsed, 1)]));
         self.advance(elapsed);
     }
 
     /// Announce a client disconnection.
     pub fn disconnect(&mut self, client_id: usize, elapsed: Duration) {
         let _ = self
             .sender
-            .send(Event::Messages(self.time, vec![(client_id, elapsed, -1)]));
+            .send(Event::Messages(timely::progress::Stamp::from_elem(self.time), vec![(client_id, elapsed, -1)]));
         self.advance(elapsed);
     }
 

differential-dataflow/examples/scc_bench.rs

@@ -0,0 +1,109 @@
+//! Timed strongly connected components over streamed rounds of edge changes.
+//!
+//! Usage: scc_bench <nodes> <edges> <rounds> [-w<workers>]
+//!
+//! All rounds are introduced before stepping, so many epochs are concurrently
+//! live and arrangements hold many incomparable capabilities in the nested
+//! iterative scopes.
+
+use std::mem;
+use std::hash::Hash;
+
+use rand::{Rng, SeedableRng, StdRng};
+
+use differential_dataflow::VecCollection;
+use differential_dataflow::input::Input;
+use differential_dataflow::operators::*;
+use differential_dataflow::lattice::Lattice;
+
+type Node = usize;
+type Edge = (Node, Node);
+
+fn main() {
+    let nodes: usize = std::env::args().nth(1).unwrap().parse().unwrap();
+    let edges: usize = std::env::args().nth(2).unwrap().parse().unwrap();
+    let rounds: usize = std::env::args().nth(3).unwrap().parse().unwrap();
+
+    timely::execute_from_args(std::env::args(), move |worker| {
+
+        let timer = std::time::Instant::now();
+        let mut probe = timely::dataflow::ProbeHandle::new();
+
+        let mut edge_input = worker.dataflow(|scope| {
+            let (edge_input, graph) = scope.new_collection();
+            _strongly_connected(graph)
+                .consolidate()
+                .probe_with(&mut probe)
+                ;
+            edge_input
+        });
+
+        let seed: &[_] = &[1, 2, 3, 4];
+        let mut rng1: StdRng = SeedableRng::from_seed(seed);
+        let mut rng2: StdRng = SeedableRng::from_seed(seed);
+
+        if worker.index() == 0 {
+            for _ in 0..edges {
+                edge_input.insert((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)));
+            }
+            for round in 1..rounds {
+                edge_input.advance_to(round);
+                edge_input.insert((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)));
+                edge_input.remove((rng2.gen_range(0, nodes), rng2.gen_range(0, nodes)));
+            }
+        }
+        edge_input.advance_to(rounds);
+        edge_input.flush();
+        edge_input.close();
+
+        while worker.step() { }
+
+        if worker.index() == 0 {
+            println!("elapsed: {:?}", timer.elapsed());
+        }
+    }).unwrap();
+}
+
+fn _strongly_connected<'scope, T>(graph: VecCollection<'scope, T, Edge>) -> VecCollection<'scope, T, Edge>
+where
+    T: timely::progress::Timestamp + Lattice + Ord + Hash,
+{
+    graph.clone().iterate(|scope, inner| {
+        let edges = graph.enter(scope);
+        let trans = edges.clone().map_in_place(|x| mem::swap(&mut x.0, &mut x.1));
+        _trim_edges(_trim_edges(inner, edges), trans)
+    })
+}
+
+fn _trim_edges<'scope, T>(cycle: VecCollection<'scope, T, Edge>, edges: VecCollection<'scope, T, Edge>) -> VecCollection<'scope, T, Edge>
+where
+    T: timely::progress::Timestamp + Lattice + Ord + Hash,
+{
+    let nodes = edges.clone()
+                     .map_in_place(|x| x.0 = x.1)
+                     .consolidate();
+
+    let labels = _reachability(cycle, nodes);
+
+    edges.consolidate()
+         .join_map(labels.clone(), |&e1,&e2,&l1| (e2,(e1,l1)))
+         .join_map(labels.clone(), |&e2,&(e1,l1),&l2| ((e1,e2),(l1,l2)))
+         .filter(|&(_,(l1,l2))| l1 == l2)
+         .map(|((x1,x2),_)| (x2,x1))
+}
+
+fn _reachability<'scope, T>(edges: VecCollection<'scope, T, Edge>, nodes: VecCollection<'scope, T, (Node, Node)>) -> VecCollection<'scope, T, Edge>
+where
+    T: timely::progress::Timestamp + Lattice + Ord + Hash,
+{
+    edges.clone()
+         .filter(|_| false)
+         .iterate(|scope, inner| {
+             let edges = edges.enter(scope);
+             let nodes = nodes.enter_at(scope, |r| 256 * (64 - (r.0 as u64).leading_zeros() as u64));
+
+             inner.join_map(edges, |_k,l,d| (*d,*l))
+                  .concat(nodes)
+                  .reduce(|_, s, t| t.push((*s[0].0, 1)))
+         })
+}

differential-dataflow/src/columnar/collection/exchange.rs

@@ -30,7 +30,7 @@ impl<U: Update, H: for<'a> FnMut(columnar::Ref<'a, U::Key>)->u64> Distributor<Re
     // so the per-group pre_lens snapshot and seal check costs O(keys × workers). Should
     // either batch keys by destination first, or detect stride-1 outer bounds and use a
     // simpler single-pass partitioning that seals once at the end.
-    fn partition<T: Clone, P: timely::communication::Push<Message<T, RecordedUpdates<U>>>>(&mut self, container: &mut RecordedUpdates<U>, time: &T, pushers: &mut [P]) {
+    fn partition<T: Clone, P: timely::communication::Push<Message<T, RecordedUpdates<U>>>>(&mut self, container: &mut RecordedUpdates<U>, stamp: &timely::progress::Stamp<T>, pushers: &mut [P]) {
         use crate::columnar::updates::child_range;
 
         let view = container.updates.view();
@@ -80,7 +80,7 @@ impl<U: Update, H: for<'a> FnMut(columnar::Ref<'a, U::Key>)->u64> Distributor<Re
             // Push the empty update to the worker that produced the original
             // values so the send stays local. Not needed for correctness, but
             // a reasonable choice.
-            Message::push_at(&mut recorded, time.clone(), &mut pushers[self.worker % pushers.len()]);
+            Message::push_at(&mut recorded, stamp.clone(), &mut pushers[self.worker % pushers.len()]);
             return;
         }
 
@@ -90,11 +90,11 @@ impl<U: Update, H: for<'a> FnMut(columnar::Ref<'a, U::Key>)->u64> Distributor<Re
                 let recorded = RecordedUpdates { updates: output.into(), records: first_records, consolidated: container.consolidated };
                 first_records = 1;
                 let mut recorded = recorded;
-                Message::push_at(&mut recorded, time.clone(), pusher);
+                Message::push_at(&mut recorded, stamp.clone(), pusher);
             }
         }
     }
-    fn flush<T: Clone, P: timely::communication::Push<Message<T, RecordedUpdates<U>>>>(&mut self, _time: &T, _pushers: &mut [P]) { }
+    fn flush<T: Clone, P: timely::communication::Push<Message<T, RecordedUpdates<U>>>>(&mut self, _stamp: &timely::progress::Stamp<T>, _pushers: &mut [P]) { }
     fn relax(&mut self) { }
 }
 

differential-dataflow/src/operators/arrange/agent.rs

@@ -118,7 +118,7 @@
             .borrow_mut()
             .trace
             .map_batches(|batch| {
-                new_queue.push_back(TraceReplayInstruction::Batch(batch.clone(), Some(Tr::Time::minimum())));
+                new_queue.push_back(TraceReplayInstruction::Batch(batch.clone(), timely::progress::Stamp::from_elem(Tr::Time::minimum())));
                 upper = Some(batch.upper().clone());
             });
 
@@ -283,7 +283,7 @@
                 let activator = scope.activator_for(Rc::clone(&info.address));
                 let queue = self.new_listener(activator);
 
                 let activator = scope.activator_for(info.address);
                 *shutdown_button_ref = Some(ShutdownButton::new(Rc::clone(&capabilities), activator));
 
                 capabilities.borrow_mut().as_mut().unwrap().insert(capability);
@@ -300,11 +300,9 @@
                                     capabilities.downgrade(&frontier.borrow()[..]);
                                 },
                                 TraceReplayInstruction::Batch(batch, hint) => {
-                                    if let Some(time) = hint {
-                                        if !batch.is_empty() {
-                                            let delayed = capabilities.delayed(&time);
-                                            output.session(&delayed).give(batch);
-                                        }
+                                    if !hint.is_empty() && !batch.is_empty() {
+                                        let delayed = capabilities.delayed_stamp(&hint);
+                                        output.session(&delayed).give(batch);
                                     }
                                 }
                             }
@@ -416,7 +414,7 @@
                 let activator = scope.activator_for(Rc::clone(&info.address));
                 let queue = self.new_listener(activator);
 
                 let activator = scope.activator_for(info.address);
                 *shutdown_button_ref = Some(ShutdownButton::new(Rc::clone(&capabilities), activator));
 
                 capabilities.borrow_mut().as_mut().unwrap().insert(capability);
@@ -441,11 +439,9 @@
                                         }
                                     },
                                     TraceReplayInstruction::Batch(batch, hint) => {
-                                        if let Some(time) = hint {
-                                            if !batch.is_empty() {
-                                                let delayed = capabilities.delayed(&time);
-                                                output.session(&delayed).give(BatchFrontier::make_from(batch, since.borrow(), until.borrow()));
-                                            }
+                                        if !hint.is_empty() && !batch.is_empty() {
+                                            let delayed = capabilities.delayed_stamp(&hint);
+                                            output.session(&delayed).give(BatchFrontier::make_from(batch, since.borrow(), until.borrow()));
                                         }
                                     }
                                 }

differential-dataflow/src/operators/arrange/arrangement.rs

@@ -25,7 +25,8 @@
 use timely::progress::Timestamp;
 use timely::progress::Antichain;
 use timely::container::{ContainerBuilder, PushInto};
-use timely::dataflow::operators::Capability;
+use timely::dataflow::operators::{Capability, CapabilitySet};
+use timely::progress::Stamp;
 
 use crate::{Data, VecCollection, AsCollection};
 use crate::difference::Semigroup;
@@ -183,7 +184,7 @@
                     while let Some(key) = cursor.get_key(batch) {
                         while let Some(val) = cursor.get_val(batch) {
                             for datum in logic(key, val) {
                                 cursor.map_times(batch, |time, diff| {
                                     session.give((datum.clone(), <BatchCursor<Tr> as Cursor>::owned_time(time), <BatchCursor<Tr> as Cursor>::owned_diff(diff)));
                                 });
                             }
@@ -392,7 +393,9 @@
             // when we realize that time intervals are complete.
 
             input.for_each(|cap, data| {
-                capabilities.insert(cap.retain(0));
+                for capability in cap.retain_stamp(0).iter() {
+                    capabilities.insert(capability.clone());
+                }
                 chunker.push_into(data);
                 while let Some(chunk) = chunker.extract() {
                     batcher.push_into(std::mem::take(chunk));
@@ -422,8 +425,7 @@
                 //
                 // 1. If any held capabilities are not in advance of the new input frontier,
                 //    we must carve out updates now in advance of the new input frontier and
-                //    transmit them as batches, which requires appropriate *single* capabilities;
-                //    Until timely dataflow supports multiple capabilities on messages, at least.
+                //    transmit them as a batch, stamped with the capabilities they retire.
                 //
                 // 2. If there are no held capabilities in advance of the new input frontier,
                 //    then there are no updates not in advance of the new input frontier and
@@ -433,37 +435,30 @@
                 // If there is at least one capability not in advance of the input frontier ...
                 if capabilities.elements().iter().any(|c| !frontier.less_equal(c.time())) {
 
-                    let mut upper = Antichain::new();   // re-used allocation for sealing batches.
-
-                    // For each capability not in advance of the input frontier ...
-                    for (index, capability) in capabilities.elements().iter().enumerate() {
-
-                        if !frontier.less_equal(capability.time()) {
-
-                            // Assemble the upper bound on times we can commit with this capabilities.
-                            // We must respect the input frontier, and *subsequent* capabilities, as
-                            // we are pretending to retire the capability changes one by one.
-                            upper.clear();
-                            for time in frontier.frontier().iter() {
-                                upper.insert(time.clone());
-                            }
-                            for other_capability in &capabilities.elements()[(index + 1) .. ] {
-                                upper.insert(other_capability.time().clone());
-                            }
-
-                            // Extract updates not in advance of `upper`.
-                            let (mut chain, description) = batcher.seal(upper.clone());
-                            let batch = Bu::seal(&mut chain, description);
-
-                            writer.insert(batch.clone(), Some(capability.time().clone()));
-
-                            // send the batch to downstream consumers, empty or not.
-                            output.session(&capabilities.elements()[index]).give(batch);
-                        }
-                    }
-
-                    // Having extracted and sent batches between each capability and the input frontier,
-                    // we should downgrade all capabilities to match the batcher's lower update frontier.
+                    // The capabilities to retire: those not in advance of the input frontier.
+                    // Each update sealed below is greater or equal to one of them, as updates
+                    // supported only by the remaining capabilities are in advance of the input
+                    // frontier and remain in the batcher.
+                    let retired = capabilities
+                        .elements()
+                        .iter()
+                        .filter(|c| !frontier.less_equal(c.time()))
+                        .cloned()
+                        .collect::<CapabilitySet<_>>();
+
+                    // Extract all updates not in advance of the input frontier, as one batch.
+                    let (mut chain, description) = batcher.seal(frontier.frontier().to_owned());
+                    let batch = Bu::seal(&mut chain, description);
+
+                    let stamp = retired.iter().map(|c| c.time().clone()).collect::<Stamp<_>>();
+                    writer.insert(batch.clone(), stamp);
+
+                    // send the batch to downstream consumers, empty or not.
+                    output.session(&retired).give(batch);
+
+                    // Having extracted and sent the batch of updates not in advance of the input
+                    // frontier, we should downgrade all capabilities to match the batcher's lower
+                    // update frontier.
                     // This may involve discarding capabilities, which is fine as any new updates arrive
                     // in messages with new capabilities.
 

differential-dataflow/src/operators/arrange/mod.rs

@@ -51,8 +51,10 @@ use crate::trace::TraceReader;
 pub enum TraceReplayInstruction<Tr: TraceReader> {
     /// Describes a frontier advance.
     Frontier(Antichain<Tr::Time>),
-    /// Describes a batch of data and a capability hint.
-    Batch(Tr::Batch, Option<Tr::Time>),
+    /// Describes a batch of data and the stamp justifying its contents.
+    ///
+    /// The stamp is empty exactly when the batch is empty.
+    Batch(Tr::Batch, timely::progress::Stamp<Tr::Time>),
 }
 
 // Short names for strongly and weakly owned activators and shared queues.

differential-dataflow/src/operators/arrange/upsert.rs

@@ -282,7 +282,7 @@ where
                                 prev_frontier.clone_from(&upper);
 
                                 // Communicate `batch` to the arrangement and the stream.
-                                writer.insert(batch.clone(), Some(capability.time().clone()));
+                                writer.insert(batch.clone(), timely::progress::Stamp::from_elem(capability.time().clone()));
                                 output.session(&capabilities.elements()[index]).give(batch);
                             }
                         }

differential-dataflow/src/operators/arrange/writer.rs

@@ -52,7 +52,7 @@ impl<Tr: Trace> TraceWriter<Tr> {
     /// The `hint` argument is either `None` in the case of an empty batch,
     /// or is `Some(time)` for a time less or equal to all updates in the
     /// batch and which is suitable for use as a capability.
-    pub fn insert(&mut self, batch: Tr::Batch, hint: Option<Tr::Time>) {
+    pub fn insert(&mut self, batch: Tr::Batch, hint: timely::progress::Stamp<Tr::Time>) {
 
         // Something is wrong if not a sequence.
         if !(&self.upper == batch.lower()) {
@@ -84,7 +84,7 @@ impl<Tr: Trace> TraceWriter<Tr> {
     /// Inserts an empty batch up to `upper`.
     pub fn seal(&mut self, upper: Antichain<Tr::Time>) {
         if self.upper != upper {
-            self.insert(Tr::Batch::empty(self.upper.clone(), upper), None);
+            self.insert(Tr::Batch::empty(self.upper.clone(), upper), timely::progress::Stamp::new());
         }
     }
 }

differential-dataflow/src/operators/count.rs

@@ -81,20 +81,20 @@ where
                 lower_limit.clear();
                 lower_limit.extend(upper_limit.borrow().iter().cloned());
 
-                let mut cap = None;
+                let mut caps = timely::dataflow::operators::CapabilitySet::new();
                 input.for_each(|capability, batches| {
-                    if cap.is_none() {                          // NB: Assumes batches are in-order
-                        cap = Some(capability.retain(0));
+                    for capability in capability.retain_stamp(0).iter() {
+                        caps.insert(capability.clone());
                     }
                     for batch in batches.drain(..) {
                         upper_limit.clone_from(batch.upper());  // NB: Assumes batches are in-order
                         batch_storage.push(batch);
                     }
                 });
 
-                if let Some(capability) = cap {
+                if !caps.is_empty() {
 
-                    let mut session = output.session(&capability);
+                    let mut session = output.session(&caps);
 
                     let (mut batch_cursor, batch_storage) = crate::trace::cursor::cursor_list(batch_storage);
                     let (mut trace_cursor, trace_storage) = crate::trace::cursor::cursor_list(trace.batches_through(lower_limit.borrow()).unwrap());