feat: add marketplace metrics, privacy features, and service registry endpoints
- Add Prometheus metrics for marketplace API throughput and error rates with new dashboard panels - Implement confidential transaction models with encryption support and access control - Add key management system with registration, rotation, and audit logging - Create services and registry routers for service discovery and management - Integrate ZK proof generation for privacy-preserving receipts - Add metrics instru
This commit is contained in:
oib
15
apps/miner-node/plugins/__init__.py
Normal file
15
apps/miner-node/plugins/__init__.py
Normal file
@ -0,0 +1,15 @@
|
||||
"""
|
||||
Miner plugin system for GPU service execution
|
||||
"""
|
||||
|
||||
from .base import ServicePlugin, PluginResult
|
||||
from .registry import PluginRegistry
|
||||
from .exceptions import PluginError, PluginNotFoundError
|
||||
|
||||
__all__ = [
|
||||
"ServicePlugin",
|
||||
"PluginResult",
|
||||
"PluginRegistry",
|
||||
"PluginError",
|
||||
"PluginNotFoundError"
|
||||
]
|
||||
111
apps/miner-node/plugins/base.py
Normal file
111
apps/miner-node/plugins/base.py
Normal file
@ -0,0 +1,111 @@
|
||||
"""
|
||||
Base plugin interface for GPU service execution
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, Any, Optional, List
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
import asyncio
|
||||
|
||||
|
||||
@dataclass
|
||||
class PluginResult:
|
||||
"""Result from plugin execution"""
|
||||
success: bool
|
||||
data: Optional[Dict[str, Any]] = None
|
||||
error: Optional[str] = None
|
||||
metrics: Optional[Dict[str, Any]] = None
|
||||
execution_time: Optional[float] = None
|
||||
|
||||
|
||||
class ServicePlugin(ABC):
|
||||
"""Base class for all service plugins"""
|
||||
|
||||
def __init__(self):
|
||||
self.service_id = None
|
||||
self.name = None
|
||||
self.version = "1.0.0"
|
||||
self.description = ""
|
||||
self.capabilities = []
|
||||
|
||||
@abstractmethod
|
||||
async def execute(self, request: Dict[str, Any]) -> PluginResult:
|
||||
"""Execute the service with given parameters"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def validate_request(self, request: Dict[str, Any]) -> List[str]:
|
||||
"""Validate request parameters, return list of errors"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_hardware_requirements(self) -> Dict[str, Any]:
|
||||
"""Get hardware requirements for this plugin"""
|
||||
pass
|
||||
|
||||
def get_metrics(self) -> Dict[str, Any]:
|
||||
"""Get plugin-specific metrics"""
|
||||
return {
|
||||
"service_id": self.service_id,
|
||||
"name": self.name,
|
||||
"version": self.version
|
||||
}
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check if plugin dependencies are available"""
|
||||
return True
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Initialize plugin resources"""
|
||||
pass
|
||||
|
||||
def cleanup(self) -> None:
|
||||
"""Cleanup plugin resources"""
|
||||
pass
|
||||
|
||||
|
||||
class GPUPlugin(ServicePlugin):
|
||||
"""Base class for GPU-accelerated plugins"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.gpu_available = False
|
||||
self.vram_gb = 0
|
||||
self.cuda_available = False
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Check GPU availability"""
|
||||
self._detect_gpu()
|
||||
|
||||
def _detect_gpu(self) -> None:
|
||||
"""Detect GPU and VRAM"""
|
||||
try:
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
self.gpu_available = True
|
||||
self.cuda_available = True
|
||||
self.vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import GPUtil
|
||||
gpus = GPUtil.getGPUs()
|
||||
if gpus:
|
||||
self.gpu_available = True
|
||||
self.vram_gb = gpus[0].memoryTotal / 1024
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
def get_hardware_requirements(self) -> Dict[str, Any]:
|
||||
"""Default GPU requirements"""
|
||||
return {
|
||||
"gpu": "any",
|
||||
"vram_gb": 4,
|
||||
"cuda": "recommended"
|
||||
}
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check GPU health"""
|
||||
return self.gpu_available
|
||||
371
apps/miner-node/plugins/blender.py
Normal file
371
apps/miner-node/plugins/blender.py
Normal file
@ -0,0 +1,371 @@
|
||||
"""
|
||||
Blender 3D rendering plugin
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
import json
|
||||
from typing import Dict, Any, List, Optional
|
||||
import time
|
||||
|
||||
from .base import GPUPlugin, PluginResult
|
||||
from .exceptions import PluginExecutionError
|
||||
|
||||
|
||||
class BlenderPlugin(GPUPlugin):
|
||||
"""Plugin for Blender 3D rendering"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.service_id = "blender"
|
||||
self.name = "Blender Rendering"
|
||||
self.version = "1.0.0"
|
||||
self.description = "Render 3D scenes using Blender"
|
||||
self.capabilities = ["render", "animation", "cycles", "eevee"]
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Initialize Blender dependencies"""
|
||||
super().setup()
|
||||
|
||||
# Check for Blender installation
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["blender", "--version"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True
|
||||
)
|
||||
self.blender_path = "blender"
|
||||
except (subprocess.CalledProcessError, FileNotFoundError):
|
||||
raise PluginExecutionError("Blender not found. Install Blender for 3D rendering")
|
||||
|
||||
# Check for bpy module (Python API)
|
||||
try:
|
||||
import bpy
|
||||
self.bpy_available = True
|
||||
except ImportError:
|
||||
self.bpy_available = False
|
||||
print("Warning: bpy module not available. Some features may be limited.")
|
||||
|
||||
def validate_request(self, request: Dict[str, Any]) -> List[str]:
|
||||
"""Validate Blender request parameters"""
|
||||
errors = []
|
||||
|
||||
# Check required parameters
|
||||
if "blend_file" not in request and "scene_data" not in request:
|
||||
errors.append("Either 'blend_file' or 'scene_data' must be provided")
|
||||
|
||||
# Validate engine
|
||||
engine = request.get("engine", "cycles")
|
||||
valid_engines = ["cycles", "eevee", "workbench"]
|
||||
if engine not in valid_engines:
|
||||
errors.append(f"Invalid engine. Must be one of: {', '.join(valid_engines)}")
|
||||
|
||||
# Validate resolution
|
||||
resolution_x = request.get("resolution_x", 1920)
|
||||
resolution_y = request.get("resolution_y", 1080)
|
||||
|
||||
if not isinstance(resolution_x, int) or resolution_x < 1 or resolution_x > 65536:
|
||||
errors.append("resolution_x must be an integer between 1 and 65536")
|
||||
if not isinstance(resolution_y, int) or resolution_y < 1 or resolution_y > 65536:
|
||||
errors.append("resolution_y must be an integer between 1 and 65536")
|
||||
|
||||
# Validate samples
|
||||
samples = request.get("samples", 128)
|
||||
if not isinstance(samples, int) or samples < 1 or samples > 10000:
|
||||
errors.append("samples must be an integer between 1 and 10000")
|
||||
|
||||
# Validate frame range for animation
|
||||
if request.get("animation", False):
|
||||
frame_start = request.get("frame_start", 1)
|
||||
frame_end = request.get("frame_end", 250)
|
||||
|
||||
if not isinstance(frame_start, int) or frame_start < 1:
|
||||
errors.append("frame_start must be >= 1")
|
||||
if not isinstance(frame_end, int) or frame_end < frame_start:
|
||||
errors.append("frame_end must be >= frame_start")
|
||||
|
||||
return errors
|
||||
|
||||
def get_hardware_requirements(self) -> Dict[str, Any]:
|
||||
"""Get hardware requirements for Blender"""
|
||||
return {
|
||||
"gpu": "recommended",
|
||||
"vram_gb": 4,
|
||||
"ram_gb": 16,
|
||||
"cuda": "recommended"
|
||||
}
|
||||
|
||||
async def execute(self, request: Dict[str, Any]) -> PluginResult:
|
||||
"""Execute Blender rendering"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Validate request
|
||||
errors = self.validate_request(request)
|
||||
if errors:
|
||||
return PluginResult(
|
||||
success=False,
|
||||
error=f"Validation failed: {'; '.join(errors)}"
|
||||
)
|
||||
|
||||
# Get parameters
|
||||
blend_file = request.get("blend_file")
|
||||
scene_data = request.get("scene_data")
|
||||
engine = request.get("engine", "cycles")
|
||||
resolution_x = request.get("resolution_x", 1920)
|
||||
resolution_y = request.get("resolution_y", 1080)
|
||||
samples = request.get("samples", 128)
|
||||
animation = request.get("animation", False)
|
||||
frame_start = request.get("frame_start", 1)
|
||||
frame_end = request.get("frame_end", 250)
|
||||
output_format = request.get("output_format", "png")
|
||||
gpu_acceleration = request.get("gpu_acceleration", self.gpu_available)
|
||||
|
||||
# Prepare input file
|
||||
input_file = await self._prepare_input_file(blend_file, scene_data)
|
||||
|
||||
# Build Blender command
|
||||
cmd = self._build_blender_command(
|
||||
input_file=input_file,
|
||||
engine=engine,
|
||||
resolution_x=resolution_x,
|
||||
resolution_y=resolution_y,
|
||||
samples=samples,
|
||||
animation=animation,
|
||||
frame_start=frame_start,
|
||||
frame_end=frame_end,
|
||||
output_format=output_format,
|
||||
gpu_acceleration=gpu_acceleration
|
||||
)
|
||||
|
||||
# Execute Blender
|
||||
output_files = await self._execute_blender(cmd, animation, frame_start, frame_end)
|
||||
|
||||
# Get render statistics
|
||||
render_stats = await self._get_render_stats(output_files[0] if output_files else None)
|
||||
|
||||
# Clean up input file if created from scene data
|
||||
if scene_data:
|
||||
os.unlink(input_file)
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
return PluginResult(
|
||||
success=True,
|
||||
data={
|
||||
"output_files": output_files,
|
||||
"count": len(output_files),
|
||||
"animation": animation,
|
||||
"parameters": {
|
||||
"engine": engine,
|
||||
"resolution": f"{resolution_x}x{resolution_y}",
|
||||
"samples": samples,
|
||||
"gpu_acceleration": gpu_acceleration
|
||||
}
|
||||
},
|
||||
metrics={
|
||||
"engine": engine,
|
||||
"frames_rendered": len(output_files),
|
||||
"render_time": execution_time,
|
||||
"time_per_frame": execution_time / len(output_files) if output_files else 0,
|
||||
"samples_per_second": (samples * len(output_files)) / execution_time if execution_time > 0 else 0,
|
||||
"render_stats": render_stats
|
||||
},
|
||||
execution_time=execution_time
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
return PluginResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
execution_time=time.time() - start_time
|
||||
)
|
||||
|
||||
async def _prepare_input_file(self, blend_file: Optional[str], scene_data: Optional[Dict]) -> str:
|
||||
"""Prepare input .blend file"""
|
||||
if blend_file:
|
||||
# Use provided file
|
||||
if not os.path.exists(blend_file):
|
||||
raise PluginExecutionError(f"Blend file not found: {blend_file}")
|
||||
return blend_file
|
||||
elif scene_data:
|
||||
# Create blend file from scene data
|
||||
if not self.bpy_available:
|
||||
raise PluginExecutionError("Cannot create scene without bpy module")
|
||||
|
||||
# Create a temporary Python script to generate the scene
|
||||
script = tempfile.mktemp(suffix=".py")
|
||||
output_blend = tempfile.mktemp(suffix=".blend")
|
||||
|
||||
with open(script, "w") as f:
|
||||
f.write(f"""
|
||||
import bpy
|
||||
import json
|
||||
|
||||
# Load scene data
|
||||
scene_data = json.loads('''{json.dumps(scene_data)}''')
|
||||
|
||||
# Clear default scene
|
||||
bpy.ops.object.select_all(action='SELECT')
|
||||
bpy.ops.object.delete()
|
||||
|
||||
# Create scene from data
|
||||
# This is a simplified example - in practice, you'd parse the scene_data
|
||||
# and create appropriate objects, materials, lights, etc.
|
||||
|
||||
# Save blend file
|
||||
bpy.ops.wm.save_as_mainfile(filepath='{output_blend}')
|
||||
""")
|
||||
|
||||
# Run Blender to create the scene
|
||||
cmd = [self.blender_path, "--background", "--python", script]
|
||||
process = await asyncio.create_subprocess_exec(*cmd)
|
||||
await process.communicate()
|
||||
|
||||
# Clean up script
|
||||
os.unlink(script)
|
||||
|
||||
return output_blend
|
||||
else:
|
||||
raise PluginExecutionError("Either blend_file or scene_data must be provided")
|
||||
|
||||
def _build_blender_command(
|
||||
self,
|
||||
input_file: str,
|
||||
engine: str,
|
||||
resolution_x: int,
|
||||
resolution_y: int,
|
||||
samples: int,
|
||||
animation: bool,
|
||||
frame_start: int,
|
||||
frame_end: int,
|
||||
output_format: str,
|
||||
gpu_acceleration: bool
|
||||
) -> List[str]:
|
||||
"""Build Blender command"""
|
||||
cmd = [
|
||||
self.blender_path,
|
||||
"--background",
|
||||
input_file,
|
||||
"--render-engine", engine,
|
||||
"--render-format", output_format.upper()
|
||||
]
|
||||
|
||||
# Add Python script for settings
|
||||
script = tempfile.mktemp(suffix=".py")
|
||||
with open(script, "w") as f:
|
||||
f.write(f"""
|
||||
import bpy
|
||||
|
||||
# Set resolution
|
||||
bpy.context.scene.render.resolution_x = {resolution_x}
|
||||
bpy.context.scene.render.resolution_y = {resolution_y}
|
||||
|
||||
# Set samples for Cycles
|
||||
if bpy.context.scene.render.engine == 'CYCLES':
|
||||
bpy.context.scene.cycles.samples = {samples}
|
||||
|
||||
# Enable GPU rendering if available
|
||||
if {str(gpu_acceleration).lower()}:
|
||||
bpy.context.scene.cycles.device = 'GPU'
|
||||
preferences = bpy.context.preferences
|
||||
cycles_preferences = preferences.addons['cycles'].preferences
|
||||
cycles_preferences.compute_device_type = 'CUDA'
|
||||
cycles_preferences.get_devices()
|
||||
for device in cycles_preferences.devices:
|
||||
device.use = True
|
||||
|
||||
# Set frame range for animation
|
||||
if {str(animation).lower()}:
|
||||
bpy.context.scene.frame_start = {frame_start}
|
||||
bpy.context.scene.frame_end = {frame_end}
|
||||
|
||||
# Set output path
|
||||
bpy.context.scene.render.filepath = '{tempfile.mkdtemp()}/render_'
|
||||
|
||||
# Save settings
|
||||
bpy.ops.wm.save_mainfile()
|
||||
""")
|
||||
|
||||
cmd.extend(["--python", script])
|
||||
|
||||
# Add render command
|
||||
if animation:
|
||||
cmd.extend(["-a"]) # Render animation
|
||||
else:
|
||||
cmd.extend(["-f", "1"]) # Render single frame
|
||||
|
||||
return cmd
|
||||
|
||||
async def _execute_blender(
|
||||
self,
|
||||
cmd: List[str],
|
||||
animation: bool,
|
||||
frame_start: int,
|
||||
frame_end: int
|
||||
) -> List[str]:
|
||||
"""Execute Blender command"""
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE
|
||||
)
|
||||
|
||||
stdout, stderr = await process.communicate()
|
||||
|
||||
if process.returncode != 0:
|
||||
error_msg = stderr.decode() if stderr else "Blender failed"
|
||||
raise PluginExecutionError(f"Blender error: {error_msg}")
|
||||
|
||||
# Find output files
|
||||
output_dir = tempfile.mkdtemp()
|
||||
output_pattern = os.path.join(output_dir, "render_*")
|
||||
|
||||
if animation:
|
||||
# Animation creates multiple files
|
||||
import glob
|
||||
output_files = glob.glob(output_pattern)
|
||||
output_files.sort() # Ensure frame order
|
||||
else:
|
||||
# Single frame
|
||||
output_files = [glob.glob(output_pattern)[0]]
|
||||
|
||||
return output_files
|
||||
|
||||
async def _get_render_stats(self, output_file: Optional[str]) -> Dict[str, Any]:
|
||||
"""Get render statistics"""
|
||||
if not output_file or not os.path.exists(output_file):
|
||||
return {}
|
||||
|
||||
# Get file size and basic info
|
||||
file_size = os.path.getsize(output_file)
|
||||
|
||||
# Try to get image dimensions
|
||||
try:
|
||||
from PIL import Image
|
||||
with Image.open(output_file) as img:
|
||||
width, height = img.size
|
||||
except:
|
||||
width = height = None
|
||||
|
||||
return {
|
||||
"file_size": file_size,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"format": os.path.splitext(output_file)[1][1:].upper()
|
||||
}
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check Blender health"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["blender", "--version"],
|
||||
capture_output=True,
|
||||
check=True
|
||||
)
|
||||
return True
|
||||
except subprocess.CalledProcessError:
|
||||
return False
|
||||
215
apps/miner-node/plugins/discovery.py
Normal file
215
apps/miner-node/plugins/discovery.py
Normal file
@ -0,0 +1,215 @@
|
||||
"""
|
||||
Plugin discovery and matching system
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Dict, List, Set, Optional
|
||||
import requests
|
||||
|
||||
from .registry import registry
|
||||
from .base import ServicePlugin
|
||||
from .exceptions import PluginNotFoundError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ServiceDiscovery:
|
||||
"""Discovers and matches services to plugins"""
|
||||
|
||||
def __init__(self, pool_hub_url: str, miner_id: str):
|
||||
self.pool_hub_url = pool_hub_url
|
||||
self.miner_id = miner_id
|
||||
self.enabled_services: Set[str] = set()
|
||||
self.service_configs: Dict[str, Dict] = {}
|
||||
self._last_update = 0
|
||||
self._update_interval = 60 # seconds
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the discovery service"""
|
||||
logger.info("Starting service discovery")
|
||||
|
||||
# Initialize plugin registry
|
||||
await registry.initialize()
|
||||
|
||||
# Initial sync
|
||||
await self.sync_services()
|
||||
|
||||
# Start background sync task
|
||||
asyncio.create_task(self._sync_loop())
|
||||
|
||||
async def sync_services(self) -> None:
|
||||
"""Sync enabled services from pool-hub"""
|
||||
try:
|
||||
# Get service configurations from pool-hub
|
||||
response = requests.get(
|
||||
f"{self.pool_hub_url}/v1/services/",
|
||||
headers={"X-Miner-ID": self.miner_id}
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
services = response.json()
|
||||
|
||||
# Update local state
|
||||
new_enabled = set()
|
||||
new_configs = {}
|
||||
|
||||
for service in services:
|
||||
if service.get("enabled", False):
|
||||
service_id = service["service_type"]
|
||||
new_enabled.add(service_id)
|
||||
new_configs[service_id] = service
|
||||
|
||||
# Find changes
|
||||
added = new_enabled - self.enabled_services
|
||||
removed = self.enabled_services - new_enabled
|
||||
updated = set()
|
||||
|
||||
for service_id in self.enabled_services & new_enabled:
|
||||
if new_configs[service_id] != self.service_configs.get(service_id):
|
||||
updated.add(service_id)
|
||||
|
||||
# Apply changes
|
||||
for service_id in removed:
|
||||
await self._disable_service(service_id)
|
||||
|
||||
for service_id in added:
|
||||
await self._enable_service(service_id, new_configs[service_id])
|
||||
|
||||
for service_id in updated:
|
||||
await self._update_service(service_id, new_configs[service_id])
|
||||
|
||||
# Update state
|
||||
self.enabled_services = new_enabled
|
||||
self.service_configs = new_configs
|
||||
self._last_update = asyncio.get_event_loop().time()
|
||||
|
||||
logger.info(f"Synced services: {len(self.enabled_services)} enabled")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to sync services: {e}")
|
||||
|
||||
async def _enable_service(self, service_id: str, config: Dict) -> None:
|
||||
"""Enable a service"""
|
||||
try:
|
||||
# Check if plugin exists
|
||||
if service_id not in registry.list_plugins():
|
||||
logger.warning(f"No plugin available for service: {service_id}")
|
||||
return
|
||||
|
||||
# Load plugin
|
||||
plugin = registry.load_plugin(service_id)
|
||||
|
||||
# Validate hardware requirements
|
||||
await self._validate_hardware_requirements(plugin, config)
|
||||
|
||||
# Configure plugin if needed
|
||||
if hasattr(plugin, 'configure'):
|
||||
await plugin.configure(config.get('config', {}))
|
||||
|
||||
logger.info(f"Enabled service: {service_id}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to enable service {service_id}: {e}")
|
||||
|
||||
async def _disable_service(self, service_id: str) -> None:
|
||||
"""Disable a service"""
|
||||
try:
|
||||
# Unload plugin to free resources
|
||||
registry.unload_plugin(service_id)
|
||||
logger.info(f"Disabled service: {service_id}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to disable service {service_id}: {e}")
|
||||
|
||||
async def _update_service(self, service_id: str, config: Dict) -> None:
|
||||
"""Update service configuration"""
|
||||
# For now, just disable and re-enable
|
||||
await self._disable_service(service_id)
|
||||
await self._enable_service(service_id, config)
|
||||
|
||||
async def _validate_hardware_requirements(self, plugin: ServicePlugin, config: Dict) -> None:
|
||||
"""Validate that miner meets plugin requirements"""
|
||||
requirements = plugin.get_hardware_requirements()
|
||||
|
||||
# This would check against actual miner hardware
|
||||
# For now, just log the requirements
|
||||
logger.debug(f"Hardware requirements for {plugin.service_id}: {requirements}")
|
||||
|
||||
async def _sync_loop(self) -> None:
|
||||
"""Background sync loop"""
|
||||
while True:
|
||||
await asyncio.sleep(self._update_interval)
|
||||
await self.sync_services()
|
||||
|
||||
async def execute_service(self, service_id: str, request: Dict) -> Dict:
|
||||
"""Execute a service request"""
|
||||
try:
|
||||
# Check if service is enabled
|
||||
if service_id not in self.enabled_services:
|
||||
raise PluginNotFoundError(f"Service {service_id} is not enabled")
|
||||
|
||||
# Get plugin
|
||||
plugin = registry.get_plugin(service_id)
|
||||
if not plugin:
|
||||
raise PluginNotFoundError(f"No plugin loaded for service: {service_id}")
|
||||
|
||||
# Execute request
|
||||
result = await plugin.execute(request)
|
||||
|
||||
# Convert result to dict
|
||||
return {
|
||||
"success": result.success,
|
||||
"data": result.data,
|
||||
"error": result.error,
|
||||
"metrics": result.metrics,
|
||||
"execution_time": result.execution_time
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to execute service {service_id}: {e}")
|
||||
return {
|
||||
"success": False,
|
||||
"error": str(e)
|
||||
}
|
||||
|
||||
def get_enabled_services(self) -> List[str]:
|
||||
"""Get list of enabled services"""
|
||||
return list(self.enabled_services)
|
||||
|
||||
def get_service_status(self) -> Dict[str, Dict]:
|
||||
"""Get status of all services"""
|
||||
status = {}
|
||||
|
||||
for service_id in registry.list_plugins():
|
||||
plugin = registry.get_plugin(service_id)
|
||||
status[service_id] = {
|
||||
"enabled": service_id in self.enabled_services,
|
||||
"loaded": plugin is not None,
|
||||
"config": self.service_configs.get(service_id, {}),
|
||||
"capabilities": plugin.capabilities if plugin else []
|
||||
}
|
||||
|
||||
return status
|
||||
|
||||
async def health_check(self) -> Dict[str, bool]:
|
||||
"""Health check all enabled services"""
|
||||
results = {}
|
||||
|
||||
for service_id in self.enabled_services:
|
||||
plugin = registry.get_plugin(service_id)
|
||||
if plugin:
|
||||
try:
|
||||
results[service_id] = await plugin.health_check()
|
||||
except Exception as e:
|
||||
logger.error(f"Health check failed for {service_id}: {e}")
|
||||
results[service_id] = False
|
||||
else:
|
||||
results[service_id] = False
|
||||
|
||||
return results
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the discovery service"""
|
||||
logger.info("Stopping service discovery")
|
||||
registry.cleanup_all()
|
||||
23
apps/miner-node/plugins/exceptions.py
Normal file
23
apps/miner-node/plugins/exceptions.py
Normal file
@ -0,0 +1,23 @@
|
||||
"""
|
||||
Plugin system exceptions
|
||||
"""
|
||||
|
||||
|
||||
class PluginError(Exception):
|
||||
"""Base exception for plugin errors"""
|
||||
pass
|
||||
|
||||
|
||||
class PluginNotFoundError(PluginError):
|
||||
"""Raised when a plugin is not found"""
|
||||
pass
|
||||
|
||||
|
||||
class PluginValidationError(PluginError):
|
||||
"""Raised when plugin validation fails"""
|
||||
pass
|
||||
|
||||
|
||||
class PluginExecutionError(PluginError):
|
||||
"""Raised when plugin execution fails"""
|
||||
pass
|
||||
318
apps/miner-node/plugins/ffmpeg.py
Normal file
318
apps/miner-node/plugins/ffmpeg.py
Normal file
@ -0,0 +1,318 @@
|
||||
"""
|
||||
FFmpeg video processing plugin
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
from typing import Dict, Any, List
|
||||
import time
|
||||
|
||||
from .base import ServicePlugin, PluginResult
|
||||
from .exceptions import PluginExecutionError
|
||||
|
||||
|
||||
class FFmpegPlugin(ServicePlugin):
|
||||
"""Plugin for FFmpeg video processing"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.service_id = "ffmpeg"
|
||||
self.name = "FFmpeg Video Processing"
|
||||
self.version = "1.0.0"
|
||||
self.description = "Transcode and process video files using FFmpeg"
|
||||
self.capabilities = ["transcode", "resize", "compress", "convert"]
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Initialize FFmpeg dependencies"""
|
||||
# Check for ffmpeg installation
|
||||
try:
|
||||
subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
|
||||
self.ffmpeg_path = "ffmpeg"
|
||||
except (subprocess.CalledProcessError, FileNotFoundError):
|
||||
raise PluginExecutionError("FFmpeg not found. Install FFmpeg for video processing")
|
||||
|
||||
# Check for NVIDIA GPU support
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["ffmpeg", "-hide_banner", "-encoders"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True
|
||||
)
|
||||
self.gpu_acceleration = "h264_nvenc" in result.stdout
|
||||
except subprocess.CalledProcessError:
|
||||
self.gpu_acceleration = False
|
||||
|
||||
def validate_request(self, request: Dict[str, Any]) -> List[str]:
|
||||
"""Validate FFmpeg request parameters"""
|
||||
errors = []
|
||||
|
||||
# Check required parameters
|
||||
if "input_url" not in request and "input_file" not in request:
|
||||
errors.append("Either 'input_url' or 'input_file' must be provided")
|
||||
|
||||
# Validate output format
|
||||
output_format = request.get("output_format", "mp4")
|
||||
valid_formats = ["mp4", "avi", "mov", "mkv", "webm", "flv"]
|
||||
if output_format not in valid_formats:
|
||||
errors.append(f"Invalid output format. Must be one of: {', '.join(valid_formats)}")
|
||||
|
||||
# Validate codec
|
||||
codec = request.get("codec", "h264")
|
||||
valid_codecs = ["h264", "h265", "vp9", "av1", "mpeg4"]
|
||||
if codec not in valid_codecs:
|
||||
errors.append(f"Invalid codec. Must be one of: {', '.join(valid_codecs)}")
|
||||
|
||||
# Validate resolution
|
||||
resolution = request.get("resolution")
|
||||
if resolution:
|
||||
valid_resolutions = ["720p", "1080p", "1440p", "4K", "8K"]
|
||||
if resolution not in valid_resolutions:
|
||||
errors.append(f"Invalid resolution. Must be one of: {', '.join(valid_resolutions)}")
|
||||
|
||||
# Validate bitrate
|
||||
bitrate = request.get("bitrate")
|
||||
if bitrate:
|
||||
if not isinstance(bitrate, str) or not bitrate.endswith(("k", "M")):
|
||||
errors.append("Bitrate must end with 'k' or 'M' (e.g., '1000k', '5M')")
|
||||
|
||||
# Validate frame rate
|
||||
fps = request.get("fps")
|
||||
if fps:
|
||||
if not isinstance(fps, (int, float)) or fps < 1 or fps > 120:
|
||||
errors.append("FPS must be between 1 and 120")
|
||||
|
||||
return errors
|
||||
|
||||
def get_hardware_requirements(self) -> Dict[str, Any]:
|
||||
"""Get hardware requirements for FFmpeg"""
|
||||
return {
|
||||
"gpu": "optional",
|
||||
"vram_gb": 2,
|
||||
"ram_gb": 8,
|
||||
"storage_gb": 10
|
||||
}
|
||||
|
||||
async def execute(self, request: Dict[str, Any]) -> PluginResult:
|
||||
"""Execute FFmpeg processing"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Validate request
|
||||
errors = self.validate_request(request)
|
||||
if errors:
|
||||
return PluginResult(
|
||||
success=False,
|
||||
error=f"Validation failed: {'; '.join(errors)}"
|
||||
)
|
||||
|
||||
# Get parameters
|
||||
input_source = request.get("input_url") or request.get("input_file")
|
||||
output_format = request.get("output_format", "mp4")
|
||||
codec = request.get("codec", "h264")
|
||||
resolution = request.get("resolution")
|
||||
bitrate = request.get("bitrate")
|
||||
fps = request.get("fps")
|
||||
gpu_acceleration = request.get("gpu_acceleration", self.gpu_acceleration)
|
||||
|
||||
# Get input file
|
||||
input_file = await self._get_input_file(input_source)
|
||||
|
||||
# Build FFmpeg command
|
||||
cmd = self._build_ffmpeg_command(
|
||||
input_file=input_file,
|
||||
output_format=output_format,
|
||||
codec=codec,
|
||||
resolution=resolution,
|
||||
bitrate=bitrate,
|
||||
fps=fps,
|
||||
gpu_acceleration=gpu_acceleration
|
||||
)
|
||||
|
||||
# Execute FFmpeg
|
||||
output_file = await self._execute_ffmpeg(cmd)
|
||||
|
||||
# Get output file info
|
||||
output_info = await self._get_video_info(output_file)
|
||||
|
||||
# Clean up input file if downloaded
|
||||
if input_source != request.get("input_file"):
|
||||
os.unlink(input_file)
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
return PluginResult(
|
||||
success=True,
|
||||
data={
|
||||
"output_file": output_file,
|
||||
"output_info": output_info,
|
||||
"parameters": {
|
||||
"codec": codec,
|
||||
"resolution": resolution,
|
||||
"bitrate": bitrate,
|
||||
"fps": fps,
|
||||
"gpu_acceleration": gpu_acceleration
|
||||
}
|
||||
},
|
||||
metrics={
|
||||
"input_size": os.path.getsize(input_file),
|
||||
"output_size": os.path.getsize(output_file),
|
||||
"compression_ratio": os.path.getsize(output_file) / os.path.getsize(input_file),
|
||||
"processing_time": execution_time,
|
||||
"real_time_factor": output_info.get("duration", 0) / execution_time if execution_time > 0 else 0
|
||||
},
|
||||
execution_time=execution_time
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
return PluginResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
execution_time=time.time() - start_time
|
||||
)
|
||||
|
||||
async def _get_input_file(self, source: str) -> str:
|
||||
"""Get input file from URL or path"""
|
||||
if source.startswith(("http://", "https://")):
|
||||
# Download from URL
|
||||
import requests
|
||||
|
||||
response = requests.get(source, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# Save to temporary file
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as f:
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
return f.name
|
||||
else:
|
||||
# Local file
|
||||
if not os.path.exists(source):
|
||||
raise PluginExecutionError(f"Input file not found: {source}")
|
||||
return source
|
||||
|
||||
def _build_ffmpeg_command(
|
||||
self,
|
||||
input_file: str,
|
||||
output_format: str,
|
||||
codec: str,
|
||||
resolution: Optional[str],
|
||||
bitrate: Optional[str],
|
||||
fps: Optional[float],
|
||||
gpu_acceleration: bool
|
||||
) -> List[str]:
|
||||
"""Build FFmpeg command"""
|
||||
cmd = [self.ffmpeg_path, "-i", input_file]
|
||||
|
||||
# Add codec
|
||||
if gpu_acceleration and codec == "h264":
|
||||
cmd.extend(["-c:v", "h264_nvenc"])
|
||||
cmd.extend(["-preset", "fast"])
|
||||
elif gpu_acceleration and codec == "h265":
|
||||
cmd.extend(["-c:v", "hevc_nvenc"])
|
||||
cmd.extend(["-preset", "fast"])
|
||||
else:
|
||||
cmd.extend(["-c:v", codec])
|
||||
|
||||
# Add resolution
|
||||
if resolution:
|
||||
resolution_map = {
|
||||
"720p": ("1280", "720"),
|
||||
"1080p": ("1920", "1080"),
|
||||
"1440p": ("2560", "1440"),
|
||||
"4K": ("3840", "2160"),
|
||||
"8K": ("7680", "4320")
|
||||
}
|
||||
width, height = resolution_map.get(resolution, (None, None))
|
||||
if width and height:
|
||||
cmd.extend(["-s", f"{width}x{height}"])
|
||||
|
||||
# Add bitrate
|
||||
if bitrate:
|
||||
cmd.extend(["-b:v", bitrate])
|
||||
cmd.extend(["-b:a", "128k"]) # Audio bitrate
|
||||
|
||||
# Add FPS
|
||||
if fps:
|
||||
cmd.extend(["-r", str(fps)])
|
||||
|
||||
# Add audio codec
|
||||
cmd.extend(["-c:a", "aac"])
|
||||
|
||||
# Output file
|
||||
output_file = tempfile.mktemp(suffix=f".{output_format}")
|
||||
cmd.append(output_file)
|
||||
|
||||
return cmd
|
||||
|
||||
async def _execute_ffmpeg(self, cmd: List[str]) -> str:
|
||||
"""Execute FFmpeg command"""
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE
|
||||
)
|
||||
|
||||
stdout, stderr = await process.communicate()
|
||||
|
||||
if process.returncode != 0:
|
||||
error_msg = stderr.decode() if stderr else "FFmpeg failed"
|
||||
raise PluginExecutionError(f"FFmpeg error: {error_msg}")
|
||||
|
||||
# Output file is the last argument
|
||||
return cmd[-1]
|
||||
|
||||
async def _get_video_info(self, video_file: str) -> Dict[str, Any]:
|
||||
"""Get video file information"""
|
||||
cmd = [
|
||||
"ffprobe",
|
||||
"-v", "quiet",
|
||||
"-print_format", "json",
|
||||
"-show_format",
|
||||
"-show_streams",
|
||||
video_file
|
||||
]
|
||||
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE
|
||||
)
|
||||
|
||||
stdout, stderr = await process.communicate()
|
||||
|
||||
if process.returncode != 0:
|
||||
return {}
|
||||
|
||||
import json
|
||||
probe_data = json.loads(stdout.decode())
|
||||
|
||||
# Extract relevant info
|
||||
video_stream = next(
|
||||
(s for s in probe_data.get("streams", []) if s.get("codec_type") == "video"),
|
||||
{}
|
||||
)
|
||||
|
||||
return {
|
||||
"duration": float(probe_data.get("format", {}).get("duration", 0)),
|
||||
"size": int(probe_data.get("format", {}).get("size", 0)),
|
||||
"width": video_stream.get("width"),
|
||||
"height": video_stream.get("height"),
|
||||
"fps": eval(video_stream.get("r_frame_rate", "0/1")),
|
||||
"codec": video_stream.get("codec_name"),
|
||||
"bitrate": int(probe_data.get("format", {}).get("bit_rate", 0))
|
||||
}
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check FFmpeg health"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["ffmpeg", "-version"],
|
||||
capture_output=True,
|
||||
check=True
|
||||
)
|
||||
return True
|
||||
except subprocess.CalledProcessError:
|
||||
return False
|
||||
321
apps/miner-node/plugins/llm_inference.py
Normal file
321
apps/miner-node/plugins/llm_inference.py
Normal file
@ -0,0 +1,321 @@
|
||||
"""
|
||||
LLM inference plugin
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import Dict, Any, List, Optional
|
||||
import time
|
||||
|
||||
from .base import GPUPlugin, PluginResult
|
||||
from .exceptions import PluginExecutionError
|
||||
|
||||
|
||||
class LLMPlugin(GPUPlugin):
|
||||
"""Plugin for Large Language Model inference"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.service_id = "llm_inference"
|
||||
self.name = "LLM Inference"
|
||||
self.version = "1.0.0"
|
||||
self.description = "Run inference on large language models"
|
||||
self.capabilities = ["generate", "stream", "chat"]
|
||||
self._model_cache = {}
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Initialize LLM dependencies"""
|
||||
super().setup()
|
||||
|
||||
# Check for transformers installation
|
||||
try:
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
|
||||
self.transformers = AutoModelForCausalLM
|
||||
self.AutoTokenizer = AutoTokenizer
|
||||
self.pipeline = pipeline
|
||||
except ImportError:
|
||||
raise PluginExecutionError("Transformers not installed. Install with: pip install transformers accelerate")
|
||||
|
||||
# Check for torch
|
||||
try:
|
||||
import torch
|
||||
self.torch = torch
|
||||
except ImportError:
|
||||
raise PluginExecutionError("PyTorch not installed. Install with: pip install torch")
|
||||
|
||||
def validate_request(self, request: Dict[str, Any]) -> List[str]:
|
||||
"""Validate LLM request parameters"""
|
||||
errors = []
|
||||
|
||||
# Check required parameters
|
||||
if "prompt" not in request:
|
||||
errors.append("'prompt' is required")
|
||||
|
||||
# Validate model
|
||||
model = request.get("model", "llama-7b")
|
||||
valid_models = [
|
||||
"llama-7b",
|
||||
"llama-13b",
|
||||
"mistral-7b",
|
||||
"mixtral-8x7b",
|
||||
"gpt-3.5-turbo",
|
||||
"gpt-4"
|
||||
]
|
||||
if model not in valid_models:
|
||||
errors.append(f"Invalid model. Must be one of: {', '.join(valid_models)}")
|
||||
|
||||
# Validate max_tokens
|
||||
max_tokens = request.get("max_tokens", 256)
|
||||
if not isinstance(max_tokens, int) or max_tokens < 1 or max_tokens > 4096:
|
||||
errors.append("max_tokens must be an integer between 1 and 4096")
|
||||
|
||||
# Validate temperature
|
||||
temperature = request.get("temperature", 0.7)
|
||||
if not isinstance(temperature, (int, float)) or temperature < 0.0 or temperature > 2.0:
|
||||
errors.append("temperature must be between 0.0 and 2.0")
|
||||
|
||||
# Validate top_p
|
||||
top_p = request.get("top_p")
|
||||
if top_p is not None and (not isinstance(top_p, (int, float)) or top_p <= 0.0 or top_p > 1.0):
|
||||
errors.append("top_p must be between 0.0 and 1.0")
|
||||
|
||||
return errors
|
||||
|
||||
def get_hardware_requirements(self) -> Dict[str, Any]:
|
||||
"""Get hardware requirements for LLM inference"""
|
||||
return {
|
||||
"gpu": "recommended",
|
||||
"vram_gb": 8,
|
||||
"ram_gb": 16,
|
||||
"cuda": "recommended"
|
||||
}
|
||||
|
||||
async def execute(self, request: Dict[str, Any]) -> PluginResult:
|
||||
"""Execute LLM inference"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Validate request
|
||||
errors = self.validate_request(request)
|
||||
if errors:
|
||||
return PluginResult(
|
||||
success=False,
|
||||
error=f"Validation failed: {'; '.join(errors)}"
|
||||
)
|
||||
|
||||
# Get parameters
|
||||
prompt = request["prompt"]
|
||||
model_name = request.get("model", "llama-7b")
|
||||
max_tokens = request.get("max_tokens", 256)
|
||||
temperature = request.get("temperature", 0.7)
|
||||
top_p = request.get("top_p", 0.9)
|
||||
do_sample = request.get("do_sample", True)
|
||||
stream = request.get("stream", False)
|
||||
|
||||
# Load model and tokenizer
|
||||
model, tokenizer = await self._load_model(model_name)
|
||||
|
||||
# Generate response
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
if stream:
|
||||
# Streaming generation
|
||||
generator = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self._generate_streaming(
|
||||
model, tokenizer, prompt, max_tokens, temperature, top_p, do_sample
|
||||
)
|
||||
)
|
||||
|
||||
# Collect all tokens
|
||||
full_response = ""
|
||||
tokens = []
|
||||
for token in generator:
|
||||
tokens.append(token)
|
||||
full_response += token
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
return PluginResult(
|
||||
success=True,
|
||||
data={
|
||||
"text": full_response,
|
||||
"tokens": tokens,
|
||||
"streamed": True
|
||||
},
|
||||
metrics={
|
||||
"model": model_name,
|
||||
"prompt_tokens": len(tokenizer.encode(prompt)),
|
||||
"generated_tokens": len(tokens),
|
||||
"tokens_per_second": len(tokens) / execution_time if execution_time > 0 else 0
|
||||
},
|
||||
execution_time=execution_time
|
||||
)
|
||||
else:
|
||||
# Regular generation
|
||||
response = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self._generate(
|
||||
model, tokenizer, prompt, max_tokens, temperature, top_p, do_sample
|
||||
)
|
||||
)
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
return PluginResult(
|
||||
success=True,
|
||||
data={
|
||||
"text": response,
|
||||
"streamed": False
|
||||
},
|
||||
metrics={
|
||||
"model": model_name,
|
||||
"prompt_tokens": len(tokenizer.encode(prompt)),
|
||||
"generated_tokens": len(tokenizer.encode(response)) - len(tokenizer.encode(prompt)),
|
||||
"tokens_per_second": (len(tokenizer.encode(response)) - len(tokenizer.encode(prompt))) / execution_time if execution_time > 0 else 0
|
||||
},
|
||||
execution_time=execution_time
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
return PluginResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
execution_time=time.time() - start_time
|
||||
)
|
||||
|
||||
async def _load_model(self, model_name: str):
|
||||
"""Load LLM model and tokenizer with caching"""
|
||||
if model_name not in self._model_cache:
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
# Map model names to HuggingFace model IDs
|
||||
model_map = {
|
||||
"llama-7b": "meta-llama/Llama-2-7b-chat-hf",
|
||||
"llama-13b": "meta-llama/Llama-2-13b-chat-hf",
|
||||
"mistral-7b": "mistralai/Mistral-7B-Instruct-v0.1",
|
||||
"mixtral-8x7b": "mistralai/Mixtral-8x7B-Instruct-v0.1",
|
||||
"gpt-3.5-turbo": "openai-gpt", # Would need OpenAI API
|
||||
"gpt-4": "openai-gpt-4" # Would need OpenAI API
|
||||
}
|
||||
|
||||
hf_model = model_map.get(model_name, model_name)
|
||||
|
||||
# Load tokenizer
|
||||
tokenizer = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self.AutoTokenizer.from_pretrained(hf_model)
|
||||
)
|
||||
|
||||
# Load model
|
||||
device = "cuda" if self.torch.cuda.is_available() else "cpu"
|
||||
model = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self.transformers.from_pretrained(
|
||||
hf_model,
|
||||
torch_dtype=self.torch.float16 if device == "cuda" else self.torch.float32,
|
||||
device_map="auto" if device == "cuda" else None,
|
||||
load_in_4bit=True if device == "cuda" and self.vram_gb < 16 else False
|
||||
)
|
||||
)
|
||||
|
||||
self._model_cache[model_name] = (model, tokenizer)
|
||||
|
||||
return self._model_cache[model_name]
|
||||
|
||||
def _generate(
|
||||
self,
|
||||
model,
|
||||
tokenizer,
|
||||
prompt: str,
|
||||
max_tokens: int,
|
||||
temperature: float,
|
||||
top_p: float,
|
||||
do_sample: bool
|
||||
) -> str:
|
||||
"""Generate text without streaming"""
|
||||
inputs = tokenizer(prompt, return_tensors="pt")
|
||||
|
||||
if self.torch.cuda.is_available():
|
||||
inputs = {k: v.cuda() for k, v in inputs.items()}
|
||||
|
||||
with self.torch.no_grad():
|
||||
outputs = model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
do_sample=do_sample,
|
||||
pad_token_id=tokenizer.eos_token_id
|
||||
)
|
||||
|
||||
# Decode only the new tokens
|
||||
new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
|
||||
response = tokenizer.decode(new_tokens, skip_special_tokens=True)
|
||||
|
||||
return response
|
||||
|
||||
def _generate_streaming(
|
||||
self,
|
||||
model,
|
||||
tokenizer,
|
||||
prompt: str,
|
||||
max_tokens: int,
|
||||
temperature: float,
|
||||
top_p: float,
|
||||
do_sample: bool
|
||||
):
|
||||
"""Generate text with streaming"""
|
||||
inputs = tokenizer(prompt, return_tensors="pt")
|
||||
|
||||
if self.torch.cuda.is_available():
|
||||
inputs = {k: v.cuda() for k, v in inputs.items()}
|
||||
|
||||
# Simple streaming implementation
|
||||
# In production, you'd use model.generate with streamer
|
||||
with self.torch.no_grad():
|
||||
for i in range(max_tokens):
|
||||
outputs = model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=1,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
do_sample=do_sample,
|
||||
pad_token_id=tokenizer.eos_token_id
|
||||
)
|
||||
|
||||
new_token = outputs[0][-1:]
|
||||
text = tokenizer.decode(new_token, skip_special_tokens=True)
|
||||
|
||||
if text == tokenizer.eos_token:
|
||||
break
|
||||
|
||||
yield text
|
||||
|
||||
# Update inputs for next iteration
|
||||
inputs["input_ids"] = self.torch.cat([inputs["input_ids"], new_token], dim=1)
|
||||
if "attention_mask" in inputs:
|
||||
inputs["attention_mask"] = self.torch.cat([
|
||||
inputs["attention_mask"],
|
||||
self.torch.ones((1, 1), device=inputs["attention_mask"].device)
|
||||
], dim=1)
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check LLM health"""
|
||||
try:
|
||||
# Try to load a small model
|
||||
await self._load_model("mistral-7b")
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def cleanup(self) -> None:
|
||||
"""Cleanup resources"""
|
||||
# Move models to CPU and clear cache
|
||||
for model, _ in self._model_cache.values():
|
||||
if hasattr(model, 'to'):
|
||||
model.to("cpu")
|
||||
self._model_cache.clear()
|
||||
|
||||
# Clear GPU cache
|
||||
if self.torch.cuda.is_available():
|
||||
self.torch.cuda.empty_cache()
|
||||
138
apps/miner-node/plugins/registry.py
Normal file
138
apps/miner-node/plugins/registry.py
Normal file
@ -0,0 +1,138 @@
|
||||
"""
|
||||
Plugin registry for managing service plugins
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Type, Optional
|
||||
import importlib
|
||||
import inspect
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
from .base import ServicePlugin
|
||||
from .exceptions import PluginError, PluginNotFoundError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PluginRegistry:
|
||||
"""Registry for managing service plugins"""
|
||||
|
||||
def __init__(self):
|
||||
self._plugins: Dict[str, ServicePlugin] = {}
|
||||
self._plugin_classes: Dict[str, Type[ServicePlugin]] = {}
|
||||
self._loaded = False
|
||||
|
||||
def register(self, plugin_class: Type[ServicePlugin]) -> None:
|
||||
"""Register a plugin class"""
|
||||
plugin_id = getattr(plugin_class, "service_id", plugin_class.__name__)
|
||||
self._plugin_classes[plugin_id] = plugin_class
|
||||
logger.info(f"Registered plugin class: {plugin_id}")
|
||||
|
||||
def load_plugin(self, service_id: str) -> ServicePlugin:
|
||||
"""Load and instantiate a plugin"""
|
||||
if service_id not in self._plugin_classes:
|
||||
raise PluginNotFoundError(f"Plugin {service_id} not found")
|
||||
|
||||
if service_id in self._plugins:
|
||||
return self._plugins[service_id]
|
||||
|
||||
try:
|
||||
plugin_class = self._plugin_classes[service_id]
|
||||
plugin = plugin_class()
|
||||
plugin.setup()
|
||||
self._plugins[service_id] = plugin
|
||||
logger.info(f"Loaded plugin: {service_id}")
|
||||
return plugin
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load plugin {service_id}: {e}")
|
||||
raise PluginError(f"Failed to load plugin {service_id}: {e}")
|
||||
|
||||
def get_plugin(self, service_id: str) -> Optional[ServicePlugin]:
|
||||
"""Get loaded plugin"""
|
||||
return self._plugins.get(service_id)
|
||||
|
||||
def unload_plugin(self, service_id: str) -> None:
|
||||
"""Unload a plugin"""
|
||||
if service_id in self._plugins:
|
||||
plugin = self._plugins[service_id]
|
||||
plugin.cleanup()
|
||||
del self._plugins[service_id]
|
||||
logger.info(f"Unloaded plugin: {service_id}")
|
||||
|
||||
def list_plugins(self) -> List[str]:
|
||||
"""List all registered plugin IDs"""
|
||||
return list(self._plugin_classes.keys())
|
||||
|
||||
def list_loaded_plugins(self) -> List[str]:
|
||||
"""List all loaded plugin IDs"""
|
||||
return list(self._plugins.keys())
|
||||
|
||||
async def load_all_from_directory(self, plugin_dir: Path) -> None:
|
||||
"""Load all plugins from a directory"""
|
||||
if not plugin_dir.exists():
|
||||
logger.warning(f"Plugin directory does not exist: {plugin_dir}")
|
||||
return
|
||||
|
||||
for plugin_file in plugin_dir.glob("*.py"):
|
||||
if plugin_file.name.startswith("_"):
|
||||
continue
|
||||
|
||||
module_name = plugin_file.stem
|
||||
try:
|
||||
# Import the module
|
||||
spec = importlib.util.spec_from_file_location(module_name, plugin_file)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
# Find plugin classes in the module
|
||||
for name, obj in inspect.getmembers(module, inspect.isclass):
|
||||
if (issubclass(obj, ServicePlugin) and
|
||||
obj != ServicePlugin and
|
||||
not name.startswith("_")):
|
||||
self.register(obj)
|
||||
logger.info(f"Auto-registered plugin from {module_name}: {name}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load plugin from {plugin_file}: {e}")
|
||||
|
||||
async def initialize(self, plugin_dir: Optional[Path] = None) -> None:
|
||||
"""Initialize the plugin registry"""
|
||||
if self._loaded:
|
||||
return
|
||||
|
||||
# Load built-in plugins
|
||||
from . import whisper, stable_diffusion, llm_inference, ffmpeg, blender
|
||||
|
||||
self.register(whisper.WhisperPlugin)
|
||||
self.register(stable_diffusion.StableDiffusionPlugin)
|
||||
self.register(llm_inference.LLMPlugin)
|
||||
self.register(ffmpeg.FFmpegPlugin)
|
||||
self.register(blender.BlenderPlugin)
|
||||
|
||||
# Load external plugins if directory provided
|
||||
if plugin_dir:
|
||||
await self.load_all_from_directory(plugin_dir)
|
||||
|
||||
self._loaded = True
|
||||
logger.info(f"Plugin registry initialized with {len(self._plugin_classes)} plugins")
|
||||
|
||||
async def health_check_all(self) -> Dict[str, bool]:
|
||||
"""Health check all loaded plugins"""
|
||||
results = {}
|
||||
for service_id, plugin in self._plugins.items():
|
||||
try:
|
||||
results[service_id] = await plugin.health_check()
|
||||
except Exception as e:
|
||||
logger.error(f"Health check failed for {service_id}: {e}")
|
||||
results[service_id] = False
|
||||
return results
|
||||
|
||||
def cleanup_all(self) -> None:
|
||||
"""Cleanup all loaded plugins"""
|
||||
for service_id in list(self._plugins.keys()):
|
||||
self.unload_plugin(service_id)
|
||||
logger.info("All plugins cleaned up")
|
||||
|
||||
|
||||
# Global registry instance
|
||||
registry = PluginRegistry()
|
||||
281
apps/miner-node/plugins/stable_diffusion.py
Normal file
281
apps/miner-node/plugins/stable_diffusion.py
Normal file
@ -0,0 +1,281 @@
|
||||
"""
|
||||
Stable Diffusion image generation plugin
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
from typing import Dict, Any, List
|
||||
import time
|
||||
import numpy as np
|
||||
|
||||
from .base import GPUPlugin, PluginResult
|
||||
from .exceptions import PluginExecutionError
|
||||
|
||||
|
||||
class StableDiffusionPlugin(GPUPlugin):
|
||||
"""Plugin for Stable Diffusion image generation"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.service_id = "stable_diffusion"
|
||||
self.name = "Stable Diffusion"
|
||||
self.version = "1.0.0"
|
||||
self.description = "Generate images from text prompts using Stable Diffusion"
|
||||
self.capabilities = ["txt2img", "img2img"]
|
||||
self._model_cache = {}
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Initialize Stable Diffusion dependencies"""
|
||||
super().setup()
|
||||
|
||||
# Check for diffusers installation
|
||||
try:
|
||||
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
|
||||
self.diffusers = StableDiffusionPipeline
|
||||
self.img2img_pipe = StableDiffusionImg2ImgPipeline
|
||||
except ImportError:
|
||||
raise PluginExecutionError("Diffusers not installed. Install with: pip install diffusers transformers accelerate")
|
||||
|
||||
# Check for torch
|
||||
try:
|
||||
import torch
|
||||
self.torch = torch
|
||||
except ImportError:
|
||||
raise PluginExecutionError("PyTorch not installed. Install with: pip install torch")
|
||||
|
||||
# Check for PIL
|
||||
try:
|
||||
from PIL import Image
|
||||
self.Image = Image
|
||||
except ImportError:
|
||||
raise PluginExecutionError("PIL not installed. Install with: pip install Pillow")
|
||||
|
||||
def validate_request(self, request: Dict[str, Any]) -> List[str]:
|
||||
"""Validate Stable Diffusion request parameters"""
|
||||
errors = []
|
||||
|
||||
# Check required parameters
|
||||
if "prompt" not in request:
|
||||
errors.append("'prompt' is required")
|
||||
|
||||
# Validate model
|
||||
model = request.get("model", "runwayml/stable-diffusion-v1-5")
|
||||
valid_models = [
|
||||
"runwayml/stable-diffusion-v1-5",
|
||||
"stabilityai/stable-diffusion-2-1",
|
||||
"stabilityai/stable-diffusion-xl-base-1.0"
|
||||
]
|
||||
if model not in valid_models:
|
||||
errors.append(f"Invalid model. Must be one of: {', '.join(valid_models)}")
|
||||
|
||||
# Validate dimensions
|
||||
width = request.get("width", 512)
|
||||
height = request.get("height", 512)
|
||||
|
||||
if not isinstance(width, int) or width < 256 or width > 1024:
|
||||
errors.append("Width must be an integer between 256 and 1024")
|
||||
if not isinstance(height, int) or height < 256 or height > 1024:
|
||||
errors.append("Height must be an integer between 256 and 1024")
|
||||
|
||||
# Validate steps
|
||||
steps = request.get("steps", 20)
|
||||
if not isinstance(steps, int) or steps < 1 or steps > 100:
|
||||
errors.append("Steps must be an integer between 1 and 100")
|
||||
|
||||
# Validate guidance scale
|
||||
guidance_scale = request.get("guidance_scale", 7.5)
|
||||
if not isinstance(guidance_scale, (int, float)) or guidance_scale < 1.0 or guidance_scale > 20.0:
|
||||
errors.append("Guidance scale must be between 1.0 and 20.0")
|
||||
|
||||
# Check img2img requirements
|
||||
if request.get("task") == "img2img":
|
||||
if "init_image" not in request:
|
||||
errors.append("'init_image' is required for img2img task")
|
||||
strength = request.get("strength", 0.8)
|
||||
if not isinstance(strength, (int, float)) or strength < 0.0 or strength > 1.0:
|
||||
errors.append("Strength must be between 0.0 and 1.0")
|
||||
|
||||
return errors
|
||||
|
||||
def get_hardware_requirements(self) -> Dict[str, Any]:
|
||||
"""Get hardware requirements for Stable Diffusion"""
|
||||
return {
|
||||
"gpu": "required",
|
||||
"vram_gb": 6,
|
||||
"ram_gb": 8,
|
||||
"cuda": "required"
|
||||
}
|
||||
|
||||
async def execute(self, request: Dict[str, Any]) -> PluginResult:
|
||||
"""Execute Stable Diffusion generation"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Validate request
|
||||
errors = self.validate_request(request)
|
||||
if errors:
|
||||
return PluginResult(
|
||||
success=False,
|
||||
error=f"Validation failed: {'; '.join(errors)}"
|
||||
)
|
||||
|
||||
# Get parameters
|
||||
prompt = request["prompt"]
|
||||
negative_prompt = request.get("negative_prompt", "")
|
||||
model_name = request.get("model", "runwayml/stable-diffusion-v1-5")
|
||||
width = request.get("width", 512)
|
||||
height = request.get("height", 512)
|
||||
steps = request.get("steps", 20)
|
||||
guidance_scale = request.get("guidance_scale", 7.5)
|
||||
num_images = request.get("num_images", 1)
|
||||
seed = request.get("seed")
|
||||
task = request.get("task", "txt2img")
|
||||
|
||||
# Load model
|
||||
pipe = await self._load_model(model_name)
|
||||
|
||||
# Generate images
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
if task == "img2img":
|
||||
# Handle img2img
|
||||
init_image_data = request["init_image"]
|
||||
init_image = self._decode_image(init_image_data)
|
||||
strength = request.get("strength", 0.8)
|
||||
|
||||
images = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: pipe(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
image=init_image,
|
||||
strength=strength,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=guidance_scale,
|
||||
num_images_per_prompt=num_images,
|
||||
generator=self._get_generator(seed)
|
||||
).images
|
||||
)
|
||||
else:
|
||||
# Handle txt2img
|
||||
images = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: pipe(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=guidance_scale,
|
||||
num_images_per_prompt=num_images,
|
||||
generator=self._get_generator(seed)
|
||||
).images
|
||||
)
|
||||
|
||||
# Encode images to base64
|
||||
encoded_images = []
|
||||
for img in images:
|
||||
buffer = io.BytesIO()
|
||||
img.save(buffer, format="PNG")
|
||||
encoded_images.append(base64.b64encode(buffer.getvalue()).decode())
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
return PluginResult(
|
||||
success=True,
|
||||
data={
|
||||
"images": encoded_images,
|
||||
"count": len(images),
|
||||
"parameters": {
|
||||
"prompt": prompt,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"steps": steps,
|
||||
"guidance_scale": guidance_scale,
|
||||
"seed": seed
|
||||
}
|
||||
},
|
||||
metrics={
|
||||
"model": model_name,
|
||||
"task": task,
|
||||
"images_generated": len(images),
|
||||
"generation_time": execution_time,
|
||||
"time_per_image": execution_time / len(images)
|
||||
},
|
||||
execution_time=execution_time
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
return PluginResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
execution_time=time.time() - start_time
|
||||
)
|
||||
|
||||
async def _load_model(self, model_name: str):
|
||||
"""Load Stable Diffusion model with caching"""
|
||||
if model_name not in self._model_cache:
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
# Determine device
|
||||
device = "cuda" if self.torch.cuda.is_available() else "cpu"
|
||||
|
||||
# Load with attention slicing for memory efficiency
|
||||
pipe = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self.diffusers.from_pretrained(
|
||||
model_name,
|
||||
torch_dtype=self.torch.float16 if device == "cuda" else self.torch.float32,
|
||||
safety_checker=None,
|
||||
requires_safety_checker=False
|
||||
)
|
||||
)
|
||||
|
||||
pipe = pipe.to(device)
|
||||
|
||||
# Enable memory optimizations
|
||||
if device == "cuda":
|
||||
pipe.enable_attention_slicing()
|
||||
if self.vram_gb < 8:
|
||||
pipe.enable_model_cpu_offload()
|
||||
|
||||
self._model_cache[model_name] = pipe
|
||||
|
||||
return self._model_cache[model_name]
|
||||
|
||||
def _decode_image(self, image_data: str) -> 'Image':
|
||||
"""Decode base64 image"""
|
||||
if image_data.startswith('data:image'):
|
||||
# Remove data URL prefix
|
||||
image_data = image_data.split(',')[1]
|
||||
|
||||
image_bytes = base64.b64decode(image_data)
|
||||
return self.Image.open(io.BytesIO(image_bytes))
|
||||
|
||||
def _get_generator(self, seed: Optional[int]):
|
||||
"""Get torch generator for reproducible results"""
|
||||
if seed is not None:
|
||||
return self.torch.Generator().manual_seed(seed)
|
||||
return None
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check Stable Diffusion health"""
|
||||
try:
|
||||
# Try to load a small model
|
||||
pipe = await self._load_model("runwayml/stable-diffusion-v1-5")
|
||||
return pipe is not None
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def cleanup(self) -> None:
|
||||
"""Cleanup resources"""
|
||||
# Move models to CPU and clear cache
|
||||
for pipe in self._model_cache.values():
|
||||
if hasattr(pipe, 'to'):
|
||||
pipe.to("cpu")
|
||||
self._model_cache.clear()
|
||||
|
||||
# Clear GPU cache
|
||||
if self.torch.cuda.is_available():
|
||||
self.torch.cuda.empty_cache()
|
||||
215
apps/miner-node/plugins/whisper.py
Normal file
215
apps/miner-node/plugins/whisper.py
Normal file
@ -0,0 +1,215 @@
|
||||
"""
|
||||
Whisper speech recognition plugin
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import tempfile
|
||||
from typing import Dict, Any, List
|
||||
import time
|
||||
|
||||
from .base import GPUPlugin, PluginResult
|
||||
from .exceptions import PluginExecutionError
|
||||
|
||||
|
||||
class WhisperPlugin(GPUPlugin):
|
||||
"""Plugin for Whisper speech recognition"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.service_id = "whisper"
|
||||
self.name = "Whisper Speech Recognition"
|
||||
self.version = "1.0.0"
|
||||
self.description = "Transcribe and translate audio files using OpenAI Whisper"
|
||||
self.capabilities = ["transcribe", "translate"]
|
||||
self._model_cache = {}
|
||||
|
||||
def setup(self) -> None:
|
||||
"""Initialize Whisper dependencies"""
|
||||
super().setup()
|
||||
|
||||
# Check for whisper installation
|
||||
try:
|
||||
import whisper
|
||||
self.whisper = whisper
|
||||
except ImportError:
|
||||
raise PluginExecutionError("Whisper not installed. Install with: pip install openai-whisper")
|
||||
|
||||
# Check for ffmpeg
|
||||
import subprocess
|
||||
try:
|
||||
subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
|
||||
except (subprocess.CalledProcessError, FileNotFoundError):
|
||||
raise PluginExecutionError("FFmpeg not found. Install FFmpeg for audio processing")
|
||||
|
||||
def validate_request(self, request: Dict[str, Any]) -> List[str]:
|
||||
"""Validate Whisper request parameters"""
|
||||
errors = []
|
||||
|
||||
# Check required parameters
|
||||
if "audio_url" not in request and "audio_file" not in request:
|
||||
errors.append("Either 'audio_url' or 'audio_file' must be provided")
|
||||
|
||||
# Validate model
|
||||
model = request.get("model", "base")
|
||||
valid_models = ["tiny", "base", "small", "medium", "large", "large-v2", "large-v3"]
|
||||
if model not in valid_models:
|
||||
errors.append(f"Invalid model. Must be one of: {', '.join(valid_models)}")
|
||||
|
||||
# Validate task
|
||||
task = request.get("task", "transcribe")
|
||||
if task not in ["transcribe", "translate"]:
|
||||
errors.append("Task must be 'transcribe' or 'translate'")
|
||||
|
||||
# Validate language
|
||||
if "language" in request:
|
||||
language = request["language"]
|
||||
if not isinstance(language, str) or len(language) != 2:
|
||||
errors.append("Language must be a 2-letter language code (e.g., 'en', 'es')")
|
||||
|
||||
return errors
|
||||
|
||||
def get_hardware_requirements(self) -> Dict[str, Any]:
|
||||
"""Get hardware requirements for Whisper"""
|
||||
return {
|
||||
"gpu": "recommended",
|
||||
"vram_gb": 2,
|
||||
"ram_gb": 4,
|
||||
"storage_gb": 1
|
||||
}
|
||||
|
||||
async def execute(self, request: Dict[str, Any]) -> PluginResult:
|
||||
"""Execute Whisper transcription"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Validate request
|
||||
errors = self.validate_request(request)
|
||||
if errors:
|
||||
return PluginResult(
|
||||
success=False,
|
||||
error=f"Validation failed: {'; '.join(errors)}"
|
||||
)
|
||||
|
||||
# Get parameters
|
||||
model_name = request.get("model", "base")
|
||||
task = request.get("task", "transcribe")
|
||||
language = request.get("language")
|
||||
temperature = request.get("temperature", 0.0)
|
||||
|
||||
# Load or get cached model
|
||||
model = await self._load_model(model_name)
|
||||
|
||||
# Get audio file
|
||||
audio_path = await self._get_audio_file(request)
|
||||
|
||||
# Transcribe
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
if task == "translate":
|
||||
result = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: model.transcribe(
|
||||
audio_path,
|
||||
task="translate",
|
||||
temperature=temperature
|
||||
)
|
||||
)
|
||||
else:
|
||||
result = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: model.transcribe(
|
||||
audio_path,
|
||||
language=language,
|
||||
temperature=temperature
|
||||
)
|
||||
)
|
||||
|
||||
# Clean up
|
||||
if audio_path != request.get("audio_file"):
|
||||
os.unlink(audio_path)
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
return PluginResult(
|
||||
success=True,
|
||||
data={
|
||||
"text": result["text"],
|
||||
"language": result.get("language"),
|
||||
"segments": result.get("segments", [])
|
||||
},
|
||||
metrics={
|
||||
"model": model_name,
|
||||
"task": task,
|
||||
"audio_duration": result.get("duration"),
|
||||
"processing_time": execution_time,
|
||||
"real_time_factor": result.get("duration", 0) / execution_time if execution_time > 0 else 0
|
||||
},
|
||||
execution_time=execution_time
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
return PluginResult(
|
||||
success=False,
|
||||
error=str(e),
|
||||
execution_time=time.time() - start_time
|
||||
)
|
||||
|
||||
async def _load_model(self, model_name: str):
|
||||
"""Load Whisper model with caching"""
|
||||
if model_name not in self._model_cache:
|
||||
loop = asyncio.get_event_loop()
|
||||
model = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self.whisper.load_model(model_name)
|
||||
)
|
||||
self._model_cache[model_name] = model
|
||||
|
||||
return self._model_cache[model_name]
|
||||
|
||||
async def _get_audio_file(self, request: Dict[str, Any]) -> str:
|
||||
"""Get audio file from URL or direct file path"""
|
||||
if "audio_file" in request:
|
||||
return request["audio_file"]
|
||||
|
||||
# Download from URL
|
||||
audio_url = request["audio_url"]
|
||||
|
||||
# Use requests to download
|
||||
import requests
|
||||
|
||||
response = requests.get(audio_url, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# Save to temporary file
|
||||
suffix = self._get_audio_suffix(audio_url)
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as f:
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
return f.name
|
||||
|
||||
def _get_audio_suffix(self, url: str) -> str:
|
||||
"""Get file extension from URL"""
|
||||
if url.endswith('.mp3'):
|
||||
return '.mp3'
|
||||
elif url.endswith('.wav'):
|
||||
return '.wav'
|
||||
elif url.endswith('.m4a'):
|
||||
return '.m4a'
|
||||
elif url.endswith('.flac'):
|
||||
return '.flac'
|
||||
else:
|
||||
return '.mp3' # Default
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check Whisper health"""
|
||||
try:
|
||||
# Check if we can load the tiny model
|
||||
await self._load_model("tiny")
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def cleanup(self) -> None:
|
||||
"""Cleanup resources"""
|
||||
self._model_cache.clear()
|
||||
Reference in New Issue
Block a user