gpu-worker 1.0.0

package/config.2gb.yaml

@@ -0,0 +1,32 @@
+# 2GB 显存配置
+name: "Worker-2GB"
+region: "asia-east"
+country: "China"
+city: "Shanghai"
+
+server:
+  url: "http://服务器A的IP:8880"  # 修改为服务器A的实际IP
+  timeout: 30
+  verify_ssl: false
+
+gpu:
+  device_id: 0
+  enable_cpu_offload: true
+
+# 2GB显存只能运行小模型
+supported_types:
+  - whisper
+  # - embedding
+
+engines:
+  whisper:
+    model_id: "openai/whisper-small"  # 约2GB显存
+  # embedding:
+  #   model_id: "BAAI/bge-small-zh-v1.5"  # 约0.5GB
+
+heartbeat_interval: 30
+poll_interval: 2
+
+load_control:
+  acceptance_rate: 1.0
+  max_concurrent_jobs: 1

package/config.8gb.yaml

@@ -0,0 +1,29 @@
+# 8GB 显存配置
+name: "Worker-8GB"
+region: "asia-east"
+country: "China"
+city: "Shanghai"
+
+server:
+  url: "http://127.0.0.1:8880"
+  timeout: 30
+  verify_ssl: false
+
+gpu:
+  device_id: 0
+  enable_cpu_offload: true  # 显存不足时使用CPU内存
+
+# 8GB显存推荐只运行一种任务
+supported_types:
+  - llm
+
+engines:
+  llm:
+    model_id: "Qwen/Qwen2.5-1.5B-Instruct"  # 约3GB显存
+
+heartbeat_interval: 30
+poll_interval: 2
+
+load_control:
+  acceptance_rate: 1.0
+  max_concurrent_jobs: 1

package/config.example.yaml

@@ -0,0 +1,72 @@
+# Worker配置示例 - 增强版
+
+# Worker标识（首次运行后会自动填充）
+worker_id: null
+token: null
+name: "My GPU Worker"
+
+# 地理信息（重要：影响任务分配）
+region: "asia-east"      # 区域代码，见下方说明
+country: "China"
+city: "Shanghai"
+timezone: "Asia/Shanghai"
+
+# 区域代码说明：
+# - asia-east: 东亚（中国、日本、韩国）
+# - asia-south: 东南亚（新加坡、泰国）
+# - europe-west: 西欧（德国、法国、英国）
+# - europe-east: 东欧
+# - america-north: 北美（美国、加拿大）
+# - america-south: 南美
+# - oceania: 大洋洲（澳大利亚）
+
+# 服务器配置
+server:
+  url: "http://localhost:8000"  # 中央服务器地址
+  timeout: 30
+  verify_ssl: true
+
+# GPU配置
+gpu:
+  enable_cpu_offload: true      # 启用CPU Offload节省显存
+  max_memory_gb: null           # 限制最大显存使用（null=不限制）
+  device_id: 0                  # 使用的GPU设备ID
+
+# 直连配置（可选，用于P2P低延迟场景）
+direct:
+  enabled: false                # 是否启用直连服务
+  host: "0.0.0.0"
+  port: 8080
+  public_url: null              # 公网可访问的URL，如 "http://your-ip:8080"
+
+# 支持的任务类型
+supported_types:
+  - "llm"
+  - "image_gen"
+  # - "whisper"
+  # - "embedding"
+
+# 引擎配置
+engines:
+  llm:
+    model_id: "Qwen/Qwen2.5-7B-Instruct"
+    # 可选配置：
+    # max_new_tokens: 2048
+    # temperature: 0.7
+
+  image_gen:
+    model_id: "black-forest-labs/FLUX.1-schnell"
+    # 可选配置：
+    # default_steps: 4
+    # default_width: 1024
+    # default_height: 1024
+
+  # whisper:
+  #   model_id: "openai/whisper-large-v3"
+
+  # embedding:
+  #   model_id: "BAAI/bge-large-zh-v1.5"
+
+# 轮询配置
+heartbeat_interval: 30   # 心跳间隔（秒）
+poll_interval: 2         # 任务轮询间隔（秒）

package/config.py

