From 25dce1ff583020f570fdbfb882794e707519da86 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 13 Oct 2023 14:03:42 +0100 Subject: [PATCH] Update readme --- README.md | 51 ++++++++++----------- examples/readme.py | 13 ++++++ examples/run_over_dataset.py | 0 permchain/pregel.py | 4 -- tests/test_pregel.py | 86 ++++++++++++++++++------------------ 5 files changed, 83 insertions(+), 71 deletions(-) create mode 100644 examples/readme.py delete mode 100644 examples/run_over_dataset.py diff --git a/README.md b/README.md index 593177dd2..594e76deb 100644 --- a/README.md +++ b/README.md @@ -7,36 +7,37 @@ ## Usage ```python -from permchain import InMemoryPubSubConnection, PubSub, Topic +from permchain import Pregel, channels -topic_one = Topic("one") -chain_one = Topic.IN.subscribe() | (lambda x: x + 'b') | topic_one.publish() -chain_two = topic_one.subscribe() | (lambda x: x + 'c') | Topic.OUT.publish() +value = channels.LastValue[str]("value") -conn = InMemoryPubSubConnection() -pubsub = PubSub(chain_one, chain_two, connection=conn) +grow_value = ( + Pregel.subscribe_to(value) + | (lambda x: x + x) + | Pregel.send_to({value: lambda x: x if len(x) < 10 else None}) +) -assert pubsub.invoke('a') == 'abc' +pubsub = Pregel(grow_value, input=value, output=value) + +assert pubsub.invoke("a") == "aaaaaaaa" ``` -Check `tests` and `examples` for more examples. +Check `examples` for more examples. ## Near-term Roadmap -- [x] Add initial retry support (pending changes in `langchain`) -- [x] Implement OUT as regular topic -- [x] Implement IN as regular topic -- [x] Add Connection.peek() to monitor past messages from all topics -- [x] Enable resuming PubSub from the "middle" of the computation -- [x] Add test for .peek() -- [x] Add "wait until topic X is done" pattern, aka. `Topic.join()` -- [ ] Move tracking of inflight processes/messages to Connection - - [ ] Use this to build retry mechanism, where any inflight messages are moved back to the respective topics when restarting - - [ ] But this would require being able to replay a message for a single listener only, which maybe requires a larger redesign of PubSub<>Connection contract than what I wanted to do here -- [ ] Detect cycles (aka. infinite loops) and throw an error - - [ ] Allow user to catch that error (by subcribing to an error topic?) -- [ ] Add example for "human in the loop" pattern, one of the two below - - [ ] Example with one permchain, which runs until it produces either 1. request for input or 2. output. The consumer code then gets the needed info, and restarts the permchain with answer, and same state id - - [ ] Allow interrupting execution by breaking out of the iterator returned by .stream() - - [ ] Build example showing a simple "human in the loop" pattern using this, ie. if a certain message asking for input is published the consumer of the iterator breaks out, does something and then restarts it -- [ ] Add Redis-backed Connection implementation +- [ ] Iterate on API + - [ ] do we want api to receive output from multiple channels in invoke() + - [ ] do we want api to send input to multiple channels in invoke() +- [ ] Implement checkpointing + - [ ] Save checkpoints at end of each step + - [ ] Load checkpoint at start of invocation + - [ ] API to specify storage backend and save key +- [ ] Add more examples + - [ ] human in the loop + - [ ] combine documents + - [ ] agent executor + - [ ] run over dataset +- [ ] Fault tolerance + - [ ] Retry individual processes in a step + - [ ] Retry entire step? diff --git a/examples/readme.py b/examples/readme.py new file mode 100644 index 000000000..87a2d7628 --- /dev/null +++ b/examples/readme.py @@ -0,0 +1,13 @@ +from permchain import Pregel, channels + +value = channels.LastValue[str]("value") + +grow_value = ( + Pregel.subscribe_to(value) + | (lambda x: x + x) + | Pregel.send_to({value: lambda x: x if len(x) < 10 else None}) +) + +pubsub = Pregel(grow_value, input=value, output=value) + +assert pubsub.invoke("a") == "aaaaaaaa" diff --git a/examples/run_over_dataset.py b/examples/run_over_dataset.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/permchain/pregel.py b/permchain/pregel.py index e3d65c7b5..903fc1e2a 100644 --- a/permchain/pregel.py +++ b/permchain/pregel.py @@ -588,7 +588,3 @@ def _apply_writes_and_prepare_next_tasks( tasks.append((proc, val)) return tasks - - -# TODO do we want api to subscribe to all channels? -# Do we want api to send input to multiple channels in invoke() diff --git a/tests/test_pregel.py b/tests/test_pregel.py index a0e8104ee..76e94ca78 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -1,5 +1,4 @@ import time -from uuid import uuid4 import pytest from pytest_mock import MockerFixture @@ -10,28 +9,28 @@ from permchain.pregel import PregelInvoke def test_invoke_single_process_in_out(mocker: MockerFixture): - input = channels.LastValue[int]("input") - output = channels.LastValue[int]("output") + input_chan = channels.LastValue[int]("input") + output_chan = channels.LastValue[int]("output") add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Pregel.subscribe_to(input) | add_one | Pregel.send_to(output) + chain = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output_chan) - pubsub = Pregel(chain, input=input, output=output) + pubsub = Pregel(chain, input=input_chan, output=output_chan) # Then invoke pubsub assert pubsub.invoke(2) == 3 def test_invoke_two_processes_in_out(mocker: MockerFixture): - input = channels.LastValue[int]("input") + input_chan = channels.LastValue[int]("input") output = channels.LastValue[int]("output") inbox = channels.Inbox[int]("inbox") add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Pregel.subscribe_to(input) | add_one | Pregel.send_to(inbox) + chain_one = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(inbox) chain_two = Pregel.subscribe_to_each(inbox) | add_one | Pregel.send_to(output) - pubsub = Pregel(chain_one, chain_two, input=input, output=output) + pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output) # Then invoke pubsub assert pubsub.invoke(2) == 4 @@ -42,13 +41,16 @@ def test_batch_two_processes_in_out(mocker: MockerFixture): time.sleep(inp / 10) return inp + 1 - input = channels.LastValue[int]("input") + input_chan = channels.LastValue[int]("input") output = channels.LastValue[int]("output") one = channels.LastValue[int]("one") - chain_one = Pregel.subscribe_to(input) | add_one_with_delay | Pregel.send_to(one) + + chain_one = ( + Pregel.subscribe_to(input_chan) | add_one_with_delay | Pregel.send_to(one) + ) chain_two = Pregel.subscribe_to(one) | add_one_with_delay | Pregel.send_to(output) - pubsub = Pregel(chain_one, chain_two, input=input, output=output) + pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output) # Then invoke pubsub assert pubsub.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] @@ -57,13 +59,13 @@ def test_batch_two_processes_in_out(mocker: MockerFixture): def test_invoke_many_processes_in_out(mocker: MockerFixture): test_size = 100 - input = channels.LastValue[int]("input") + input_chan = channels.LastValue[int]("input") output = channels.LastValue[int]("output") topics: list[channels.Channel] = [channels.LastValue[int]("zero")] add_one = mocker.Mock(side_effect=lambda x: x + 1) chains: list[PregelInvoke] = [ - Pregel.subscribe_to(input) | add_one | Pregel.send_to(topics[0]) + Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(topics[0]) ] for i in range(test_size - 2): topics.append(channels.LastValue[int](str(i))) @@ -72,7 +74,7 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture): ) chains.append(Pregel.subscribe_to(topics[-1]) | add_one | Pregel.send_to(output)) - pubsub = Pregel(*chains, input=input, output=output) + pubsub = Pregel(*chains, input=input_chan, output=output) for _ in range(10): # Then invoke pubsub @@ -82,13 +84,13 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture): def test_batch_many_processes_in_out(mocker: MockerFixture): test_size = 100 - input = channels.LastValue[int]("input") + input_chan = channels.LastValue[int]("input") output = channels.LastValue[int]("output") topics: list[channels.Channel] = [channels.LastValue[int]("zero")] add_one = mocker.Mock(side_effect=lambda x: x + 1) chains: list[PregelInvoke] = [ - Pregel.subscribe_to(input) | add_one | Pregel.send_to(topics[0]) + Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(topics[0]) ] for i in range(test_size - 2): topics.append(channels.LastValue[int](str(i))) @@ -97,7 +99,7 @@ def test_batch_many_processes_in_out(mocker: MockerFixture): ) chains.append(Pregel.subscribe_to(topics[-1]) | add_one | Pregel.send_to(output)) - pubsub = Pregel(*chains, input=input, output=output) + pubsub = Pregel(*chains, input=input_chan, output=output) for _ in range(10): # Then invoke pubsub @@ -113,13 +115,13 @@ def test_batch_many_processes_in_out(mocker: MockerFixture): def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture): add_one = mocker.Mock(side_effect=lambda x: x + 1) - input = channels.LastValue[int]("input") + input_chan = channels.LastValue[int]("input") output = channels.LastValue[int]("output") - chain_one = Pregel.subscribe_to(input) | add_one | Pregel.send_to(output) - chain_two = Pregel.subscribe_to(input) | add_one | Pregel.send_to(output) + chain_one = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output) + chain_two = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output) - pubsub = Pregel(chain_one, chain_two, input=input, output=output) + pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output) with pytest.raises(channels.InvalidUpdateError): # LastValue channels can only be updated once per iteration @@ -129,13 +131,13 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture): def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture): add_one = mocker.Mock(side_effect=lambda x: x + 1) - input = channels.LastValue[int]("input") + input_chan = channels.LastValue[int]("input") output = channels.Inbox[int]("output") - chain_one = Pregel.subscribe_to(input) | add_one | Pregel.send_to(output) - chain_two = Pregel.subscribe_to(input) | add_one | Pregel.send_to(output) + chain_one = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output) + chain_two = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output) - pubsub = Pregel(chain_one, chain_two, input=input, output=output) + pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output) # An Inbox channel accumulates updates into a sequence assert pubsub.invoke(2) == (3, 3) @@ -145,15 +147,15 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture): add_one = mocker.Mock(side_effect=lambda x: x + 1) add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x)) - input = channels.LastValue[int]("input") + input_chan = channels.LastValue[int]("input") output = channels.LastValue[int]("output") inbox = channels.Inbox[int]("inbox") - chain_one = Pregel.subscribe_to(input) | add_one | Pregel.send_to(inbox) - chain_three = Pregel.subscribe_to(input) | add_one | Pregel.send_to(inbox) + chain_one = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(inbox) + chain_three = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(inbox) chain_four = Pregel.subscribe_to(inbox) | add_10_each | Pregel.send_to(output) - pubsub = Pregel(chain_one, chain_three, chain_four, input=input, output=output) + pubsub = Pregel(chain_one, chain_three, chain_four, input=input_chan, output=output) # Then invoke pubsub # We get a single array result as chain_four waits for all publishers to finish @@ -166,12 +168,12 @@ def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture): add_one = mocker.Mock(side_effect=lambda x: x + 1) add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) - input = channels.LastValue[int]("input") + input_chan = channels.LastValue[int]("input") output = channels.LastValue[int]("output") inner_pubsub = Pregel( - Pregel.subscribe_to(input) | add_one | Pregel.send_to(output), - input=input, + Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output), + input=input_chan, output=output, ) @@ -179,7 +181,7 @@ def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture): outbox_one = channels.LastValue[int]("outbox_one") chain_one = ( - Pregel.subscribe_to(input) | add_10_each | Pregel.send_to(inbox_one).map() + Pregel.subscribe_to(input_chan) | add_10_each | Pregel.send_to(inbox_one).map() ) chain_two = ( Pregel.subscribe_to(inbox_one) @@ -189,7 +191,7 @@ def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture): ) chain_three = Pregel.subscribe_to(outbox_one) | sum | Pregel.send_to(output) - pubsub = Pregel(chain_one, chain_two, chain_three, input=input, output=output) + pubsub = Pregel(chain_one, chain_two, chain_three, input=input_chan, output=output) # Then invoke pubsub for _ in range(10): @@ -199,12 +201,12 @@ def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture): def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture): add_one = mocker.Mock(side_effect=lambda x: x + 1) - input = channels.LastValue[int]("input") + input_chan = channels.LastValue[int]("input") output = channels.LastValue[int]("output") between = channels.LastValue[int]("between") chain_one = ( - Pregel.subscribe_to(input) + Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to( {output: RunnablePassthrough(), between: RunnablePassthrough()} @@ -212,22 +214,22 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture): ) chain_two = Pregel.subscribe_to(between) | add_one | Pregel.send_to(output) - pubsub = Pregel(chain_one, chain_two, input=input, output=output) + pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output) # Then invoke pubsub assert [c for c in pubsub.stream(2)] == [3, 4] def test_invoke_two_processes_no_out(mocker: MockerFixture): - input = channels.LastValue[int]("input") + input_chan = channels.LastValue[int]("input") output = channels.LastValue[int]("output") between = channels.LastValue[int]("between") add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain_one = Pregel.subscribe_to(input) | add_one | Pregel.send_to(between) + chain_one = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(between) chain_two = Pregel.subscribe_to(between) | add_one - pubsub = Pregel(chain_one, chain_two, input=input, output=output) + pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output) # Then invoke pubsub # It finishes executing (once no more messages being published) @@ -236,7 +238,7 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture): def test_invoke_two_processes_no_in(mocker: MockerFixture): - input = channels.LastValue[int]("input") + input_chan = channels.LastValue[int]("input") output = channels.LastValue[int]("output") between = channels.LastValue[int]("between") @@ -244,7 +246,7 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture): chain_one = Pregel.subscribe_to(between) | add_one | Pregel.send_to(output) chain_two = Pregel.subscribe_to(between) | add_one - pubsub = Pregel(chain_one, chain_two, input=input, output=output) + pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output) with pytest.raises(ValueError): assert pubsub.invoke(2) is None