pcm-coordinator/internal/svc/servicecontext.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/participant"
	"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
	Ai              *participant.Ai
}

func NewServiceContext(c config.Config) *ServiceContext {
	promClient, err := tracker.NewPrometheus(c.Monitoring.PromUrl)

	if err != nil {
		logx.Errorf("InitPrometheus err: %v", err)
		panic("InitPrometheus err")
	}
	httpClient := resty.New()
	httpClient.SetTimeout(1 * time.Second)
	alertClient := tracker.NewAlertClient(c.Monitoring.AlertUrl)

	//添加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}
	hpcStorage := &database.HpcStorage{DbEngin: dbEngin}
	cache := make(map[string]interface{}, 0)
	aiService, err := service.NewAiService(&c, storage, cache)
	if err != nil {
		logx.Error(err.Error())
		panic(err)
	}
	hpcService, err := service.NewHpcService(&c, hpcStorage, cache)
	if err != nil {
		logx.Error(err.Error())
		panic(err)
	}
	scheduler := scheduler.NewSchdlr(aiService, storage, hpcStorage, hpcService)
	ai, err := participant.New(storage, c.Participant.AdapterId)
	if err != nil {
		logx.Error(err.Error())
		panic(err)
	}
	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,
		Ai:              ai,
	}
}