class WorkRun(LightningWork):
def __init__(self, tmpdir):
super().__init__()
self.var_a = None
self.tmpdir = tmpdir
def run(self):
self.var_a = Payload("something")
assert self.var_a.name == "var_a"
assert self.var_a._origin == "root.a"
assert self.var_a.hash == "9bd514ad51fc33d895c50657acd0f0582301cf3e"
source_path = pathlib.Path(self.tmpdir, self.var_a.name)
assert not source_path.exists()
response = self.var_a._handle_get_request(
self,
GetRequest(
name="var_a",
hash=self.var_a.hash,
source="root.a",
path=str(source_path),
destination="root",
),
)
assert source_path.exists()
assert self.var_a.load(str(source_path)) == "something"
assert not response.exception
def test_path_attach_queues():
path = Payload(None)
request_queue = Mock()
response_queue = Mock()
path._attach_queues(request_queue=request_queue,
response_queue=response_queue)
assert path._request_queue is request_queue
assert path._response_queue is response_queue
def test_payload_pickable():
payload = Payload("MyObject")
payload._origin = "root.x.y.z"
payload._consumer = "root.p.q.r"
payload._name = "var_a"
loaded = pickle.loads(pickle.dumps(payload))
assert isinstance(loaded, Payload)
assert loaded._origin == payload._origin
assert loaded._consumer == payload._consumer
assert loaded._name == payload._name
assert loaded._request_queue is None
assert loaded._response_queue is None
def test_payload_copy():
"""Test that Payload creates an exact copy when passing a Payload instance to the constructor."""
payload = Payload(None)
payload._origin = "origin"
payload._consumer = "consumer"
payload._request_queue = "MockQueue"
payload._response_queue = "MockQueue"
payload_copy = deepcopy(payload)
assert payload_copy._origin == payload._origin
assert payload_copy._consumer == payload._consumer
assert payload_copy._request_queue == payload._request_queue
assert payload_copy._response_queue == payload._response_queue
def __init__(self):
super().__init__()
self.var_a = Payload(None)
def run(self):
self.value_all = Payload(["A", "B", "C"])
self.value_b = Payload("B")
self.value_c = Payload("C")
def on_after_run(self, res: Any):
"""Called after the python script is executed."""
for name in self.outputs:
setattr(self, name, Payload(res[name]))
def run(self):
self.payload = Payload(2)