pcm-coordinator/internal/svc/servicecontext.go

139 lines
4.6 KiB
Go

/*
Copyright (c) [2023] [pcm]
[pcm-coordinator] is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2.
You may obtain a copy of Mulan PSL v2 at:
http://license.coscl.org.cn/MulanPSL2
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPaRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
See the Mulan PSL v2 for more details.
*/
package svc
import (
"github.com/go-redis/redis/v8"
"github.com/go-resty/resty/v2"
alert "github.com/prometheus/alertmanager/api/v2/client"
"github.com/robfig/cron/v3"
"github.com/zeromicro/go-zero/core/logx"
"github.com/zeromicro/go-zero/zrpc"
"gitlink.org.cn/JointCloud/pcm-ac/hpcacclient"
"gitlink.org.cn/JointCloud/pcm-coordinator/internal/config"
"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler"
"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/database"
"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/service"
"gitlink.org.cn/JointCloud/pcm-coordinator/pkg/tracker"
"gitlink.org.cn/JointCloud/pcm-coordinator/pkg/utils"
"gitlink.org.cn/JointCloud/pcm-modelarts/client/imagesservice"
"gitlink.org.cn/JointCloud/pcm-modelarts/client/modelartsservice"
"gitlink.org.cn/JointCloud/pcm-octopus/octopusclient"
"gitlink.org.cn/JointCloud/pcm-openstack/openstackclient"
slurmClient "gitlink.org.cn/JointCloud/pcm-slurm/slurmclient"
"gorm.io/driver/mysql"
"gorm.io/gorm"
"gorm.io/gorm/logger"
"gorm.io/gorm/schema"
"time"
)
type ServiceContext struct {
Config config.Config
RedisClient *redis.Client
Cron *cron.Cron
ModelArtsRpc modelartsservice.ModelArtsService
ModelArtsImgRpc imagesservice.ImagesService
DbEngin *gorm.DB
ACRpc hpcacclient.HpcAC
THRpc slurmClient.Slurm
OctopusRpc octopusclient.Octopus
OpenstackRpc openstackclient.Openstack
MonitorClient map[int64]tracker.Prometheus
PromClient tracker.Prometheus
AlertClient *alert.AlertmanagerAPI
HttpClient *resty.Client
Scheduler *scheduler.Scheduler
}
func NewServiceContext(c config.Config) *ServiceContext {
promClient, err := tracker.NewPrometheus(c.Monitoring.PromUrl)
if err != nil {
logx.Errorf("InitPrometheus err: %v", err)
panic("InitSnowflake err")
}
httpClient := resty.New()
httpClient.SetTimeout(1 * time.Second)
alertClient := tracker.NewAlertClient(c.Monitoring.AlertUrl)
if err != nil {
logx.Errorf("InitPrometheus err: %v", err)
panic("InitSnowflake err")
}
//添加snowflake支持
err = utils.InitSnowflake(c.SnowflakeConf.MachineId)
if err != nil {
logx.Errorf("InitSnowflake err: %v", err)
panic("InitSnowflake err")
}
//启动Gorm支持
dbEngin, err := gorm.Open(mysql.Open(c.DB.DataSource), &gorm.Config{
NamingStrategy: schema.NamingStrategy{
SingularTable: true, // 使用单数表名,启用该选项,此时,`User` 的表名应该是 `t_user`
},
Logger: logger.Default.LogMode(logger.Warn),
})
if err != nil {
logx.Errorf("数据库连接失败, err%v", err)
panic(err)
}
sqlDB, err := dbEngin.DB()
// SetMaxIdleConns 设置空闲连接池中连接的最大数量
sqlDB.SetMaxIdleConns(10)
// SetMaxOpenConns 设置打开数据库连接的最大数量。
sqlDB.SetMaxOpenConns(50)
// SetConnMaxLifetime 设置了连接可复用的最大时间。
sqlDB.SetConnMaxLifetime(time.Hour)
if err != nil {
logx.Error(err.Error())
return nil
}
redisClient := redis.NewClient(&redis.Options{
Addr: c.Redis.Host,
Password: c.Redis.Pass,
})
// scheduler
storage := &database.AiStorage{DbEngin: dbEngin}
aiService, err := service.NewAiService(&c, storage)
if err != nil {
logx.Error(err.Error())
return nil
}
scheduler := scheduler.NewSchdlr(aiService, storage)
return &ServiceContext{
DbEngin: dbEngin,
Cron: cron.New(cron.WithSeconds()),
Config: c,
RedisClient: redisClient,
ModelArtsRpc: modelartsservice.NewModelArtsService(zrpc.MustNewClient(c.ModelArtsRpcConf)),
ModelArtsImgRpc: imagesservice.NewImagesService(zrpc.MustNewClient(c.ModelArtsImgRpcConf)),
ACRpc: hpcacclient.NewHpcAC(zrpc.MustNewClient(c.ACRpcConf)),
OctopusRpc: octopusclient.NewOctopus(zrpc.MustNewClient(c.OctopusRpcConf)),
OpenstackRpc: openstackclient.NewOpenstack(zrpc.MustNewClient(c.OpenstackRpcConf)),
MonitorClient: make(map[int64]tracker.Prometheus),
PromClient: promClient,
AlertClient: alertClient,
HttpClient: httpClient,
Scheduler: scheduler,
}
}