「我正在参加Trae「超级体验官」创意实践征文，」 「本文所使用的 Trae 免费下载链接：」 「www.trae.ai/?utm_source…」
背景
我们实现了一个k8s任务分发的项目，项目之前会从普罗米修斯获取每个节点的资源使用情况，后来这种实现有个问题，经常把资源过度使用，导致系统组件出现异常，所以想在前端进行一个资源管理，给系统留20%的资源。
安装 Trae
安装过程很丝滑：「www.trae.ai/?utm_source…」
不过记得科学上网，登录。
Prompt
在已有项目上新增一个资源管控功能 1、使用go语言实现 2、设计一个表，字段有节点名称，卫星名称，CPU总量，CPU余量，GPU总量，GPU余量，内存总量，内存余量 3、数据库使用 sqlite orm.go 已经实现了部分功能
DB, err = gorm.Open(sqlite.Open("/mnt/data/task-manager/task-manager.db"), &gorm.Config{
Logger: logger.Default.LogMode(logger.Silent),
})
if err != nil {
fmt.Println("connect fail:", err)
return
}
fmt.Println("connect success:", DB != nil)
// Migrate the schema
_ = DB.AutoMigrate(&DBTaskRequest{})
_ = DB.AutoMigrate(&DBSubTask{})
_ = DB.AutoMigrate(&VolSubTask{})
4、实现表数据的初始化，kubectl get nodes ,获取节点的资源信息 client.go 已经实现了
// NewKubernetesClient 创建一个新的 Kubernetes 客户端
func NewKubernetesClient() (*KubernetesClient, error) {
// 使用集群内配置
config, err := rest.InClusterConfig()
if err != nil {
return nil, fmt.Errorf("error creating in-cluster config: %v", err)
}
clientset, err := kubernetes.NewForConfig(config)
if err != nil {
return nil, fmt.Errorf("error creating clientset: %v", err)
}
return &KubernetesClient{clientset: clientset}, nil
}
// ListNodes 获取所有节点信息
func (k *KubernetesClient) ListNodes() (*corev1.NodeList, error) {
return k.clientset.CoreV1().Nodes().List(context.Background(), metav1.ListOptions{})
}
5、实现更新接口,查询接口 6、调用更新接口，实现一个函数来判断：这个任务是否调度成功到集群中，资源是否占用。
使用 Trae
1、引用文件
为了保持代码一致性，我们经常使用引用项目中的原代码，这里可以是代码片段，也可以是文件。而且有自动检索功能。第一次体验有点惊艳。
2、代码生成
依据生产代码，速度很快。
我来帮你实现资源管控功能。
- 首先，让我们在 orm.go 中添加新的数据结构：
// ... existing code ...
// NodeResource 节点资源信息表
type NodeResource struct {
gorm.Model
NodeName string `json:"node_name" gorm:"uniqueIndex"`
SatelliteName string `json:"satellite_name"`
TotalCPU float32 `json:"total_cpu"`
FreeCPU float32 `json:"free_cpu"`
TotalGPU float32 `json:"total_gpu"`
FreeGPU float32 `json:"free_gpu"`
TotalMemory float32 `json:"total_memory"` // 单位：GB
FreeMemory float32 `json:"free_memory"` // 单位：GB
}
func Connect() {
// ... existing code ...
// 添加 NodeResource 表的迁移
_ = DB.AutoMigrate(&NodeResource{})
}
// 新增资源管理相关的函数
func InitNodeResources(nodes []NodeResource) error {
return DB.Create(&nodes).Error
}
func UpdateNodeResource(nodeName string, updates map[string]interface{}) error {
return DB.Model(&NodeResource{}).Where("node_name = ?", nodeName).Updates(updates).Error
}
func GetNodeResource(nodeName string) (*NodeResource, error) {
var resource NodeResource
err := DB.Where("node_name = ?", nodeName).First(&resource).Error
return &resource, err
}
func ListNodeResources() ([]NodeResource, error) {
var resources []NodeResource
err := DB.Find(&resources).Error
return resources, err
}
// CheckResourceAvailability 检查任务是否可以被调度
func CheckResourceAvailability(nodeName string, reqCPU, reqGPU, reqMemory float32) (bool, error) {
var resource NodeResource
err := DB.Where("node_name = ?", nodeName).First(&resource).Error
if err != nil {
return false, err
}
if resource.FreeCPU >= reqCPU && resource.FreeGPU >= reqGPU && resource.FreeMemory >= reqMemory {
return true, nil
}
return false, nil
}
2. 在 kubernetes 包中添加资源初始化和更新的功能：
// ... existing code ...
// GetNodeResources 获取节点的资源总量
func (k *KubernetesClient) GetNodeResources(nodeName string) (cpu float32, memory float32, err error) {
node, err := k.GetNode(nodeName)
if err != nil {
return 0, 0, err
}
cpuQuantity := node.Status.Capacity[corev1.ResourceCPU]
memoryQuantity := node.Status.Capacity[corev1.ResourceMemory]
cpuFloat, _ := cpuQuantity.AsInt64()
memoryFloat := float32(memoryQuantity.Value()) / (1024 * 1024 * 1024) // 转换为 GB
return float32(cpuFloat), memoryFloat, nil
}
// GetNodeAllocatableResources 获取节点的可分配资源
func (k *KubernetesClient) GetNodeAllocatableResources(nodeName string) (cpu float32, memory float32, err error) {
node, err := k.GetNode(nodeName)
if err != nil {
return 0, 0, err
}
cpuQuantity := node.Status.Allocatable[corev1.ResourceCPU]
memoryQuantity := node.Status.Allocatable[corev1.ResourceMemory]
cpuFloat, _ := cpuQuantity.AsInt64()
memoryFloat := float32(memoryQuantity.Value()) / (1024 * 1024 * 1024) // 转换为 GB
return float32(cpuFloat), memoryFloat, nil
}
3. 创建一个新的文件来实现资源管理的主要逻辑：
package resource
import (
"fmt"
"your-project/pkg/kubernetes"
"your-project/pkg/mysql"
)
type ResourceManager struct {
k8sClient *kubernetes.KubernetesClient
}
func NewResourceManager() (*ResourceManager, error) {
client, err := kubernetes.NewKubernetesClient()
if err != nil {
return nil, err
}
return &ResourceManager{k8sClient: client}, nil
}
// InitializeNodeResources 初始化节点资源信息
func (rm *ResourceManager) InitializeNodeResources() error {
nodes, err := rm.k8sClient.ListNodes()
if err != nil {
return err
}
var nodeResources []mysql.NodeResource
for _, node := range nodes.Items {
totalCPU, totalMemory, err := rm.k8sClient.GetNodeResources(node.Name)
if err != nil {
continue
}
allocatedCPU, allocatedMemory, err := rm.k8sClient.NodeAllocatedResources(node.Name)
if err != nil {
continue
}
// 转换已分配资源的字符串为 float32
allocCPUFloat := float32(0)
fmt.Sscanf(allocatedCPU, "%f", &allocCPUFloat)
allocMemFloat := float32(0)
fmt.Sscanf(allocatedMemory, "%fGi", &allocMemFloat)
nodeResources = append(nodeResources, mysql.NodeResource{
NodeName: node.Name,
SatelliteName: node.Labels["satellite-name"], // 假设有这个标签
TotalCPU: totalCPU,
FreeCPU: totalCPU - allocCPUFloat,
TotalGPU: 0, // GPU 信息需要从节点标签或其他方式获取
FreeGPU: 0,
TotalMemory: totalMemory,
FreeMemory: totalMemory - allocMemFloat,
})
}
return mysql.InitNodeResources(nodeResources)
}
// UpdateNodeResources 更新节点资源信息
func (rm *ResourceManager) UpdateNodeResources(nodeName string) error {
totalCPU, totalMemory, err := rm.k8sClient.GetNodeResources(nodeName)
if err != nil {
return err
}
allocatedCPU, allocatedMemory, err := rm.k8sClient.NodeAllocatedResources(nodeName)
if err != nil {
return err
}
// 转换已分配资源的字符串为 float32
allocCPUFloat := float32(0)
fmt.Sscanf(allocatedCPU, "%f", &allocCPUFloat)
allocMemFloat := float32(0)
fmt.Sscanf(allocatedMemory, "%fGi", &allocMemFloat)
updates := map[string]interface{}{
"total_cpu": totalCPU,
"free_cpu": totalCPU - allocCPUFloat,
"total_memory": totalMemory,
"free_memory": totalMemory - allocMemFloat,
}
return mysql.UpdateNodeResource(nodeName, updates)
}
使用说明：
- 初始化资源信息：
resourceManager, err := resource.NewResourceManager()
if err != nil {
// 处理错误
}
err = resourceManager.InitializeNodeResources()
if err != nil {
// 处理错误
}
2. 更新节点资源信息：
err = resourceManager.UpdateNodeResources("node-name")
if err != nil {
// 处理错误
}
3. 检查任务是否可以调度：
canSchedule, err := mysql.CheckResourceAvailability("node-name", reqCPU, reqGPU, reqMemory)
if err != nil {
// 处理错误
}
if !canSchedule {
// 资源不足，无法调度
}
这个实现提供了：
- 节点资源信息的数据库表结构
- 资源信息的初始化功能
- 资源信息的更新接口
- 资源查询接口
- 任务调度可行性检查功能
你可以根据实际需求调整字段和功能。例如，如果需要添加更多的资源类型或者其他特性，可以相应地修改 NodeResource 结构体和相关函数。
3、检视代码并验证
- 看了下数据结构生成的都非常正确，然后点击应用按钮
- 新增代码会以不同颜色显示，点击全部接受
第一个功能完美实现，在数据库新增了一个表结构来实现资源的管理。并把实现了数据初始化，从k8s集群获取资源然后写入数据库里面。
第二个功能，是遍历获取k8s集群资源的数据，提供给数据库初始化函数。
第三个功能，Trae 帮我新建了一个模块，专门用于资源管理的接口实现。这没个5.6年的开发经验，写不出这样的代码。
4、环境验证&微调
在主逻辑里面调用生成的函数，再生成一些单元测试。
总结
我使用过 cursor、windsurf ,Trae不仅继承了他们的优点，并在此基础上进一步优化了体验。
