Add optional 'persist' query to PUT/PATCH array (bluesky#1213)

* Add optional 'persist' query to PUT/PATCH array

* Add CHANGELOG entry for `persist` query parameter.

* Add 'persist' option to array client write()/patch()

* Test 'persist' option for put array/full

* Test 'persist' option for patch array/full

* Receive all test_websockets messages

* Apply linter recommendations

* 'persist=None' for shorter url query

* Unindent asserts of websockets 'persist' tests

* Abstract websockets test steps

* Set DaskArrayClient default persist=True

* Apply linter fixes to test_websockets

* Websocket persist tests without extend array

* Raise Conflict for PATCH array with extend but no persist

* Template Conflicts to ValueError message conversion

* Test for ValueError in append_array with persist=False

* Handle test_updates_append_without_persist() separately

* Add 'Non-persistent Updates' to streaming tutorial

* Show Container.new() for non-persistent streaming

* Serialize NDSlice when streaming

* Test chunked writes in websocket stream

* Report array persist,extend bad args

Author: padraic-shafer

CHANGELOG.md

@@ -5,6 +5,11 @@ Write the date in place of the "Unreleased" in the case a new version is release
 
 ## Unreleased
 
+### Added
+
+- Optional `persist` query parameter to PUT and PATCH /array/... routes, and
+  the corresponding DaskArrayClient methods: `write`, `write_block`, `patch`.
+
 ### Changed
 
 - The public demo hosted by NSLS2 has moved from `tiled-demo.blueskyproject.io`

docs/source/tutorials/streaming.md

@@ -49,6 +49,7 @@ If using a config file, add this section:
     ```
 ````
 
+(write-stream-data)=
 ## Write and Stream Data
 
 ```py
@@ -175,10 +176,82 @@ dormant. These intervals are configurable via the settings `data_ttl` and
 `seq_ttl` respectively under `streaming_cach` configuration.  (In the `tiled
 serve` CLI they are `--cache-data-ttl` and `--cache-seq-ttl`.)
 
+## Performance Optimization: Non-persistent Updates
+
+To rapidly stream large arrays of transitory data, it might be unnecessary
+or even undesirable to persist the intermediate arrays to storage.
+In this case, pass the parameter `persist=False` to `ArrayClient.write()`
+or `ArrayClient.patch()`.
+
+This could be useful for serving reconstructed images to an on-screen display
+while an iterative algorithm is refining that reconstruction. Only the final
+image needs to be persisted. The intermediate images can be discarded as each
+new image is received.
+
+To try this yourself, modify the code in [](#write-stream-data) to include the
+`persist=False` when sending updates.
+
+```py
+# Intial state: x == [1, 2, 3]
+
+# PUT (write) new values
+x.write(numpy.array([4, 5, 6]), persist=False)
+  # x == [1, 2, 3], update.data() == [4, 5, 6]
+x.write(numpy.array([7, 8, 9]), persist=False)
+  # x == [1, 2, 3], update.data() == [7, 8, 9]
+
+# Persist the final array
+x.write(numpy.array([11, 12, 13]))  # x == [11, 12, 13], persist == True
+```
+
+```py
+# Intial state: x == [1, 2, 3]
+
+# PATCH (patch) new values
+x.patch(numpy.array([11], offset=(0,)), persist=False)
+  # x == [1, 2, 3], update.data() == [11], update.offset() == (0,)
+x.patch(numpy.array([12], offset=(1,)), persist=False)
+  # x == [1, 2, 3], update.data() == [12], update.offset() == (1,)
+x.patch(numpy.array([13], offset=(2,)), persist=False)
+  # x == [1, 2, 3], update.data() == [13], update.offset() == (2,)
+
+# Persist the final array
+x.write(numpy.array([11, 12, 13]))  # x == [11, 12, 13], persist == True
+```
+
+The array could initially be left empty to further reduce the writes to disk.
+
+```py
+client = from_uri('http://localhost:8000', api_key='secret')
+x = client.new(
+    structure_family=StructureFamily.array,
+    data_sources=[tiff_data_source],  # DataSource details omitted for brevity
+    data_sources=[
+        # In-memory DataSource details omitted for brevity
+        DataSource(structure=structure, structure_family=StructureFamily.array),
+        # Or use an external data source, such as a TIFF file
+        tiff_data_source,  # DataSource details omitted for brevity
+    ],
+    key='x',
+)
+# Intial state: x == [?, ?, ?]
+
+# PUT (write) new values
+x.write(numpy.array([4, 5, 6]), persist=False)
+  # x == [?, ?, ?], update.data() == [4, 5, 6]
+x.write(numpy.array([7, 8, 9]), persist=False)
+  # x == [?, ?, ?], update.data() == [7, 8, 9]
+
+# Persist the final array
+x.write(numpy.array([11, 12, 13]))  # x == [11, 12, 13], persist == True
+```
+
 ## Limitations
 
 This feature is in a very early preview stage.
 
 - External registered tabular data is not yet supported.
 - Other data structures (sparse, awkward) are not yet supported.
+- Arrays can be patched with either `extend=True` or `persist=False`,
+  but not both.
 - Deletions of nodes are not yet visible to subscribers.

tiled/_tests/test_websockets.py

@@ -1,5 +1,6 @@
 import sys
 
+import dask.array
 import msgpack
 import numpy as np
 import pytest
@@ -12,6 +13,71 @@
 )
 
 
+def send_ws_updates(client, update_func, start=1, count=1, persist=None):
+    """Helper to send updates via Tiled client in websocket tests."""
+    for i in range(start, start + count):
+        new_arr = np.arange(10) + i
+        update_func(client, new_arr, i, persist=persist)
+
+
+# An update_func for send_ws_updates
+def overwrite_array(client, new_arr, seq_num, persist=None):
+    _ = seq_num  # seq_num is unused for these updates
+    client.write(new_arr, persist=persist)
+
+
+# An update_func for send_ws_updates
+def write_array_block(client, new_arr, seq_num, persist=None):
+    client.write_block(new_arr, block=(seq_num - 1, 0), persist=persist)
+
+
+# An update_func for send_ws_updates
+def patch_array(client, new_arr, seq_num, persist=None):
+    _ = seq_num  # seq_num is unused for these updates
+    client.patch(new_arr, offset=(0,), persist=persist)
+
+
+# An update_func for send_ws_updates
+def append_array(client, new_arr, seq_num, persist=None):
+    client.patch(new_arr, offset=(10 * seq_num,), extend=True, persist=persist)
+
+
+def receive_ws_updates(websocket, count=1):
+    """Helper to receive updates in websocket tests."""
+    # Receive all updates
+    received = []
+    for _ in range(count + 1):  # +1 for schema
+        msg_bytes = websocket.receive_bytes()
+        msg = msgpack.unpackb(msg_bytes)
+        received.append(msg)
+
+    # Verify all messages received (schema + n updates)
+    assert len(received) == count + 1
+
+    return received
+
+
+def verify_ws_updates(received, start=1, chunked=False):
+    """Verify that we received messages with the expected data"""
+    for i, msg in enumerate(received):
+        if i == 0:  # schema
+            assert "type" in msg
+            assert "version" in msg
+        else:
+            assert "type" in msg
+            assert "timestamp" in msg
+            assert "payload" in msg
+            if chunked:
+                assert msg["shape"] == [1, 10]
+            else:
+                assert msg["shape"] == [10]
+
+            # Verify payload contains the expected array data
+            payload_array = np.frombuffer(msg["payload"], dtype=np.int64)
+            expected_array = np.arange(10) + (start - 1) + i
+            np.testing.assert_array_equal(payload_array, expected_array)
+
+
 def test_subscribe_immediately_after_creation_websockets(tiled_websocket_context):
     context = tiled_websocket_context
     client = from_context(context)
@@ -26,36 +92,12 @@ def test_subscribe_immediately_after_creation_websockets(tiled_websocket_context
         "/api/v1/stream/single/test_stream_immediate?envelope_format=msgpack",
         headers={"Authorization": "Apikey secret"},
     ) as websocket:
-        # Write updates using Tiled client
-        for i in range(1, 4):
-            new_arr = np.arange(10) + i
-            streaming_node.write(new_arr)
-
-        # Receive all updates
-        received = []
-        for _ in range(3):
-            msg_bytes = websocket.receive_bytes()
-            msg = msgpack.unpackb(msg_bytes)
-            received.append(msg)
-
-        # Verify all updates received in order
-        assert len(received) == 3
+        # Send 3 updates using Tiled client that overwrite the array
+        send_ws_updates(streaming_node, overwrite_array, count=3)
 
-        # Check that we received messages with the expected data
-        for i, msg in enumerate(received):
-            if i == 0:  # schema
-                assert "type" in msg
-                assert "version" in msg
-            else:
-                assert "type" in msg
-                assert "timestamp" in msg
-                assert "payload" in msg
-                assert msg["shape"] == [10]
-
-                # Verify payload contains the expected array data
-                payload_array = np.frombuffer(msg["payload"], dtype=np.int64)
-                expected_array = np.arange(10) + i
-                np.testing.assert_array_equal(payload_array, expected_array)
+        # Receive and validate all updates
+        received = receive_ws_updates(websocket, count=3)
+        verify_ws_updates(received)
 
 
 def test_websocket_connection_to_non_existent_node(tiled_websocket_context):
@@ -93,38 +135,13 @@ def test_subscribe_after_first_update_websockets(tiled_websocket_context):
         "/api/v1/stream/single/test_stream_after_update?envelope_format=msgpack",
         headers={"Authorization": "Apikey secret"},
     ) as websocket:
