diff --git a/toolbox/core/bst_containers.m b/toolbox/core/bst_containers.m index bce5d82d49..a2b1a3d671 100644 --- a/toolbox/core/bst_containers.m +++ b/toolbox/core/bst_containers.m @@ -5,7 +5,7 @@ % [errMsg, engineName] = bst_containers('GetEngine') % [errMsg, imageList] = bst_containers('GetImages') % [errMsg, imageSha] = bst_containers('ImportImage', imageSource, [imageTag]) -% errMsg = bst_containers('RunContainer', containerName, imageSha, [volumes], [isDaemon]) +% errMsg = bst_containers('RunContainer', containerName, imageSha, [volumes], [isDaemon], [containerArgs]) % [errMsg, cmdout] = bst_containers('ExecInContainer', containerName, cmdStr) % [errMsg, containerInfo] = bst_containers('GetContainerInfo', containerName) % errMsg = bst_containers('StopContainer', containerName, [isForced=0]) @@ -248,9 +248,12 @@ %% ===== RUN CONTAINER AS DAEMON ===== -function errMsg = RunContainer(containerName, imageSha, volumes, isDaemon) -% USAGE: errMsg = bst_containers('RunContainer', containerName, imageSha, volumes, isDaemon) +function errMsg = RunContainer(containerName, imageSha, volumes, isDaemon, containerArgs) +% USAGE: errMsg = bst_containers('RunContainer', containerName, imageSha, volumes, isDaemon, containerArgs) % Validate inputs + if nargin < 5 || isempty(containerArgs) + containerArgs = ''; + end if nargin < 4 || isempty(isDaemon) isDaemon = 0; end @@ -275,15 +278,21 @@ volumesStr = strjoin(pairs, ' '); end + % Use GPU with container engine + gpuStr = ''; + if bst_get('ContainerUseGpu') && system('which nvidia-smi') == 0 + gpuStr = '--gpus all'; + end + % Run container switch engineName case 'docker' if ~isDaemon % Run ENTRYPOINT - cmdStr = sprintf('docker run --rm --name %s %s %s', containerName, volumesStr, imageSha); + cmdStr = sprintf('docker run --rm --name %s %s %s %s %s', containerName, gpuStr, volumesStr, imageSha, containerArgs); else % Replace ENTRYPOINT (if any) with `sleep infinity` - cmdStr = sprintf('docker run -d --name %s %s --entrypoint sleep %s infinity', containerName, volumesStr, imageSha); + cmdStr = sprintf('docker run -d --name %s %s %s --entrypoint sleep %s infinity', containerName, gpuStr, volumesStr, imageSha); end [status, cmdout] = system(cmdStr); end @@ -308,6 +317,10 @@ return end + % Flag to track interruption + processState = containers.Map({'isInterruptCleanup'}, {1}); + % Clean up on function end, errors or Ctrl+C is pressed + cleanupObj = onCleanup(@() ProcessInterrupted(containerName, processState)); % Run command switch engineName case 'docker' @@ -316,11 +329,15 @@ else commandWrapper = ''''; % Single quote end - [status, cmdout] = system(['docker exec ' containerName ' sh -c ' commandWrapper cmdStr commandWrapper]); + % Execute the running container + commandExec = ['docker exec ' containerName ' sh -c ' commandWrapper cmdStr commandWrapper]; + [status, cmdout] = system(commandExec, '-echo'); if status ~= 0 errMsg = strtrim(cmdout); end end + % Code in container ended normally + processState('isInterruptCleanup') = 0; end @@ -428,6 +445,15 @@ end +%% ===== PROCESS INTERRUPTED ===== +function ProcessInterrupted(containerName, processState) + if processState('isInterruptCleanup') + bst_plugin('Unload', regexprep(containerName, '^bst_', '')); + bst_error('The process running in the container was interrupted', 'Container', 0); + end +end + + %% ===== GET ONLINE MANIFEST DIGEST ===== function [errMsg, manifestSha] = GetOnlineManifest(imageSource) manifestSha = ''; diff --git a/toolbox/core/bst_plugin.m b/toolbox/core/bst_plugin.m index dea4bd8403..875867855e 100644 --- a/toolbox/core/bst_plugin.m +++ b/toolbox/core/bst_plugin.m @@ -209,15 +209,9 @@ PlugDesc(end+1) = GetStruct('resection-identification'); PlugDesc(end).Version = 'latest'; PlugDesc(end).Category = 'Anatomy'; - PlugDesc(end).AutoUpdate = 1; - PlugDesc(end).URLzip = ['https://neuroimage.usc.edu/bst/getupdate.php?d=bst_resection_identification_' OsType '.zip']; - PlugDesc(end).TestFile = 'resection_identification'; - if strcmp(OsType, 'win64') - PlugDesc(end).TestFile = [PlugDesc(end).TestFile, '.bat']; - end - PlugDesc(end).URLinfo = 'https://github.com/ajoshiusc/auto_resection_mask/tree/brainstorm-plugin'; + PlugDesc(end).URLinfo = 'https://github.com/ajoshiusc/auto_resection_mask/tree/brainstorm-container'; + PlugDesc(end).ImageSource = ['docker.io/brainstormtools/auto-resection-mask:' PlugDesc(end).Version]; PlugDesc(end).CompiledStatus = 1; - PlugDesc(end).LoadFolders = {'bin'}; % === ANATOMY: ROAST === PlugDesc(end+1) = GetStruct('roast'); diff --git a/toolbox/process/functions/process_resection_identification.m b/toolbox/process/functions/process_resection_identification.m index 0d4d7099d9..dbf469f0cb 100644 --- a/toolbox/process/functions/process_resection_identification.m +++ b/toolbox/process/functions/process_resection_identification.m @@ -130,18 +130,22 @@ errMsg = 'The fiducials (NAS, LPA, RPA) are missing in the pre-op (default) MRI. Set them first before proceeding.'; return; end - - % Install/load resection-identification plugin - [isOk, errInstall, PlugDesc] = bst_plugin('Install', 'resection-identification', isInteractive); - if ~isOk - errMsg = [errMsg, errInstall]; - return; - end - + + % === CALL RESECTION-IDENTIFICATION PIPELINE === + % Container plugin name + plugName = 'resection-identification'; + bst_progress('text', ['Calling ' plugName]); + tic; + % Ensure container plugin: Installs and/or Loads + % Install container plugin === Import image into container engine + % Load container plugin === Run container in standby (name ['bst_' plugName]) + ensureRes = bst_plugin('Ensure', plugName); + % Retrieve info of container + [errMsg, containerInfo] = bst_containers('GetContainerInfo', ['bst_' plugName]); + % === SAVE BOTH MRI AS NIfTI === - bst_progress('text', 'Exporting pre- and post-op MRI...'); - % Create temporary folder - TmpDir = bst_get('BrainstormTmpDir', 0, 'resection_identification'); + % Get temporary folder from container info + TmpDir = containerInfo.volumes{1,1}; % Save pre-op MRI preOpNii = bst_fullfile(TmpDir, 'preop.nii'); out_mri_nii(sMriPreOp, preOpNii); @@ -149,20 +153,19 @@ postOpNii = bst_fullfile(TmpDir, 'postop.nii'); sMriPostOp = in_mri_bst(MriFilePostOp); out_mri_nii(sMriPostOp, postOpNii); - - % === CALL RESECTION-IDENTIFICATION PIPELINE === - bst_progress('text', 'Calling resection-identification...'); - % Get resection-identification executable - ResecExe = bst_fullfile(PlugDesc.Path, PlugDesc.SubFolder, PlugDesc.TestFile); - % Call resection-identification - strCall = ['"' ResecExe '"' ' ' '"' preOpNii '"' ' ' '"' postOpNii '"' ' ' '"' TmpDir '"']; - disp(['RESEC_ID > System call: ' strCall]); - tic; - status = system(strCall); - if (status ~= 0) - errMsg = 'Error during resection-identification, see logs in the command window.'; - bst_progress('stop'); - return; + % Make dir for processed files + mkdir(bst_fullfile(TmpDir, 'temp_dir_resection')); + + % === RUN COMMAND IN CONTAINER ===== + if isempty(errMsg) && containerInfo.isRunning + dataPath = containerInfo.volumes{1,2}; + command = [' python3 auto_resection_mask.py ' dataPath '/preop.nii ' dataPath '/postop.nii']; + errMsg = bst_containers('ExecInContainer', containerInfo.name, command); + else + errMsg = 'Container is not running'; + end + if ~isempty(errMsg) + return end disp(['RESEC_ID > Computation completed in: ' num2str(round(toc)) ' s']); @@ -182,8 +185,10 @@ Post2PreOpNii = bst_fullfile(TmpDir, 'postop.nonlin.post2pre.nii.gz'); MriFilePost2PreOp = import_mri(iSubject, Post2PreOpNii, 'Nifti1', 0, 1, 'postop_coreg_preop'); - % Delete the temporary files - file_delete(TmpDir, 1, 1); + % Unload container plugin === Stop container and Delete bind files + if ensureRes > 0 + bst_plugin('Unload', plugName); + end % Return success isOk = 1; end