ei-pipe-sdk 0.1.5__tar.gz

ei_pipe_sdk-0.1.5/PKG-INFO

@@ -0,0 +1,97 @@
+Metadata-Version: 2.4
+Name: ei-pipe-sdk
+Version: 0.1.5
+Summary: Minimal inference SDK for algorithm engineers on the ei-pipeline ei_infer node
+Requires-Python: <3.15,>=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: arq>=0.26
+Requires-Dist: redis>=8.0.1
+
+# ei-pipe-sdk
+
+面向算法同学的极简模型接入 SDK。把你的模型写成一个 `EiInfer` 的两个方法，ei-pipeline
+就能调度它推理：**你不用关心数据怎么进来、结果怎么提交，也不用关心部署**。
+
+## 安装
+
+```bash
+pip install ei-pipe-sdk
+# 或先在仓库 ei_pipe_sdk/ 下本地构建
+uv build && pip install dist/ei_pipe_sdk-*.whl
+```
+
+## 三步接入
+
+```python
+# my_model.py
+from ei_infer import EiInfer, register
+
+
+class MyModel(EiInfer):
+    def init(self, model_cfg: dict | None = None) -> None:
+        # 进程启动时调用一次：加载权重、tokenizer、选设备。失败会让 worker 启动即崩溃。
+        self.model = load_weights(...)
+
+    def inference(self, input_path: str, output_path: str, node_info: dict) -> dict:
+        # input_path / output_path: worker 已恢复好的绝对路径
+        #   = {EI_INFER_DATA_ROOT}/{time_bucket}/{pipeline_id}/<相对路径>
+        # input_path 是输入（文件或目录）；output_path 是允许写结果文件的目录（可能为空）。
+        # 是否落盘由模型自己决定；返回的 dict 会被 worker 回传给 pipeline。
+        return {"detections": self.model(open(input_path))}
+
+
+register("my-model", MyModel())   # 注册名 = 管道节点 params.name 的值
+```
+
+```bash
+pip install ei-pipe-sdk
+# 在算法部署里跑（pod 内需挂载 PFS data_root 到 EI_INFER_DATA_ROOT）
+EI_INFER_DATA_ROOT=/lpai/pvc/ei-autolabel-bd-ga-infer python -m ei_infer.worker
+```
+
+## 协议
+
+调度端（ei-pipeline 的 `ei_infer` 节点）会按统一协议给你发任务：
+
+| 字段 | 类型 | 说明 |
+|---|---|---|
+| `name` | str | 模型注册名，由你 `register("name", ...)` 时命名 |
+| `input_path` | str | 输入路径（相对 pipe 共享目录，如 `data/`），worker 恢复为绝对路径 |
+| `output_path` | str? | 允许写结果文件的目录（相对 pipe 共享目录，可选） |
+| `version` | str? | 版本提示，透传保留 |
+
+worker 用 `node_info.time_bucket + pipeline_id + EI_INFER_DATA_ROOT` 把 `input_path` /
+`output_path` 恢复成绝对路径，调用 `inference(input_path, output_path, node_info)`。
+**是否把结果落盘由模型自己决定**，worker 不替模型落盘，只把模型返回的 dict 原样回传。
+整个链路里**你的函数只看绝对路径 input_path / output_path 和 node_info**，其余（数据就位、
+并发、轮询）全部由 SDK 兜底。
+
+## 配置（环境变量）
+
+| 变量 | 默认 | 说明 |
+|---|---|---|
+| `EI_INFER_REDIS_URL` | 回退 `EI_ARQ_REDIS_URL` → `EI_REDIS_URL` → `redis://localhost:6379/0` | 消费队列的 Redis |
+| `EI_INFER_QUEUE_NAME` | `arq:{env}:ei:model-infer` | 消费的任务队列，默认与 ei-pipeline `ei_infer` 节点指向一致 |
+| `EI_INFER_MAX_JOBS` | `4` | 每 worker 并发任务数（arq `max_jobs`） |
+| `EI_INFER_MAX_THREADS` | `max(4, max_jobs)` | 跑同步推理的线程池大小 |
+| `EI_INFER_DATA_ROOT` | `/lpai/pvc/ei-autolabel-bd-ga-infer` | PFS data_root 在 pod 内的挂载根 |
+
+## 多模型
+
+多个模型可在同一个 worker 进程里 `register` 多个名字；每个模型各自实现自己的一套动作。
+一个队列可由多个模型 worker 消费，任务按 `name` 路由到注册了同名模型的进程。
+
+## 失败模式
+
+- `inference` 返回非 dict → 任务失败（错误信息含模型名）。
+- `name` 未注册 → 任务失败，节点错误信息会指出未知模型名。
+- `init` 抛异常 → worker 启动即退出，方便在你自己的日志里早发现。
+
+## 本地联调
+
+```bash
+# 起一个本地 redis 后（redis-server）
+python -m ei_infer.worker           # 默认队列 arq:local:ei:model-infer
+```
+
+然后在 ei-pipeline 用 `ei_infer` 节点提交一条任务即可走通。

