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

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

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

1. arena submit pytorchjob \
2.   --name=bloom-sft \
3.   --gpus=1 \
4.   --image=registry.cn-hangzhou.aliyuncs.com/acs/deepspeed:v0.9.0-chat \
5.   --data=training-data:/model \
6.   --tensorboard \
7.   --logdir=/model/logs \
8.   --config test_ac.yaml
9.   "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，

1. var submitTmplStr string = `arena submit pytorchjob
2.   --name. {[.Name}]
3.   --workers={{.Workers}}
4.   --image={{.Image}}
5.   {{if .RDMA}}
6.   --rdma
7.   {{end}}
8.   --retry {{.Retry}}
9.   --gpus {{.RequestGPUs}}
10.   --memory [{.RequestMemory}}
11.   --cpu {{.RequestCPUs}}
12.   --share-memory {{.RequestShareMemory}}
13.   -n {{.BelongTenantId}}
14.   --selector {[.ResourceSpec}}
15.   -l kueue.x-k8s.io/queue-name={{.BelongTenantId}}-local-queue
16.   -e TZ=Asia/Shanghai
17.   {{ range $key,$value :=.Ants}} -a {{$key}}={{$value}} {{end}}
18.   {{ range .Data}} --data {{.}} {{end}}
19.   --config {[.KubeConfig}}
20.   --loglevel=info
21.   {{ .Command}}

复制代码

golang text template使用{{ }}作为占位符，支持if/else、 foreach， 渲染格式对于开发者友好，后期也很容易定位问题。
这里还有个小技巧, 如果你模板内容很长（还要求是一个紧凑的字符串），模板中的cli参数之间可使用\n\t来间隔，渲染完成后自行去掉即可。
产生渲染文本：

1. tmpl, err := template.New("submitArgs").Parse(submitTmplStr)
2. if err != nil {
3.   zap.L().Error("parse submit template failure", zap.Error(err)〕
4.   return err
5. }
6. var buf bytes.Buffer
7. err = tmpl.Execute(&buf, ja)   // 产出渲染文本
8. if err != nil {
9.   zap.L().Error("execute submit template failure", zap.Error(err)〕
10.   return err
11. }
13. submitStr := buf.String()
14. replacer := strings.NewReplacer("\n","","\t","") // 移除模板中插入的\n\t
15. submitStr = replacer.Replace(submitStr)

复制代码

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

1. env:= os.Environ()
2. cmd:= exec.Command("sh","-c",submitStr)
3. cmd.Env= env
4. by,err := cmd.CombinedOutput()
5. if err != nil  {
6.    // ...
7. }
8. return nil

复制代码

理论上还有另外一种启动子进程的方式：
exec.Command("arena",args...)， args 是参数切片，或者叫不定长参数，但是我们已经用模板字符串了，所以这里使用sh来启动子进程并自行解释arena参数比较合适。
That's All 本文记录了在实践【云原生算力平台】时用到的两个通用技能点， 后期我会持续记录本项目实践中遇到的挑战，并适时给出整个架构的宏观全貌和设计考量。

来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！