initial commit

Author: winsomeglint

asr-worker/asr_worker/constants.py

@@ -8,8 +8,4 @@
 
 PARAKEET = "parakeet"
 
-DEFAULT_TEMPORAL_ADDRESS = "temporal:7233"
-
-RESPONSE_SUCCESS = "success"
-
-RESPONSE_ERROR = "error"
+ASR_WORKFLOW_NAME = "asr-workflow"

asr-worker/asr_worker/models.py asr-worker/asr_worker/objects.pyasr-worker/asr_worker/models.py renamed to asr-worker/asr_worker/objects.py

@@ -1,7 +1,8 @@
 from dataclasses import dataclass, field, asdict
-from typing import Optional
 
-from asr_worker.constants import PARAKEET
+from .constants import PARAKEET
+
+from datashare_python.objects import WorkerResponse
 
 
 @dataclass
@@ -45,10 +46,7 @@ class ASRInputs:
     pipeline: ASRPipelineConfig
 
 
-@dataclass
-class ASRResponse:
+class ASRResponse(WorkerResponse):
     """ASR workflow response"""
 
-    status: str
     transcriptions: list[dict] = field(default_factory=list)
-    error: Optional[str] = None

asr-worker/asr_worker/worker.py

@@ -9,15 +9,15 @@
 
 from temporalio import workflow
 
