Metadata-Version: 2.4
Name: factory-queue
Version: 0.2.3
Summary: 流水线工厂模块，支持多节点、多工位、资源控制、磁盘溢出
Home-page: https://github.com/yourusername/factory-queue
Author: wyf
Author-email: Your Name <your.email@example.com>
License: MIT
Project-URL: Homepage, https://github.com/yourusername/factory-queue
Project-URL: Bug Reports, https://github.com/yourusername/factory-queue/issues
Project-URL: Source, https://github.com/yourusername/factory-queue
Keywords: pipeline,factory,queue,node,multi-threading,disk-overflow
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: License :: OSI Approved :: MIT License
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-Python: >=3.7
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: psutil>=5.8.0
Dynamic: author
Dynamic: home-page
Dynamic: license-file
Dynamic: requires-python

# Factory Queue

[![PyPI version](https://badge.fury.io/py/factory-queue.svg)](https://badge.fury.io/py/factory-queue)
[![Python](https://img.shields.io/pypi/pyversions/factory-queue.svg)](https://pypi.org/project/factory-queue/)
[![License](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT)

基于生产者-消费者模式的 Python 流水线工具，支持多节点、多工位并行、多分支、资源控制、磁盘溢出、完成回调。

适用场景：**采集 → 解析 → 入库**、**批量文件处理**、**ETL 流水线**、**任何需要并行 + 流式的数据处理任务**。

---

## 目录

- [功能特性](#功能特性)
- [安装](#安装)
- [快速开始](#快速开始)
- [核心 API](#核心-api)
  - [Factory](#factory)
  - [Node](#node)
  - [ResourceConfig](#resourceconfig)
  - [FeedCounters](#feedcounters)
- [处理函数返回值规则](#处理函数返回值规则)
- [多分支流水线](#多分支流水线)
- [高级功能](#高级功能)
  - [finish_func 完成回调](#finish_func-完成回调)
  - [setup / teardown](#setup--teardown)
  - [batch_size 批量叶子节点](#batch_size-批量叶子节点)
  - [磁盘溢出](#磁盘溢出)
  - [内存比例配置](#内存比例配置)
  - [监控与日志](#监控与日志)
- [常见陷阱与故障排查](#常见陷阱与故障排查)
- [性能调优建议](#性能调优建议)
- [完整示例](#完整示例)
- [模块结构](#模块结构)
- [许可证](#许可证)

---

## 功能特性

- ✅ **流水线设计** — 链式 `head → create_node → create_node` 自动管理依赖
- ✅ **多工位并行** — 每个节点 `node_num=N` 启动 N 个工作线程
- ✅ **多分支流水线** — 一个节点可以输出多种产物给不同下游 (`feed=0/1/2...`)
- ✅ **完成回调 (finish_func)** — 某个输入触发的所有产物被下游消费完后回调，可拿消费计数
- ✅ **资源控制** — 内存上限 (`max_memory_mb`) + 队列长度上限 (`max_queue_size`)
- ✅ **磁盘溢出** — 内存不足时自动写盘，防止 OOM；也可禁用让队列阻塞等待
- ✅ **批量叶子节点** — `batch_size>1` 时攒批后一次调用，适合数据库批量 upsert
- ✅ **实时监控** — 定时输出各节点 / 队列状态
- ✅ **彩色日志** — 不同级别日志使用不同颜色
- ✅ **优雅退出** — `with Factory(...) as f:` 自动 `close()`

---

## 安装

```bash
pip install factory-queue
```

或从本地源码安装：

```bash
git clone <your-repo-url>
cd factory-queue
pip install -e .
```

---

## 快速开始

### 1. 线性流水线

最简形式：头节点 → 中间节点 → 叶子节点。

```python
import logging
from factory_queue import Factory, ResourceConfig

logging.basicConfig(level=logging.INFO)


def fetch(data_id):
    """头节点：拉取原始数据"""
    return {"id": data_id, "value": data_id * 10}


def process(data):
    """中间节点：加工"""
    data["processed"] = data["value"] * 2
    return data


def save(batch):
    """叶子节点：批量入库 (batch_size=100)"""
    print(f"写入 {len(batch)} 条 -> {batch[0]['id']}~{batch[-1]['id']}")


config = ResourceConfig(max_memory_mb=512, max_queue_size=10000)

with Factory(resource_config=config) as factory:
    head = factory.head(func=fetch, node_num=2, name="拉取")
    mid = head.create_node(func=process, node_num=4, name="加工")
    mid.create_node(func=save, node_num=2, batch_size=100, name="入库")

    factory.start()
    for i in range(1000):
        factory.feed(i)
    factory.end_feed()
    factory.wait_complete()
```

### 2. 分支流水线

一个节点产出多种数据，分发到不同下游：

```python
def classify(data_id):
    """按奇偶分到两个分支"""
    raw = {"id": data_id, "value": data_id * 10}
    # 返回 dict: key=feed 索引, value=给该 feed 的数据
    if data_id % 2 == 0:
        return {0: raw}   # feed 0: 偶数分支
    else:
        return {1: raw}   # feed 1: 奇数分支


with Factory(resource_config=ResourceConfig(max_memory_mb=512)) as factory:
    head = factory.head(func=classify, node_num=2, name="分类")

    # 分支 1 (feed=0)
    even = head.create_node(func=lambda d: {**d, "tag": "even"},
                            node_num=2, feed=0, name="偶数处理")
    even.create_node(func=print, node_num=1, name="偶数保存")

    # 分支 2 (feed=1)
    odd = head.create_node(func=lambda d: {**d, "tag": "odd"},
                           node_num=1, feed=1, name="奇数处理")
    odd.create_node(func=print, node_num=1, name="奇数保存")

    factory.start()
    for i in range(100):
        factory.feed(i)
    factory.end_feed()
    factory.wait_complete()
```

---

## 核心 API

### `Factory`

流水线主类，管理所有节点、线程、队列。

#### 构造参数

```python
Factory(
    resource_config: ResourceConfig = None,    # 资源配置 (不传则用默认)
    enable_monitor: bool = True,               # 是否启用状态监控
    monitor_interval: float = 10.0,            # 监控输出间隔 (秒)
    max_memory_percent: float = None,          # 比例配置内存上限 (0~1), 覆盖 resource_config
    disable_disk_mode: bool = False,           # True=内存满时阻塞, 不写盘
    hide_technical_terms: bool = True,         # 日志里隐藏底层技术词汇
)
```

#### 主要方法

| 方法 | 说明 |
|---|---|
| `head(func, node_num, name, ...)` | 创建头节点（流水线起点） |
| `start()` | 启动所有线程，必须先调用 |
| `feed(data)` | 启动后逐个投喂数据 |
| `feed_batch(list)` | 启动后批量投喂 |
| `bulk_feed(list)` | **启动前**批量投喂（启动后不能用） |
| `end_feed()` | 通知所有头节点投喂结束 |
| `wait_complete(timeout=None)` | 阻塞等待所有数据处理完成 |
| `stop()` | 强制停止 |
| `close()` | 清理资源（with 语句自动调用） |
| `stats` | 拿到各队列统计字典 |

---

### `Node`

流水线节点。通过 `factory.head(...)` 创建头节点，`node.create_node(...)` 创建下游节点。

#### `head()` / `create_node()` 公共参数

```python
Node(
    func: Callable,                # 处理函数: func(data, **kwargs) 或 func(batch, **kwargs)
    args: Dict = None,             # 传给 func 的固定 kwargs
    node_num: int = 1,             # 工位数量 (并行线程数)
    name: str = None,              # 节点名 (队列名会用到)
    save_result: bool = False,     # 是否把返回值保存到文件
    print_processing: bool = True, # 是否打印处理日志
    batch_size: int = 1,           # 批量大小 (仅叶子节点有效)
    setup_func: Callable = None,   # 每个工作线程启动时调用一次 (仅叶子节点)
    teardown_func: Callable = None,# 每个工作线程结束时调用一次 (仅叶子节点)
    output_queue_size: int = None, # 输出队列大小 (不传用全局)
    finish_func: Callable = None,  # 完成回调 (见下文)
    finish_args: Dict = None,      # 传给 finish_func 的固定 kwargs
    log_finish: bool = True,       # 是否打印完成日志
)
```

`create_node()` 额外有：

```python
feed: Union[int, str] = 0,         # 使用上游哪个产物 (0=第一个, 1=第二个, ...)
```

#### 节点类型

- **头节点**：`factory.head(...)` 创建，接收外部 `feed()` 的数据
- **中间节点**：有下游节点，用 `_build_process_func` 包装，返回值会被分发到下游
- **叶子节点**：没有 `create_node` 过，仅消费数据，用 `_build_consume_func` 包装

---

### `ResourceConfig`

```python
from dataclasses import dataclass
from typing import Optional

@dataclass
class ResourceConfig:
    max_memory_mb: int = 1024              # 内存上限 (超过则写盘)
    max_queue_size: int = 10000            # 队列长度上限 (满了阻塞生产者)
    temp_dir: Optional[str] = None         # 磁盘溢出临时目录 (默认系统 temp)
    check_interval: float = 1.0            # 资源检查间隔 (秒)
```

示例：

```python
# 固定 2GB 内存
config = ResourceConfig(max_memory_mb=2048, max_queue_size=50000)

# 自定义 temp 目录 (建议放最大可用盘)
import os
config = ResourceConfig(max_memory_mb=2048, temp_dir=r'D:\factory_tmp')
```

---

### `FeedCounters`

`finish_func` 的第三个参数，用来查询某个输入产生的各 feed 消费数。

```python
def on_done(original_data, worker, feed_counters, **finish_args):
    consumed_feed_0 = feed_counters.get_consumed(0)   # feed 0 消费数
    consumed_feed_1 = feed_counters.get_consumed(1)   # feed 1 消费数
    produced_feed_0 = feed_counters.get_produced(0)   # feed 0 生产数
    stats = feed_counters.get_all_stats()             # {feed_idx: {produced, consumed}}
```

---

## 处理函数返回值规则

处理函数 `func(data, **kwargs)` 的返回值会被自动分发到下游：

| 返回类型 | 行为 |
|---|---|
| `None` | 不输出，不入队 |
| `dict` `{0: x, 1: y}` | 按键匹配 feed：`x` 进 feed_0，`y` 进 feed_1 |
| `tuple` `(x, y)` | **按位置**分发：`x` 进 feed_0，`y` 进 feed_1 |
| `list` `[a, b, c]` | 整个 list 作为**一个元素**入 feed_0（不会展开） |
| 单个值（非上述类型） | 直接入 feed_0 |
| **Generator** | 逐 yield 分发；每个 yield 值按上述规则再判定 |

举例：

```python
def gen_node(data):
    """生成器节点：一次产多个"""
    for item in data["items"]:
        yield item                # 每条单独入 feed_0
    yield {1: data["summary"]}    # 最后给 feed_1 发一个汇总
```

**注意**：`dict` / `tuple` 分发时，值为 `None` 的 feed 会被跳过（不入队、不计入生产）。

---

## 多分支流水线

一个节点可以创建多个下游分支：

```python
head = factory.head(func=fetch, node_num=2, name="拉取")

# 分支 A: feed=0, 入 PG
branch_a = head.create_node(func=to_pg_row, node_num=3, feed=0, name="入库")
branch_a.create_node(func=batch_upsert, node_num=1, batch_size=500, name="写入PG")

# 分支 B: feed=1, 上传 MinIO
branch_b = head.create_node(func=upload_minio, node_num=5, feed=1, name="上传对象")
branch_b.create_node(func=record_url, node_num=1, name="记录URL")
```

上游节点 yield `(item_for_feed_0, item_for_feed_1)` 元组时按位置分发：

```python
def fetch(item):
    pg_row = {...}
    minio_obj = {...}
    return pg_row, minio_obj    # → feed_0 + feed_1
```

### 队列命名规则

队列名 = **上游节点名** + `_feed_` + feed 索引

例：
- `拉取_feed_0`、`拉取_feed_1`（head 节点名为"拉取"）
- `入库_feed_0`（中间节点名为"入库"）

`factory.stats` 查询时**必须用上游节点名**：

```python
stats = factory.stats
queues = stats.get('queues', {})
q = queues.get('拉取_feed_0', {})   # ✅ 用上游名
q = queues.get('入库_feed_0', {})   # ❌ 错, 这是下游叶子节点名
print(q.get('total_get', 0))         # 该队列总共消费数
```

---

## 高级功能

### finish_func 完成回调

**触发时机**：某条输入触发的所有产物（可能跨多个 feed）都被下游叶子节点成功消费后。

**典型用途**：处理完一个 zip 后写 Redis 标记完成。

```python
def on_zip_done(zip_path, worker, feed_counters, **finish_args):
    """一个 zip 处理完所有专利后回调"""
    patent_count = feed_counters.get_consumed(0)   # feed_0 消费数 = 解析出的专利数
    redis_ = DataTool().redis_tool()
    redis_.hset('patent_zip_done', os.path.basename(zip_path), patent_count)


head = factory.head(func=parse_zip, node_num=3, name="解析",
                    finish_func=on_zip_done)
head.create_node(func=upload_pg, node_num=10, batch_size=1000, name="上传")
```

#### finish_func 签名

```python
def finish_func(original_data, worker, feed_counters, **finish_args):
    ...
```

- `original_data`：传入该节点的原始输入（即上游 yield 它的那条数据）
- `worker`：该节点的 Producer 对象（一般不用）
- `feed_counters`：`FeedCounters` 对象，查消费数
- `finish_args`：`Node(finish_args={...})` 传入的额外参数

#### ⚠ 触发条件（重要）

`finish_func` 要求**所有 feed 的 consumed == produced**：
- 上游 yield 后 `tracker.add_feed(feed_idx, count)` 记录生产数
- 下游叶子消费成功后 `tracker.mark_consumed(feed_idx, 1)`
- 任何一个下游抛异常 → `mark_consumed` 不执行 → **finish_func 永远不触发**
- 用户的 `func` 里 `return`（包括跳过）不算异常，`mark_consumed` 照常执行

**排错**：finish_func 不触发时，检查下游叶子节点有没有吞异常、有没有某条数据卡在队列。

---

### setup / teardown

每个工作线程启动 / 结束时各调用一次，用来初始化线程局部资源（PG 连接、Redis 客户端等）。

```python
def setup_upload(worker):
    """每个上传线程启动时建立自己的 PG 连接"""
    tool = DataTool()
    worker.set_attr('pg_client', tool.pg_patent_db())

def teardown_upload(worker):
    """线程结束时关闭连接"""
    pg = worker.get_attr('pg_client')
    if pg:
        pg.close()

# 叶子节点用:
node.create_node(func=save_batch, node_num=10, batch_size=1000,
                 setup_func=setup_upload, teardown_func=teardown_upload,
                 name="上传")
```

**重要**：
- `worker.set_attr / get_attr` 是**线程局部**存储，不同节点、不同线程之间不共享
- **finish_func 里不能用 worker.get_attr**（worker 是 Producer，不是 Consumer，没设过那些属性）。finish_func 里需要资源时自己用 `DataTool()` 新建

---

### batch_size 批量叶子节点

叶子节点的 `batch_size > 1` 时，会攒够一批再调用 `func(batch_list, ...)`，整个 batch 作为 list 传入。

```python
def batch_upsert(batch, **kwargs):
    """batch 是 list, 长度 = batch_size (或 flush 时的剩余)"""
    df = pd.DataFrame(batch)
    pg.db_upsert(df, table='xxx', p_key=['id'])

node.create_node(func=batch_upsert, node_num=5, batch_size=1000)
```

**适用场景**：数据库批量 upsert（比单条快几十倍）、批量 HTTP 请求、批量文件写。

中间节点没有 batch_size 概念，每次处理一个输入。

---

### 磁盘溢出

内存达到 `max_memory_mb` 时，新入队的数据自动写盘 (`temp_dir`)，消费时再读回。

```python
config = ResourceConfig(
    max_memory_mb=2048,             # 内存上限 2GB
    max_queue_size=100000,          # 单队列最长 10 万
    temp_dir=r'D:\spill',           # 溢出目录 (建议放最大可用盘)
)
```

**禁用磁盘模式**（让队列在内存满时阻塞而非写盘）：

```python
Factory(resource_config=config, disable_disk_mode=True)
```

---

### 内存比例配置

按系统总内存的比例配置上限（替代 `max_memory_mb`）：

```python
# 最多用 50% 系统内存
Factory(max_memory_percent=0.5)

# 最多用 30% 系统内存, 同时启用监控
Factory(max_memory_percent=0.3, enable_monitor=True, monitor_interval=5.0)
```

`max_memory_percent` 会覆盖 `resource_config.max_memory_mb`。

---

### 监控与日志

启用监控后，每隔 `monitor_interval` 秒打印各队列状态：

```python
Factory(enable_monitor=True, monitor_interval=10.0)
```

输出示例：
```
2026-06-23 10:30:00 - Factory - INFO - [监控] 解析_feed_0: size=234/50000, get=1200/5000
2026-06-23 10:30:00 - Factory - INFO - [监控] 入库_feed_0: size=0/50000, get=1100/5000
```

**关闭某节点的处理日志**（如下载/上传大量小数据时太吵）：

```python
node = head.create_node(..., print_processing=False)
```

**关闭 finish 日志**：

```python
node = head.create_node(..., finish_func=fn, log_finish=False)
```

**屏蔽 Factory 的"批量处理完成"日志**（与 `[上传]` 业务日志重复）：

```python
import logging
class _SkipBatchDoneLog(logging.Filter):
    def filter(self, record):
        return '批量处理完成' not in record.getMessage()
logging.getLogger('Factory').addFilter(_SkipBatchDoneLog())
```

---

## 常见陷阱与故障排查

### 1. `finish_func` 永远不触发

**原因**：某个下游叶子节点的 `func` 抛了异常，`mark_consumed` 没执行。

**解决**：
- 在叶子节点的 `func` 里 try/except 吞异常，或确保不会抛
- 看日志有没有 `[节点 xxx] 处理函数执行失败`

### 2. `finish_func` 里取不到 worker 属性

**原因**：finish_func 的 `worker` 是当前节点的 Producer，不是下游 Consumer；`setup_func` 设的属性在 Consumer 上。

**解决**：finish_func 里自己 `DataTool().xxx()` 新建连接。

### 3. 队列名查不到

**原因**：`factory.stats['queues']` 里的 key 用的是**上游节点名 + `_feed_N`**，不是下游节点名。

```python
# 假设 head 名为 "解析", 下游名为 "上传"
stats['queues']['解析_feed_0']   # ✅
stats['queues']['上传_feed_0']   # ❌ 不存在
```

### 4. 多分支时某个 feed 数据丢失

**原因**：上游 yield 的 tuple/dict 里该位置是 `None`，会被跳过（不入队、不 add_feed）。

**解决**：确保该 feed 有数据时返回真实值，没数据时也 yield 占位符（如空 list）。

### 5. `RuntimeError: 请先调用 start()`

**原因**：`feed()` / `feed_batch()` 必须在 `factory.start()` 之后调用。

```python
factory.start()         # ✅ 先启动
for x in data:
    factory.feed(x)
factory.end_feed()
factory.wait_complete()
```

### 6. `bulk_feed` 报错 "流水线已启动"

**原因**：`bulk_feed` 只能在 `start()` **之前**调用，用来预加载大批数据。

```python
factory.bulk_feed(huge_list)   # ✅ start 前
factory.start()
factory.end_feed()
factory.wait_complete()
```

---

## 性能调优建议

### 节点工位数 (`node_num`)

| 节点类型 | 建议 |
|---|---|
| **CPU 密集型**（解析、转换） | `node_num = CPU 核心数` 或略多 |
| **I/O 密集型**（HTTP、DB、S3） | `node_num = 5 ~ 30`，看下游服务抗压 |
| **FTP 下载**（WIPO FTP 这种限并发的） | `node_num = 3 ~ 5`，太大反而被服务端拒连 |

### batch_size

数据库 upsert 强烈建议用 `batch_size`：
- PG upsert: `batch_size=1000` 通常比 `batch_size=1` 快 30~50 倍
- 但也不能太大，单批 > 10000 可能撑爆 PG work_mem

### 内存配置

- 小数据量：`max_memory_mb=512` ~ `1024`
- 大批量 + 大对象（专利全文、图像 bytes）：`max_memory_mb=4096` ~ `8192`
- 配合 `disable_disk_mode=True` 用，让队列满时阻塞，不写盘（避免磁盘 I/O）

### 死锁排查

流水线卡住时：
1. 看监控日志各队列 size 是否一直在增长
2. 看线程栈（py-spy / `threading.enumerate()`）
3. 常见原因：上游比下游快太多导致队列满阻塞；`finish_func` 不触发导致 `wait_complete` 死等

---

## 完整示例

完整的"采集 → 解析 → 入库"流水线，演示所有主要特性：

```python
import logging
import os
import zipfile
import pandas as pd
from factory_queue import Factory, ResourceConfig
from tool.database import DataTool

logging.basicConfig(level=logging.INFO)


# ===== 处理函数 =====
def parse_zip(zip_path, **kwargs):
    """中间节点 (有下游): 流式解析 zip, yield 每条记录"""
    with zipfile.ZipFile(zip_path) as zf:
        for name in zf.namelist():
            if not name.endswith('.xml'):
                continue
            with zf.open(name) as f:
                # ... 解析逻辑 ...
                yield {"id": name, "content": f.read().decode('utf-8')}


def upsert_db(batch, **kwargs):
    """叶子节点: 攒批 upsert 到 PG"""
    worker = kwargs.get('worker')
    pg = worker.get_attr('pg_client')
    df = pd.DataFrame(batch)
    pg.db_upsert(df, table='patent.data', p_key=['id'])


# ===== 资源管理 =====
def setup_db(worker):
    worker.set_attr('pg_client', DataTool().pg_patent_db())

def teardown_db(worker):
    pg = worker.get_attr('pg_client')
    if pg:
        pg.close()


# ===== 完成回调 =====
def on_zip_done(zip_path, worker, feed_counters, **kwargs):
    redis_ = DataTool().redis_tool()
    count = feed_counters.get_consumed(0)
    redis_.hset('patent_zip_done', os.path.basename(zip_path), count)
    print(f'[完成] {zip_path} -> {count} 条')


# ===== 流水线 =====
def main():
    config = ResourceConfig(max_memory_mb=4096, max_queue_size=50000)
    factory = Factory(resource_config=config, enable_monitor=True, monitor_interval=10.0)

    head = factory.head(func=parse_zip, node_num=3, name="解析",
                        finish_func=on_zip_done, log_finish=True)
    head.create_node(
        func=upsert_db, node_num=10, batch_size=1000,
        setup_func=setup_db, teardown_func=teardown_db,
        name="上传", print_processing=False,
    )

    factory.start()

    for zip_path in ['/data/zip1.zip', '/data/zip2.zip', '/data/zip3.zip']:
        factory.feed(zip_path)
    factory.end_feed()
    factory.wait_complete()
    factory.close()


if __name__ == '__main__':
    main()
```

---

## 模块结构

```
factory-queue/
├── factory_queue/
│   ├── __init__.py       # 对外导出 Factory / Node / ResourceConfig 等
│   ├── core.py           # 底层实现: BaseFactory / DiskBackedQueue / Producer / Consumer
│   └── factory.py        # 高层接口: Factory / Node / FeedCounters / GroupTracker
├── tests/
│   ├── test_framework.py
│   └── test_monitor_log.py
├── pyproject.toml
├── setup.py
├── README.md
└── LICENSE
```

### 底层组件（core.py，一般用户不直接用）

| 类 | 作用 |
|---|---|
| `BaseFactory` | 底层工厂，管理线程和队列 |
| `DiskBackedQueue` | 内存 + 磁盘溢出队列 |
| `Producer` | 头节点 / 中间节点的工作线程 |
| `Consumer` | 叶子节点的工作线程 |
| `ProducerGroup` / `ConsumerGroup` | 线程组管理 |
| `ColoredFormatter` | 彩色日志格式化器 |

---

## 许可证

MIT License

## 作者

wyf (`__author__` in `__init__.py`)

## 贡献

欢迎提交 Issue 和 Pull Request！