-        # Write more updates
-        for i in range(2, 4):
-            new_arr = np.arange(10) + i
-            streaming_node.write(new_arr)
+        # Send 2 more updates that overwrite the array
+        send_ws_updates(streaming_node, overwrite_array, start=2, count=2)
 
         # Should only receive the 2 new updates
-        received = []
-        for _ in range(2):
-            msg_bytes = websocket.receive_bytes()
-            msg = msgpack.unpackb(msg_bytes)
-            received.append(msg)
-
-        # Verify only new updates received
-        assert len(received) == 2
-
-        # Check that we received messages with the expected data
-        for i, msg in enumerate(received):
-            if i == 0:  # schema
-                assert "type" in msg
-                assert "version" in msg
-            else:
-                assert "type" in msg
-                assert "timestamp" in msg
-                assert "payload" in msg
-                assert msg["shape"] == [10]
-
-                # Verify payload contains the expected array data
-                payload_array = np.frombuffer(msg["payload"], dtype=np.int64)
-                expected_array = np.arange(10) + (
-                    i + 1
-                )  # i+2 because we start from update 1
-                np.testing.assert_array_equal(payload_array, expected_array)
+        received = receive_ws_updates(websocket, count=2)
+        # Content starts with update #2
+        verify_ws_updates(received, start=2)
 
 
 def test_subscribe_after_first_update_from_beginning_websockets(
@@ -198,6 +215,121 @@ def test_subscribe_after_first_update_from_beginning_websockets(
             np.testing.assert_array_equal(payload_array, expected_array)
 
 
+@pytest.mark.parametrize("write_op", (overwrite_array, patch_array))
+@pytest.mark.parametrize("persist", (None, True, False))
+def test_updates_persist_write(tiled_websocket_context, write_op, persist):
+    context = tiled_websocket_context
+    client = from_context(context)
+    test_client = context.http_client
+
+    # Create streaming array node using Tiled client
+    arr = np.arange(10)
+    streaming_node = client.write_array(arr, key="test_stream_immediate")
+
+    # Connect WebSocket using TestClient with msgpack format and authorization
+    with test_client.websocket_connect(
+        "/api/v1/stream/single/test_stream_immediate?envelope_format=msgpack",
+        headers={"Authorization": "Apikey secret"},
+    ) as websocket:
+        # Send 3 updates using Tiled client that write values into the array
+        send_ws_updates(streaming_node, write_op, count=3, persist=persist)
+
+        # Receive and validate all updates
+        received = receive_ws_updates(websocket, count=3)
+        verify_ws_updates(received)
+
+    # Verify values of persisted data
+    if persist or persist is None:
+        expected_persisted = np.arange(10) + 3  # Final sent values
+    else:
+        expected_persisted = arr  # Original values
+    persisted_data = streaming_node.read()
+    np.testing.assert_array_equal(persisted_data, expected_persisted)
+
+
+@pytest.mark.parametrize("persist", (None, True, False))
+def test_updates_persist_write_block(tiled_websocket_context, persist):
+    context = tiled_websocket_context
+    client = from_context(context)
+    test_client = context.http_client
+
+    # Create a streaming chunked array node using Tiled client
+    _arr = np.array([np.arange(10) for _ in range(3)])
+    arr = dask.array.from_array(_arr, chunks=(1, 10))  # Chunk along first axis
+    streaming_node = client.write_array(arr, key="test_stream_immediate")
+
+    # Connect WebSocket using TestClient with msgpack format and authorization
+    with test_client.websocket_connect(
+        "/api/v1/stream/single/test_stream_immediate?envelope_format=msgpack",
+        headers={"Authorization": "Apikey secret"},
+    ) as websocket:
+        # Send 3 updates using Tiled client that write values into the array
+        send_ws_updates(streaming_node, write_array_block, count=3, persist=persist)
+
+        # Receive and validate all updates
+        received = receive_ws_updates(websocket, count=3)
+        verify_ws_updates(received, chunked=True)
+
+    # Verify values of persisted data
+    if persist or persist is None:
+        # Combined effect of all sent values
+        expected_persisted = np.array([np.arange(10) + i for i in range(1, 4)])
+    else:
+        # Original values
+        expected_persisted = arr
+    persisted_data = streaming_node.read()
+    np.testing.assert_array_equal(persisted_data, expected_persisted)
+
+
+# Extending an array with persist=False is not yet supported
+@pytest.mark.parametrize("persist", (None, True))
+def test_updates_persist_append(tiled_websocket_context, persist):
+    context = tiled_websocket_context
+    client = from_context(context)
+    test_client = context.http_client
+
+    # Create streaming array node using Tiled client
+    arr = np.arange(10)
+    streaming_node = client.write_array(arr, key="test_stream_immediate")
+
+    # Connect WebSocket using TestClient with msgpack format and authorization
+    with test_client.websocket_connect(
+        "/api/v1/stream/single/test_stream_immediate?envelope_format=msgpack",
+        headers={"Authorization": "Apikey secret"},
+    ) as websocket:
+        # Send 3 updates using Tiled client that append to the array
+        send_ws_updates(streaming_node, append_array, count=3, persist=persist)
+
+        # Receive and validate all updates
+        received = receive_ws_updates(websocket, count=3)
+        verify_ws_updates(received)
+
+    # Verify values of persisted data
+    if persist or persist is None:
+        # Combined effect of all sent values
+        expected_persisted = np.array(
+            [np.arange(10) + i for i in range(0, 4)]
+        ).flatten()
+    else:
+        # Original values
+        expected_persisted = arr
+    persisted_data = streaming_node.read()
+    np.testing.assert_array_equal(persisted_data, expected_persisted)
+
+
+def test_updates_append_without_persist(tiled_websocket_context):
+    context = tiled_websocket_context
+    client = from_context(context)
+
+    # Create streaming array node using Tiled client
+    arr = np.arange(10)
+    streaming_node = client.write_array(arr, key="test_stream_immediate")
+
+    with pytest.raises(ValueError, match="Cannot PATCH an array with both parameters"):
+        # Extending an array with persist=False is not yet supported
+        send_ws_updates(streaming_node, append_array, count=1, persist=False)
+
+
 def test_close_stream_success(tiled_websocket_context):
     """Test successful close of an existing stream."""
     context = tiled_websocket_context