-from asr_worker.constants import (
+from datashare_python.constants import DEFAULT_TEMPORAL_ADDRESS
+from .constants import (
     ASR_TASK_QUEUE,
     ASR_WORKER_NAME,
-    DEFAULT_TEMPORAL_ADDRESS,
 )
-from asr_worker.workflow import ASRWorkflow
+from .workflow import ASRWorkflow
 
 with workflow.unsafe.imports_passed_through():
-    from asr_worker.activities import ASRActivities
+    from .activities import ASRActivities
 
 LOGGER = logging.getLogger(__name__)
 

asr-worker/asr_worker/workflow.py

@@ -5,16 +5,16 @@
 from more_itertools import flatten
 from temporalio import workflow
 
-from asr_worker.constants import _TEN_MINUTES, RESPONSE_ERROR, RESPONSE_SUCCESS
-from asr_worker.models import ASRInputs, ASRResponse
+from asr_worker.models import ASRResponse, ASRInputs
+from asr_worker.constants import _TEN_MINUTES, RESPONSE_SUCCESS, RESPONSE_ERROR
 
 with workflow.unsafe.imports_passed_through():
     from asr_worker.activities import ASRActivities
 
 
 # TODO: Figure out which modules are violating sandbox restrictions
 #  and grant a limited passthrough
-@workflow.defn(name="asr.transcription", sandboxed=False)
+@workflow.defn(sandboxed=False)
 class ASRWorkflow:
     """ASR workflow definition"""
 
@@ -28,6 +28,7 @@ async def run(self, inputs: ASRInputs) -> ASRResponse:
         :param inputs: ASRInputs
         :return: ASRResponse
         """
+
         try:
             # Preprocessing
             preprocessed_batches = await gather(
@@ -98,6 +99,3 @@ async def run(self, inputs: ASRInputs) -> ASRResponse:
         except ValueError as e:
             workflow.logger.exception(e)
             return ASRResponse(status=RESPONSE_ERROR, error=str(e))
-
-
-WORKFLOWS = [ASRWorkflow]

asr-worker/pyproject.toml

@@ -21,6 +21,7 @@ authors = [
 readme = "README.md"
 requires-python = ">=3.11.0, <3.14"
 dependencies = [
+    "datashare-python",
     "temporalio>=1.22.0",
     "caul==0.1.5",
     "pyarrow<21.0.0",
@@ -37,6 +38,8 @@ workflows = "asr_worker.workflows:REGISTRY"
 [project.entry-points."datashare.activities"]
 activities = "asr_worker.activities:REGISTRY"
 
+
+
 [[tool.uv.index]]
 name = "pytorch-cpu"
 url = "https://download.pytorch.org/whl/cpu"
@@ -52,6 +55,7 @@ torchaudio = [
 torchcodec = [
   { index = "pytorch-cpu" }
 ]
+datashare-python = { path = "../datashare-python", editable = true }
 
 [tool.uv]
 package = true

…rker/tests/e2e/asr/test_transcription.py …r-worker/tests/e2e/test_transcription.pyasr-worker/tests/e2e/asr/test_transcription.py renamed to asr-worker/tests/e2e/test_transcription.py asr-worker/tests/e2e/asr/test_transcription.py renamed to asr-worker/tests/e2e/test_transcription.py

@@ -11,3 +11,15 @@
 DEFAULT_DS_ADDRESS = "http://localhost:8080"
 
 DEFAULT_NAMESPACE = "datashare-default"
+
+WORKER_RESPONSE_SUCCESS = "success"
+
+WORKER_RESPONSE_ERROR = "error"
+
+CPU = "cpu"
+
+CUDA = "cuda"
+
+MPS = "mps"
+
+MKL = "mkl"

datashare-python/datashare_python/constants.py

@@ -3,10 +3,12 @@
 from dataclasses import dataclass
 from datetime import UTC, datetime
 from enum import StrEnum, unique
-from typing import Any, Literal, Self
+from typing import Any, Literal, Self, Optional, TypeVar, Awaitable
 
 from temporalio import workflow
 
+from datashare_python.constants import WORKER_RESPONSE_SUCCESS, WORKER_RESPONSE_ERROR
+
 with workflow.unsafe.imports_passed_through():
     from icij_common.es import DOC_CONTENT, DOC_LANGUAGE, DOC_ROOT_ID, ID_, SOURCE
 
@@ -23,14 +25,28 @@
 logger = logging.getLogger(__name__)
 
 
+T = TypeVar("T")
+
+
+Predicate = Callable[[T], bool] | Callable[[T], Awaitable[bool]]
+
+
 class BaseModel(_BaseModel):
     model_config = merge_configs(icij_config(), no_enum_values_config())
 
 
+class BasePayload(_BaseModel):
+    model_config = icij_config()
+
+
 class DatashareModel(BaseModel):
     model_config = merge_configs(BaseModel.model_config, lowercamel_case_config())
 
 
+class LowerCamelCaseModel(_BaseModel):
+    model_config = merge_configs(icij_config(), lowercamel_case_config())
+
+
 @unique
 class TaskState(StrEnum):
     CREATED = "CREATED"
@@ -142,3 +158,16 @@ def from_es(cls, es_doc: dict) -> Self:
             root_document=sources[DOC_ROOT_ID],
             tags=sources.get("tags", []),
         )
+
+
+# Temporal objects
+class WorkerResponse(BasePayload):
+    """Generic worker response"""
+
+    status: str
+    error: Optional[str] = None
+
+
+class WorkerResponseStatus(StrEnum):
+    SUCCESS = WORKER_RESPONSE_SUCCESS
+    ERROR = WORKER_RESPONSE_ERROR

datashare-python/datashare_python/objects.py

@@ -6,9 +6,19 @@
 from dataclasses import dataclass
 from functools import partial, wraps
 from inspect import signature
-from typing import ParamSpec, TypeVar
+from itertools import islice
+from typing import (
+    ParamSpec,
+    TypeVar,
+    AsyncIterator,
+    AsyncIterable,
+    Awaitable,
+    Generator,
+    Iterable,
+)
 
 import nest_asyncio
+import torch
 from icij_common.logging_utils import (
     DATE_FMT,
     STREAM_HANDLER_FMT,
@@ -23,6 +33,8 @@
 from temporalio.common import SearchAttributeKey
 from temporalio.exceptions import ApplicationError
 
+from .constants import CPU, MPS, MKL, CUDA
+from .objects import Predicate
 from .types_ import ProgressRateHandler, RawProgressHandler
 
 DependencyLabel = str | None
@@ -201,6 +213,7 @@ async def wrapper(*args, **kwargs) -> T:
             # recreate kwargs from pargs
             new_args, new_kwargs = _unpack_positional_args(args, keyword_only, params)
             return await activity_fn(*new_args, **new_kwargs, **kwargs)
+
     else:
 
         @wraps(activity_fn)
@@ -211,9 +224,11 @@ def wrapper(*args, **kwargs) -> T:
 
     # Update the decorated function signature to appear as p-args only
     new_params = [
-        p.replace(kind=inspect.Parameter.POSITIONAL_OR_KEYWORD)
-        if p.kind == inspect.Parameter.KEYWORD_ONLY
-        else p
+        (
+            p.replace(kind=inspect.Parameter.POSITIONAL_OR_KEYWORD)
+            if p.kind == inspect.Parameter.KEYWORD_ONLY
+            else p
+        )
         for p in params
     ]
     wrapper.__signature__ = sig.replace(parameters=new_params)
@@ -252,6 +267,7 @@ async def wrapper(*args, **kwargs) -> T:
                     raise
                 except Exception as e:
                     raise fatal_error_from_exception(e) from e
+
         else:
 
             @wraps(activity_fn)
@@ -374,3 +390,83 @@ def _handlers(
                 handler.addFilter(worker_id_filter)
             handler.setLevel(log_level)
         return handlers
+
+
+# Temporal utils
+async def async_batches(
+    iterable: AsyncIterable[T], batch_size: int
+) -> AsyncIterator[tuple[T]]:
+    it = aiter(iterable)
+    if batch_size < 1:
+        raise ValueError("n must be at least one")
+    while True:
+        batch = []
+        while len(batch) < batch_size:
+            try:
+                batch.append(await anext(it))
+            except StopAsyncIteration:
+                if batch:
+                    yield tuple(batch)
+                return
+        yield tuple(batch)
+
+
+def batches(
+    iterable: Iterable[T], batch_size: int
+) -> Generator[tuple[T, ...], None, None]:
+    if batch_size < 1:
+        raise ValueError("n must be at least one")
+    it = iter(iterable)
+    while batch := tuple(islice(it, batch_size)):
+        yield batch
+
+
+async def maybe_await(maybe_awaitable: Awaitable[T] | T) -> T:
+    if inspect.isawaitable(maybe_awaitable):
+        return await maybe_awaitable
+    return maybe_awaitable
+
+
+async def once(item: T) -> AsyncIterator[T]:
+    yield item
+
+
+def before_and_after(
+    iterable: AsyncIterable[T], predicate: Predicate[T]
+) -> tuple[AsyncIterable[T], AsyncIterable[T]]:
+    transition = asyncio.get_event_loop().create_future()
+
+    async def true_iterator() -> AsyncIterator[T]:
+        async for elem in iterable:
+            if await maybe_await(predicate(elem)):
+                yield elem
+            else:
+                transition.set_result(elem)
+                return
+        transition.set_exception(StopAsyncIteration)
+
+    async def remainder_iterator() -> AsyncIterator[T]:
+        try:
+            yield await transition
+        except StopAsyncIteration:
+            return
+        async for elm in iterable:
+            yield elm
+
+    return true_iterator(), remainder_iterator()
+
+
+# Torch utils
+def find_device(device_name: str = CPU) -> torch.Device:
+    """Find a device by name if available
+
+    :param device_name: Device name
+    :return: torch.Device
+    """
+    if (
+        hasattr(torch.backends, device_name)
+        and getattr(torch.backends, device_name).is_available()
+    ):
+        return torch.device(device_name)
+
+    return torch.device(CPU)

datashare-python/datashare_python/utils.py

@@ -1,4 +1,8 @@
-from datashare_python.utils import positional_args_only
+from typing import AsyncGenerator, AsyncIterable, AsyncIterator
+
+from aiostream.stream import chain
+
+from datashare_python.utils import positional_args_only, before_and_after, once
 from temporalio import activity
 
 
@@ -15,3 +19,29 @@ def test_keyword_safe_activity() -> None:
         raise AssertionError(
             "couldn't create activity from keyword only function "
         ) from e
+
+
+async def _num_gen() -> AsyncGenerator[int, None]:
+    for i in range(10):
+        yield i // 3
+
+
+async def test_before_and_after() -> None:
+    # Given
+    async def group_by_iterator(
+        items: AsyncIterable[int],
+    ) -> AsyncIterator[AsyncIterator[int]]:
+        while True:
+            try:
+                next_item = await anext(aiter(items))
+            except StopAsyncIteration:
+                return
+            gr, items = before_and_after(items, lambda x, next_i=next_item: x == next_i)
+            yield chain(once(next_item), gr)
+
+    # When
+    grouped = []
+    async for group in group_by_iterator(_num_gen()):
+        group = [item async for item in group]  # noqa: PLW2901
+        grouped.append(group)
+    assert grouped == [[0, 0, 0], [1, 1, 1], [2, 2, 2], [3]]