Database Integration

DBManager Overview

The DBManager class (smocs/db/mysql_api_v0.py) provides all database operations for agents. It handles connection management, data storage, querying, and sampling.

Key capabilities:

• Store sensor data with timestamps
• Store predictions/actions
• Store SARSA tuples for RL agents
• Sample batches for training
• Register agent metadata
• Query database size

Initialization

In Agent Threads

Each thread creates its own DBManager instance:

from smocs.db.mysql_api_v0 import DBManager

def _setup_db_connection(self) -> DBManager:
    """Setup database connection for this thread."""
    db_config = {
        'agent_id': self.agent_id,
        'host': os.environ.get('MYSQL_HOST', 'localhost'),
        'port': int(os.environ.get('MYSQL_PORT', 3306)),
        'user': os.environ.get('MYSQL_USER', 'root'),
        'pwd': os.environ['MYSQL_ROOT_PASSWORD'],
        'database': os.environ.get('MYSQL_DATABASE', 'agentdb')
    }
    return DBManager(db_config)

Usage in thread:

self.db_manager = self._setup_db_connection()

Configuration

Required environment variables:

• MYSQL_HOST: Database host (usually 'localhost' in container)
• MYSQL_PORT: Database port (3306)
• MYSQL_USER: Database user ('root')
• MYSQL_ROOT_PASSWORD: Password (from .env)
• MYSQL_DATABASE: Database name ('agentdb')

Database Tables

agent_information

Stores agent metadata and registration:

CREATE TABLE agent_information (
    id INT AUTO_INCREMENT PRIMARY KEY,
    registered_id VARCHAR(50) NOT NULL,
    agent_name VARCHAR(50),
    config BLOB,
    info BLOB
);

Usage: Managed by AgentBase, not directly by threads

agent_inferences

Stores sensor data and predictions:

CREATE TABLE agent_inferences (
    id INT AUTO_INCREMENT PRIMARY KEY,
    state_source_timestamp DATETIME(6) NOT NULL,
    state_received_timestamp DATETIME(6) NOT NULL,
    state BLOB NOT NULL,
    prediction_timestamp DATETIME(6),
    prediction BLOB
);

Fields:

• state_source_timestamp: Timestamp from original sensor message
• state_received_timestamp: When agent received the data
• state: Sensor values as numpy array (serialized to BLOB)
• prediction: Model output as numpy array (optional)

agent_replay

Stores SARSA tuples for RL agents:

CREATE TABLE agent_replay (
    id INT AUTO_INCREMENT PRIMARY KEY,
    state_id INT NOT NULL,
    action_success BOOL,
    reward BLOB NOT NULL,
    next_state_source_timestamp DATETIME(6) NOT NULL,
    next_state_received_timestamp DATETIME(6) NOT NULL,
    next_state BLOB NOT NULL,
    terminate BOOL NOT NULL,
    truncate BOOL NOT NULL,
    info BLOB,
    FOREIGN KEY (state_id) REFERENCES agent_inferences(id)
);

Relationship: Links to agent_inferences via state_id

Core Operations

Store Sensor Data

Used by: Data Ingest Thread

def record_sensor_data(self, data: Dict[str, Any]) -> int:
    """
    Store sensor readings to agent_inferences table.
    
    Args:
        data: Dictionary with keys:
            - 'state_source_timestamp': str, format 'YYYY-MM-DD HH:MM:SS.ffffff'
            - 'state_received_timestamp': str, format 'YYYY-MM-DD HH:MM:SS.ffffff'
            - 'state': numpy array of sensor values
            
    Returns:
        0 on success, 1 on failure
    """

Example:

from datetime import datetime
import numpy as np

timestamp = datetime.fromtimestamp(message_data['timestamp'])
sensor_values = np.array([1.23, 4.56, 7.89], dtype=np.float32)

data = {
    'state_source_timestamp': timestamp.strftime('%Y-%m-%d %H:%M:%S.%f'),
    'state_received_timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f'),
    'state': sensor_values
}

status = self.db_manager.record_sensor_data(data)
if status == 0:
    logging.info("Data stored successfully")

Important:

• Timestamps must be strings in exact format
• state must be numpy array (converted to BLOB automatically)
• Returns 0 for success, 1 for failure

Store Prediction

Used by: ML Inference Thread

