Skip to content

Commit 22608d4

Browse files
Deimos28Deimos28
authored
Add optional parameter to force payload offloading (#608)
When sending a gRPC payload to Temporal, if the payload is too large, we offload some of it to Amazon S3: this commit allows you to do this even if the payload is small. This is useful when uploading several things at once (eg registering multiple child workflows), where any single payload is small but their concatenation is too big and no offload happens.
1 parent dd6a743 commit 22608d4

2 files changed

Lines changed: 89 additions & 3 deletions

File tree

‎src/mistralai/extra/tests/test_workflow_encoding.py‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -760,3 +760,83 @@ async def test_workflow_encoding_hook_handles_gzipped_response():
760760
assert isinstance(result, httpx.Response)
761761
response_body = json.loads(result.content)
762762
assert response_body["result"] == original_data
763+
764+
765+
@pytest.mark.asyncio
766+
async def test_encode_payload_content_does_not_offload_below_threshold(monkeypatch):
767+
storage = InMemoryBlobStorage()
768+
monkeypatch.setattr(
769+
"mistralai.extra.workflows.encoding.payload_encoder.get_blob_storage",
770+
lambda _: storage,
771+
)
772+
config = WorkflowEncodingConfig(
773+
payload_offloading=PayloadOffloadingConfig(
774+
min_size_bytes=1024,
775+
storage_config=BlobStorageConfig(
776+
storage_provider=StorageProvider.S3,
777+
bucket_name="test-bucket",
778+
),
779+
),
780+
)
781+
encoder = PayloadEncoder(encoding_config=config)
782+
context = WorkflowContext(namespace="test", execution_id="exec")
783+
784+
data, options = await encoder.encode_payload_content(b"tiny", context)
785+
786+
assert options == []
787+
assert storage.blobs == {}
788+
assert data == b"tiny"
789+
790+
791+
@pytest.mark.asyncio
792+
async def test_encode_payload_content_force_offload_bypasses_threshold(monkeypatch):
793+
storage = InMemoryBlobStorage()
794+
monkeypatch.setattr(
795+
"mistralai.extra.workflows.encoding.payload_encoder.get_blob_storage",
796+
lambda _: storage,
797+
)
798+
config = WorkflowEncodingConfig(
799+
payload_offloading=PayloadOffloadingConfig(
800+
min_size_bytes=1024,
801+
storage_config=BlobStorageConfig(
802+
storage_provider=StorageProvider.S3,
803+
bucket_name="test-bucket",
804+
),
805+
),
806+
)
807+
encoder = PayloadEncoder(encoding_config=config)
808+
context = WorkflowContext(namespace="test", execution_id="exec")
809+
810+
data, options = await encoder.encode_payload_content(
811+
b"tiny", context, force_offload=True
812+
)
813+
814+
assert EncodedPayloadOptions.OFFLOADED in options
815+
assert len(storage.blobs) == 1
816+
ref = json.loads(data)
817+
assert storage.blobs[ref["key"]] == b"tiny"
818+
819+
820+
@pytest.mark.asyncio
821+
async def test_encode_payload_content_force_offload_is_idempotent(monkeypatch):
822+
storage = InMemoryBlobStorage()
823+
monkeypatch.setattr(
824+
"mistralai.extra.workflows.encoding.payload_encoder.get_blob_storage",
825+
lambda _: storage,
826+
)
827+
config = WorkflowEncodingConfig(
828+
payload_offloading=PayloadOffloadingConfig(
829+
min_size_bytes=1024,
830+
storage_config=BlobStorageConfig(
831+
storage_provider=StorageProvider.S3,
832+
bucket_name="test-bucket",
833+
),
834+
),
835+
)
836+
encoder = PayloadEncoder(encoding_config=config)
837+
context = WorkflowContext(namespace="test", execution_id="exec")
838+
839+
await encoder.encode_payload_content(b"same", context, force_offload=True)
840+
await encoder.encode_payload_content(b"same", context, force_offload=True)
841+
842+
assert len(storage.blobs) == 1

‎src/mistralai/extra/workflows/encoding/payload_encoder.py‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -213,7 +213,10 @@ def _decrypt(self, data: bytes) -> bytes:
213213
) from main_exc
214214

215215
async def _handle_offloading(
216-
self, data: bytes, context: Optional[WorkflowContext]
216+
self,
217+
data: bytes,
218+
context: Optional[WorkflowContext],
219+
force: bool = False,
217220
) -> tuple[bytes, bool]:
218221
if (
219222
self.offloading_config is None
@@ -223,7 +226,7 @@ async def _handle_offloading(
223226
"You must configure payload offloading storage"
224227
)
225228

226-
if len(data) < self.offloading_config.min_size_bytes:
229+
if not force and len(data) < self.offloading_config.min_size_bytes:
227230
return data, False
228231

229232
if not context:
@@ -325,6 +328,7 @@ async def encode_payload_content(
325328
context: Optional[WorkflowContext] = None,
326329
*,
327330
allow_offloading: bool = True,
331+
force_offload: bool = False,
328332
) -> tuple[bytes, list[EncodedPayloadOptions]]:
329333
"""Handle payload encoding.
330334
@@ -353,7 +357,9 @@ async def encode_payload_content(
353357
encoding_options.append(EncodedPayloadOptions.COMPRESSED)
354358

355359
if allow_offloading and self.offloading_config is not None:
356-
data, offloaded = await self._handle_offloading(data, context)
360+
data, offloaded = await self._handle_offloading(
361+
data, context, force=force_offload
362+
)
357363
if offloaded:
358364
encoding_options.append(EncodedPayloadOptions.OFFLOADED)
359365

0 commit comments

Comments
 (0)