分布式训练实验平台自动化管理：任务调度与日志聚合
大家好，今天我们来探讨如何搭建一个分布式训练实验平台，并自动化管理训练任务和日志结果。在深度学习领域，模型训练的计算量日益增长，单机训练往往耗时过长。分布式训练应运而生，能够显著缩短训练时间，但也带来了任务管理和结果分析上的挑战。一个好的实验平台能够简化这些流程，提高研发效率。
本次讲座将分为以下几个部分：
- 架构设计： 平台整体架构的设计思路，包括各个模块的职责和交互。
- 任务调度： 如何将训练任务分配到不同的计算节点，并进行有效的资源管理。
- 日志聚合： 如何从各个计算节点收集训练日志，并进行统一的存储和分析。
- 结果管理： 如何管理训练结果，包括模型文件、评估指标等。
- 代码示例： 使用Python和相关工具，演示关键模块的实现。
1. 架构设计
一个分布式训练实验平台的核心目标是简化训练流程，提高资源利用率，并方便结果分析。 我们可以将平台划分为以下几个核心模块：
- 任务管理模块 (Task Management)： 负责接收用户提交的训练任务，并将任务信息存储到数据库中。任务信息包括模型配置、数据集路径、训练参数、资源需求等。
- 调度器模块 (Scheduler)： 负责根据任务的资源需求和集群的资源状态，将任务分配到合适的计算节点。调度器需要考虑资源利用率、任务优先级等因素。
- 计算节点模块 (Compute Node)： 负责接收调度器分配的任务，执行训练脚本，并将训练日志和结果上传到指定的存储位置。
- 日志聚合模块 (Log Aggregation)： 负责从各个计算节点收集训练日志，并将日志存储到统一的日志存储系统中。
- 结果管理模块 (Result Management)： 负责管理训练结果，包括模型文件、评估指标等。结果管理模块可以将结果存储到数据库或对象存储中。
- 监控模块 (Monitoring)： 负责监控集群的资源状态和任务的运行状态，并提供可视化界面。
各个模块之间的交互关系如下图所示：
用户 --> 任务管理模块 --> 调度器模块 --> 计算节点模块
|
V
日志聚合模块 & 结果管理模块 & 监控模块2. 任务调度
任务调度的核心是根据任务的资源需求和集群的资源状态，将任务分配到合适的计算节点。常见的调度算法包括：
- 先来先服务 (FCFS)： 按照任务提交的顺序进行调度。
- 最短作业优先 (SJF)： 优先调度运行时间最短的任务。
- 优先级调度： 按照任务的优先级进行调度。
- 资源公平调度 (Fair Scheduling)： 确保每个用户或队列都能获得公平的资源分配。
这里我们实现一个简单的基于优先级的调度器，使用Python的 threading 模块模拟并发环境：
import threading
import time
import random
class Task:
def __init__(self, task_id, priority, resource_request):
self.task_id = task_id
self.priority = priority
self.resource_request = resource_request
self.status = "pending" # pending, running, completed, failed
def run(self):
print(f"Task {self.task_id} running with priority {self.priority}")
self.status = "running"
time.sleep(random.randint(1, 5)) # Simulate task execution time
self.status = "completed"
print(f"Task {self.task_id} completed")
class ComputeNode:
def __init__(self, node_id, total_resources):
self.node_id = node_id
self.total_resources = total_resources
self.available_resources = total_resources
self.running_task = None
def allocate_resources(self, task):
if self.available_resources >= task.resource_request:
self.available_resources -= task.resource_request
self.running_task = task
return True
else:
return False
def release_resources(self):
if self.running_task:
self.available_resources += self.running_task.resource_request
self.running_task = None
class Scheduler:
def __init__(self, compute_nodes):
self.compute_nodes = compute_nodes
self.task_queue = []
self.lock = threading.Lock() # Protect shared resources
def add_task(self, task):
with self.lock:
self.task_queue.append(task)
self.task_queue.sort(key=lambda x: x.priority) # Sort by priority (lower is higher priority)
def schedule(self):
while True:
with self.lock:
if self.task_queue:
task = self.task_queue.pop(0) # Get highest priority task
for node in self.compute_nodes:
if node.allocate_resources(task):
print(f"Task {task.task_id} scheduled on Node {node.node_id}")
threading.Thread(target=self.run_task, args=(task, node)).start()
break # Task assigned, move to next task
else:
# No node available, put task back in queue (optional)
print(f"No available nodes for Task {task.task_id}, requeuing.")
self.task_queue.insert(0, task) # Put back at the front to retry later
time.sleep(1) # Check for new tasks periodically
def run_task(self, task, node):
task.run()
node.release_resources()
with self.lock:
print(f"Task {task.task_id} finished on Node {node.node_id}. Resources released.")
# Example Usage
if __name__ == "__main__":
node1 = ComputeNode("Node1", 10)
node2 = ComputeNode("Node2", 5)
scheduler = Scheduler([node1, node2])
# Create some tasks with different priorities and resource requests
task1 = Task("Task1", 1, 3) # High priority, small resource request
task2 = Task("Task2", 3, 7) # Low priority, large resource request
task3 = Task("Task3", 2, 2) # Medium priority, small resource request
task4 = Task("Task4", 1, 1) # High priority, very small resource request
scheduler.add_task(task1)
scheduler.add_task(task2)
scheduler.add_task(task3)
scheduler.add_task(task4)
# Start the scheduler in a separate thread
scheduler_thread = threading.Thread(target=scheduler.schedule)
scheduler_thread.daemon = True # Allow the main thread to exit
scheduler_thread.start()
# Let the scheduler run for a while
time.sleep(10)
print("Main thread exiting.")解释：
Task类定义了任务的属性，包括任务ID、优先级、资源需求和状态。ComputeNode类定义了计算节点的属性，包括节点ID、总资源数和可用资源数。Scheduler类实现了调度器，维护一个任务队列，并根据任务的优先级和节点的资源状态进行调度。add_task方法将任务添加到任务队列中，并按照优先级进行排序。schedule方法不断从任务队列中取出优先级最高的任务，并尝试将其分配到可用的计算节点上。run_task方法模拟任务的执行过程，并释放计算节点上的资源。- 使用
threading.Lock()确保对共享资源（task_queue，节点资源）的并发访问是线程安全的。 scheduler_thread.daemon = True允许主线程在所有非守护线程完成后退出。
3. 日志聚合
训练过程中会产生大量的日志信息，包括训练进度、损失函数值、评估指标等。为了方便分析，需要将这些日志信息从各个计算节点收集起来，并存储到统一的日志存储系统中。常见的日志存储系统包括：
- Elasticsearch： 一个开源的分布式搜索和分析引擎，可以存储和搜索大量的日志数据。
- Fluentd： 一个开源的数据收集器，可以将日志数据从不同的来源收集起来，并发送到不同的目的地。
- Kafka： 一个分布式流处理平台，可以用于构建实时数据管道。
这里我们使用 Python 的 logging 模块和 filebeat (或其他日志收集工具) 模拟日志聚合的过程。 首先，在每个计算节点上的训练脚本中，使用 logging 模块记录日志：
import logging
# Configure logging
logging.basicConfig(filename='training.log', level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s')
def train_model():
logging.info("Starting training...")
for epoch in range(10):
logging.info(f"Epoch {epoch+1}: Loss = {epoch * 0.1}")
# Simulate training process
# ...
logging.info("Training finished.")
if __name__ == "__main__":
train_model()解释：
logging.basicConfig函数配置日志的输出位置、日志级别和日志格式。logging.info函数用于记录信息级别的日志。
然后，使用 filebeat 将各个计算节点上的 training.log 文件收集起来，并发送到 Elasticsearch 或其他日志存储系统中。 filebeat 的配置文件（filebeat.yml）如下所示：
filebeat.inputs:
- type: log
enabled: true
paths:
- /path/to/training.log # Change this to the actual path
output.elasticsearch:
hosts: ["localhost:9200"] # Change this to your Elasticsearch host解释：
filebeat.inputs部分配置日志的输入来源，这里指定了training.log文件。output.elasticsearch部分配置日志的输出目的地，这里指定了 Elasticsearch 的地址。
最后，可以使用 Kibana 或其他可视化工具，对 Elasticsearch 中存储的日志数据进行分析和可视化。
4. 结果管理
训练完成后，需要将训练结果（包括模型文件、评估指标等）保存起来，以便后续使用。常见的存储方式包括：
- 文件系统： 将模型文件和评估指标保存到文件系统中。
- 对象存储： 将模型文件和评估指标保存到对象存储服务中，例如 Amazon S3、Google Cloud Storage 或 Azure Blob Storage。
- 数据库： 将评估指标保存到数据库中。
这里我们使用 Python 的 pickle 模块将模型文件保存到文件系统中，并将评估指标保存到 CSV 文件中：
import pickle
import csv
def save_model(model, filepath):
with open(filepath, 'wb') as f:
pickle.dump(model, f)
print(f"Model saved to {filepath}")
def save_metrics(metrics, filepath):
with open(filepath, 'w', newline='') as f:
writer = csv.writer(f)
writer.writerow(['Epoch', 'Loss', 'Accuracy']) # Write header
for epoch, (loss, accuracy) in enumerate(metrics):
writer.writerow([epoch + 1, loss, accuracy])
print(f"Metrics saved to {filepath}")
# Example usage
if __name__ == "__main__":
# Simulate a trained model
class SimpleModel:
def __init__(self):
self.weights = [1.0, 2.0, 3.0]
model = SimpleModel()
save_model(model, 'model.pkl')
# Simulate training metrics
metrics = [(0.5, 0.8), (0.4, 0.9), (0.3, 0.95)]
save_metrics(metrics, 'metrics.csv')解释：
save_model函数使用pickle.dump函数将模型对象保存到文件中。save_metrics函数使用csv.writer函数将评估指标保存到 CSV 文件中。
为了更好地管理这些结果，可以考虑使用数据库来存储模型的元数据（例如模型名称、训练时间、数据集、评估指标等），并使用对象存储来存储模型文件。 可以使用 Python 的 sqlite3 模块连接到 SQLite 数据库，并将模型元数据保存到数据库中：
import sqlite3
def create_table(db_path):
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
cursor.execute('''
CREATE TABLE IF NOT EXISTS models (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT,
training_time DATETIME,
dataset TEXT,
loss REAL,
accuracy REAL,
model_path TEXT
)
''')
conn.commit()
conn.close()
def insert_model_metadata(db_path, name, training_time, dataset, loss, accuracy, model_path):
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
cursor.execute('''
INSERT INTO models (name, training_time, dataset, loss, accuracy, model_path)
VALUES (?, ?, ?, ?, ?, ?)
''', (name, training_time, dataset, loss, accuracy, model_path))
conn.commit()
conn.close()
# Example usage
if __name__ == "__main__":
db_path = 'models.db'
create_table(db_path)
# Simulate model metadata
name = 'MyModel'
training_time = '2023-10-27 10:00:00'
dataset = 'MNIST'
loss = 0.3
accuracy = 0.95
model_path = 'model.pkl'
insert_model_metadata(db_path, name, training_time, dataset, loss, accuracy, model_path)
print("Model metadata saved to database.")解释：
create_table函数创建models表，用于存储模型元数据。insert_model_metadata函数将模型元数据插入到models表中。
5. 代码示例：整合各模块
下面将上述各个模块的代码片段整合起来，构成一个简单的分布式训练实验平台的示例。 为了简化，这里只演示了任务提交、调度和日志记录的过程。
# (Task, ComputeNode, Scheduler classes - as defined previously)
import threading
import time
import random
import logging
import sqlite3
import pickle
import csv
import datetime
import os
# Configure logging
logging.basicConfig(filename='scheduler.log', level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s')
class Task:
def __init__(self, task_id, priority, resource_request, model_config, dataset_path, output_dir):
self.task_id = task_id
self.priority = priority
self.resource_request = resource_request
self.model_config = model_config
self.dataset_path = dataset_path
self.output_dir = output_dir
self.status = "pending" # pending, running, completed, failed
def run(self):
logging.info(f"Task {self.task_id} running with priority {self.priority}")
self.status = "running"
# Simulate model training based on config
try:
# Simulate training loop
for epoch in range(self.model_config.get('epochs', 10)):
loss = random.uniform(0.1, 1.0) # Simulate loss
accuracy = random.uniform(0.7, 0.99) # Simulate accuracy
logging.info(f"Task {self.task_id} - Epoch {epoch+1}: Loss = {loss:.4f}, Accuracy = {accuracy:.4f}")
time.sleep(random.uniform(0.1, 0.5)) # Simulate training time
# Simulate saving the model
model_filename = os.path.join(self.output_dir, f"model_{self.task_id}.pkl")
with open(model_filename, 'wb') as f:
pickle.dump({"weights": [random.random() for _ in range(10)]}, f) # Simulate model weights
logging.info(f"Task {self.task_id} - Model saved to {model_filename}")
# Simulate saving metrics
metrics_filename = os.path.join(self.output_dir, f"metrics_{self.task_id}.csv")
with open(metrics_filename, 'w', newline='') as f:
writer = csv.writer(f)
writer.writerow(['Epoch', 'Loss', 'Accuracy'])
for epoch in range(self.model_config.get('epochs', 10)):
loss = random.uniform(0.1, 1.0)
accuracy = random.uniform(0.7, 0.99)
writer.writerow([epoch + 1, loss, accuracy])
logging.info(f"Task {self.task_id} - Metrics saved to {metrics_filename}")
self.status = "completed"
logging.info(f"Task {self.task_id} completed successfully.")
except Exception as e:
self.status = "failed"
logging.error(f"Task {self.task_id} failed: {e}")
class ComputeNode:
def __init__(self, node_id, total_resources, db_path):
self.node_id = node_id
self.total_resources = total_resources
self.available_resources = total_resources
self.running_task = None
self.db_path = db_path
def allocate_resources(self, task):
if self.available_resources >= task.resource_request:
self.available_resources -= task.resource_request
self.running_task = task
return True
else:
return False
def release_resources(self):
if self.running_task:
self.available_resources += self.running_task.resource_request
self.running_task = None
def insert_model_metadata(self, name, training_time, dataset, loss, accuracy, model_path):
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
cursor.execute('''
INSERT INTO models (name, training_time, dataset, loss, accuracy, model_path)
VALUES (?, ?, ?, ?, ?, ?)
''', (name, training_time, dataset, loss, accuracy, model_path))
conn.commit()
conn.close()
class Scheduler:
def __init__(self, compute_nodes, db_path):
self.compute_nodes = compute_nodes
self.task_queue = []
self.lock = threading.Lock()
self.db_path = db_path
self.create_table()
def create_table(self):
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
cursor.execute('''
CREATE TABLE IF NOT EXISTS models (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT,
training_time DATETIME,
dataset TEXT,
loss REAL,
accuracy REAL,
model_path TEXT
)
''')
conn.commit()
conn.close()
def add_task(self, task):
with self.lock:
self.task_queue.append(task)
self.task_queue.sort(key=lambda x: x.priority)
def schedule(self):
while True:
with self.lock:
if self.task_queue:
task = self.task_queue.pop(0)
for node in self.compute_nodes:
if node.allocate_resources(task):
logging.info(f"Task {task.task_id} scheduled on Node {node.node_id}")
threading.Thread(target=self.run_task, args=(task, node)).start()
break
else:
logging.warning(f"No available nodes for Task {task.task_id}, requeuing.")
self.task_queue.insert(0, task)
time.sleep(1)
def run_task(self, task, node):
task.run()
node.release_resources()
with self.lock:
logging.info(f"Task {task.task_id} finished on Node {node.node_id}. Resources released.")
# Simulate getting final loss and accuracy (replace with actual values)
final_loss = random.uniform(0.1, 0.5)
final_accuracy = random.uniform(0.8, 0.99)
# Save model metadata to database
now = datetime.datetime.now().isoformat()
model_filename = os.path.join(task.output_dir, f"model_{task.task_id}.pkl")
node.insert_model_metadata(f"Model_{task.task_id}", now, task.dataset_path, final_loss, final_accuracy, model_filename)
logging.info(f"Model metadata for Task {task.task_id} saved to database.")
if __name__ == "__main__":
db_path = 'models.db'
node1 = ComputeNode("Node1", 10, db_path)
node2 = ComputeNode("Node2", 5, db_path)
scheduler = Scheduler([node1, node2], db_path)
# Define some tasks with model configurations, dataset paths and output directories
task1 = Task("Task1", 1, 3, {'epochs': 5}, "MNIST", "output/task1")
task2 = Task("Task2", 3, 7, {'epochs': 10}, "CIFAR10", "output/task2")
task3 = Task("Task3", 2, 2, {'epochs': 7}, "ImageNet", "output/task3")
task4 = Task("Task4", 1, 1, {'epochs': 3}, "CustomData", "output/task4")
# Create output directories if they don't exist
for task in [task1, task2, task3, task4]:
os.makedirs(task.output_dir, exist_ok=True)
scheduler.add_task(task1)
scheduler.add_task(task2)
scheduler.add_task(task3)
scheduler.add_task(task4)
# Start the scheduler in a separate thread
scheduler_thread = threading.Thread(target=scheduler.schedule)
scheduler_thread.daemon = True
scheduler_thread.start()
# Let the scheduler run for a while
time.sleep(20)
logging.info("Main thread exiting.")
print("Main thread exiting.")解释：
- 这个例子将之前的
Task,ComputeNode, 和Scheduler类整合在了一起。 - 每个
Task对象现在包含了model_config,dataset_path, 和output_dir属性，用于更真实地模拟一个训练任务。 Task.run方法现在模拟了基于model_config的训练过程，包括训练循环和模型/指标的保存。 它使用logging记录训练过程。ComputeNode对象现在包含了db_path属性，并负责将模型元数据插入到数据库中。Scheduler对象现在负责创建数据库表，并在run_task完成后将模型元数据保存到数据库中。- 主程序创建了一些
Task对象，并将它们添加到调度器中。 每个任务都有一个output_dir，用于保存模型和指标文件。
这个例子展示了一个简单的分布式训练实验平台的核心组件是如何协同工作的。 实际的平台会更加复杂，但基本原理是相同的。
训练平台的核心机制
整个训练平台的核心机制包括任务调度，日志聚合和结果管理，通过python多线程和数据库操作保证并发安全和数据持久化。
进一步扩展的方向
为了使平台更加完善，可以考虑以下扩展方向：
- Web界面： 提供一个Web界面，方便用户提交任务、查看任务状态和下载训练结果。
- 自动化部署： 使用 Docker 或其他容器化技术，自动化部署平台到不同的计算节点上。
- 资源管理： 集成 Kubernetes 或其他资源管理系统，实现更灵活的资源分配。
- 模型版本管理： 使用 MLflow 或其他模型版本管理工具，管理不同版本的模型。
- 监控和告警： 使用 Prometheus 和 Grafana 或其他监控工具，监控集群的资源状态和任务的运行状态，并发送告警信息。