最近在倒騰“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 本文記錄了在實踐【云原生算力平臺】時用到的兩個通用技能點， 后期我會持續記錄本項目實踐中遇到的挑戰，并適時給出整個架構的宏觀全貌和設計考量。