100天跟着CP学PostgreSQL+AI,第20天 : 分布式训练架构：PostgreSQL 如何支持多 GPU 协同

import horovod.torch as hvd
import psycopg2

# 初始化Horovod
hvd.init()
rank = hvd.rank()
world_size = hvd.size()

# 连接PostgreSQL数据库
conn = psycopg2.connect(
    dbname="training_data",
    user="user",
    password="password",
    host="postgresql-host"
)
cur = conn.cursor()

# 计算每个GPU的数据范围
total_samples = cur.execute("SELECT COUNT(*) FROM samples").fetchone()[0]
samples_per_gpu = total_samples // world_size
offset = rank * samples_per_gpu

# 获取当前GPU的数据
cur.execute(f"SELECT data, label FROM samples LIMIT {samples_per_gpu} OFFSET {offset}")
data, labels = cur.fetchall()

# 转换为PyTorch张量
data = torch.tensor(data)
labels = torch.tensor(labels)

CREATE EXTENSION pg_lockman;

CREATE TABLE parameters (
    param_id SERIAL PRIMARY KEY,
    name VARCHAR(50) UNIQUE,
    value JSONB
);

CREATE TABLE locks (
    lock_name VARCHAR(50) PRIMARY KEY,
    owner VARCHAR(50),
    acquired_at TIMESTAMP
);

import psycopg2
from pg_lockman import LockManager

lock_manager = LockManager(conn)

# 获取锁
lock_manager.acquire_lock("param_lock")

# 更新参数
cur.execute("UPDATE parameters SET value = %s WHERE name = %s", (new_value, param_name))
conn.commit()

# 释放锁
lock_manager.release_lock("param_lock")

PostgreSQL Deployment 和 Service

apiVersion: apps/v1
kind: Deployment
metadata:
  name: postgresql
spec:
  replicas: 1
  selector:
    matchLabels:
      app: postgresql
  template:
    metadata:
      labels:
        app: postgresql
    spec:
      containers:
      - name: postgresql
        image: postgres:13
        env:
        - name: POSTGRES_USER
          value: "user"
        - name: POSTGRES_PASSWORD
          value: "password"
        - name: POSTGRES_DB
          value: "training_data"
        ports:
        - containerPort: 5432

---
apiVersion: v1
kind: Service
metadata:
  name: postgresql-service
spec:
  selector:
    app: postgresql
  ports:
  - protocol: TCP
  port: 5432
  targetPort: 5432

训练任务 Deployment

apiVersion: apps/v1
kind: Deployment
metadata:
  name: training-job
spec:
  replicas: 4  # 4个GPU节点
  selector:
    matchLabels:
      app: training
  template:
    metadata:
      labels:
        app: training
    spec:
      containers:
      - name: training-container
        image: your-docker-image:latest
        env:
        - name: HOROVOD_RANK
          valueFrom:
            fieldRef:
              fieldPath: metadata.labels['horovod/worker']
        - name: HOROVOD_SIZE
          value: "4"
        - name: POSTGRES_HOST
          value: "postgresql-service"
        - name: POSTGRES_USER
          value: "user"
        - name: POSTGRES_PASSWORD
          value: "password"
        - name: POSTGRES_DB
          value: "training_data"
        resources:
          limits:
            nvidia.com/gpu: 1  # 每个容器分配1个GPU

import psycopg2
import torch

def load_data_from_postgresql(rank, world_size):
    conn = psycopg2.connect(
        dbname="training_data",
        user="user",
        password="password",
        host="postgresql-service"
    )
    cur = conn.cursor()
    total_samples = cur.execute("SELECT COUNT(*) FROM samples").fetchone()[0]
    samples_per_gpu = total_samples // world_size
    offset = rank * samples_per_gpu
    cur.execute(f"SELECT data, label FROM samples LIMIT {samples_per_gpu} OFFSET {offset}")
    data, labels = cur.fetchall()
    return torch.tensor(data), torch.tensor(labels)

模型定义

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(100, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)
        return x

训练模块

import horovod.torch as hvd
from model import SimpleModel
from data_loader import load_data_from_postgresql

hvd.init()
torch.cuda.set_device(hvd.local_rank())

model = SimpleModel().to(hvd.local_rank())
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 广播模型参数到所有GPU
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
hvd.broadcast_optimizer_state(optimizer, root_rank=0)

data, labels = load_data_from_postgresql(hvd.rank(), hvd.size())
data = data.to(hvd.local_rank())
labels = labels.to(hvd.local_rank())

for epoch in range(10):
    optimizer.zero_grad()
    outputs = model(data)
    loss = nn.CrossEntropyLoss()(outputs, labels)
    loss.backward()
    hvd.allreduce_gradients(model.parameters(), op=hvd.AverageOp)
    optimizer.step()
    if hvd.rank() == 0:
        print(f"Epoch {epoch}, Loss: {loss.item()}")

参数服务器模块（使用 pg_lockman）

from pg_lockman import LockManager
import psycopg2

class ParameterServer:
    def __init__(self):
        self.conn = psycopg2.connect(
            dbname="training_data",
            user="user",
            password="password",
            host="postgresql-service"
        )
        self.lock_manager = LockManager(self.conn)

    def get_parameter(self, name):
        cur = self.conn.cursor()
        cur.execute("SELECT value FROM parameters WHERE name = %s", (name,))
        return cur.fetchone()[0]

    def update_parameter(self, name, new_value):
        self.lock_manager.acquire_lock(f"param_{name}_lock")
        cur = self.conn.cursor()
        cur.execute("UPDATE parameters SET value = %s WHERE name = %s", (new_value, name))
        self.conn.commit()
        self.lock_manager.release_lock(f"param_{name}_lock")