dwave-sona-core 0.2.7__tar.gz

dwave_sona_core-0.2.7/PKG-INFO

@@ -0,0 +1,145 @@
+Metadata-Version: 2.1
+Name: dwave-sona-core
+Version: 0.2.7
+Summary: 
+Author: dwave-dev
+Requires-Python: >=3.7,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.7
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Requires-Dist: boto3 (>=1.26.85,<2.0.0)
+Requires-Dist: confluent-kafka (>=2.1.1,<3.0.0)
+Requires-Dist: loguru (>=0.6.0,<0.7.0)
+Requires-Dist: pydantic[dotenv] (>=1.10.5,<2.0.0)
+Requires-Dist: redis (>=4.5.4,<5.0.0)
+Requires-Dist: typer (>=0.7.0,<0.8.0)
+Description-Content-Type: text/markdown
+
+# Dwave SONA Core
+
+迪威智能 SONA 服務專用核心開發套件
+
+## 已使用模組
+### Inferencers
+| 名稱    | 連結                                                                                                                       |
+| ------ | ----------------------------------------------------------------------------------------------------------------------    |
+| 轉檔工具 | [dwave-sona-inferencer-media-tools](https://github.com/DeepWaveInc/dwave-sona-inferencer-media-tools)                     |
+| 字幕產生 | [dwave-sona-inferencer-captioner](https://github.com/DeepWaveInc/dwave-sona-inferencer-captioner)                         |
+| 人聲降噪 | [dwave-sona-inferencer-denoise](https://github.com/DeepWaveInc/dwave-sona-inferencer-denoise)                             |
+| 語者分離 | [dwave-sona-inferencer-vad](https://github.com/DeepWaveInc/dwave-sona-inferencer-vad)                                     |
+| 音軌分離 | [dwave-sona-inferencer-svs](https://github.com/DeepWaveInc/dwave-sona-inferencer-svs)                                     |
+| 全曲轉譜 | [dwave-sona-inferencer-singing-transcription](https://github.com/DeepWaveInc/dwave-sona-inferencer-singing-transcription) |
+| 字幕摘要 | [dwave-sona-inferencer-summarizer](https://github.com/DeepWaveInc/dwave-sona-inferencer-summarizer)                       |
+
+### Workers
+| 名稱    | 連結                                                                                        |
+| ------ | ------------------------------------------------------------------------------------------ |
+| 控制中心 |[dwave-sona-worker-controller](https://github.com/DeepWaveInc/dwave-sona-worker-controller) |
+| 郵件寄送 |[dwave-sona-worker-mail](https://github.com/DeepWaveInc/dwave-sona-worker-mail)             |
+
+
+### Cronjobs
+| 名稱     | 連結                                                                              |
+| ------- | -------------------------------------------------------------------------------- |
+| 數發部排程 |[dwave-sona-cronjob-moda](https://github.com/DeepWaveInc/dwave-sona-cronjob-moda)|
+
+
+## 安裝與使用
+
+### 開發環境需求
+
+- Python 3.8 或更新版本
+- poetry
+
+### 安裝與使用
+
+1. 環境建構
+
+```sh
+$ pip install poetry
+```
+
+2. 下載與安裝
+
+```sh
+$ export PYTHON_KEYRING_BACKEND=keyring.backends.null.Keyring
+$ poetry add git+ssh://git@github.com/DeepWaveInc/dwave-sona-core.git
+```
+
+3. 撰寫 Inferencer 模組
+
+```python
+# example/basic.py
+from pathlib import Path
+from typing import Dict, List
+
+from loguru import logger
+from sona.core.messages import Context, File, Job, Result
+from sona.inferencers import InferencerBas
+
+
+class BasicExample(InferencerBase):
+    inferencer = "basic"  # Inferencer 名稱
+
+    def on_load(self) -> None:
+        """
+        載入函式，Worker 啟動時呼叫
+        """
+        logger.info(f"Download {self.__class__.__name__} models...")
+
+    def inference(self, params: Dict, files: List[File]) -> Result:
+        """
+        訊息處理函式，Worker 接受到新訊息後呼叫
+        :param params: 處理參數
+        :param files: 處理檔案
+        :return: 處理結果
+        """
+        logger.info(f"Get params {params}")
+        logger.info(f"Get files {files}")
+
+        filname = "output.wav"
+        Path(filname).touch(exist_ok=True)
+        return Result(
+            files=[File(label="output", path=filname)],
+            data={"data_key": "data_val"},
+        )
+
+    def context_example(self) -> Context:
+        """
+        範例訊息，供 Worker 測試及 API 開發時參考使用
+        :return: 範例訊息
+        """
+        filname = "input.wav"
+        Path(filname).touch(exist_ok=True)
+
+        params = {"param_key": "param_val"}
+        files = [File(label="input", path=filname)]
+        return Context(
+            jobs=[
+                Job(
+                    name="basic_job",
+                    topic=self.get_topic(),
+                    params=params,
+                    files=files,
+                )
+            ]
+        )
+
+```
+
+4. Worker 測試
+
+```sh
+$ poetry run python -m sona inferencer test inferencer.basic.BasicExample
+2023-03-22 02:52:27.392 | INFO     | sona.workers.inferencer:on_load:33 - Loading inferencer: basic
+2023-03-22 02:52:27.392 | INFO     | inferencer.basic:on_load:13 - Download BasicExample models...
+2023-03-22 02:52:27.392 | INFO     | sona.workers.inferencer:on_load:36 - Susbcribe on sona.worker.inferencer.basic(MockConsumer)
+2023-03-22 02:52:27.392 | INFO     | sona.workers.inferencer:on_context:42 - [sona.worker.inferencer.basic] recv: {"id": "5fe59e8bcd4b4efb84462cdbcad4a3b4", "header": {}, "jobs": [{"name": "basic_job", "topic": "sona.worker.inferencer.basic", "params": {"param_key": "param_val"}, "files": [{"label": "input", "path": "input.wav"}], "extra_params": {}, "extra_files": {}}], "fallbacks": [], "results": {}, "states": []}
+2023-03-22 02:52:27.392 | INFO     | inferencer.basic:inference:16 - Get params {'param_key': 'param_val'}
+2023-03-22 02:52:27.392 | INFO     | inferencer.basic:inference:17 - Get files [File(label='input', path='input.wav')]
+2023-03-22 02:52:27.393 | INFO     | sona.workers.inferencer:on_context:59 - [sona.worker.inferencer.basic] success: {"id": "5fe59e8bcd4b4efb84462cdbcad4a3b4", "header": {}, "jobs": [{"name": "basic_job", "topic": "sona.worker.inferencer.basic", "params": {"param_key": "param_val"}, "files": [{"label": "input", "path": "input.wav"}], "extra_params": {}, "extra_files": {}}], "fallbacks": [], "results": {"basic_job": {"files": [{"label": "output", "path": "output.wav"}], "data": {"data_key": "data_val"}}}, "states": [{"job_name": "basic_job", "exec_time": 0.00029676100000000015, "exception": {}}]}
+```
+

dwave_sona_core-0.2.7/README.md

@@ -0,0 +1,124 @@
+# Dwave SONA Core
+
+迪威智能 SONA 服務專用核心開發套件
+
+## 已使用模組
+### Inferencers
+| 名稱    | 連結                                                                                                                       |
+| ------ | ----------------------------------------------------------------------------------------------------------------------    |
+| 轉檔工具 | [dwave-sona-inferencer-media-tools](https://github.com/DeepWaveInc/dwave-sona-inferencer-media-tools)                     |
+| 字幕產生 | [dwave-sona-inferencer-captioner](https://github.com/DeepWaveInc/dwave-sona-inferencer-captioner)                         |
+| 人聲降噪 | [dwave-sona-inferencer-denoise](https://github.com/DeepWaveInc/dwave-sona-inferencer-denoise)                             |
+| 語者分離 | [dwave-sona-inferencer-vad](https://github.com/DeepWaveInc/dwave-sona-inferencer-vad)                                     |
+| 音軌分離 | [dwave-sona-inferencer-svs](https://github.com/DeepWaveInc/dwave-sona-inferencer-svs)                                     |
+| 全曲轉譜 | [dwave-sona-inferencer-singing-transcription](https://github.com/DeepWaveInc/dwave-sona-inferencer-singing-transcription) |
+| 字幕摘要 | [dwave-sona-inferencer-summarizer](https://github.com/DeepWaveInc/dwave-sona-inferencer-summarizer)                       |
+
+### Workers
+| 名稱    | 連結                                                                                        |
+| ------ | ------------------------------------------------------------------------------------------ |
+| 控制中心 |[dwave-sona-worker-controller](https://github.com/DeepWaveInc/dwave-sona-worker-controller) |
+| 郵件寄送 |[dwave-sona-worker-mail](https://github.com/DeepWaveInc/dwave-sona-worker-mail)             |
+
+
+### Cronjobs
+| 名稱     | 連結                                                                              |
+| ------- | -------------------------------------------------------------------------------- |
+| 數發部排程 |[dwave-sona-cronjob-moda](https://github.com/DeepWaveInc/dwave-sona-cronjob-moda)|
+
+
+## 安裝與使用
+
+### 開發環境需求
+
+- Python 3.8 或更新版本
+- poetry
+
+### 安裝與使用
+
+1. 環境建構
+
+```sh
+$ pip install poetry
+```
+
+2. 下載與安裝
+
+```sh
+$ export PYTHON_KEYRING_BACKEND=keyring.backends.null.Keyring
+$ poetry add git+ssh://git@github.com/DeepWaveInc/dwave-sona-core.git
+```
+
+3. 撰寫 Inferencer 模組
+
+```python
+# example/basic.py
+from pathlib import Path
+from typing import Dict, List
+
+from loguru import logger
+from sona.core.messages import Context, File, Job, Result
+from sona.inferencers import InferencerBas
+
+
+class BasicExample(InferencerBase):
+    inferencer = "basic"  # Inferencer 名稱
+
+    def on_load(self) -> None:
+        """
+        載入函式，Worker 啟動時呼叫
+        """
+        logger.info(f"Download {self.__class__.__name__} models...")
+
+    def inference(self, params: Dict, files: List[File]) -> Result:
+        """
+        訊息處理函式，Worker 接受到新訊息後呼叫
+        :param params: 處理參數
+        :param files: 處理檔案
+        :return: 處理結果
+        """
+        logger.info(f"Get params {params}")
+        logger.info(f"Get files {files}")
+
+        filname = "output.wav"
+        Path(filname).touch(exist_ok=True)
+        return Result(
+            files=[File(label="output", path=filname)],
+            data={"data_key": "data_val"},
+        )
+
+    def context_example(self) -> Context:
+        """
+        範例訊息，供 Worker 測試及 API 開發時參考使用
+        :return: 範例訊息
+        """
+        filname = "input.wav"
+        Path(filname).touch(exist_ok=True)
+
+        params = {"param_key": "param_val"}
+        files = [File(label="input", path=filname)]
+        return Context(
+            jobs=[
+                Job(
+                    name="basic_job",
+                    topic=self.get_topic(),
+                    params=params,
+                    files=files,
+                )
+            ]
+        )
+
+```
+
+4. Worker 測試
+
+```sh
+$ poetry run python -m sona inferencer test inferencer.basic.BasicExample
+2023-03-22 02:52:27.392 | INFO     | sona.workers.inferencer:on_load:33 - Loading inferencer: basic
+2023-03-22 02:52:27.392 | INFO     | inferencer.basic:on_load:13 - Download BasicExample models...
+2023-03-22 02:52:27.392 | INFO     | sona.workers.inferencer:on_load:36 - Susbcribe on sona.worker.inferencer.basic(MockConsumer)
+2023-03-22 02:52:27.392 | INFO     | sona.workers.inferencer:on_context:42 - [sona.worker.inferencer.basic] recv: {"id": "5fe59e8bcd4b4efb84462cdbcad4a3b4", "header": {}, "jobs": [{"name": "basic_job", "topic": "sona.worker.inferencer.basic", "params": {"param_key": "param_val"}, "files": [{"label": "input", "path": "input.wav"}], "extra_params": {}, "extra_files": {}}], "fallbacks": [], "results": {}, "states": []}
+2023-03-22 02:52:27.392 | INFO     | inferencer.basic:inference:16 - Get params {'param_key': 'param_val'}
+2023-03-22 02:52:27.392 | INFO     | inferencer.basic:inference:17 - Get files [File(label='input', path='input.wav')]
+2023-03-22 02:52:27.393 | INFO     | sona.workers.inferencer:on_context:59 - [sona.worker.inferencer.basic] success: {"id": "5fe59e8bcd4b4efb84462cdbcad4a3b4", "header": {}, "jobs": [{"name": "basic_job", "topic": "sona.worker.inferencer.basic", "params": {"param_key": "param_val"}, "files": [{"label": "input", "path": "input.wav"}], "extra_params": {}, "extra_files": {}}], "fallbacks": [], "results": {"basic_job": {"files": [{"label": "output", "path": "output.wav"}], "data": {"data_key": "data_val"}}}, "states": [{"job_name": "basic_job", "exec_time": 0.00029676100000000015, "exception": {}}]}
+```

dwave_sona_core-0.2.7/pyproject.toml

@@ -0,0 +1,24 @@
+[tool.poetry]
+name = "dwave-sona-core"
+version = "0.2.7"
+description = ""
+authors = ["dwave-dev"]
+license = ""
+readme = "README.md"
+packages = [{include = "sona"}]
+
+[tool.poetry.dependencies]
+python = "^3.7"
+typer = "^0.7.0"
+loguru = "^0.6.0"
+pydantic = {extras = ["dotenv"], version = "^1.10.5"}
+boto3 = "^1.26.85"
+redis = "^4.5.4"
+confluent-kafka = "^2.1.1"
+
+[tool.poetry.group.dev.dependencies]
+black = "^23.1.0"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

dwave_sona_core-0.2.7/sona/__init__.py

@@ -0,0 +1,62 @@
+import asyncio
+
+import typer
+from loguru import logger
+
+from sona.core.consumers import create_consumer
+from sona.core.messages.context import Context
+from sona.core.producers import MockProducer, create_producer
+from sona.core.storages import LocalStorage, ShareStorageBase, create_storage
+from sona.inferencers import InferencerBase
+from sona.settings import settings
+from sona.workers import *
+
+worker_app = typer.Typer()
+inference_app = typer.Typer()
+
+
+@worker_app.command()
+def run(
+    inferencer: str = settings.SONA_INFERENCER,
+    worker: str = settings.SONA_WORKER,
+):
+    try:
+        worker_cls = WorkerBase.load_class(worker)
+        if worker_cls == InferencerWorker:
+            worker: WorkerBase = worker_cls(InferencerBase.load_class(inferencer)())
+        else:
+            worker: WorkerBase = worker_cls()
+        worker.set_storage(create_storage())
+        worker.set_producer(create_producer())
+        worker.set_consumer(create_consumer())
+        asyncio.run(worker.start())
+    except Exception as e:
+        logger.exception(e)
+
+
+@worker_app.command("test")
+def worker_test(worker_class: str, file: str = None):
+    worker: WorkerBase = WorkerBase.load_class(worker_class)()
+    _test(worker, file)
+
+
+@inference_app.command("test")
+def inferencer_test(inferencer_class: str, file: str = None):
+    worker: WorkerBase = InferencerWorker(InferencerBase.load_class(inferencer_class)())
+    _test(worker, file)
+
+
+def _test(worker: WorkerBase, file: str):
+    storage: ShareStorageBase = LocalStorage()
+    if file:
+        context: Context = Context.parse_file(file)
+    else:
+        context: Context = worker.context_example()
+    worker.set_storage(storage)
+    worker.set_producer(MockProducer())
+    asyncio.run(worker.test(context))
+
+
+app = typer.Typer()
+app.add_typer(worker_app, name="worker")
+app.add_typer(inference_app, name="inferencer")

dwave_sona_core-0.2.7/sona/__main__.py

@@ -0,0 +1,4 @@
+from sona import app
+
+if __name__ == "__main__":
+    app()

dwave_sona_core-0.2.7/sona/core/consumers/__init__.py

@@ -0,0 +1,18 @@
+from sona.settings import settings
+
+from .base import ConsumerBase
+from .kafka import KafkaConsumer
+from .redis import RedisConsumer
+from .sqs import SQSConsumer
+
+__all__ = ["ConsumerBase", "KafkaConsumer", "RedisConsumer", "SQSConsumer"]
+
+
+def create_consumer():
+    if settings.SONA_CONSUMER_REDIS_URL:
+        return RedisConsumer()
+    if settings.SONA_CONSUMER_KAFKA_SETTING:
+        return KafkaConsumer()
+    raise Exception(
+        "Consumer settings not found, please set SONA_CONSUMER_KAFKA_SETTING or SONA_CONSUMER_REDIS_URL"
+    )

dwave_sona_core-0.2.7/sona/core/consumers/base.py

@@ -0,0 +1,11 @@
+import abc
+
+
+class ConsumerBase:
+    @abc.abstractmethod
+    def subscribe(self, topic: str) -> None:
+        return NotImplemented
+
+    @abc.abstractmethod
+    async def consume(self) -> str:
+        return NotImplemented

dwave_sona_core-0.2.7/sona/core/consumers/kafka.py

@@ -0,0 +1,30 @@
+import asyncio
+
+from confluent_kafka import Consumer
+from loguru import logger
+
+from sona.settings import settings
+
+from .base import ConsumerBase
+
+CONSUMER_SETTING = settings.SONA_CONSUMER_KAFKA_SETTING
+
+
+class KafkaConsumer(ConsumerBase):
+    def __init__(self, configs=CONSUMER_SETTING):
+        self.consumer = Consumer(configs)
+
+    def subscribe(self, topic):
+        self.consumer.subscribe([topic])
+
+    async def consume(self):
+        loop = asyncio.get_running_loop()
+        while True:
+            msg = await loop.run_in_executor(None, self.consumer.poll, 1)
+            if not msg:
+                continue
+            if msg.error():
+                logger.warning(f"kafka error: {msg}")
+                continue
+            self.consumer.commit()
+            yield msg.value()

dwave_sona_core-0.2.7/sona/core/consumers/redis.py

@@ -0,0 +1,36 @@
+import uuid
+from loguru import logger
+import redis.asyncio as redis
+from redis.exceptions import ResponseError
+
+from sona.settings import settings
+
+from .base import ConsumerBase
+
+REDIS_URL = settings.SONA_CONSUMER_REDIS_URL
+CONSUMER_GROUP= settings.SONA_CONSUMER_REDIS_GROUP
+
+
+class RedisConsumer(ConsumerBase):
+    def __init__(self, url=REDIS_URL):
+        self.topics = []
+        self.redis = redis.from_url(url)
+        self.client_id = str(uuid.uuid4())
+
+    def subscribe(self, topic):
+        self.topics += [topic]
+
+    async def consume(self):
+        for topic in self.topics:
+            try:
+                await self.redis.xgroup_create(topic, CONSUMER_GROUP, id='0', mkstream=True)
+            except ResponseError as e:
+                logger.warning(e)
+
+        while True:
+            stream_list = {topic: '>' for topic in self.topics}
+            streams = await self.redis.xreadgroup(CONSUMER_GROUP, self.client_id, stream_list, block=1000 * 60, noack=True)
+            for _stream_key, stream in streams:
+                for _id, message in stream:
+                    if message is not None:
+                        yield message[b"data"].decode()

dwave_sona_core-0.2.7/sona/core/consumers/sqs.py

@@ -0,0 +1,34 @@
+import asyncio
+import functools
+
+import boto3
+
+from .base import ConsumerBase
+
+
+class SQSConsumer(ConsumerBase):
+    def __init__(self):
+        self.sqs = boto3.resource("sqs")
+        self.queue = None
+
+    def subscribe(self, topic):
+        self.queue = self.sqs.get_queue_by_name(QueueName=topic)
+
+    async def consume(self):
+        loop = asyncio.get_running_loop()
+        while True:
+            messages = await loop.run_in_executor(
+                None,
+                functools.partial(
+                    self.queue.receive_messages,
+                    AttributeNames=["ApproximateReceiveCount"],
+                    MaxNumberOfMessages=1,
+                    WaitTimeSeconds=20,
+                ),
+            )
+            if len(messages) <= 0:
+                continue
+            for message in messages:
+                body = message.body
+                message.delete()
+                yield body

dwave_sona_core-0.2.7/sona/core/messages/__init__.py

@@ -0,0 +1,8 @@
+from .base import MessageBase
+from .context import Context
+from .file import File
+from .job import Job
+from .result import Result
+from .state import State
+
+__all__ = ["MessageBase", "Context", "File", "Job", "Result", "State"]

dwave_sona_core-0.2.7/sona/core/messages/base.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+
+class MessageBase(BaseModel):
+    def mutate(self, **kwargs) -> MessageBase:
+        kwargs = {**self.dict(), **kwargs}
+        return self.__class__(**kwargs)
+
+    def to_message(self) -> str:
+        return self.json()
+
+    class Config:
+        allow_mutation = False

dwave_sona_core-0.2.7/sona/core/messages/context.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import time
+import uuid
+from typing import Dict, List
+
+from pydantic import Field
+
+from .base import MessageBase
+from .job import Job
+from .result import Result
+from .state import State
+
+
+class Context(MessageBase):
+    id: uuid.UUID = Field(default_factory=uuid.uuid4)
+    start_time: int = Field(default_factory=time.time_ns)
+    reporters: List[str] = []
+    fallbacks: List[str] = []
+    headers: Dict = {}
+    jobs: List[Job]
+    results: Dict[str, Result] = {}
+    states: List[State] = []
+
+    @property
+    def duration(self):
+        return float(time.time_ns() - self.start_time) * (0.1**9)
+
+    @property
+    def current_job(self) -> Job:
+        if len(self.states) == len(self.jobs):
+            return None
+        return self.jobs[len(self.states)]
+
+    def find_job(self, job_name):
+        for job in self.jobs:
+            if job.name == job_name:
+                return job
+        return None
+
+    def next_context(self, state, result=None) -> Context:
+        current_job = self.current_job
+        context = self.mutate(states=self.states + [state])
+        if result:
+            context = context.mutate(results={**self.results, current_job.name: result})
+        return context

dwave_sona_core-0.2.7/sona/core/messages/file.py

@@ -0,0 +1,9 @@
+from typing import Dict
+
+from .base import MessageBase
+
+
+class File(MessageBase):
+    label: str
+    path: str
+    metadata: Dict = {}

dwave_sona_core-0.2.7/sona/core/messages/job.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from typing import Dict, List
+
+from loguru import logger
+
+from sona.core.messages.result import Result
+from sona.core.utils.dict_utils import find_value_from_nested_keys
+
+from .base import MessageBase
+from .file import File
+
+
+class Job(MessageBase):
+    name: str
+    topic: str
+    params: Dict = {}
+    files: List[File] = []
+    extra_params: Dict[str, str] = {}  # Will be deprecated in 1.0.0
+    extra_files: Dict[str, str] = {}  # Will be deprecated in 1.0.0
+    required_result_params: Dict[str, str] = {}
+    required_result_files: Dict[str, str] = {}
+
+    @property
+    def required_params(self):
+        if self.extra_params:
+            logger.warning("extra_params will be deprecated in version 1.0.0")
+        return {**self.extra_params, **self.required_result_params}
+
+    @property
+    def required_files(self):
+        if self.extra_files:
+            logger.warning("extra_files will be deprecated in version 1.0.0")
+        return {**self.extra_files, **self.required_result_files}
+
+    def prepare_params(self, results: Dict[str, Result]):
+        params = {**self.params}
+        for key, target in self.required_params.items():
+            targets = target.split("__")
+            job_name, keys = targets[0], targets[1:]
+            params[key] = find_value_from_nested_keys(keys, results[job_name].data)
+        return params
+
+    def prepare_files(self, results: Dict[str, Result]):
+        file_map = {file.label: file for file in self.files}
+        for label, target in self.required_files.items():
+            job, files_label = target.split("__")
+            file = results[job].find_file(files_label).mutate(label=label)
+            file_map[label] = file
+        return list(file_map.values())

dwave_sona_core-0.2.7/sona/core/messages/result.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from typing import Dict, List
+
+from .base import MessageBase
+from .file import File
+
+
+class Result(MessageBase):
+    files: List[File] = []
+    data: Dict = {}
+
+    def find_file(self, label) -> File:
+        for file in self.files:
+            if file.label == label:
+                return file
+        return None

dwave_sona_core-0.2.7/sona/core/messages/state.py

@@ -0,0 +1,33 @@
+import socket
+import time
+import traceback
+from typing import Dict
+
+from pydantic import Field
+
+from .base import MessageBase
+
+
+class State(MessageBase):
+    job_name: str
+    node_name: str = socket.gethostname()
+    timestamp: int = Field(default_factory=time.time_ns)
+    exec_time: float = 0
+    exception: Dict = {}
+
+    @property
+    def duration(self):
+        return float(time.time_ns() - self.timestamp) * (0.1**9)
+
+    @classmethod
+    def start(cls, job_name):
+        return State(job_name=job_name)
+
+    def complete(self):
+        return self.mutate(exec_time=self.duration)
+
+    def fail(self, exception):
+        return self.mutate(
+            exec_time=self.duration,
+            exception={"message": str(exception), "traceback": traceback.format_exc()},
+        )

dwave_sona_core-0.2.7/sona/core/producers/__init__.py

@@ -0,0 +1,25 @@
+from sona.settings import settings
+
+from .base import ProducerBase
+from .kafka import KafkaProducer
+from .mock import MockProducer
+from .redis import RedisProducer
+from .sqs import SQSProducer
+
+__all__ = [
+    "ProducerBase",
+    "KafkaProducer",
+    "RedisProducer",
+    "SQSProducer",
+    "MockProducer",
+]
+
+
+def create_producer():
+    if settings.SONA_PRODUCER_REDIS_URL:
+        return RedisProducer()
+    if settings.SONA_PRODUCER_KAFKA_SETTING:
+        return KafkaProducer()
+    raise Exception(
+        "Producer settings not found, please set SONA_PRODUCER_KAFKA_SETTING or SONA_PRODUCER_REDIS_URL"
+    )

dwave_sona_core-0.2.7/sona/core/producers/base.py

@@ -0,0 +1,7 @@
+import abc
+
+
+class ProducerBase:
+    @abc.abstractmethod
+    def emit(self, topic, message) -> None:
+        return NotImplemented

dwave_sona_core-0.2.7/sona/core/producers/kafka.py

@@ -0,0 +1,23 @@
+from confluent_kafka import Producer
+
+from sona.settings import settings
+
+from .base import ProducerBase
+
+PRODUCER_SETTING = settings.SONA_PRODUCER_KAFKA_SETTING
+
+
+class KafkaProducer(ProducerBase):
+    def __init__(self, configs=PRODUCER_SETTING):
+        self.producer = Producer(configs)
+
+    def emit(self, topic, message):
+        self.producer.poll(0)
+        self.producer.produce(
+            topic, message.encode("utf-8"), callback=self.__delivery_report
+        )
+        self.producer.flush()
+
+    def __delivery_report(self, err, msg):
+        if err:
+            raise Exception(msg.error())

dwave_sona_core-0.2.7/sona/core/producers/mock.py

@@ -0,0 +1,8 @@
+from loguru import logger
+
+from .base import ProducerBase
+
+
+class MockProducer(ProducerBase):
+    def emit(self, topic, message) -> None:
+        logger.info(f"emit [{topic}] {message}")

dwave_sona_core-0.2.7/sona/core/producers/redis.py

@@ -0,0 +1,15 @@
+import redis
+
+from sona.settings import settings
+
+from .base import ProducerBase
+
+REDIS_URL = settings.SONA_PRODUCER_REDIS_URL
+
+
+class RedisProducer(ProducerBase):
+    def __init__(self, url=REDIS_URL):
+        self.redis = redis.from_url(url)
+
+    def emit(self, topic, message):
+        self.redis.xadd(topic, {"data": message})

dwave_sona_core-0.2.7/sona/core/producers/sqs.py

@@ -0,0 +1,12 @@
+import boto3
+
+from .base import ProducerBase
+
+
+class SQSProducer(ProducerBase):
+    def __init__(self):
+        self.sqs = boto3.resource("sqs")
+
+    def emit(self, topic, message):
+        queue = self.sqs.get_queue_by_name(QueueName=topic)
+        queue.send_message(MessageBody=message)

dwave_sona_core-0.2.7/sona/core/storages/__init__.py

@@ -0,0 +1,13 @@
+from sona.settings import settings
+
+from .base import ShareStorageBase
+from .local import LocalStorage
+from .s3 import S3Storage
+
+__all__ = ["ShareStorageBase", "LocalStorage", "S3Storage"]
+
+
+def create_storage():
+    if settings.SONA_STORAGE_SETTING:
+        return S3Storage()
+    return LocalStorage()

dwave_sona_core-0.2.7/sona/core/storages/base.py

@@ -0,0 +1,51 @@
+from pathlib import Path
+from typing import List
+
+from sona.core.messages.file import File
+
+
+class ShareStorageBase:
+    def __init__(self):
+        self.cached_path = []
+
+    def pull_all(self, files: List[File]) -> List[File]:
+        return [self.pull(file) for file in files]
+
+    def pull(self, file: File) -> File:
+        if Path(file.path).is_file():
+            return file
+        if not self.is_valid(file.path):
+            raise Exception("Invalid share storage path")
+        path = self.on_pull(file.path)
+        self.cached_path.append(path)
+        return file.mutate(path=path)
+
+    def push_all(self, files: List[File]) -> List[File]:
+        return [self.push(file) for file in files]
+
+    def push(self, file: File) -> File:
+        if self.is_valid(file.path):
+            return file
+        if not Path(file.path).is_file():
+            raise Exception(f"Missing file: {file}")
+        self.cached_path.append(file.path)
+        return file.mutate(path=self.on_push(file.path))
+
+    def clean(self) -> None:
+        for path in self.cached_path:
+            file_path = Path(path)
+            if file_path.exists():
+                file_path.unlink()
+
+    def create_storage(self):
+        return self.__class__()
+
+    # Callbacks
+    def is_valid(self, path) -> bool:
+        return NotImplemented
+
+    def on_pull(self, path) -> str:
+        return NotImplemented
+
+    def on_push(self, path) -> str:
+        return NotImplemented

dwave_sona_core-0.2.7/sona/core/storages/local.py

@@ -0,0 +1,32 @@
+import os
+import shutil
+from pathlib import Path
+
+from sona.core.storages.base import ShareStorageBase
+from sona.settings import settings
+
+
+class LocalStorage(ShareStorageBase):
+    def __init__(self,
+                 local_root=settings.SONA_STORAGE_LOCAL_ROOT,
+                 storage_dir=settings.SONA_STORAGE_DIR):
+        super().__init__()
+        self.local_root = local_root
+        self.storage_dir = storage_dir
+
+    def is_valid(self, path: str) -> bool:
+        return path.startswith(self.local_root)
+
+    def on_pull(self, path: str) -> str:
+        tmp_path = os.path.relpath(path, self.local_root)
+        Path(tmp_path).parent.mkdir(parents=True, exist_ok=True)
+        with open(path, "rb") as f_in, open(tmp_path, "wb") as f_out:
+            shutil.copyfileobj(f_in, f_out)
+        return str(tmp_path)
+
+    def on_push(self, path: str) -> str:
+        local_path = Path(self.local_root) / self.storage_dir / Path(path).name
+        local_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(path, "rb") as f_in, open(local_path, "wb") as f_out:
+            shutil.copyfileobj(f_in, f_out)
+        return str(local_path)

dwave_sona_core-0.2.7/sona/core/storages/s3.py

@@ -0,0 +1,53 @@
+import datetime
+import hashlib
+import re
+from pathlib import Path
+
+import boto3
+from botocore.client import Config
+
+from sona.core.storages.base import ShareStorageBase
+from sona.settings import settings
+
+
+class S3Storage(ShareStorageBase):
+    def __init__(
+        self,
+        bucket=settings.SONA_STORAGE_BUCKET,
+        upload_dir=settings.SONA_STORAGE_DIR,
+        configs=settings.SONA_STORAGE_SETTING,
+    ):
+        super().__init__()
+        self.bucket = bucket
+        self.upload_dir = upload_dir
+        self.configs = configs
+
+    @property
+    def client(self):
+        configs = self.configs or {}
+        configs.update({"config": Config(signature_version="s3v4")})
+        return boto3.resource("s3", **configs).meta.client
+
+    def is_valid(self, path: str) -> bool:
+        return re.match(r"^[Ss]3://.*", path)
+
+    def on_pull(self, path: str) -> str:
+        match = re.match(r"[Ss]3://([-_A-Za-z0-9]+)/(.+)", path)
+        bucket = match.group(1)
+        obj_key = match.group(2)
+        filepath = Path(obj_key)
+        filepath.parent.mkdir(parents=True, exist_ok=True)
+        self.client.download_file(bucket, obj_key, str(filepath))
+        return str(filepath)
+
+    def on_push(self, path: str) -> str:
+        md5 = hashlib.md5()
+        with open(path, "rb") as f:
+            for chunk in iter(lambda: f.read(4096), b""):
+                md5.update(chunk)
+        obj_key = f"{md5.hexdigest()}{''.join(Path(path).suffixes)}"
+        obj_key = (
+            Path(self.upload_dir) / datetime.date.today().strftime("%Y%m%d") / obj_key
+        )
+        self.client.upload_file(path, self.bucket, str(obj_key))
+        return f"S3://{self.bucket}/{obj_key}"

dwave_sona_core-0.2.7/sona/core/utils/cls_utils.py

@@ -0,0 +1,13 @@
+import sys
+import traceback
+
+
+def import_class(import_str):
+    mod_str, _sep, class_str = import_str.rpartition(".")
+    __import__(mod_str)
+    try:
+        return getattr(sys.modules[mod_str], class_str)
+    except AttributeError:
+        raise ImportError(
+            f"Class {class_str} cannot be found ({traceback.format_exception(*sys.exc_info())})"
+        )

dwave_sona_core-0.2.7/sona/core/utils/dict_utils.py

@@ -0,0 +1,7 @@
+def find_value_from_nested_keys(keys, dict_):
+    dict_value = dict_
+    for key in keys:
+        dict_value = dict_value.get(key)
+        if not dict_value:
+            return None
+    return dict_value

dwave_sona_core-0.2.7/sona/inferencers/__init__.py

@@ -0,0 +1,3 @@
+from .base import InferencerBase
+
+__all__ = ["InferencerBase"]

dwave_sona_core-0.2.7/sona/inferencers/base.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import abc
+from typing import Dict, List
+
+from sona.core.messages import Context, File, Result
+from sona.core.utils.cls_utils import import_class
+from sona.settings import settings
+
+TOPIC_PREFIX = settings.SONA_INFERENCER_TOPIC_PREFIX
+
+
+class InferencerBase:
+    inferencer = NotImplemented
+
+    def on_load(self) -> None:
+        return
+
+    def context_example(self) -> Context:
+        return None
+
+    @classmethod
+    def get_topic(cls):
+        return f"{TOPIC_PREFIX}.{cls.inferencer}"
+
+    @classmethod
+    def load_class(cls, import_str):
+        inferencer_cls = import_class(import_str)
+        if inferencer_cls not in cls.__subclasses__():
+            raise Exception(f"Unknown inferencer class: {import_str}")
+        return inferencer_cls
+
+    @abc.abstractmethod
+    def inference(self, params: Dict, files: List[File]) -> Result:
+        return NotImplemented

dwave_sona_core-0.2.7/sona/inferencers/mock.py

@@ -0,0 +1,64 @@
+from pathlib import Path
+from typing import Dict, List
+
+from loguru import logger
+
+from sona.core.messages import Context, File, Job, Result
+from sona.inferencers import InferencerBase
+from sona.settings import settings
+
+
+class MockInferencer(InferencerBase):
+    inferencer = "mock"
+
+    def on_load(self) -> None:
+        logger.info(f"Download {self.__class__.__name__} models...")
+
+    def inference(self, params: Dict, files: List[File]) -> Result:
+        logger.info(f"Get params {params}")
+        logger.info(f"Get files {files}")
+        filepath = "output.wav"
+        Path(filepath).touch(exist_ok=True)
+        return Result(
+            files=[File(label="output", path=filepath)],
+            data={"output_key": "output_val"},
+        )
+
+    def context_example(self) -> Context:
+        storage = settings.SONA_STORAGE_LOCAL_ROOT
+        filepath1 = f"{storage}/input1.wav"
+        filepath2 = f"{storage}/input2.wav"
+        Path(storage).mkdir(exist_ok=True, parents=True)
+        Path(filepath1).touch(exist_ok=True)
+        Path(filepath2).touch(exist_ok=True)
+
+        return Context(
+            jobs=[
+                Job(
+                    name="input1",
+                    topic=self.get_topic(),
+                    params={"input1_key": "input1_val"},
+                    files=[
+                        File(
+                            label="input1",
+                            path=filepath1,
+                        )
+                    ],
+                    required_result_params={
+                        "input2_key": "input2__input2_key1__input2_key2"
+                    },
+                    required_result_files={"input2": "input2__input2_r"},
+                )
+            ],
+            results={
+                "input2": Result(
+                    files=[
+                        File(
+                            label="input2_r",
+                            path=filepath2,
+                        )
+                    ],
+                    data={"input2_key1": {"input2_key2": "input2_val"}},
+                )
+            },
+        )

dwave_sona_core-0.2.7/sona/settings.py

@@ -0,0 +1,36 @@
+from typing import Any, Dict, Optional
+
+from pydantic import BaseSettings, RedisDsn, root_validator
+
+
+class Settings(BaseSettings):
+    # Consumer settings
+    SONA_CONSUMER_KAFKA_SETTING: Optional[Dict] = None
+    SONA_CONSUMER_REDIS_URL: Optional[RedisDsn] = None
+    SONA_CONSUMER_REDIS_GROUP: Optional[str] = "dwave.anonymous"
+
+    # Producer settings
+    SONA_PRODUCER_KAFKA_SETTING: Optional[Dict] = None
+    SONA_PRODUCER_REDIS_URL: Optional[RedisDsn] = None
+
+    # Storage settings
+    SONA_STORAGE_DIR: str = "_tmp"
+    SONA_STORAGE_SETTING: Dict = None
+    SONA_STORAGE_BUCKET: str = "sona"
+    SONA_STORAGE_LOCAL_ROOT: str = "_share"
+
+    # Inferencer settings
+    SONA_WORKER: str = "sona.workers.InferencerWorker"
+    SONA_INFERENCER: str = None
+    SONA_INFERENCER_TOPIC_PREFIX: str = "dwave.inferencer"
+
+    @root_validator(pre=False)
+    def load_settings(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+        return values
+
+    class Config:
+        env_file = ".env"
+        env_file_encoding = "utf-8"
+
+
+settings = Settings()

dwave_sona_core-0.2.7/sona/workers/__init__.py

@@ -0,0 +1,4 @@
+from .base import WorkerBase
+from .inferencer import InferencerWorker
+
+__all__ = ["WorkerBase", "InferencerWorker"]

dwave_sona_core-0.2.7/sona/workers/base.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import abc
+
+from loguru import logger
+
+from sona.core.consumers import ConsumerBase
+from sona.core.messages import Context
+from sona.core.producers import ProducerBase
+from sona.core.storages.base import ShareStorageBase
+from sona.core.utils.cls_utils import import_class
+
+
+class WorkerBase:
+    topic: str = "dummy"
+
+    def set_consumer(self, consumer: ConsumerBase):
+        self.consumer = consumer
+
+    def set_producer(self, producer: ProducerBase):
+        self.producer = producer
+
+    def set_storage(self, storage: ShareStorageBase):
+        self.storage = storage
+
+    async def start(self):
+        await self.on_load()
+        self.topic = self.get_topic()
+        logger.info(f"Susbcribe on {self.topic}({self.consumer.__class__.__name__})")
+        self.consumer.subscribe(self.topic)
+        async for message in self.consumer.consume():
+            try:
+                self.storage = self.storage.create_storage()
+                context = Context.parse_raw(message)
+                await self.on_context(context)
+            except Exception as e:
+                logger.warning(f"[{self.topic}] error: {e}, msg: {message}")
+            finally:
+                self.storage.clean()
+
+    async def test(self, context: Context):
+        await self.on_load()
+        await self.on_context(context)
+
+    @classmethod
+    def get_topic(cls) -> str:
+        return cls.topic
+
+    @classmethod
+    def load_class(cls, import_str):
+        worker_cls = import_class(import_str)
+        if worker_cls not in cls.__subclasses__():
+            raise Exception(f"Unknown worker class: {import_str}")
+        return worker_cls
+
+    # Callbacks
+    @abc.abstractmethod
+    async def on_load(self) -> None:
+        return NotImplemented
+
+    @abc.abstractmethod
+    async def on_context(self, message: Context) -> Context:
+        return NotImplemented
+
+    def context_example(self) -> Context:
+        return None

dwave_sona_core-0.2.7/sona/workers/inferencer.py

@@ -0,0 +1,73 @@
+import asyncio
+
+from loguru import logger
+
+from sona.core.messages import Context, Job, State
+from sona.inferencers import InferencerBase
+from sona.settings import settings
+
+from .base import WorkerBase
+
+TOPIC_PREFIX = settings.SONA_INFERENCER_TOPIC_PREFIX
+INFERENCER_CLASS = settings.SONA_INFERENCER
+
+
+class InferencerWorker(WorkerBase):
+    def __init__(self, inferencer: InferencerBase):
+        super().__init__()
+        self.inferencer = inferencer
+
+    async def on_load(self):
+        logger.info(f"Loading inferencer: {self.inferencer.inferencer}")
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(None, self.inferencer.on_load)
+
+    async def on_context(self, context: Context):
+        try:
+            logger.info(f"[{self.topic}] recv: {context.to_message()}")
+
+            # Prepare process data
+            current_job: Job = context.current_job
+            current_state: State = State.start(current_job.name)
+            params = current_job.prepare_params(context.results)
+            files = current_job.prepare_files(context.results)
+            files = self.storage.pull_all(files)
+
+            # Process
+            # TODO: Make process cancelable
+            loop = asyncio.get_running_loop()
+            result = await loop.run_in_executor(
+                None, self.inferencer.inference, params, files
+            )
+            result = result.mutate(files=self.storage.push_all(result.files))
+
+            # Create success context
+            current_state = current_state.complete()
+            next_context = context.next_context(current_state, result)
+            logger.info(f"[{self.topic}] success: {next_context.to_message()}")
+
+            # Emit message
+            next_job = next_context.current_job
+            if next_job:
+                self.producer.emit(next_job.topic, next_context.to_message())
+            else:
+                for topic in next_context.reporters:
+                    self.producer.emit(topic, next_context.to_message())
+            return next_context
+
+        except Exception as e:
+            # Create fail context
+            current_state = current_state.fail(e)
+            next_context = context.next_context(current_state)
+            logger.exception(f"[{self.topic}] error: {next_context.to_message()}")
+
+            # Emit message
+            for topic in next_context.fallbacks:
+                self.producer.emit(topic, next_context.to_message())
+            return next_context
+
+    def context_example(self):
+        return self.inferencer.context_example()
+
+    def get_topic(self):
+        return self.inferencer.get_topic()