refactor: worker configuration dependency

Author: ClemDoum

.github/workflows/test_asr_worker.yml

@@ -50,7 +50,7 @@ jobs:
       - name: Run tests
         run: |
           cd asr-worker
-          uv sync --frozen --extra preprocessing --extra inference --dev
+          uv sync --frozen --extra preprocessing --extra inference --extra cpu --dev
           uv run --frozen python -m pytest --timeout=180 -vvv --cache-clear --show-capture=all -r A
 
 concurrency:

asr-worker/asr_worker/activities.py

@@ -2,7 +2,7 @@
 from contextlib import AbstractContextManager, contextmanager
 from itertools import tee
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
 
 from caul.objects import ASRResult, PreprocessedInput
 from caul.tasks import (
@@ -12,6 +12,7 @@
     Postprocessor,
     Preprocessor,
 )
+from datashare_python.dependencies import lifespan_worker_config
 from datashare_python.types_ import ProgressRateHandler
 from datashare_python.utils import (
     ActivityWithProgress,
@@ -67,7 +68,7 @@ async def search_audios(
         self, project: str, query: dict[str, Any], batch_size: int
     ) -> list[Path]:
         es_client = lifespan_es_client()
-        worker_config = ASRWorkerConfig()
+        worker_config = cast(ASRWorkerConfig, lifespan_worker_config())
         batch_dir_name = activity_contextual_id()
         workdir = worker_config.workdir
         batch_root = workdir / batch_dir_name
@@ -87,7 +88,7 @@ def preprocess(
         self, paths: list[Path] | Path, config: ParakeetPreprocessorConfig
     ) -> list[Path]:
         # TODO: this shouldn't be necessary, fix this bug
-        worker_config = ASRWorkerConfig()
+        worker_config = cast(ASRWorkerConfig, lifespan_worker_config())
         audio_root = worker_config.audios_root
         workdir = worker_config.workdir
         # TODO: implement caching
@@ -120,7 +121,7 @@ def infer(
     ) -> list[Path]:
         # TODO: fix this temporal by, we shouldn't have to reload
         config = _INFERENCE_CONFIG_TYPE_ADAPTER.validate_python(config)
-        worker_config = ASRWorkerConfig()
+        worker_config = cast(ASRWorkerConfig, lifespan_worker_config())
         workdir = worker_config.workdir
         preprocessed_inputs = _LIST_OF_PATH_ADAPTER.validate_python(preprocessed_inputs)
         if progress is not None:
@@ -168,7 +169,7 @@ def postprocess(
         # TODO: this shouldn't be necessary, fix this bug
         input_paths = _LIST_OF_PATH_ADAPTER.validate_python(input_paths)
         config = ParakeetPostprocessorConfig.model_validate(config)
-        worker_config = ASRWorkerConfig()
+        worker_config = cast(ASRWorkerConfig, lifespan_worker_config())
         workdir = worker_config.workdir
         artifacts_root = worker_config.artifacts_root
         inference_results = _LIST_OF_PATH_ADAPTER.validate_python(inference_results)

asr-worker/asr_worker/dependencies.py

@@ -5,7 +5,9 @@
 
 from datashare_python.dependencies import (  # noqa: F401
     lifespan_es_client,
+    lifespan_worker_config,
     set_es_client,
+    set_worker_config,
 )
 
 
@@ -23,6 +25,7 @@ def set_multiprocessing_start_method(**_) -> Generator[None, None, None]:
 
 
 REGISTRY = {
-    "inference": [set_multiprocessing_start_method],
-    "io": [set_es_client],
+    "inference": [set_worker_config, set_multiprocessing_start_method],
+    "io": [set_worker_config, set_es_client],
+    "preprocessing": [set_worker_config],
 }

asr-worker/entrypoints/inference_gpu_worker.sh

@@ -1,5 +1,6 @@
 #!/bin/bash
 
 uv run --no-sync datashare-python worker start \
+    --dependencies inference \
     --queue worker-template.classify-gpu \
     --activities asr.transcription.infer

asr-worker/entrypoints/io_worker.sh

@@ -2,7 +2,7 @@
 set -e
 
 uv run --no-sync datashare-python worker start \
+    --dependencies io \
     --queue asr.io \
     --activities asr.transcription.search-audios \
-    --workflows asr.transcription \
-    --dependencies io
+    --workflows asr.transcription

asr-worker/entrypoints/preprocessing_worker.sh

@@ -2,5 +2,6 @@
 set -e
 
 uv run --no-sync datashare-python worker start \
+    --dependencies preprocessing \
     --queue asr.preprocessing.cpu \
     --activities asr.transcription.preprocess

asr-worker/pyproject.toml

@@ -16,21 +16,24 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-preprocessing = [
-    "torchaudio==2.11.0",
-    "torchcodec==0.11.0",
+cpu = [ # just a placeholder to create an extra marker used to switch to torch gpu build
+  "torch==2.11.0",
+  "torchaudio==2.11.0",
+  "torchcodec==0.11.0",
+]
+gpu = [ # just a placeholder to create an extra marker used to switch to torch gpu build
+  "torch==2.11.0+cu128",
+  "torchaudio==2.11.0",
+  "torchcodec==0.11.0",
 ]
 inference = [
     "caul[nemo]==0.2.16",
     "kaldialign==0.9.3",
     "ml-dtypes==0.5.4",
     "numpy==2.3.0",
     "pyarrow==20.0.0",
-    "torchaudio==2.11.0",
-    "torch==2.11.0",
-    "torchcodec==0.11.0",
 ]
-gpu = [ # just a placeholder to create an extra marker used to switch to torch gpu build
+preprocessing = [
 ]
 
 [project.entry-points."datashare.workflows"]
@@ -69,22 +72,28 @@ explicit = true
 
 [tool.uv.sources]
 torch = [
-  { index = "pytorch-gpu", marker = "extra == 'gpu'" },
-  { index = "pytorch-cpu", marker = "extra != 'gpu'"},
+  { index = "pytorch-gpu", extra = "gpu" },
+  { index = "pytorch-cpu", extra = "cpu"},
 ]
 torchaudio = [
-  { index = "pytorch-gpu", marker = "extra == 'gpu'" },
-  { index = "pytorch-cpu", marker = "extra != 'gpu'"},
+  { index = "pytorch-gpu", extra = "gpu" },
+  { index = "pytorch-cpu", extra = "cpu"},
 ]
 torchcodec = [
-  { index = "pytorch-gpu", marker = "extra == 'gpu'" },
-  { index = "pytorch-cpu", marker = "extra != 'gpu'"},
+  { index = "pytorch-gpu", extra = "gpu" },
+  { index = "pytorch-cpu", extra = "cpu"},
 ]
 datashare-python = { path = "../datashare-python", editable = true }
 
 [tool.uv]
 package = true
 override-dependencies = ["kaldialign>=0.9.3"]
+conflicts = [
+    [
+      { extra = "cpu" },
+      { extra = "gpu" },
+    ],
+]
 
 [tool.uv.pip]
 prerelease = "if-necessary"

asr-worker/tests/conftest.py

@@ -1,9 +1,8 @@
-import os
-from pathlib import Path
-
 import pytest
+from _pytest.tmpdir import TempPathFactory
 from asr_worker.config import ASRWorkerConfig
 from asr_worker.dependencies import set_multiprocessing_start_method
+from datashare_python.config import DatashareClientConfig, TemporalClientConfig
 from datashare_python.conftest import (  # noqa: F401
     TEST_PROJECT,
     doc_3,
@@ -20,26 +19,30 @@
 from datashare_python.objects import Document
 from datashare_python.types_ import ContextManagerFactory
 from icij_common.es import ESClient
-from icij_common.test_utils import reset_env  # noqa: F401
 
 
 @pytest.fixture(scope="session")
 def test_deps() -> list[ContextManagerFactory]:
     return [set_temporal_client, set_es_client, set_multiprocessing_start_method]
 
 
-@pytest.fixture
-def mocked_worker_config_in_env(reset_env, tmp_path: Path) -> ASRWorkerConfig:  # noqa: ANN001, ARG001, F811
-    audios_path = tmp_path / "audios"
-    audios_path.mkdir()
-    os.environ["DS_WORKER_AUDIOS_ROOT"] = str(audios_path)
-    artifacts_path = tmp_path / "artifacts"
-    artifacts_path.mkdir()
-    os.environ["DS_WORKER_ARTIFACTS_ROOT"] = str(artifacts_path)
+@pytest.fixture(scope="session")
+def test_worker_config(tmp_path_factory: TempPathFactory) -> ASRWorkerConfig:  # noqa: ANN001, ARG001, F811
+    tmp_path = tmp_path_factory.mktemp("test-")
+    audios_root = tmp_path / "audios"
+    audios_root.mkdir()
+    artifacts_root = tmp_path / "artifacts"
+    artifacts_root.mkdir()
     workdir = tmp_path / "workdir"
     workdir.mkdir()
-    os.environ["DS_WORKER_WORKDIR"] = str(workdir)
-    return ASRWorkerConfig()
+    return ASRWorkerConfig(
+        log_level="DEBUG",
+        datashare=DatashareClientConfig(url="http://localhost:8080"),
+        temporal=TemporalClientConfig(host="localhost:7233"),
+        audios_root=audios_root,
+        artifacts_root=artifacts_root,
+        workdir=workdir,
+    )
 
 
 @pytest.fixture(scope="session")