@@ -0,0 +1,213 @@
+"""
+Worker配置 - 支持环境变量和YAML配置
+优先级: 环境变量 > config.yaml > 默认值
+"""
+from pydantic import BaseModel, Field
+from typing import List, Optional, Dict, Any
+import yaml
+from pathlib import Path
+import os
+
+
+def get_env(key: str, default: Any = None, cast: type = str) -> Any:
+    """获取环境变量并转换类型"""
+    value = os.getenv(key)
+    if value is None:
+        return default
+
+    if cast == bool:
+        return value.lower() in ('true', '1', 'yes', 'on')
+    elif cast == list:
+        return [x.strip() for x in value.split(',') if x.strip()]
+
+    try:
+        return cast(value)
+    except (ValueError, TypeError):
+        return default
+
+
+class ServerConfig(BaseModel):
+    """服务器配置"""
+    url: str = Field(default_factory=lambda: get_env('GPU_SERVER_URL', 'http://localhost:8000'))
+    timeout: int = Field(default_factory=lambda: get_env('GPU_SERVER_TIMEOUT', 30, int))
+    verify_ssl: bool = Field(default_factory=lambda: get_env('GPU_SERVER_VERIFY_SSL', True, bool))
+
+
+class GPUConfig(BaseModel):
+    """GPU配置"""
+    enable_cpu_offload: bool = Field(default_factory=lambda: get_env('GPU_ENABLE_CPU_OFFLOAD', True, bool))
+    max_memory_gb: Optional[float] = Field(default_factory=lambda: get_env('GPU_MAX_MEMORY_GB', None, float))
+    device_id: int = Field(default_factory=lambda: get_env('GPU_DEVICE_ID', 0, int))
+
+
+class DirectConfig(BaseModel):
+    """直连配置"""
+    enabled: bool = Field(default_factory=lambda: get_env('GPU_DIRECT_ENABLED', False, bool))
+    host: str = Field(default_factory=lambda: get_env('GPU_DIRECT_HOST', '0.0.0.0'))
+    port: int = Field(default_factory=lambda: get_env('GPU_DIRECT_PORT', 8080, int))
+    public_url: Optional[str] = Field(default_factory=lambda: get_env('GPU_DIRECT_PUBLIC_URL', None))
+
+
+class LoadControlConfig(BaseModel):
+    """负载控制配置"""
+    acceptance_rate: float = Field(default_factory=lambda: get_env('GPU_ACCEPTANCE_RATE', 1.0, float))
+    max_concurrent_jobs: int = Field(default_factory=lambda: get_env('GPU_MAX_CONCURRENT_JOBS', 1, int))
+    max_jobs_per_hour: int = Field(default_factory=lambda: get_env('GPU_MAX_JOBS_PER_HOUR', 0, int))
+    working_hours_start: Optional[int] = Field(default_factory=lambda: get_env('GPU_WORKING_HOURS_START', None, int))
+    working_hours_end: Optional[int] = Field(default_factory=lambda: get_env('GPU_WORKING_HOURS_END', None, int))
+
+
+class WorkerConfig(BaseModel):
+    """Worker配置"""
+
+    # Worker标识（首次运行后自动填充）
+    worker_id: Optional[str] = Field(default_factory=lambda: get_env('GPU_WORKER_ID', None))
+    token: Optional[str] = Field(default_factory=lambda: get_env('GPU_WORKER_TOKEN', None))
+    name: Optional[str] = Field(default_factory=lambda: get_env('GPU_WORKER_NAME', None))
+
+    # 地理信息
+    region: str = Field(default_factory=lambda: get_env('GPU_REGION', 'asia-east'))
+    country: Optional[str] = Field(default_factory=lambda: get_env('GPU_COUNTRY', None))
+    city: Optional[str] = Field(default_factory=lambda: get_env('GPU_CITY', None))
+    timezone: Optional[str] = Field(default_factory=lambda: get_env('GPU_TIMEZONE', None))
+
+    # 服务器配置
+    server: ServerConfig = Field(default_factory=ServerConfig)
+
+    # GPU配置
+    gpu: GPUConfig = Field(default_factory=GPUConfig)
+
+    # 直连配置
+    direct: DirectConfig = Field(default_factory=DirectConfig)
+
+    # 负载控制
+    load_control: LoadControlConfig = Field(default_factory=LoadControlConfig)
+
+    # 支持的任务类型
+    supported_types: List[str] = Field(
+        default_factory=lambda: get_env('GPU_SUPPORTED_TYPES', ['llm', 'image_gen'], list)
+    )
+
+    # 引擎配置
+    engines: Dict[str, Dict[str, Any]] = Field(default_factory=dict)
+
+    # 轮询配置
+    heartbeat_interval: int = Field(default_factory=lambda: get_env('GPU_HEARTBEAT_INTERVAL', 30, int))
+    poll_interval: int = Field(default_factory=lambda: get_env('GPU_POLL_INTERVAL', 2, int))
+
+    def save(self, path: str = "config.yaml"):
+        """保存配置到YAML文件"""
+        data = self.model_dump()
+        with open(path, "w", encoding="utf-8") as f:
+            yaml.dump(data, f, default_flow_style=False, allow_unicode=True)
+
+    @classmethod
+    def from_env(cls) -> 'WorkerConfig':
+        """从环境变量创建配置"""
+        return cls()
+
+
+def load_dotenv(path: str = ".env"):
+    """加载.env文件到环境变量"""
+    env_path = Path(path)
+    if not env_path.exists():
+        return
+
+    with open(env_path, encoding='utf-8') as f:
+        for line in f:
+            line = line.strip()
+            # 跳过空行和注释
+            if not line or line.startswith('#'):
+                continue
+
+            # 解析 KEY=VALUE
+            if '=' in line:
+                key, _, value = line.partition('=')
+                key = key.strip()
+                value = value.strip()
+
+                # 移除引号
+                if value and value[0] in ('"', "'") and value[-1] == value[0]:
+                    value = value[1:-1]
+
+                # 只设置未定义的环境变量
+                if key and key not in os.environ:
+                    os.environ[key] = value
+
+
+def load_config(path: str = "config.yaml") -> WorkerConfig:
+    """
+    加载配置
+    优先级: 环境变量 > config.yaml > 默认值
+    """
+    # 首先加载.env文件
+    load_dotenv()
+
+    config_path = Path(path)
+
+    if config_path.exists():
+        with open(config_path, encoding="utf-8") as f:
+            data = yaml.safe_load(f) or {}
+
+        # 处理嵌套配置
+        if "server" in data and isinstance(data["server"], dict):
+            data["server"] = ServerConfig(**data["server"])
+        if "gpu" in data and isinstance(data["gpu"], dict):
+            data["gpu"] = GPUConfig(**data["gpu"])
+        if "direct" in data and isinstance(data["direct"], dict):
+            data["direct"] = DirectConfig(**data["direct"])
+        if "load_control" in data and isinstance(data["load_control"], dict):
+            data["load_control"] = LoadControlConfig(**data["load_control"])
+
+        config = WorkerConfig(**data)
+    else:
+        # 仅从环境变量创建配置
+        config = WorkerConfig()
+
+    # 从环境变量加载引擎配置
+    _load_engine_configs_from_env(config)
+
+    return config
+
+
+def _load_engine_configs_from_env(config: WorkerConfig):
+    """从环境变量加载引擎配置"""
+    env_models = {
+        'llm': get_env('GPU_LLM_MODEL'),
+        'image_gen': get_env('GPU_IMAGE_MODEL'),
+        'vision': get_env('GPU_VISION_MODEL'),
+        'whisper': get_env('GPU_WHISPER_MODEL'),
+        'embedding': get_env('GPU_EMBEDDING_MODEL'),
+    }
+
+    for engine_type, model_id in env_models.items():
+        if model_id:
+            if engine_type not in config.engines:
+                config.engines[engine_type] = {}
+            config.engines[engine_type]['model_id'] = model_id
+
+
+# 默认引擎配置
+DEFAULT_ENGINE_CONFIGS = {
+    "llm": {
+        "model_id": "Qwen/Qwen2.5-7B-Instruct",
+        "max_new_tokens": 2048,
+        "temperature": 0.7,
+    },
+    "image_gen": {
+        "model_id": "Zhihu-ai/Z-Image-Turbo",
+        "default_steps": 4,
+        "default_width": 1024,
+        "default_height": 1024,
+    },
+    "vision": {
+        "model_id": "THUDM/glm-4v-9b",
+        "max_new_tokens": 1024,
+    },
+    "whisper": {
+        "model_id": "openai/whisper-large-v3",
+    },
+    "embedding": {
+        "model_id": "BAAI/bge-large-zh-v1.5",
+    }
+}