ei_pipe_sdk-0.1.5/README.md

@@ -0,0 +1,88 @@
+# ei-pipe-sdk
+
+面向算法同学的极简模型接入 SDK。把你的模型写成一个 `EiInfer` 的两个方法，ei-pipeline
+就能调度它推理：**你不用关心数据怎么进来、结果怎么提交，也不用关心部署**。
+
+## 安装
+
+```bash
+pip install ei-pipe-sdk
+# 或先在仓库 ei_pipe_sdk/ 下本地构建
+uv build && pip install dist/ei_pipe_sdk-*.whl
+```
+
+## 三步接入
+
+```python
+# my_model.py
+from ei_infer import EiInfer, register
+
+
+class MyModel(EiInfer):
+    def init(self, model_cfg: dict | None = None) -> None:
+        # 进程启动时调用一次：加载权重、tokenizer、选设备。失败会让 worker 启动即崩溃。
+        self.model = load_weights(...)
+
+    def inference(self, input_path: str, output_path: str, node_info: dict) -> dict:
+        # input_path / output_path: worker 已恢复好的绝对路径
+        #   = {EI_INFER_DATA_ROOT}/{time_bucket}/{pipeline_id}/<相对路径>
+        # input_path 是输入（文件或目录）；output_path 是允许写结果文件的目录（可能为空）。
+        # 是否落盘由模型自己决定；返回的 dict 会被 worker 回传给 pipeline。
+        return {"detections": self.model(open(input_path))}
+
+
+register("my-model", MyModel())   # 注册名 = 管道节点 params.name 的值
+```
+
+```bash
+pip install ei-pipe-sdk
+# 在算法部署里跑（pod 内需挂载 PFS data_root 到 EI_INFER_DATA_ROOT）
+EI_INFER_DATA_ROOT=/lpai/pvc/ei-autolabel-bd-ga-infer python -m ei_infer.worker
+```
+
+## 协议
+
+调度端（ei-pipeline 的 `ei_infer` 节点）会按统一协议给你发任务：
+
+| 字段 | 类型 | 说明 |
+|---|---|---|
+| `name` | str | 模型注册名，由你 `register("name", ...)` 时命名 |
+| `input_path` | str | 输入路径（相对 pipe 共享目录，如 `data/`），worker 恢复为绝对路径 |
+| `output_path` | str? | 允许写结果文件的目录（相对 pipe 共享目录，可选） |
+| `version` | str? | 版本提示，透传保留 |
+
+worker 用 `node_info.time_bucket + pipeline_id + EI_INFER_DATA_ROOT` 把 `input_path` /
+`output_path` 恢复成绝对路径，调用 `inference(input_path, output_path, node_info)`。
+**是否把结果落盘由模型自己决定**，worker 不替模型落盘，只把模型返回的 dict 原样回传。
+整个链路里**你的函数只看绝对路径 input_path / output_path 和 node_info**，其余（数据就位、
+并发、轮询）全部由 SDK 兜底。
+
+## 配置（环境变量）
+
+| 变量 | 默认 | 说明 |
+|---|---|---|
+| `EI_INFER_REDIS_URL` | 回退 `EI_ARQ_REDIS_URL` → `EI_REDIS_URL` → `redis://localhost:6379/0` | 消费队列的 Redis |
+| `EI_INFER_QUEUE_NAME` | `arq:{env}:ei:model-infer` | 消费的任务队列，默认与 ei-pipeline `ei_infer` 节点指向一致 |
+| `EI_INFER_MAX_JOBS` | `4` | 每 worker 并发任务数（arq `max_jobs`） |
+| `EI_INFER_MAX_THREADS` | `max(4, max_jobs)` | 跑同步推理的线程池大小 |
+| `EI_INFER_DATA_ROOT` | `/lpai/pvc/ei-autolabel-bd-ga-infer` | PFS data_root 在 pod 内的挂载根 |
+
+## 多模型
+
+多个模型可在同一个 worker 进程里 `register` 多个名字；每个模型各自实现自己的一套动作。
+一个队列可由多个模型 worker 消费，任务按 `name` 路由到注册了同名模型的进程。
+
+## 失败模式
+
+- `inference` 返回非 dict → 任务失败（错误信息含模型名）。
+- `name` 未注册 → 任务失败，节点错误信息会指出未知模型名。
+- `init` 抛异常 → worker 启动即退出，方便在你自己的日志里早发现。
+
+## 本地联调
+
+```bash
+# 起一个本地 redis 后（redis-server）
+python -m ei_infer.worker           # 默认队列 arq:local:ei:model-infer
+```
+
+然后在 ei-pipeline 用 `ei_infer` 节点提交一条任务即可走通。

