Golang 文本模板，你指定没用过

最近在倒腾“AI大模型基础设施”，宏观目标是做一个基于云原生的AI算力平台，目前因公司隐私暂不能公开宏观背景和技术方案，姑且记录实践中遇到的一些技能点。

Arena是阿里云开源的命令行工具，帮助算法开发者和数据科学家，在无需云原生技能的前提下，提交训练和评测任务、管控模型, 现在已经捐赠给kubeflow社区。

提交一个pytorchjob微调训练模型的姿势如下，具体的请参考官网。

arena submit pytorchjob \

  --name=bloom-sft \

  --gpus=1 \

  --image=registry.cn-hangzhou.aliyuncs.com/acs/deepspeed:v0.9.0-chat \

  --data=training-data:/model \

  --tensorboard \

  --logdir=/model/logs \

  --config test_ac.yaml

  "cd /model/DeepSpeedExamples/applications/DeepSpeed-Chat/training/step1_supervised_finetuning && bash training_scripts/other_language/run_chinese.sh /model/bloom-560m-sft"

Arena将应用层的行为动作抽象为cli参数，渲染为helm的应用模板，使用helm创建云原生机器学习任务。

Arena有golang、python、java sdk, 早期我们尝试使用Arena sdk提交训练任务，后面发现这个工具截止1.4.2版本时依旧是只支持单用户，那么在用户粒度上使用该sdk就会有线程安全问题，这个二开的话改动还挺大。

基于对平台客户和并发度的分析，我们认为可以采用Arena子进程的方式规避这个问题。

于是问题就变成了在web后端进程启动Arena子进程来完成动作，同时我们发现Arena cli底层是将cli参数渲染进helm模板，于是我们借葫芦画瓢也可以利用golang text template将saas前端的内容渲染进arena cli参数。

1. 准备text template，产生需要渲染的Arena Cli参数字符串。

新建一个golang模板文件tmpl.go，下面是提交训练任务的golang text template，

var submitTmplStr string = `arena submit pytorchjob

  --name. {[.Name}]

  --workers={{.Workers}}

  --image={{.Image}}

  {{if .RDMA}}

  --rdma

  {{end}}

  --retry {{.Retry}}

  --gpus {{.RequestGPUs}}

  --memory [{.RequestMemory}}

  --cpu {{.RequestCPUs}}

  --share-memory {{.RequestShareMemory}}

  -n {{.BelongTenantId}}

  --selector {[.ResourceSpec}}

  -l kueue.x-k8s.io/queue-name={{.BelongTenantId}}-local-queue

  -e TZ=Asia/Shanghai

  {{ range $key,$value :=.Ants}} -a {{$key}}={{$value}} {{end}}

  {{ range .Data}} --data {{.}} {{end}}

  --config {[.KubeConfig}}

  --loglevel=info

  {{ .Command}}

golang text template使用{{ }}作为占位符，支持if/else、 foreach，渲染格式对于开发者友好，后期也很容易定位问题。

这里还有个小技巧, 如果你模板内容很长（还要求是一个紧凑的字符串），模板中的cli参数之间可使用\n\t来间隔，渲染完成后自行去掉即可。

产生渲染文本：

tmpl, err := template.New("submitArgs").Parse(submitTmplStr)

if err != nil {

  zap.L().Error("parse submit template failure", zap.Error(err)〕

  return err

}

var buf bytes.Buffer

err = tmpl.Execute(&buf, ja)   // 产出渲染文本

if err != nil {

  zap.L().Error("execute submit template failure", zap.Error(err)〕

  return err

}

submitStr := buf.String()

replacer := strings.NewReplacer("\n","","\t","") // 移除模板中插入的\n\t

submitStr = replacer.Replace(submitStr)

2. web后端启动Arena子进程，加载参数，提交训练任务。

env:= os.Environ()

cmd:= exec.Command("sh","-c",submitStr)

cmd.Env= env

by,err := cmd.CombinedOutput()

if err != nil  {

   // ...

}

return nil

理论上还有另外一种启动子进程的方式：

exec.Command("arena",args...)， args 是参数切片，或者叫不定长参数，但是我们已经用模板字符串了，所以这里使用sh来启动子进程并自行解释arena参数比较合适。

That's All 本文记录了在实践【云原生算力平台】时用到的两个通用技能点，后期我会持续记录本项目实践中遇到的挑战，并适时给出整个架构的宏观全貌和设计考量。