def record_prediction(self, prediction: np.ndarray, prediction_timestamp: str, 
                     key_value: str, key: str = "state_source_timestamp") -> int:
    """
    Update agent_inferences with prediction result.
    
    Args:
        prediction: Numpy array of model output
        prediction_timestamp: str, format 'YYYY-MM-DD HH:MM:SS.ffffff'
        key_value: Value to match (timestamp or state_id)
        key: Column to match on ('state_source_timestamp' or 'state_id')
        
    Returns:
        0 on success, 1 on failure
    """

Example:

# After performing inference
reconstructed = self.model.predict(input_data)

status = self.db_manager.record_prediction(
    prediction=reconstructed,
    prediction_timestamp=datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f'),
    key_value=source_timestamp_str,
    key='state_source_timestamp'
)

Use cases:

• Store autoencoder reconstructions
• Store classifier predictions
• Store any model output

Get State ID

Used by: RL Control Agents

def get_state_id(self, source_timestamp: str) -> Optional[int]:
    """
    Retrieve state ID from timestamp.
    
    Args:
        source_timestamp: Timestamp string
        
    Returns:
        Integer ID or None if not found
    """

Example:

# Get ID for storing SARSA tuple
state_id = self.db_manager.get_state_id(source_timestamp)
if state_id is None:
    logging.error("State not found in database")
    return

Store SARSA Tuple

Used by: RL Control Agent Data Ingest Thread

def record_controls_tuple(self, data: Dict[str, Any], state_id: int) -> int:
    """
    Store RL transition to agent_replay table.
    
    Args:
        data: Dictionary with keys:
            - 'next_state': numpy array
            - 'reward': numpy array (can be scalar wrapped in array)
            - 'terminate': bool
            - 'truncate': bool
            - 'info': dict (optional)
        state_id: Foreign key to agent_inferences
        
    Returns:
        0 on success, 1 on failure
    """

Example:

# Parse SARSA from Kafka
state, action, reward, next_state, terminated, truncated = sarsa_tuple

# Get state_id for current state
state_id = self.db_manager.get_state_id(state_timestamp)

# Store SARSA tuple
sarsa_data = {
    'next_state': next_state,
    'reward': np.array([reward], dtype=np.float32),
    'terminate': terminated,
    'truncate': truncated,
    'info': {}
}

status = self.db_manager.record_controls_tuple(sarsa_data, state_id)

Query Operations

Get Database Size

Used by: All training threads

def get_size(self, table_name: str) -> int:
    """
    Get row count for table.
    
    Args:
        table_name: 'agent_inferences', 'agent_replay', or 'agent_information'
        
    Returns:
        Number of rows in table
    """

Example:

total_samples = self.db_manager.get_size("agent_inferences")
logging.info(f"Database contains {total_samples} samples")

if total_samples < self.min_training_samples:
    return None  # Not enough data to train

Sample Training Batch

Used by: ML Training Thread

def sample_batch(self, batch_size: int, segment_length: int, 
                 agent_type: str, mode: str = "random") -> Optional[Dict]:
    """
    Sample sequences for training.
    
    Args:
        batch_size: Number of sequences to sample
        segment_length: Length of each sequence
        agent_type: 'diagnostics' or 'controls'
        mode: 'random' or 'latest'
        
    Returns:
        Dictionary with keys:
            - 'state_source_timestamp': List of timestamps
            - 'state': List of state sequences
            - 'prediction': List of prediction sequences (if available)
            - 'next_state': List (controls only)
            - 'reward': List (controls only)
            - 'terminate': List (controls only)
            - 'truncate': List (controls only)
    """

For diagnostic agents:

batch_data = self.db_manager.sample_batch(
    batch_size=32,
    segment_length=10,  # 10 timesteps per sequence
    agent_type="diagnostics",
    mode="random"
)

if batch_data is None:
    return None

# Extract state sequences
states = batch_data['state']  # List of length batch_size
# Each element is a list of segment_length states
# states[0] = [state_t0, state_t1, ..., state_t9]

For RL agents:

batch_data = self.db_manager.sample_batch(
    batch_size=256,
    segment_length=1,  # Single transitions
    agent_type="controls",
    mode="random"
)

# Contains state, next_state, reward, terminate, truncate

Sampling modes:

