diff --git a/permchain/pregel.py b/permchain/pregel.py index fde0346e1..d20ac9429 100644 --- a/permchain/pregel.py +++ b/permchain/pregel.py @@ -250,6 +250,9 @@ class Pregel(Generic[Output], RunnableSerializable[dict[str, Any] | Any, Output] chains_flat.extend(chain) else: chains_flat.append(chain) + + validate_chains_channels(chains_flat, channels, input, output) + super().__init__( chains=chains_flat, channels=channels, @@ -584,3 +587,32 @@ def _apply_writes_and_prepare_next_tasks( tasks.append((proc, val)) return tasks + + +def validate_chains_channels( + chains: Sequence[PregelInvoke | PregelBatch], + channels: Mapping[str, Channel], + input: str | None, + output: str | Sequence[str], +) -> None: + subscribed_channels = set() + for chain in chains: + if isinstance(chain, PregelInvoke): + subscribed_channels.update(chain.channels.values()) + elif isinstance(chain, PregelBatch): + subscribed_channels.add(chain.channel) + else: + raise TypeError( + f"Invalid chain type {type(chain)}, expected Pregel.subscribe_to() or Pregel.subscribe_to_each()" + ) + + if input is not None and input not in subscribed_channels: + raise ValueError(f"Input channel {input} is not subscribed to by any chain") + + for chan in subscribed_channels: + if chan not in channels: + raise ValueError(f"Channel {chan} is subscribed to, but not initialized") + + if isinstance(output, str): + if output not in channels: + raise ValueError(f"Output channel {output} is not initialized") diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 76e94ca78..3c0e04585 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -9,48 +9,66 @@ from permchain.pregel import PregelInvoke def test_invoke_single_process_in_out(mocker: MockerFixture): - 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_chan) | add_one | Pregel.send_to(output_chan) + chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") - pubsub = Pregel(chain, input=input_chan, output=output_chan) + app = Pregel( + (chain,), + channels={ + "input": channels.LastValue[int](), + "output": channels.LastValue[int](), + }, + input="input", + output="output", + ) # Then invoke pubsub - assert pubsub.invoke(2) == 3 + assert app.invoke(2) == 3 def test_invoke_two_processes_in_out(mocker: MockerFixture): - 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_chan) | add_one | Pregel.send_to(inbox) - chain_two = Pregel.subscribe_to_each(inbox) | add_one | Pregel.send_to(output) + chain_one = Pregel.subscribe_to("input") | 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_chan, output=output) + pubsub = Pregel( + [chain_one, chain_two], + channels={ + "input": channels.LastValue[int](), + "output": channels.LastValue[int](), + "inbox": channels.Inbox[int](), + }, + input="input", + output="output", + ) # Then invoke pubsub assert pubsub.invoke(2) == 4 -def test_batch_two_processes_in_out(mocker: MockerFixture): +def test_batch_two_processes_in_out(): def add_one_with_delay(inp: int) -> int: time.sleep(inp / 10) return inp + 1 - input_chan = channels.LastValue[int]("input") - output = channels.LastValue[int]("output") - one = channels.LastValue[int]("one") - chain_one = ( - Pregel.subscribe_to(input_chan) | add_one_with_delay | Pregel.send_to(one) + Pregel.subscribe_to("input") | add_one_with_delay | Pregel.send_to("one") + ) + chain_two = ( + Pregel.subscribe_to("one") | add_one_with_delay | Pregel.send_to("output") ) - chain_two = Pregel.subscribe_to(one) | add_one_with_delay | Pregel.send_to(output) - pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output) + pubsub = Pregel( + chain_one, + chain_two, + channels={ + "input": channels.LastValue[int](), + "output": channels.LastValue[int](), + "one": channels.LastValue[int](), + }, + input="input", + output="output", + ) # Then invoke pubsub assert pubsub.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] @@ -58,52 +76,53 @@ def test_batch_two_processes_in_out(mocker: MockerFixture): def test_invoke_many_processes_in_out(mocker: MockerFixture): test_size = 100 - - 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) + + chans = { + "input": channels.LastValue[int](), + "output": channels.LastValue[int](), + "-1": channels.LastValue[int](), + } chains: list[PregelInvoke] = [ - Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(topics[0]) + Pregel.subscribe_to("input") | add_one | Pregel.send_to("-1") ] for i in range(test_size - 2): - topics.append(channels.LastValue[int](str(i))) + chans[str(i)] = channels.LastValue[int]() chains.append( - Pregel.subscribe_to(topics[-2]) | add_one | Pregel.send_to(topics[-1]) + Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i)) ) - chains.append(Pregel.subscribe_to(topics[-1]) | add_one | Pregel.send_to(output)) + chains.append(Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output")) - pubsub = Pregel(*chains, input=input_chan, output=output) + app = Pregel(*chains, channels=chans, input="input", output="output") for _ in range(10): - # Then invoke pubsub - assert pubsub.invoke(2, {"recursion_limit": test_size}) == 2 + test_size + assert app.invoke(2, {"recursion_limit": test_size}) == 2 + test_size def test_batch_many_processes_in_out(mocker: MockerFixture): test_size = 100 - - 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) + + chans = { + "input": channels.LastValue[int](), + "output": channels.LastValue[int](), + "-1": channels.LastValue[int](), + } chains: list[PregelInvoke] = [ - Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(topics[0]) + Pregel.subscribe_to("input") | add_one | Pregel.send_to("-1") ] for i in range(test_size - 2): - topics.append(channels.LastValue[int](str(i))) + chans[str(i)] = channels.LastValue[int]() chains.append( - Pregel.subscribe_to(topics[-2]) | add_one | Pregel.send_to(topics[-1]) + Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i)) ) - chains.append(Pregel.subscribe_to(topics[-1]) | add_one | Pregel.send_to(output)) + chains.append(Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output")) - pubsub = Pregel(*chains, input=input_chan, output=output) + app = Pregel(*chains, channels=chans, input="input", output="output") for _ in range(10): # Then invoke pubsub - assert pubsub.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [ + assert app.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [ 2 + test_size, 1 + test_size, 3 + test_size, @@ -115,138 +134,182 @@ 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_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_chan, output=output) + app = Pregel( + chain_one, + chain_two, + channels={ + "input": channels.LastValue[int](), + "output": channels.LastValue[int](), + }, + input="input", + output="output", + ) with pytest.raises(channels.InvalidUpdateError): # LastValue channels can only be updated once per iteration - pubsub.invoke(2) + app.invoke(2) def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture): add_one = mocker.Mock(side_effect=lambda x: x + 1) - 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_chan, output=output) + app = Pregel( + chain_one, + chain_two, + channels={ + "input": channels.LastValue[int](), + "output": channels.Inbox[int](), + }, + input="input", + output="output", + ) # An Inbox channel accumulates updates into a sequence - assert pubsub.invoke(2) == (3, 3) + assert app.invoke(2) == (3, 3) 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_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_four = Pregel.subscribe_to("inbox") | add_10_each | Pregel.send_to("output") - 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) + app = Pregel( + chain_one, + chain_three, + chain_four, + channels={ + "input": channels.LastValue[int](), + "output": channels.LastValue[int](), + "inbox": channels.Inbox[int](), + }, + input="input", + output="output", + ) - pubsub = Pregel(chain_one, chain_three, chain_four, input=input_chan, output=output) - - # Then invoke pubsub + # Then invoke app # We get a single array result as chain_four waits for all publishers to finish # before operating on all elements published to topic_two as an array for _ in range(100): - assert pubsub.invoke(2) == [13, 13] + assert app.invoke(2) == [13, 13] 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_chan = channels.LastValue[int]("input") - output = channels.LastValue[int]("output") - - inner_pubsub = Pregel( - Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output), - input=input_chan, - output=output, + inner_app = Pregel( + Pregel.subscribe_to("input") | add_one | Pregel.send_to("output"), + channels={ + "input": channels.LastValue[int](), + "output": channels.LastValue[int](), + }, + input="input", + output="output", ) - inbox_one = channels.Inbox[int]("inbox_one") - outbox_one = channels.LastValue[int]("outbox_one") - chain_one = ( - Pregel.subscribe_to(input_chan) | add_10_each | Pregel.send_to(inbox_one).map() + Pregel.subscribe_to("input") | add_10_each | Pregel.send_to("inbox_one").map() ) chain_two = ( - Pregel.subscribe_to(inbox_one) - | inner_pubsub.map() + Pregel.subscribe_to("inbox_one") + | inner_app.map() | sorted - | Pregel.send_to(outbox_one) + | Pregel.send_to("outbox_one") ) - chain_three = Pregel.subscribe_to(outbox_one) | sum | Pregel.send_to(output) + chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.send_to("output") - pubsub = Pregel(chain_one, chain_two, chain_three, input=input_chan, output=output) + app = Pregel( + chain_one, + chain_two, + chain_three, + channels={ + "input": channels.LastValue[int](), + "output": channels.LastValue[int](), + "inbox_one": channels.Inbox[int](), + "outbox_one": channels.LastValue[int](), + }, + input="input", + output="output", + ) # Then invoke pubsub for _ in range(10): - assert pubsub.invoke([2, 3]) == 27 + assert app.invoke([2, 3]) == 27 def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture): add_one = mocker.Mock(side_effect=lambda x: x + 1) - input_chan = channels.LastValue[int]("input") - output = channels.LastValue[int]("output") - between = channels.LastValue[int]("between") - chain_one = ( - Pregel.subscribe_to(input_chan) + Pregel.subscribe_to("input") | add_one - | Pregel.send_to( - {output: RunnablePassthrough(), between: RunnablePassthrough()} - ) + | Pregel.send_to(output=RunnablePassthrough(), between=RunnablePassthrough()) ) - chain_two = Pregel.subscribe_to(between) | add_one | Pregel.send_to(output) + chain_two = Pregel.subscribe_to("between") | add_one | Pregel.send_to("output") - pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output) + app = Pregel( + chain_one, + chain_two, + channels={ + "input": channels.LastValue[int](), + "output": channels.LastValue[int](), + "between": channels.LastValue[int](), + }, + input="input", + output="output", + ) # Then invoke pubsub - assert [c for c in pubsub.stream(2)] == [3, 4] + assert [c for c in app.stream(2)] == [3, 4] def test_invoke_two_processes_no_out(mocker: MockerFixture): - 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_chan) | add_one | Pregel.send_to(between) - chain_two = Pregel.subscribe_to(between) | add_one + chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("between") + chain_two = Pregel.subscribe_to("between") | add_one - pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output) + app = Pregel( + chain_one, + chain_two, + channels={ + "input": channels.LastValue[int](), + "output": channels.LastValue[int](), + "between": channels.LastValue[int](), + }, + input="input", + output="output", + ) # Then invoke pubsub # It finishes executing (once no more messages being published) # but returns nothing, as nothing was published to OUT topic - assert pubsub.invoke(2) is None + assert app.invoke(2) is None def test_invoke_two_processes_no_in(mocker: MockerFixture): - 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(between) | add_one | Pregel.send_to(output) - chain_two = Pregel.subscribe_to(between) | add_one - pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output) + chain_one = Pregel.subscribe_to("between") | add_one | Pregel.send_to("output") + chain_two = Pregel.subscribe_to("between") | add_one with pytest.raises(ValueError): - assert pubsub.invoke(2) is None + app = Pregel( + chain_one, + chain_two, + channels={ + "input": channels.LastValue[int](), + "output": channels.LastValue[int](), + "between": channels.LastValue[int](), + }, + input="input", + output="output", + )