Done setup and integration redis and postgresql
This commit is contained in:
parent
37c2e2a4d4
commit
a1d358aead
3 changed files with 189 additions and 18 deletions
|
@ -12,11 +12,40 @@ import time
|
|||
import uuid
|
||||
from ultralytics import YOLO
|
||||
from urllib.parse import urlparse
|
||||
from .database import DatabaseManager
|
||||
|
||||
# Create a logger specifically for this module
|
||||
logger = logging.getLogger("detector_worker.pympta")
|
||||
|
||||
def load_pipeline_node(node_config: dict, mpta_dir: str, redis_client) -> dict:
|
||||
def validate_redis_config(redis_config: dict) -> bool:
|
||||
"""Validate Redis configuration parameters."""
|
||||
required_fields = ["host", "port"]
|
||||
for field in required_fields:
|
||||
if field not in redis_config:
|
||||
logger.error(f"Missing required Redis config field: {field}")
|
||||
return False
|
||||
|
||||
if not isinstance(redis_config["port"], int) or redis_config["port"] <= 0:
|
||||
logger.error(f"Invalid Redis port: {redis_config['port']}")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def validate_postgresql_config(pg_config: dict) -> bool:
|
||||
"""Validate PostgreSQL configuration parameters."""
|
||||
required_fields = ["host", "port", "database", "username", "password"]
|
||||
for field in required_fields:
|
||||
if field not in pg_config:
|
||||
logger.error(f"Missing required PostgreSQL config field: {field}")
|
||||
return False
|
||||
|
||||
if not isinstance(pg_config["port"], int) or pg_config["port"] <= 0:
|
||||
logger.error(f"Invalid PostgreSQL port: {pg_config['port']}")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def load_pipeline_node(node_config: dict, mpta_dir: str, redis_client, db_manager=None) -> dict:
|
||||
# Recursively load a model node from configuration.
|
||||
model_path = os.path.join(mpta_dir, node_config["modelFile"])
|
||||
if not os.path.exists(model_path):
|
||||
|
@ -46,16 +75,22 @@ def load_pipeline_node(node_config: dict, mpta_dir: str, redis_client) -> dict:
|
|||
"triggerClasses": trigger_classes,
|
||||
"triggerClassIndices": trigger_class_indices,
|
||||
"crop": node_config.get("crop", False),
|
||||
"cropClass": node_config.get("cropClass"),
|
||||
"minConfidence": node_config.get("minConfidence", None),
|
||||
"multiClass": node_config.get("multiClass", False),
|
||||
"expectedClasses": node_config.get("expectedClasses", []),
|
||||
"parallel": node_config.get("parallel", False),
|
||||
"actions": node_config.get("actions", []),
|
||||
"parallelActions": node_config.get("parallelActions", []),
|
||||
"model": model,
|
||||
"branches": [],
|
||||
"redis_client": redis_client
|
||||
"redis_client": redis_client,
|
||||
"db_manager": db_manager
|
||||
}
|
||||
logger.debug(f"Configured node {node_config['modelId']} with trigger classes: {node['triggerClasses']}")
|
||||
for child in node_config.get("branches", []):
|
||||
logger.debug(f"Loading branch for parent node {node_config['modelId']}")
|
||||
node["branches"].append(load_pipeline_node(child, mpta_dir, redis_client))
|
||||
node["branches"].append(load_pipeline_node(child, mpta_dir, redis_client, db_manager))
|
||||
return node
|
||||
|
||||
def load_pipeline_from_zip(zip_source: str, target_dir: str) -> dict:
|
||||
|
@ -168,21 +203,42 @@ def load_pipeline_from_zip(zip_source: str, target_dir: str) -> dict:
|
|||
redis_client = None
|
||||
if "redis" in pipeline_config:
|
||||
redis_config = pipeline_config["redis"]
|
||||
try:
|
||||
redis_client = redis.Redis(
|
||||
host=redis_config["host"],
|
||||
port=redis_config["port"],
|
||||
password=redis_config.get("password"),
|
||||
db=redis_config.get("db", 0),
|
||||
decode_responses=True
|
||||
)
|
||||
redis_client.ping()
|
||||
logger.info(f"Successfully connected to Redis at {redis_config['host']}:{redis_config['port']}")
|
||||
except redis.exceptions.ConnectionError as e:
|
||||
logger.error(f"Failed to connect to Redis: {e}")
|
||||
redis_client = None
|
||||
if not validate_redis_config(redis_config):
|
||||
logger.error("Invalid Redis configuration, skipping Redis connection")
|
||||
else:
|
||||
try:
|
||||
redis_client = redis.Redis(
|
||||
host=redis_config["host"],
|
||||
port=redis_config["port"],
|
||||
password=redis_config.get("password"),
|
||||
db=redis_config.get("db", 0),
|
||||
decode_responses=True
|
||||
)
|
||||
redis_client.ping()
|
||||
logger.info(f"Successfully connected to Redis at {redis_config['host']}:{redis_config['port']}")
|
||||
except redis.exceptions.ConnectionError as e:
|
||||
logger.error(f"Failed to connect to Redis: {e}")
|
||||
redis_client = None
|
||||
|
||||
return load_pipeline_node(pipeline_config["pipeline"], mpta_dir, redis_client)
|
||||
# Establish PostgreSQL connection if configured
|
||||
db_manager = None
|
||||
if "postgresql" in pipeline_config:
|
||||
pg_config = pipeline_config["postgresql"]
|
||||
if not validate_postgresql_config(pg_config):
|
||||
logger.error("Invalid PostgreSQL configuration, skipping database connection")
|
||||
else:
|
||||
try:
|
||||
db_manager = DatabaseManager(pg_config)
|
||||
if db_manager.connect():
|
||||
logger.info(f"Successfully connected to PostgreSQL at {pg_config['host']}:{pg_config['port']}")
|
||||
else:
|
||||
logger.error("Failed to connect to PostgreSQL")
|
||||
db_manager = None
|
||||
except Exception as e:
|
||||
logger.error(f"Error initializing PostgreSQL connection: {e}")
|
||||
db_manager = None
|
||||
|
||||
return load_pipeline_node(pipeline_config["pipeline"], mpta_dir, redis_client, db_manager)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"Error parsing pipeline.json: {str(e)}", exc_info=True)
|
||||
return None
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue