zishuo/ascend-docker-runtime
1
0
Fork 0
mirror of https://github.com/Ascend/ascend-docker-runtime.git synced 2025-10-06 03:56:49 +08:00
Code Issues Packages Projects Releases Wiki Activity

Match-id-736351bcbab9825232a29001080a89b7d32c4ff3

Browse Source
This commit is contained in:
BianTanggui
2022-06-13 17:45:20 +08:00
parent 56a513d839
commit 32b45190bf
1 changed files with 16 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
runtime/main.go
View File

@@ -28,6 +28,7 @@ const (
hookDefaultFilePath = "/usr/local/bin/ascend-docker-hook"
dockerRuncFile = "docker-runc"
runcFile = "runc"
envLength = 2
)
var (
@@ -108,6 +109,19 @@ func addHook(spec *specs.Spec) error {
})
}
hasVirtualFlag := false
for _, line := range spec.Process.Env {
words := strings.Split(line, "=")
if len(words) == envLength && strings.TrimSpace(words[0]) == "ASCEND_RUNTIME_OPTIONS" {
if strings.Contains(words[1], "VIRTUAL") {
hasVirtualFlag = true
}
}
}
if hasVirtualFlag {
return nil
}
vdevice, err := dcmi.CreateVDevice(&dcmi.NpuWorker{}, spec)
if err != nil {
@@ -126,12 +140,11 @@ func updateEnvAndPostHook(spec *specs.Spec, vdevice dcmi.VDeviceInfo) {
needAddVirtualFlag := true
for _, line := range spec.Process.Env {
words := strings.Split(line, "=")
const LENGTH int = 2
if len(words) == LENGTH && strings.TrimSpace(words[0]) == "ASCEND_VISIBLE_DEVICES" {
if len(words) == envLength && strings.TrimSpace(words[0]) == "ASCEND_VISIBLE_DEVICES" {
newEnv = append(newEnv, fmt.Sprintf("ASCEND_VISIBLE_DEVICES=%d", vdevice.VdeviceID))
continue
}
if len(words) == LENGTH && strings.TrimSpace(words[0]) == "ASCEND_RUNTIME_OPTIONS" {
if len(words) == envLength && strings.TrimSpace(words[0]) == "ASCEND_RUNTIME_OPTIONS" {
needAddVirtualFlag = false
if strings.Contains(words[1], "VIRTUAL") {
newEnv = append(newEnv, line)