package/direct_server.py

@@ -0,0 +1,140 @@
+"""
+直连服务器 - 允许客户端直接与Worker通信
+跳过中央服务器，降低延迟
+"""
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import Optional, Dict, Any
+import uvicorn
+import logging
+from threading import Thread
+
+logger = logging.getLogger(__name__)
+
+
+class DirectInferenceRequest(BaseModel):
+    """直连推理请求"""
+    type: str
+    params: Dict[str, Any]
+    timeout_seconds: int = 300
+
+
+class DirectInferenceResponse(BaseModel):
+    """直连推理响应"""
+    success: bool
+    result: Optional[Dict[str, Any]] = None
+    error: Optional[str] = None
+    processing_time_ms: int = 0
+
+
+class DirectServer:
+    """
+    直连服务器
+
+    允许客户端绕过中央服务器，直接与Worker通信
+    适用于低延迟场景
+    """
+
+    def __init__(self, worker, host: str = "0.0.0.0", port: int = 8080):
+        self.worker = worker
+        self.host = host
+        self.port = port
+        self.app = FastAPI(title="Worker Direct API")
+        self._setup_routes()
+        self.server = None
+
+    def _setup_routes(self):
+        """设置路由"""
+
+        @self.app.get("/health")
+        async def health():
+            return {
+                "status": "healthy",
+                "worker_id": self.worker.worker_id,
+                "worker_status": self.worker.status,
+                "supported_types": list(self.worker.engines.keys())
+            }
+
+        @self.app.get("/status")
+        async def status():
+            gpu_info = self.worker._get_gpu_info()
+            return {
+                "worker_id": self.worker.worker_id,
+                "status": self.worker.status,
+                "current_job": self.worker.current_job_id,
+                "supported_types": list(self.worker.engines.keys()),
+                "gpu_info": gpu_info,
+                "accepting_jobs": self.worker.accepting_jobs
+            }
+
+        @self.app.post("/inference", response_model=DirectInferenceResponse)
+        async def direct_inference(request: DirectInferenceRequest):
+            """
+            直连推理接口
+
+            客户端可以直接调用此接口进行推理，跳过中央服务器
+            """
+            import time
+
+            # 检查是否接受任务
+            if not self.worker.accepting_jobs:
+                raise HTTPException(503, "Worker is going offline")
+
+            # 检查是否空闲
+            if self.worker.status != "idle":
+                raise HTTPException(503, "Worker is busy")
+
+            # 检查引擎
+            engine = self.worker.engines.get(request.type)
+            if not engine:
+                raise HTTPException(
+                    400,
+                    f"Unsupported type: {request.type}. "
+                    f"Supported: {list(self.worker.engines.keys())}"
+                )
+
+            # 标记为忙碌
+            self.worker.status = "busy"
+
+            try:
+                start_time = time.time()
+                result = engine.inference(request.params)
+                processing_time_ms = int((time.time() - start_time) * 1000)
+
+                return DirectInferenceResponse(
+                    success=True,
+                    result=result,
+                    processing_time_ms=processing_time_ms
+                )
+
+            except Exception as e:
+                logger.error(f"Direct inference error: {e}")
+                return DirectInferenceResponse(
+                    success=False,
+                    error=str(e)
+                )
+
+            finally:
+                self.worker.status = "idle"
+
+    def start(self):
+        """启动服务器（阻塞）"""
+        config = uvicorn.Config(
+            self.app,
+            host=self.host,
+            port=self.port,
+            log_level="warning"
+        )
+        self.server = uvicorn.Server(config)
+        self.server.run()
+
+    def start_background(self):
+        """后台启动服务器"""
+        thread = Thread(target=self.start, daemon=True)
+        thread.start()
+        return thread
+
+    def stop(self):
+        """停止服务器"""
+        if self.server:
+            self.server.should_exit = True

package/distributed/__init__.py

@@ -0,0 +1,35 @@
+"""分布式推理组件
+
+实现跨 Worker 的模型分片推理，参考 Petals 项目设计：
+- DistributedInferenceSession: 分布式推理会话管理
+- WorkerSession: Worker 级别会话
+- ModelShard: 模型分片加载器
+- GRPCServer: Worker 间 P2P 通信
+"""
+from .session import (
+    DistributedInferenceSession,
+    WorkerSession,
+    SessionManager,
+)
+from .model_shard import (
+    ModelShard,
+    ShardedModelLoader,
+    get_layer_range_for_worker,
+)
+from .kv_cache import (
+    DistributedKVCacheManager,
+    PagedKVCache,
+    KVCachePool,
+)
+
+__all__ = [
+    "DistributedInferenceSession",
+    "WorkerSession",
+    "SessionManager",
+    "ModelShard",
+    "ShardedModelLoader",
+    "get_layer_range_for_worker",
+    "DistributedKVCacheManager",
+    "PagedKVCache",
+    "KVCachePool",
+]