ei_pipe_sdk-0.1.5/pyproject.toml

@@ -0,0 +1,48 @@
+[project]
+name = "ei-pipe-sdk"
+version = "0.1.5"
+description = "Minimal inference SDK for algorithm engineers on the ei-pipeline ei_infer node"
+readme = "README.md"
+requires-python = ">=3.11,<3.15"
+dependencies = [
+    "arq>=0.26",
+    # arq caps redis<6, but it is runtime-compatible with redis-py 8; aligned with
+    # the ei-pipeline override so both projects share one Redis client version.
+    "redis>=8.0.1",
+]
+
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/ei_infer", "src/ei_path"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+addopts = "--strict-markers --strict-config"
+asyncio_mode = "auto"
+pythonpath = [".", "src"]
+
+[dependency-groups]
+dev = [
+    "pytest>=8.4.1",
+    "pytest-asyncio>=1.1.0",
+    "fakeredis[lua]>=2.37.0",
+    "pytest-xdist>=3.8.0",
+]
+
+[[tool.uv.index]]
+name = "liauto-pypi-l5"
+url = "https://artifactory.ep.chehejia.com/artifactory/api/pypi/liauto-pypi-l5/simple"
+default = true
+
+[tool.uv]
+# Use only the system CPython; never download a managed (uv-managed) interpreter.
+# The pipeline pod/ci image provides python 3.11+ itself, so this avoids hitting
+# the network-to-fetch-CPython when the image lacks a matching interpreter.
+python-preference = "only-system"
+# arq caps redis<6, but it is runtime-compatible with redis-py 8; override so
+# arq can coexist with redis>=8.0.1 (mirrors the ei-pipeline override).
+override-dependencies = ["redis>=8.0.1"]

ei_pipe_sdk-0.1.5/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

ei_pipe_sdk-0.1.5/src/ei_infer/__init__.py