• "random": Random sampling across all data
• "latest": Most recent data first

Agent Registration

Used by: AgentBase (automatic)

def register_agent(self, agent_id: str, agent_name: str, 
                   config: Dict = None, info: Dict = None) -> int:
    """
    Register agent in database.
    
    Args:
        agent_id: Unique UUID string
        agent_name: Human-readable name
        config: Configuration dictionary
        info: Additional info dictionary
        
    Returns:
        0 on success, 1 on failure
    """

Called automatically by AgentBase.start(), no manual intervention needed.

Update Agent Info

def update_agent_info(self, agent_id: str, info_updates: Dict) -> int:
    """
    Update agent status/information.
    
    Args:
        agent_id: Agent UUID
        info_updates: Dictionary of updates to merge
        
    Returns:
        0 on success, 1 on failure
    """

Example:

self.update_agent_status({
    'status': 'training',
    'last_training_time': time.time()
})

Retrieve Agent Info

def get_agent_info(self, agent_id: str) -> Optional[Dict]:
    """
    Retrieve agent information.
    
    Args:
        agent_id: Agent UUID
        
    Returns:
        Dictionary with 'registered_id', 'agent_name', 'config', 'info' or None
    """

Data Serialization

Numpy Arrays to BLOB

Automatic conversion:

# This numpy array...
state = np.array([1.0, 2.0, 3.0], dtype=np.float32)

# ...is automatically converted to bytes for storage
data = {
    'state': state  # DBManager handles conversion
}
self.db_manager.record_sensor_data(data)

Retrieving BLOB Data

Automatic conversion:

batch_data = self.db_manager.sample_batch(...)

# States are returned as numpy arrays
states = batch_data['state']
# states[0] is already a numpy array, ready to use

Dictionary Serialization

Pickle used for dictionaries:

config = {'learning_rate': 0.001, 'epochs': 50}

# Stored as pickled BLOB
self.db_manager.register_agent(
    agent_id=self.agent_id,
    agent_name='MyAgent',
    config=config  # Automatically pickled
)

# Retrieved as dictionary
agent_info = self.db_manager.get_agent_info(self.agent_id)
config = agent_info['config']  # Automatically unpickled

Database Cleanup

Close Connection

def close(self):
    """Close database connection and cursor."""

Automatic cleanup: Called by thread cleanup methods

Common Patterns

Check Data Before Training

def get_training_data(self):
    # Check total samples
    total_samples = self.db_manager.get_size("agent_inferences")
    
    if total_samples < self.min_training_samples:
        logging.info(f"Need {self.min_training_samples}, have {total_samples}")
        return None
    
    # Check for new data
    if total_samples <= self.last_training_count:
        logging.debug("No new data")
        return None
    
    # Sample batch
    batch_data = self.db_manager.sample_batch(...)
    
    # Update count
    self.last_training_count = total_samples
    
    return batch_data

Store with Error Handling

def store_message(self, message_data, topic, partition, offset):
    try:
        # Prepare data
        data = {...}
        
        # Store to database
        status = self.db_manager.record_sensor_data(data)
        
        if status == 0:
            return True
        else:
            logging.error(f"Storage failed with status {status}")
            return False
            
    except Exception as e:
        logging.error(f"Error storing data: {e}")
        return False

Sample with Validation

def get_training_data(self):
    batch_data = self.db_manager.sample_batch(
        batch_size=32,
        segment_length=10,
        agent_type="diagnostics",
        mode="random"
    )
    
    if batch_data is None or len(batch_data['state']) == 0:
        logging.warning("No data returned from sample_batch")
        return None
    
    # Validate data shape
    states = np.array(batch_data['state'])
    logging.info(f"Sampled {states.shape} for training")
    
    return states

Best Practices

Connection management:

• One DBManager per thread
• Let threads handle cleanup automatically
• Don't share DBManager between threads

Error handling:

• Always check return status (0 = success)
• Catch exceptions around database operations
• Log failures with context

Data types:

• Use numpy arrays for sensor data
• Use proper datetime format for timestamps
• Convert to float32 for consistency

Sampling strategy:

• Start with small batch sizes
• Use "latest" mode for debugging
• Use "random" mode for training diversity

Performance:

• Sample in batches, not individual records
• Track last_training_count to avoid reprocessing
• Use appropriate segment_length for your model

---