updated ai octopus

Former-commit-id: e1760e20de1bab6e6d9f4e5366672750be8fcc94
This commit is contained in:
tzwang 2024-02-26 00:36:35 +08:00
parent c7597592f9
commit efdf09649d
1 changed files with 45 additions and 5 deletions

View File

@ -49,6 +49,14 @@ const (
GCU = "GCU"
EnflameT20 = 128
BASE_TOPS = 128
CAMBRICON = "cambricon"
)
var (
cardAliasMap = map[string]string{
MLU: CAMBRICON,
GCU: GCU,
}
)
func NewOctopusLink(ctx context.Context, svcCtx *svc.ServiceContext, name string, id int64) *OctopusLink {
@ -327,7 +335,7 @@ func (o *OctopusLink) generateResourceId(option *option.AiOption) error {
}
if option.ResourceType == CARD {
err = setResourceIdByCard(option, specResp, GCU)
err = setResourceIdByCard(option, specResp, MLU)
if err != nil {
return err
}
@ -398,8 +406,20 @@ func setResourceIdByCard(option *option.AiOption, specs *octopus.GetResourceSpec
}
case MLU: // 1, 2, 4
if cardNum*2 == s {
if cardNum/2 == s {
option.ResourceId = spec.Id
option.ComputeCard = computeCard
return nil
}
if 1 < cardNum/2 && cardNum/2 <= 2 && s == 2 {
option.ResourceId = spec.Id
option.ComputeCard = computeCard
return nil
}
if 2 < cardNum/2 && s == 4 {
option.ResourceId = spec.Id
option.ComputeCard = computeCard
return nil
}
}
}
@ -434,8 +454,28 @@ func (o *OctopusLink) generateDatasetsId(option *option.AiOption) error {
}
func (o *OctopusLink) generateImageId(option *option.AiOption) error {
if option.TaskType == "" {
return errors.New("TaskType not set")
}
req := &octopus.GetPresetImageListReq{
Platform: o.platform,
PageIndex: o.pageIndex,
PageSize: o.pageSize,
}
resp, err := o.svcCtx.OctopusRpc.GetPresetImageList(o.ctx, req)
if err != nil {
return err
}
if !resp.Success {
return errors.New("failed to get PresetImages")
}
for _, image := range resp.Payload.Images {
if strings.Contains(image.ImageName, option.TaskType) && strings.Contains(image.ImageName, cardAliasMap[option.ComputeCard]) {
option.ImageId = image.Id
return nil
}
}
return errors.New("failed to get ImageId")
}
func (o *OctopusLink) generateCmd(option *option.AiOption) error {