diff --git a/src/db_connector.py b/src/db_connector.py
index 08b16c1..9d62166 100644
--- a/src/db_connector.py
+++ b/src/db_connector.py
@@ -33,25 +33,30 @@ class DBConnector:
         pool_min: int = 5,
         pool_recycle: int = 3600,
     ):
-        self._pool = PooledDB(
-            creator=pymysql,
-            maxconnections=pool_size,
-            mincached=pool_min,
-            maxcached=pool_size,
-            maxusage=200,
-            blocking=True,
-            max_idle_time=pool_recycle,
-            connection_timeout=10,
-            charset="utf8mb4",
-            cursorclass=pymysql.cursors.DictCursor,
-            host=host,
-            port=port,
-            database=database,
-            user=user,
-            password=password,
-            read_timeout=30,
-            write_timeout=30,
-        )
+        # PooledDB parameters
+        pool_config = {
+            "creator": pymysql,
+            "maxconnections": pool_size,
+            "mincached": pool_min,
+            "maxcached": pool_size,
+            "maxusage": 200,
+            "blocking": True,
+        }
+        
+        # PyMySQL connection parameters  
+        connection_params = {
+            "host": host,
+            "port": port,
+            "database": database,
+            "user": user,
+            "password": password,
+            "charset": "utf8mb4",
+            "cursorclass": pymysql.cursors.DictCursor,
+            "read_timeout": 30,
+            "write_timeout": 30,
+        }
+        
+        self._pool = PooledDB(**pool_config, **connection_params)
         logger.info(
             "DB pool initialized: host=%s db=%s pool_size=%d min=%d",
             host, database, pool_size, pool_min,
diff --git a/src/logging_config.py b/src/logging_config.py
index 050f212..93d6cf6 100644
--- a/src/logging_config.py
+++ b/src/logging_config.py
@@ -14,8 +14,6 @@ import logging.handlers
 import sys
 from datetime import datetime, timezone
 
-from python_json_logger import json_formatter
-
 
 def setup_logging(
     level: str = "INFO",
diff --git a/worker/Dockerfile b/worker/Dockerfile
index 3ba6fa5..5f35ff6 100644
--- a/worker/Dockerfile
+++ b/worker/Dockerfile
@@ -1,6 +1,5 @@
 # Base image: CUDA 11.8 runtime on Ubuntu 22.04
-#FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04
-FROM nvidia/cuda:13.3.1-cudnn-runtime-ubuntu22.04
+FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04
 
 # Avoid interactive prompts during build
 ENV DEBIAN_FRONTEND=noninteractive
@@ -36,7 +35,8 @@ RUN groupadd -g 1000 appuser && \
 WORKDIR /app
 
 # Create necessary directories
-RUN mkdir -p /scratch /models /data/training /data/output /logs
+RUN mkdir -p /scratch /models /data/training /data/output /logs && \
+    chown -R appuser:appuser /app /scratch /models /data /logs
 
 # Copy requirements first for better caching
 COPY worker/requirements.txt /app/requirements.txt