@@ -0,0 +1,16 @@
+"""ei-infer: minimal inference SDK for algorithm models on ei-pipeline.
+
+Public surface for an algorithm engineer:
+
+- ``EiInfer``      base class to implement (``init`` + ``inference``)
+- ``EI_REGISTER``  decorator to put a model into the in-process registry
+
+Everything else (registry internals, task runner, worker entry) lives in
+submodules and is not re-exported here on purpose.
+"""
+
+from .model import EiInfer
+from .registry import EI_REGISTER
+from .worker import run_loop
+
+__all__ = ["EiInfer", "EI_REGISTER", "run_loop"]

ei_pipe_sdk-0.1.5/src/ei_infer/env.py

@@ -0,0 +1,80 @@
+"""SDK environment helpers: EI_INFER_* with fallback to the shared Redis vars."""
+
+import os
+from urllib.parse import unquote, urlparse
+
+from arq.connections import RedisSettings
+
+DEFAULT_DATA_ROOT = "/lpai/pvc/ei-autolabel-bd-ga-infer"
+
+
+def resolve_env() -> str:
+    """Resolve active environment, mirroring ei-pipeline's envget_core_env."""
+    return os.getenv("EI_ENV") or os.getenv("LI_ENV") or "local"
+
+
+def infer_redis_url() -> str:
+    """Redis the worker consumes from.
+
+    EI_INFER_REDIS_URL overrides; otherwise fall back to EI_ARQ_REDIS_URL, then
+    EI_REDIS_URL, then a local dev default. Allows a zero-config local run and
+    reuse of the ei-pipeline cluster in production.
+    """
+    return (
+        os.getenv("EI_INFER_REDIS_URL")
+        or os.getenv("EI_ARQ_REDIS_URL")
+        or os.getenv("EI_REDIS_URL")
+        or "redis://localhost:6379/0"
+    )
+
+
+def redis_settings_from_url(url: str) -> RedisSettings:
+    """Parse a Redis URL into RedisSettings, decoding percent-encoded auth.
+
+    arq's ``RedisSettings.from_dsn`` forwards ``urlparse``'s raw username and
+    password through to the AUTH command, so a URL like
+    ``redis://user:p%40ss@host`` would authenticate with the literal bytes
+    ``p%40ss`` instead of ``p@ss``. Decode both fields with ``unquote`` (not
+    ``unquote_plus``) so a literal ``+`` in a password is preserved.
+    """
+    settings = RedisSettings.from_dsn(url)
+    parsed = urlparse(url)
+    if parsed.username is not None:
+        settings.username = unquote(parsed.username)
+    if parsed.password is not None:
+        settings.password = unquote(parsed.password)
+    return settings
+
+
+def infer_queue_name() -> str:
+    """Task queue this worker consumes.
+
+    Default matches the ei-pipeline ei_infer node's build_default_queue_name()
+    (env segment hashed when > 32 chars), so a zero-config pipe and worker land
+    on the same queue.
+    """
+    return os.getenv("EI_INFER_QUEUE_NAME") or f"arq:{resolve_env()}:ei:model-infer"
+
+
+def infer_max_jobs() -> int:
+    """Concurrent jobs per worker (arq max_jobs). Default: 4. Must be positive."""
+    value = int(os.getenv("EI_INFER_MAX_JOBS") or "4")
+    if value <= 0:
+        raise ValueError("EI_INFER_MAX_JOBS must be positive")
+    return value
+
+
+def infer_max_threads() -> int:
+    """Thread pool size for sync inference, at least one per concurrent job."""
+    return max(4, infer_max_jobs())
+
+
+def infer_data_root() -> str:
+    """PFS data_root mount point inside the pod.
+
+    The pipeline sends an input_path relative to this mount, and the worker
+    forwards it unchanged to ``model.inference(input_path, node_info)``; the
+    algorithm joins this root to read the input file. Default matches the
+    ei-pipeline PFS data root.
+    """
+    return os.getenv("EI_INFER_DATA_ROOT") or DEFAULT_DATA_ROOT

ei_pipe_sdk-0.1.5/src/ei_infer/model.py

@@ -0,0 +1,46 @@
+"""Algorithm model abstraction: implement exactly two methods to join the pipeline."""
+
+from abc import ABC, abstractmethod
+
+
+class EiInfer(ABC):
+    """What an algorithm engineer implements to plug a model into ei-pipeline.
+
+    Convention
+    ----------
+    - ``init`` runs once per worker process (model warm-up). Fail fast here:
+      a model that cannot initialize must crash the worker on boot, not the
+      first job. The current worker calls ``init()`` with no arguments;
+      ``model_cfg`` is reserved for a future config loader that would pass a
+      parsed dict in.
+    - ``inference`` receives the absolute input path, an optional absolute
+      output directory, and the caller's ``node_info``; it **returns a dict**.
+      The worker returns that dict to the pipeline and does NOT persist it -
+      the model decides whether to write result files into ``output_path``.
+    """
+
+    @abstractmethod
+    def init(self, model_cfg: dict | None = None) -> None:
+        """Load model artifacts (weights, tokenizer, device).
+
+        Called once per worker process before the first inference. ``model_cfg``
+        is reserved for a future caller that would pass parsed config (e.g. a
+        yaml); today the worker calls ``init()`` with no arguments, so
+        implementations that need per-model config should read it themselves
+        (env var, local file, etc.). Raise to abort worker startup.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def inference(self, input_path: str, output_path: str, node_info: dict) -> dict:
+        """Run one inference and return the result dict.
+
+        ``input_path`` is the absolute path of the input the pipeline staged
+        (a file or a directory of material). ``output_path`` is the absolute
+        directory (may be empty) where the model may persist result files when
+        it chooses to; whether to write is the model's decision. ``node_info``
+        carries caller context (pipeline_id, node_name, params, ...), may be an
+        empty dict but is always provided. The returned dict must be
+        JSON-serializable; it becomes the pipeline node's result.
+        """
+        raise NotImplementedError

ei_pipe_sdk-0.1.5/src/ei_infer/registry.py

@@ -0,0 +1,75 @@
+"""Process-local, thread-safe model registry: name -> EiInfer instance."""
+
+import threading
+from collections.abc import Callable
+from typing import TypeVar
+
+from .model import EiInfer
+
+T = TypeVar("T", bound=EiInfer)
+
+_lock = threading.RLock()
+_registry: dict[str, EiInfer] = {}
+
+
+def _register_instance(name: str, model: EiInfer) -> None:
+    """Register ``model`` under ``name`` (internal, used by the decorator).
+
+    Re-registering the *same* instance under the same name is idempotent;
+    registering a *different* instance under an existing name raises, so a
+    running worker can never silently swap a model mid-flight.
+    """
+    if not isinstance(name, str) or not name.strip():
+        raise ValueError("model name must be a non-empty string")
+    if not isinstance(model, EiInfer):
+        raise TypeError("model must be an EiInfer instance")
+    with _lock:
+        existing = _registry.get(name)
+        if existing is not None and existing is not model:
+            raise ValueError(f"model '{name}' already registered")
+        _registry[name] = model
+
+
+def EI_REGISTER(model_name: str) -> Callable[[type[T]], type[T]]:
+    """Register an EiInfer subclass as a decorator.
+
+    Usage::
+
+        @EI_REGISTER("example-model")
+        class ExampleModel(EiInfer):
+            ...
+
+    The model is instantiated with the model name as its sole constructor
+    argument (``Model(model_name)``) and put into the registry at class
+    definition time; the worker pre-loads it via ``init()`` on startup, so
+    load model artifacts in ``init()`` rather than ``__init__``. Registering a
+    different class under an existing name raises.
+    """
+    if not isinstance(model_name, str) or not model_name.strip():
+        raise ValueError("model name must be a non-empty string")
+
+    def decorator(model_cls: type[T]) -> type[T]:
+        _register_instance(model_name, model_cls(model_name))  # type: ignore
+        return model_cls
+
+    return decorator
+
+
+def get_model(name: str) -> EiInfer:
+    """Return the registered model, raising KeyError when unknown."""
+    with _lock:
+        if name not in _registry:
+            raise KeyError(f"model '{name}' is not registered")
+        return _registry[name]
+
+
+def get_registered_names() -> list[str]:
+    """Return sorted registered model names (for startup pre-loading)."""
+    with _lock:
+        return sorted(_registry)
+
+
+def clear_registry() -> None:
+    """Drop all registrations. Test / restart hook only."""
+    with _lock:
+        _registry.clear()

ei_pipe_sdk-0.1.5/src/ei_infer/task.py

@@ -0,0 +1,124 @@
+"""arq task that turns a pipeline job into a call into the model registry."""
+
+import asyncio
+from concurrent.futures.thread import ThreadPoolExecutor
+
+from ei_path.paths import get_shared_dirpath, normalize_relative_path
+
+from .env import infer_data_root
+from .registry import get_model
+
+
+class ModelTaskError(Exception):
+    """A deterministic job failure (bad payload / unregistered model / bad result)."""
+
+
+def resolve_absolute_paths(
+    data: dict,
+    node_info: dict | None,
+    data_root: str,
+) -> tuple[str, str]:
+    """Resolve input/output to absolute paths on the shared PFS root.
+
+    The pipeline sends paths relative to the pipe shared dir; this joins them
+    onto {EI_INFER_DATA_ROOT}/{time_bucket}/{pipeline_id} with traversal
+    checks (mirrors the pipeline's pfs_storage.paths helpers). Returns
+    ``(absolute_input, absolute_output)``; output is "" when data has no
+    output_path.
+    """
+    raw_input = str(data.get("input_path") or "").strip()
+    if not raw_input:
+        raise ModelTaskError("job data requires 'input_path'")
+    normalized_input = normalize_relative_path(raw_input)
+
+    raw_output = data.get("output_path")
+    if raw_output is None:
+        normalized_output = ""
+    else:
+        stripped = str(raw_output).strip()
+        if not stripped:
+            normalized_output = ""
+        else:
+            normalized_output = normalize_relative_path(stripped)
+
+    base = get_shared_dirpath(data_root, node_info)  # {root}/{time_bucket}/{pipe_id}
+    absolute_input = str(base / normalized_input)
+    absolute_output = (
+        str(base / normalized_output).rstrip("/") if normalized_output else ""
+    )
+    return absolute_input, absolute_output
+
+
+class _WorkerExecutor:
+    """Process-wide, lazily created thread pool for running sync inference.
+
+    arq runs each job as an async task; model.inference() is synchronous and
+    must not block the event loop, so it is executed in this pool. Created once
+    and shut down at the very end of the process.
+    """
+
+    _pool: ThreadPoolExecutor | None = None
+
+    @classmethod
+    def get(cls) -> ThreadPoolExecutor:
+        if cls._pool is None:
+            from .env import infer_max_threads
+
+            cls._pool = ThreadPoolExecutor(
+                max_workers=infer_max_threads(), thread_name_prefix="model-infer"
+            )
+        return cls._pool
+
+    @classmethod
+    def shutdown(cls) -> None:
+        if cls._pool is not None:
+            cls._pool.shutdown(wait=True)
+            cls._pool = None
+
+
+async def run_inference(ctx: dict, data: dict, node_info: dict | None = None) -> dict:
+    """arq task served by ``ei_infer.worker``.
+
+    Protocol (shared with the ei-pipeline ``ei_infer`` node):
+      data = {"name": str, "input_path": str, "output_path": str, ...}
+      - ``name``       model registered in the in-process registry
+      - ``input_path`` PFS path relative to the pipe shared dir
+                       {time_bucket}/{pipe_id}/... (may be a file or a dir)
+      - ``output_path``(optional) PFS dir the model may write result files to
+      - ``version``    optional, passed through untouched (reserved)
+
+    The worker resolves absolute paths from node_info + EI_INFER_DATA_ROOT:
+      absolute = {EI_INFER_DATA_ROOT}/{time_bucket}/{pipeline_id}/<path>
+    then calls ``model.inference(absolute_input, absolute_output, node_info)``.
+    The result dict is returned to the pipeline; the worker does NOT persist it
+    - writing result files is up to the model (into output_path if given).
+
+    Returns ``{"status": "success", "result": <inference dict>, "name": ...}``.
+    """
+    name = data.get("name") or ""
+    if not name:
+        raise ModelTaskError("job data requires 'name'")
+
+    model = get_model(name)  # KeyError when unknown -> job fails cleanly
+
+    try:
+        absolute_input, absolute_output = resolve_absolute_paths(
+            data, node_info, infer_data_root()
+        )
+    except ValueError as error:
+        raise ModelTaskError(f"invalid data path (name={name}): {error}") from None
+
+    loop = asyncio.get_running_loop()
+    result = await loop.run_in_executor(
+        _WorkerExecutor.get(),
+        model.inference,
+        absolute_input,
+        absolute_output,
+        node_info or {},
+    )
+    if not isinstance(result, dict):
+        raise ModelTaskError(
+            f"model '{name}' inference returned {type(result).__name__}, expected dict"
+        )
+
+    return {"status": "success", "result": result, "name": name}

ei_pipe_sdk-0.1.5/src/ei_infer/worker.py

@@ -0,0 +1,78 @@
+"""The ei-infer arq worker entry: ``python -m ei_infer.worker``.
+
+The worker consumes the ``arq:{env}:ei:model-infer`` queue (override with
+EI_INFER_QUEUE_NAME), pre-loads every model registered in the process on
+startup, and dispatches each job to the matching model via a thread pool.
+
+All tuning knobs come from the environment (EI_INFER_*); there are no CLI
+arguments on purpose so containers only need env vars.
+"""
+
+from arq import cron  # noqa: F401  (documentation convenience)
+
+from .env import (
+    infer_max_jobs,
+    infer_queue_name,
+    infer_redis_url,
+    redis_settings_from_url,
+)
+from .registry import get_registered_names
+from .task import ModelTaskError, _WorkerExecutor, run_inference  # noqa: F401
+
+# Re-export so a project entry script can do `from ei_infer import worker`.
+__all__ = ["build_worker_settings", "run_loop"]
+
+
+async def startup(ctx: dict) -> None:
+    """Pre-load every registered model. Any init() failure crashes the worker
+    on boot instead of surfacing on the first job."""
+    from .registry import get_model
+
+    for name in get_registered_names():
+        model = get_model(name)
+        model.init()
+
+
+async def shutdown(ctx: dict) -> None:
+    _WorkerExecutor.shutdown()
+
+
+def build_worker_settings() -> type:
+    """Build WorkerSettings at call time so env (EI_INFER_*) is read fresh.
+
+    The worker config must track env set at boot, not import-time snapshots.
+    """
+
+    class _WorkerSettings:
+        functions = [run_inference]
+        queue_name = infer_queue_name()
+        redis_settings = redis_settings_from_url(infer_redis_url())
+        on_startup = startup
+        on_shutdown = shutdown
+        max_jobs = infer_max_jobs()
+
+    return _WorkerSettings
+
+
+def run_loop() -> None:
+    """Run the worker with config read from EI_INFER_* env vars.
+
+    Entry for ``python -m ei_infer.worker``; blocks until interrupted.
+    """
+    settings = build_worker_settings()
+
+    from arq.worker import Worker
+
+    worker = Worker(
+        functions=settings.functions,
+        queue_name=settings.queue_name,
+        redis_settings=settings.redis_settings,
+        on_startup=settings.on_startup,
+        on_shutdown=settings.on_shutdown,
+        max_jobs=settings.max_jobs,
+        handle_signals=False,  # let the container runtime own signal handling
+    )
+    try:
+        worker.run()
+    finally:
+        _WorkerExecutor.shutdown()

ei_pipe_sdk-0.1.5/src/ei_path/__init__.py

@@ -0,0 +1 @@
+"""Shared-dir path helpers for the algorithm worker (available as ``ei_path``)."""