nvidia-docker

nodemanager/core/RemoteExecutor.cpp

@@ -148,11 +148,12 @@ pplx::task<json::value> RemoteExecutor::StartTask(StartTaskArgs&& args, std::str
     {
         Logger::Info(args.JobId, args.TaskId, args.StartInfo.TaskRequeueCount, "MPI non-master task found, skip creating the process.");
         std::string dockerImage = args.StartInfo.EnvironmentVariables["CCP_DOCKER_IMAGE"];
+        std::string isNvidiaDocker = args.StartInfo.EnvironmentVariables["CCP_DOCKER_NVIDIA"];
         if (!dockerImage.empty())
         {
             taskInfo->IsPrimaryTask = false;
             std::string output;
-            if (0 != System::ExecuteCommandOut(output, "/bin/bash", "StartMpiContainer.sh", taskInfo->TaskId, userName, dockerImage))
+            if (0 != System::ExecuteCommandOut(output, "/bin/bash", "StartMpiContainer.sh", taskInfo->TaskId, userName, dockerImage, isNvidiaDocker))
             {
                 Logger::Info(taskInfo->JobId, taskInfo->TaskId, taskInfo->GetTaskRequeueCount(), "Start MPI container successfully.");
             }

nodemanager/scripts/PrepareTask.sh

@@ -30,16 +30,17 @@ if [ "$isDockerTask" == "1" ]; then
 	additionalOption=$(GetDockerAdditionalOption $taskFolder)
 	envFile=$(GetDockerTaskEnvFile $taskFolder)
 	containerIdFile=$(GetContainerIdFile $taskFolder)
-	docker run -id \
-			--name $containerName \
-			--cpuset-cpus $affinity \
-			--env-file $envFile \
-			--cidfile $containerIdFile \
-			-v $taskFolder:$taskFolder:z \
-			$volumeOption \
-			$mpiContainerStartOption \
-			$additionalOption \
-			$dockerImage $ContainerPlaceholderCommand 2>&1
+	dockerEngine=$(GetDockerEngine $taskFolder)
+	$dockerEngine run -id \
+				--name $containerName \
+				--cpuset-cpus $affinity \
+				--env-file $envFile \
+				--cidfile $containerIdFile \
+				-v $taskFolder:$taskFolder:z \
+				$volumeOption \
+				$mpiContainerStartOption \
+				$additionalOption \
+				$dockerImage $ContainerPlaceholderCommand 2>&1
 	
 	ec=$?
 	if [ $ec -ne 0 ]

nodemanager/scripts/StartMpiContainer.sh

@@ -9,14 +9,16 @@
 taskId=$1
 userName=$2
 dockerImage=$3
+nvidiaOption=$4
 
 containerName=$(GetContainerName "$taskId_$MpiContainerSuffix")
 mpiContainerStartOption=$(GetMpiContainerStartOption $userName)
+dockerEngine=$(GetDockerEngine $nvidiaOption)
 
-docker run -id \
-        --name $containerName \
-        $mpiContainerStartOption \
-        $dockerImage $ContainerPlaceholderCommand 2>&1
+$dockerEngine run -id \
+            --name $containerName \
+            $mpiContainerStartOption \
+            $dockerImage $ContainerPlaceholderCommand 2>&1
 
 if [ $? -ne 0 ]
 then

nodemanager/scripts/common.sh

@@ -154,6 +154,22 @@ function CheckDockerDebugMode
 	fi	
 }
 
+function GetDockerEngine
+{
+	local nvidiaOption=$1
+	if [ -z $nvidiaOption ] || [ "$nvidiaOption" == "0" ]; then
+		echo "docker"
+	else
+		local taskFolder=$1
+		local nvidiaOption=$(cat $(GetDockerTaskEnvFile $taskFolder) | grep "CCP_DOCKER_NVIDIA=" | cut -d '=' -f 2)
+		if [ -z $nvidiaOption ] || [ "$nvidiaOption" == "0" ]; then
+			echo "docker"
+		else
+			echo "nvidia-docker"
+		fi
+	fi	
+}
+
 function GetDockerImageName
 {
 	local taskFolder=$1