diff --git a/.vscode/opsqueue.code-workspace b/.vscode/opsqueue.code-workspace new file mode 100644 index 00000000..dc0b7ae7 --- /dev/null +++ b/.vscode/opsqueue.code-workspace @@ -0,0 +1,13 @@ +{ + "folders": [ + { + "path": "..", + "name": "opsqueue", + }, + { + "path": "../libs/opsqueue_python", + "name": "opsqueue_python", + } + ], + "settings": {} +} diff --git a/libs/opsqueue_python/python/opsqueue/producer.py b/libs/opsqueue_python/python/opsqueue/producer.py index 82a877e3..97ff9ee9 100644 --- a/libs/opsqueue_python/python/opsqueue/producer.py +++ b/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], diff --git a/libs/opsqueue_python/src/errors.rs b/libs/opsqueue_python/src/errors.rs index 76b41725..339a50ec 100644 --- a/libs/opsqueue_python/src/errors.rs +++ b/libs/opsqueue_python/src/errors.rs @@ -153,6 +153,7 @@ impl From> for PyErr { } } +#[derive(Debug)] pub struct SubmissionFailed( pub crate::common::SubmissionFailed, pub crate::common::ChunkFailed, diff --git a/libs/opsqueue_python/src/producer.rs b/libs/opsqueue_python/src/producer.rs index e3167a9b..0899d608 100644 --- a/libs/opsqueue_python/src/producer.rs +++ b/libs/opsqueue_python/src/producer.rs @@ -427,6 +427,148 @@ impl ProducerClient { Ok(res) } + /// Stream output chunks as soon as each consumer has completed them. + #[must_use] + pub fn stream_submission_chunks(&self, submission_id: SubmissionId) -> PyChunksIter { + self.streaming_submission_chunks(submission_id) + } + + #[allow(clippy::too_many_lines)] + 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(Box::new( + 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 +599,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> { + 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 +731,41 @@ impl ProducerClient { } } -pub type ChunksStream = BoxStream<'static, CPyResult, ChunkRetrievalError>>; +#[derive(Debug)] +enum StreamingChunkError { + Retrieval(ChunkRetrievalError), + Internal(InternalProducerClientError), + Failed(Box), + 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> for PyErr { + fn from(value: CError) -> 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, StreamingChunkError>>; #[pyclass(module = "opsqueue")] pub struct PyChunksIter { @@ -574,16 +774,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 +797,7 @@ impl PyChunksIter { slf } - fn __next__(&self, py: Python<'_>) -> Option, ChunkRetrievalError>> { + fn __next__(&self, py: Python<'_>) -> Option, 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 || { diff --git a/libs/opsqueue_python/tests/test_roundtrip.py b/libs/opsqueue_python/tests/test_roundtrip.py index 23d9a149..8cc5fa00 100644 --- a/libs/opsqueue_python/tests/test_roundtrip.py +++ b/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]"]