From dbe9293a3674db5463104b1ec7d8d6c759bcd135 Mon Sep 17 00:00:00 2001 From: Liam Griffiths Date: Mon, 24 Jun 2024 13:34:25 -0400 Subject: [PATCH 1/4] Add `depends` option to set edges without future inputs --- src/Node.ts | 17 ++++++++++++++++ src/Substrate.ts | 13 +++++++++++- tests/Node.test.ts | 11 ++++++----- tests/Substrate.test.ts | 44 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 79 insertions(+), 6 deletions(-) diff --git a/src/Node.ts b/src/Node.ts index b49ebe3..0cfa285 100644 --- a/src/Node.ts +++ b/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,9 @@ export abstract class Node { this.cache_age = opts?.cache_age; this.cache_keys = opts?.cache_keys; this.max_retries = opts?.max_retries; + this.id = opts?.id ?? generator(this.node); + this.hide = opts?.hide ?? false; + this.depends = opts?.depends ?? []; } /** @@ -138,6 +145,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) { diff --git a/src/Substrate.ts b/src/Substrate.ts index a103ffb..ab70438 100644 --- a/src/Substrate.ts +++ b/src/Substrate.ts @@ -198,10 +198,21 @@ export class Substrate { } } + const allEdges: Record> = {}; + for (let n of allNodes) { + allEdges[n.id] = new Set(); + 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); + return fromIds.map((fromId: string) => [fromId, toId]); + }), initial_args: {}, // @deprecated }; } diff --git a/tests/Node.test.ts b/tests/Node.test.ts index 6e912a8..ae2bf8c 100644 --- a/tests/Node.test.ts +++ b/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])); }); }); diff --git a/tests/Substrate.test.ts b/tests/Substrate.test.ts index b99bc5d..a21c6b9 100644 --- a/tests/Substrate.test.ts +++ b/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] }); + + 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: [], + }); + }); }); }); From a37fbd28b92b65d828b2a4860fa32747ec59fe2d Mon Sep 17 00:00:00 2001 From: Liam Griffiths Date: Mon, 24 Jun 2024 13:38:35 -0400 Subject: [PATCH 2/4] Add a duplicate `depends` in test to make sure we do not make duplicate edges --- tests/Substrate.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/Substrate.test.ts b/tests/Substrate.test.ts index a21c6b9..f835f1d 100644 --- a/tests/Substrate.test.ts +++ b/tests/Substrate.test.ts @@ -148,7 +148,7 @@ 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] }); + const c = new FooNode({ c: 789 }, { id: "c", depends: [a, b, b] }); // intentionally using b twice here const result = Substrate.serialize(a, b, c); From aef48882f86f2e577a0db779ecf64a878990368c Mon Sep 17 00:00:00 2001 From: Liam Griffiths Date: Mon, 24 Jun 2024 13:38:55 -0400 Subject: [PATCH 3/4] Formatting --- src/Node.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Node.ts b/src/Node.ts index 0cfa285..6067e09 100644 --- a/src/Node.ts +++ b/src/Node.ts @@ -154,7 +154,7 @@ export abstract class Node { futures.add(future); } } - + const collectFutures = (obj: any) => { if (Array.isArray(obj)) { for (let item of obj) { From e2d3c35c47d7fc50874276d33a7a5d65efb4080a Mon Sep 17 00:00:00 2001 From: Liam Griffiths Date: Wed, 3 Jul 2024 11:08:01 -0400 Subject: [PATCH 4/4] Handle null/undefined cases, include third edge arg, add example --- examples/explicit-edges.ts | 35 +++++++++++++++++++++++++++++++++++ src/Node.ts | 6 ++---- src/Substrate.ts | 2 +- tests/Substrate.test.ts | 6 +++--- 4 files changed, 41 insertions(+), 8 deletions(-) create mode 100755 examples/explicit-edges.ts diff --git a/examples/explicit-edges.ts b/examples/explicit-edges.ts new file mode 100755 index 0000000..c8dad37 --- /dev/null +++ b/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(); diff --git a/src/Node.ts b/src/Node.ts index 6067e09..e642348 100644 --- a/src/Node.ts +++ b/src/Node.ts @@ -54,8 +54,6 @@ export abstract class Node { this.cache_age = opts?.cache_age; this.cache_keys = opts?.cache_keys; this.max_retries = opts?.max_retries; - this.id = opts?.id ?? generator(this.node); - this.hide = opts?.hide ?? false; this.depends = opts?.depends ?? []; } @@ -114,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; @@ -172,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]); } diff --git a/src/Substrate.ts b/src/Substrate.ts index ab70438..17cacff 100644 --- a/src/Substrate.ts +++ b/src/Substrate.ts @@ -211,7 +211,7 @@ export class Substrate { futures: Array.from(allFutures).map((future) => future.toJSON()), edges: Object.keys(allEdges).flatMap((toId: string) => { let fromIds: string[] = Array.from(allEdges[toId] as Set); - return fromIds.map((fromId: string) => [fromId, toId]); + return fromIds.map((fromId: string) => [fromId, toId, {}]); }), initial_args: {}, // @deprecated }; diff --git a/tests/Substrate.test.ts b/tests/Substrate.test.ts index f835f1d..dc08d00 100644 --- a/tests/Substrate.test.ts +++ b/tests/Substrate.test.ts @@ -154,9 +154,9 @@ describe("Substrate", () => { expect(result).toEqual({ edges: [ - ["a", "b"], - ["a", "c"], - ["b", "c"], + ["a", "b", {}], + ["a", "c", {}], + ["b", "c", {}], ], initial_args: {}, nodes: [