Feat: stream completed chunks

Author: vincevannoort

.vscode/opsqueue.code-workspace

@@ -0,0 +1,13 @@
+{
+	"folders": [
+		{
+			"path": "..",
+			"name": "opsqueue",
+		},
+		{
+			"path": "../libs/opsqueue_python",
+			"name": "opsqueue_python",
+		}
+	],
+	"settings": {}
+}

libs/opsqueue_python/python/opsqueue/producer.py

@@ -231,6 +231,14 @@ def run_submission_chunks(
         )
         return self.blocking_stream_completed_submission_chunks(submission_id)
 
+    def stream_submission_chunks(self, submission_id: SubmissionId) -> Iterator[bytes]:
+        return self.inner.stream_submission_chunks(submission_id)  # type: ignore[no-any-return]
+
+    async def async_stream_submission_chunks(
+        self, submission_id: SubmissionId
+    ) -> AsyncIterator[bytes]:
+        return await self.inner.async_stream_submission_chunks(submission_id)  # type: ignore[no-any-return]
+
     async def async_run_submission_chunks(
         self,
         chunk_contents: Iterable[bytes],

libs/opsqueue_python/src/errors.rs

@@ -153,6 +153,7 @@ impl From<CError<TooManyMatchingSubmissions>> for PyErr {
     }
 }
 
+#[derive(Debug)]
 pub struct SubmissionFailed(
     pub crate::common::SubmissionFailed,
     pub crate::common::ChunkFailed,

libs/opsqueue_python/src/producer.rs

@@ -427,6 +427,146 @@ impl ProducerClient {
         Ok(res)
     }
 
+    /// Stream output chunks as soon as each consumer has completed them.
+    pub fn stream_submission_chunks(&self, submission_id: SubmissionId) -> PyChunksIter {
+        self.streaming_submission_chunks(submission_id)
+    }
+
+    fn streaming_submission_chunks(&self, submission_id: SubmissionId) -> PyChunksIter {
+        let client = self.client.clone();
+        let object_store_client = self.object_store_client.clone();
+        let stream = futures::stream::unfold(
+            (
+                client,
+                object_store_client,
+                submission_id,
+                u63::new(0),
+                None,
+                Duration::from_millis(10),
+            ),
+            |(client, object_store_client, submission_id, index, prefix, interval)| async move {
+                let mut interval = interval;
+                loop {
+                    let status = match client.get_submission(submission_id.into()).await {
+                        Ok(Some(status)) => status,
+                        Ok(None) => {
+                            return Some((
+                                Err(StreamingChunkError::SubmissionNotFound),
+                                (
+                                    client,
+                                    object_store_client,
+                                    submission_id,
+                                    index,
+                                    prefix,
+                                    interval,
+                                ),
+                            ));
+                        }
+                        Err(error) => {
+                            return Some((
+                                Err(StreamingChunkError::Internal(error)),
+                                (
+                                    client,
+                                    object_store_client,
+                                    submission_id,
+                                    index,
+                                    prefix,
+                                    interval,
+                                ),
+                            ));
+                        }
+                    };
+
+                    match status {
+                        submission::SubmissionStatus::InProgress(submission) => {
+                            let prefix = prefix.clone().or(submission.prefix);
+                            if index < submission.chunks_done.into() {
+                                let prefix = prefix
+                                    .expect("in-progress submissions have an object-store prefix");
+                                let result = object_store_client
+                                    .retrieve_chunk(&prefix, index.into(), ChunkType::Output)
+                                    .await
+                                    .map_err(StreamingChunkError::Retrieval);
+                                return Some((
+                                    result,
+                                    (
+                                        client,
+                                        object_store_client,
+                                        submission_id,
+                                        index + u63::new(1),
+                                        Some(prefix),
+                                        interval,
+                                    ),
+                                ));
+                            }
+                        }
+                        submission::SubmissionStatus::Completed(submission) => {
+                            let prefix = prefix.clone().or(submission.prefix);
+                            if index < submission.chunks_total.into() {
+                                let prefix = prefix
+                                    .expect("completed submissions have an object-store prefix");
+                                let result = object_store_client
+                                    .retrieve_chunk(&prefix, index.into(), ChunkType::Output)
+                                    .await
+                                    .map_err(StreamingChunkError::Retrieval);
+                                return Some((
+                                    result,
+                                    (
+                                        client,
+                                        object_store_client,
+                                        submission_id,
+                                        index + u63::new(1),
+                                        Some(prefix),
+                                        interval,
+                                    ),
+                                ));
+                            }
+                            return None;
+                        }
+                        submission::SubmissionStatus::Failed(submission, chunk) => {
+                            let failure =
+                                crate::common::ChunkFailed::from_internal(chunk, &submission);
+                            return Some((
+                                Err(StreamingChunkError::Failed(
+                                    crate::errors::SubmissionFailed(submission.into(), failure),
+                                )),
+                                (
+                                    client,
+                                    object_store_client,
+                                    submission_id,
+                                    index,
+                                    prefix,
+                                    interval,
+                                ),
+                            ));
+                        }
+                        submission::SubmissionStatus::Cancelled(_) => {
+                            return Some((
+                                Err(StreamingChunkError::Cancelled),
+                                (
+                                    client,
+                                    object_store_client,
+                                    submission_id,
+                                    index,
+                                    prefix,
+                                    interval,
+                                ),
+                            ));
+                        }
+                    }
+
+                    tokio::time::sleep(interval).await;
+                    if interval < SUBMISSION_POLLING_INTERVAL {
+                        interval = (interval * 2).min(SUBMISSION_POLLING_INTERVAL);
+                    }
+                }
+            },
+        )
+        .map(|item| item.map_err(CError))
+        .boxed();
+        PyChunksIter::from_stream(self, stream)
+    }
+
     /// Blocks (and short-polls) until the submission is completed.
     ///
     /// We start with a small short-polling interval
@@ -457,6 +597,30 @@ impl ProducerClient {
         })
     }
 
