Add depends option to set edges without future inputs (#91)

* Add `depends` option to set edges without future inputs

* Add a duplicate `depends` in test to make sure we do not make duplicate edges

* Formatting

* Handle null/undefined cases, include third edge arg, add example

Author: liamgriffiths

examples/explicit-edges.ts

@@ -0,0 +1,35 @@
+#!/usr/bin/env -S npx ts-node --transpileOnly
+
+import { Substrate, Box } from "substrate";
+
+async function main() {
+  const SUBSTRATE_API_KEY = process.env["SUBSTRATE_API_KEY"];
+
+  const substrate = new Substrate({ apiKey: SUBSTRATE_API_KEY });
+
+  // One way to see that the edges determine the node run order is to use the RunPython node to print the
+  // timestamp of when the node was run. Because the RunPython node isn't available in the TypeScript SDK
+  // I've taken a pickled python function that does a `print(time.time())` and am overriding the nodes
+  // here to be RunPython nodes instead.
+  const toRunPython = (node: Box) => {
+    node.node = "RunPython";
+    node.args = {
+      pkl_function:
+        "gAWV5QEAAAAAAACMF2Nsb3VkcGlja2xlLmNsb3VkcGlja2xllIwOX21ha2VfZnVuY3Rpb26Uk5QoaACMDV9idWlsdGluX3R5cGWUk5SMCENvZGVUeXBllIWUUpQoSwBLAEsASwFLAktDQxBkAWQAbAB9AHwAoAChAFMAlE5LAIaUjAR0aW1llIWUaAuMOC9Vc2Vycy9saWFtL3dvcmsvc3Vic3RyYXRlLXB5dGhvbi9leGFtcGxlcy9ydW5fcHl0aG9uLnB5lIwKcHJpbnRfdGltZZRLKEMECAEIAZQpKXSUUpR9lCiMC19fcGFja2FnZV9flE6MCF9fbmFtZV9flIwIX19tYWluX1+UjAhfX2ZpbGVfX5RoDHVOTk50lFKUaACMEl9mdW5jdGlvbl9zZXRzdGF0ZZSTlGgXfZR9lChoE2gNjAxfX3F1YWxuYW1lX1+UaA2MD19fYW5ub3RhdGlvbnNfX5R9lIwOX19rd2RlZmF1bHRzX1+UTowMX19kZWZhdWx0c19flE6MCl9fbW9kdWxlX1+UjAhfX21haW5fX5SMB19fZG9jX1+UTowLX19jbG9zdXJlX1+UTowXX2Nsb3VkcGlja2xlX3N1Ym1vZHVsZXOUXZSMC19fZ2xvYmFsc19flH2UdYaUhlIwLg==",
+      kwargs: {},
+      pip_install: null,
+      python_version: "3.10.13",
+    };
+    return node;
+  };
+
+  const a = toRunPython(new Box({ value: "" }, { id: "a" }));
+
+  const b = toRunPython(new Box({ value: "" }, { id: "b", depends: [a] }));
+
+  const c = toRunPython(new Box({ value: "" }, { id: "c", depends: [b] }));
+
+  const res = await substrate.run(c, a, b);
+  console.log(res.json);
+}
+main();

src/Node.ts

@@ -19,6 +19,8 @@ export type Options = {
   cache_keys?: string[];
   /** Max number of times to retry this node if it fails. Default: null means no retries */
   max_retries?: number;
+  /** Specify nodes that this node depends on. */
+  depends?: Node[];
 };
 
 export abstract class Node {
@@ -38,6 +40,8 @@ export abstract class Node {
   cache_keys?: string[];
   /** Max number of times to retry this node if it fails. Default: null means no retries */
   max_retries?: number;
+  /** Specify nodes that this node depends on. */
+  depends: Node[];
 
   /** TODO this field stores the last response, but it's just temporary until the internals are refactored */
   protected _response: SubstrateResponse | undefined;
@@ -50,6 +54,7 @@ export abstract class Node {
     this.cache_age = opts?.cache_age;
     this.cache_keys = opts?.cache_keys;
     this.max_retries = opts?.max_retries;
+    this.depends = opts?.depends ?? [];
   }
 
   /**
@@ -107,7 +112,7 @@ export abstract class Node {
         return obj.toPlaceholder();
       }
 
-      if (typeof obj === "object") {
+      if (obj && typeof obj === "object") {
         return Object.keys(obj).reduce((acc: any, k: any) => {
           acc[k] = withPlaceholders(obj[k]);
           return acc;
@@ -138,6 +143,16 @@ export abstract class Node {
 
     nodes.add(this);
 
+    for (let node of this.depends) {
+      const references = node.references();
+      for (let node of references.nodes) {
+        nodes.add(node);
+      }
+      for (let future of references.futures) {
+        futures.add(future);
+      }
+    }
+
     const collectFutures = (obj: any) => {
       if (Array.isArray(obj)) {
         for (let item of obj) {
@@ -155,7 +170,7 @@ export abstract class Node {
         return;
       }
 
-      if (typeof obj === "object") {
+      if (obj && typeof obj === "object") {
         for (let key of Object.keys(obj)) {
           collectFutures(obj[key]);
         }

src/Substrate.ts

@@ -198,10 +198,21 @@ export class Substrate {
       }
     }
 
+    const allEdges: Record<string, Set<string>> = {};
+    for (let n of allNodes) {
+      allEdges[n.id] = new Set<string>();
+      for (let d of n.depends) {
+        allEdges[n.id]!.add(d.id);
+      }
+    }
+
     return {
       nodes: Array.from(allNodes).map((node) => node.toJSON()),
       futures: Array.from(allFutures).map((future) => future.toJSON()),
-      edges: [], // @deprecated
+      edges: Object.keys(allEdges).flatMap((toId: string) => {
+        let fromIds: string[] = Array.from(allEdges[toId] as Set<string>);
+        return fromIds.map((fromId: string) => [fromId, toId, {}]);
+      }),
       initial_args: {}, // @deprecated
     };
   }

tests/Node.test.ts

@@ -31,17 +31,18 @@ describe("Node", () => {
   });
 
   test(".references", () => {
-    const a = new FooNode({ x: "x" });
+    const a = new FooNode({ x: "x" }, { id: "a" });
     const f1 = a.future.get("x");
     const f2 = a.future.get("y");
-    const b = new FooNode({ x: f1, z: f2 });
+    const b = new FooNode({ x: f1, z: f2 }, { id: "b" });
     const f3 = b.future.get("x");
-    const c = new FooNode({ x: f3 });
+    const c = new FooNode({ x: f3 }, { id: "c" });
+    const d = new FooNode({}, { id: "d", depends: [c] });
 
     // @ts-ignore (protected)
-    const { nodes, futures } = c.references();
+    const { nodes, futures } = d.references();
 
-    expect(nodes).toEqual(new Set([a, b, c]));
+    expect(nodes).toEqual(new Set([a, b, c, d]));
     expect(futures).toEqual(new Set([f1, f2, f3]));
   });
 });

tests/Substrate.test.ts

@@ -144,5 +144,49 @@ describe("Substrate", () => {
         ],
       });
     });
+
+    test("when there are nodes and we use the `depends` key", () => {
+      const a = new FooNode({ a: 123 }, { id: "a" });
+      const b = new FooNode({ b: 456 }, { id: "b", depends: [a] });
+      const c = new FooNode({ c: 789 }, { id: "c", depends: [a, b, b] }); // intentionally using b twice here
+
+      const result = Substrate.serialize(a, b, c);
+
+      expect(result).toEqual({
+        edges: [
+          ["a", "b", {}],
+          ["a", "c", {}],
+          ["b", "c", {}],
+        ],
+        initial_args: {},
+        nodes: [
+          {
+            node: "FooNode",
+            id: a.id,
+            args: {
+              a: 123,
+            },
+            _should_output_globally: true,
+          },
+          {
+            node: "FooNode",
+            id: b.id,
+            args: {
+              b: 456,
+            },
+            _should_output_globally: true,
+          },
+          {
+            node: "FooNode",
+            id: c.id,
+            args: {
+              c: 789,
+            },
+            _should_output_globally: true,
+          },
+        ],
+        futures: [],
+      });
+    });
   });
 });