Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions src/digitaltwin/runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -1516,8 +1516,12 @@ async def _run_component(
note_flow_task(selecting)
answer_ms = await selecting

# answer is an investigator id.
if isinstance(answer_ms, tuple) and len(answer_ms) == 2:
# answer is an investigator id. A selector that ran as a
# remote function task gets its (id, kwargs) pair back as a
# JSON list -- the wire has no tuples -- so both spellings
# mean the pair.
if (isinstance(answer_ms, (tuple, list))
and len(answer_ms) == 2):
i_select, model_kwargs = answer_ms
else:
i_select = answer_ms
Expand Down
10 changes: 8 additions & 2 deletions src/digitaltwin/streaming.py
Original file line number Diff line number Diff line change
Expand Up @@ -849,9 +849,15 @@ def resolve(
sub_addr: Optional[str] = None,
) -> "PubSubConfig":
"""Describe the configured broker: explicit addresses, else the
environment, else the loopback defaults (see `config`)."""
environment, else the loopback defaults (see `config`).

return cls(namespace, *stream_addresses(pub_addr, sub_addr))
The backend kind comes from `DT_STREAM_BACKEND` like everywhere
else -- an external producer on an orbit deployment needs only
the same environment the deployment itself runs with.
"""

return cls(namespace, *stream_addresses(pub_addr, sub_addr),
kind=stream_backend())

async def connect_backend(self, timeout: Optional[float] = None) -> PubSubBackend:
"""Open the transport alone, without namespace semantics.
Expand Down
66 changes: 66 additions & 0 deletions test/unit/test_agent_selection.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
"""A SciAgent's selection answer, in both spellings the wire produces.

A selector registered as a remote function task returns its
``(investigator_id, model_kwargs)`` pair as a JSON *list* -- the wire
has no tuples. Both spellings must select.
"""

import pytest

from digitaltwin import (
DataType,
DTRuntime,
ModelInvestigator,
SciAgent,
TypedData,
)

NO_FLOW = None

X = DataType("sel-in")
Y = DataType("sel-out")


class Inv(ModelInvestigator):
async def main_loop(self, runtime):
async def infer(in_data, k=1.0):
return TypedData(Y, k * in_data.data)

runtime.set_inference_task(infer)
runtime.publish_new_model({"k": 2.0}, {})


class Agent(SciAgent):
"""One investigator; the selector's answer shape is injected."""

def __init__(self, flow, answer_shape):
super().__init__(flow)
self.inv = Inv(flow)
self.answer_shape = answer_shape

async def main_loop(self, runtime):
runtime.start_investigator(self.inv)

async def select(in_data, i_id=0):
return self.answer_shape(i_id)

runtime.set_model_selection_task(select)
runtime.update_model_selector(i_id=self.inv.get_id())


@pytest.mark.parametrize("shape", [
lambda i: (i, {"k": 3.0}), # in-process selector: a tuple
lambda i: [i, {"k": 3.0}], # remote function task: JSON made it a list
], ids=["tuple", "list"])
async def test_a_selection_pair_selects_in_both_spellings(
shape, stream_clients):
runtime = DTRuntime(NO_FLOW, await stream_clients("agent-sel"))
runtime.add_agent(Agent(NO_FLOW, shape), X, Y)
runtime.start()

answer = await runtime.get_inference(TypedData(X, 7.0), Y)

assert answer.data == 21.0
assert runtime.state == "running"

await runtime.stop()
11 changes: 11 additions & 0 deletions test/unit/test_streaming.py
Original file line number Diff line number Diff line change
Expand Up @@ -201,3 +201,14 @@ async def test_close_releases_task_sockets_and_context(broker, no_task_leaks):
await twin.close()
with pytest.raises(RuntimeError):
await twin.publish(SENSOR, "nope")


def test_resolve_takes_the_backend_from_the_environment(monkeypatch):
"""An external producer on an orbit deployment needs only the same
environment the deployment runs with (`DT_STREAM_BACKEND`)."""

monkeypatch.delenv("DT_STREAM_BACKEND", raising=False)
assert PubSubConfig.resolve().kind == "zmq"

monkeypatch.setenv("DT_STREAM_BACKEND", "orbit")
assert PubSubConfig.resolve().kind == "orbit"
Loading