+    /// Return an awaitable that resolves immediately to an async iterator of output chunks.
+    ///
+    /// The iterator polls submission progress and yields each output chunk as soon as it is ready.
+    ///
+    /// # Errors
+    ///
+    /// Returns a Python error if creating the awaitable fails.
+    pub fn async_stream_submission_chunks<'p>(
+        &self,
+        py: Python<'p>,
+        submission_id: SubmissionId,
+    ) -> PyResult<Bound<'p, PyAny>> {
+        let me = self.clone();
+        let _tokio_active_runtime_guard = me.runtime.enter();
+        async_util::future_into_py(
+            py,
+            async_util::async_detach(Box::pin(async move {
+                Ok(PyChunksAsyncIter::from(
+                    me.streaming_submission_chunks(submission_id),
+                ))
+            })),
+        )
+    }
+
     /// Return an awaitable that resolves to an async iterator of output chunks.
     ///
     /// # Errors
@@ -565,7 +729,41 @@ impl ProducerClient {
     }
 }
 
-pub type ChunksStream = BoxStream<'static, CPyResult<Vec<u8>, ChunkRetrievalError>>;
+#[derive(Debug)]
+enum StreamingChunkError {
+    Retrieval(ChunkRetrievalError),
+    Internal(InternalProducerClientError),
+    Failed(crate::errors::SubmissionFailed),
+    SubmissionNotFound,
+    Cancelled,
+}
+
+impl std::fmt::Display for StreamingChunkError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Retrieval(error) => error.fmt(f),
+            Self::Internal(error) => error.fmt(f),
+            Self::Failed(_) => write!(f, "Submission failed"),
+            Self::SubmissionNotFound => write!(f, "Submission not found"),
+            Self::Cancelled => write!(f, "Submission cancelled"),
+        }
+    }
+}
+
+impl std::error::Error for StreamingChunkError {}
+
+impl From<CError<StreamingChunkError>> for PyErr {
+    fn from(value: CError<StreamingChunkError>) -> Self {
+        match value.0 {
+            StreamingChunkError::Retrieval(error) => CError(error).into(),
+            StreamingChunkError::Internal(error) => CError(error).into(),
+            StreamingChunkError::Failed(error) => CError(error).into(),
+            error => PyException::new_err(error.to_string()),
+        }
+    }
+}
+
+type ChunksStream = BoxStream<'static, CPyResult<Vec<u8>, StreamingChunkError>>;
 
 #[pyclass(module = "opsqueue")]
 pub struct PyChunksIter {
@@ -574,16 +772,20 @@ pub struct PyChunksIter {
 }
 
 impl PyChunksIter {
+    fn from_stream(client: &ProducerClient, stream: ChunksStream) -> Self {
+        Self {
+            stream: Arc::new(tokio::sync::Mutex::new(stream)),
+            runtime: client.runtime.clone(),
+        }
+    }
+
     pub(crate) fn new(client: &ProducerClient, prefix: String, chunks_total: u63) -> Self {
         let stream = client
             .object_store_client
             .retrieve_chunks(prefix, chunks_total, ChunkType::Output)
-            .map_err(CError)
+            .map_err(|error| CError(StreamingChunkError::Retrieval(error)))
             .boxed();
-        Self {
-            stream: Arc::new(tokio::sync::Mutex::new(stream)),
-            runtime: client.runtime.clone(),
-        }
+        Self::from_stream(client, stream)
     }
 }
 
@@ -593,7 +795,7 @@ impl PyChunksIter {
         slf
     }
 
-    fn __next__(&self, py: Python<'_>) -> Option<CPyResult<Vec<u8>, ChunkRetrievalError>> {
+    fn __next__(&self, py: Python<'_>) -> Option<CPyResult<Vec<u8>, StreamingChunkError>> {
         // The only time we need the GIL is when turning the result back.
         // By unlocking here, we reduce the chance of deadlocks.
         py.detach(move || {

libs/opsqueue_python/tests/test_roundtrip.py

@@ -3,6 +3,8 @@
 # - use `RUST_LOG="opsqueue=info"` (or `opsqueue=debug` or `debug` for even more verbosity), together with to the pytest option `-s` AKA `--capture=no`, to debug the opsqueue binary itself.
 
 from collections.abc import Iterator, Sequence
+import asyncio
+import time
 from opsqueue.producer import (
     SubmissionId,
     ProducerClient,
@@ -667,3 +669,95 @@ def test_lookup_too_many_submission_ids_by_strategic_metadata() -> None:
             )
         assert exc.type is TooManyMatchingSubmissionsError
         assert exc.value.max_submissions == max_
+
+
+def test_streams_completed_chunks_before_submission_finishes(
+    opsqueue: OpsqueueProcess,
+    any_consumer_strategy: StrategyDescription,
+) -> None:
+    url = "file:///tmp/opsqueue/test_streaming_results"
+    producer_client = ProducerClient(f"localhost:{opsqueue.port}", url)
+    submission_id = producer_client.insert_submission_chunks(
+        [b"[1]", b"[2]"], chunk_size=1
+    )
+
+    def complete_chunks(
+        _submission_id_value: int,
+        strategy: StrategyDescription,
+    ) -> None:
+        consumer_client = ConsumerClient(f"localhost:{opsqueue.port}", url)
+        chunks = sorted(
+            consumer_client.reserve_chunks(
+                max=2,
+                strategy=strategy_from_description(strategy),
+            ),
+            key=lambda chunk: chunk.chunk_index,
+        )
+        consumer_client.complete_chunk(
+            chunks[0].submission_id,
+            chunks[0].submission_prefix,
+            chunks[0].chunk_index,
+            chunks[0].input_content,
+        )
+        time.sleep(0.25)
+        consumer_client.complete_chunk(
+            chunks[1].submission_id,
+            chunks[1].submission_prefix,
+            chunks[1].chunk_index,
+            chunks[1].input_content,
+        )
+
+    with background_process(
+        complete_chunks,
+        args=(submission_id.id, any_consumer_strategy),
+    ):
+        results = producer_client.stream_submission_chunks(submission_id)
+        assert next(results) == b"[1]"
+        assert next(results) == b"[2]"
+
+
+def test_async_streams_completed_chunks_before_submission_finishes(
+    opsqueue: OpsqueueProcess,
+    any_consumer_strategy: StrategyDescription,
+) -> None:
+    url = "file:///tmp/opsqueue/test_async_streaming_results"
+    producer_client = ProducerClient(f"localhost:{opsqueue.port}", url)
+    submission_id = producer_client.insert_submission_chunks(
+        [b"[1]", b"[2]"], chunk_size=1
+    )
+
+    def complete_chunks(
+        _submission_id_value: int,
+        strategy: StrategyDescription,
+    ) -> None:
+        consumer_client = ConsumerClient(f"localhost:{opsqueue.port}", url)
+        chunks = sorted(
+            consumer_client.reserve_chunks(
+                max=2,
+                strategy=strategy_from_description(strategy),
+            ),
+            key=lambda chunk: chunk.chunk_index,
+        )
+        consumer_client.complete_chunk(
+            chunks[0].submission_id,
+            chunks[0].submission_prefix,
+            chunks[0].chunk_index,
+            chunks[0].input_content,
+        )
+        time.sleep(0.25)
+        consumer_client.complete_chunk(
+            chunks[1].submission_id,
+            chunks[1].submission_prefix,
+            chunks[1].chunk_index,
+            chunks[1].input_content,
+        )
+
+    async def collect() -> list[bytes]:
+        results = await producer_client.async_stream_submission_chunks(submission_id)
+        return [chunk async for chunk in results]
+
+    with background_process(
+        complete_chunks,
+        args=(submission_id.id, any_consumer_strategy),
+    ):
+        assert asyncio.run(collect()) == [b"[1]", b"[2]"]