3737import java .io .IOException ;
3838import java .nio .charset .StandardCharsets ;
3939import java .util .Arrays ;
40+ import java .util .ArrayList ;
41+ import java .util .List ;
4042import java .util .regex .Matcher ;
4143import java .util .regex .Pattern ;
4244
@@ -78,7 +80,7 @@ public void addCreateArgs(AbstractDockerLauncher launcher,
7880 String value ;
7981 if ("executor" .equals (getValue ())) {
8082 String executorNum = launcher .getEnvironment ().get ("EXECUTOR_NUMBER" );
81- String nvidiasmiOutput = executeWithOutput (launcher .getInner (), "nvidia-smi" , "-L" );
83+ String nvidiasmiOutput = executeWithOutput (launcher .getInner (), "nvidia-smi" , "-L" );
8284 if (isMIG (nvidiasmiOutput )) {
8385 value = getMIG (nvidiasmiOutput , executorNum );
8486 } else {
@@ -134,7 +136,7 @@ private String executeWithOutput(Launcher launcher, String... args) {
134136 }
135137
136138 private boolean isMIG (String output ) {
137- Pattern pattern = Pattern .compile ("(MIG-GPU-[a-f0-9\- \/ ]+)" );
139+ Pattern pattern = Pattern .compile ("(MIG-GPU-[a-f0-9\\ - \ \ /]+)" );
138140 Matcher m = pattern .matcher (output );
139141
140142 if (m .find ()) {
@@ -144,13 +146,14 @@ private boolean isMIG(String output) {
144146 }
145147
146148 private String getMIG (String output , String executor ) {
149+ int executorNum = Integer .parseInt (executor );
147150 List <String > uuids = new ArrayList <String >();
148- Pattern pattern = Pattern .compile ("(MIG-GPU-[a-f0-9\- \/ ]+)" );
151+ Pattern pattern = Pattern .compile ("(MIG-GPU-[a-f0-9\\ - \ \ /]+)" );
149152 Matcher m = pattern .matcher (output );
150153
151154 while (m .find ()) {
152155 uuids .add (m .group ());
153156 }
154- return uuids [ executor ] ;
157+ return uuids . get ( executorNum ) ;
155158 }
156159}
0 commit comments