tomwer 1.2.0a1__py3-none-any.whl → 1.2.0a3__py3-none-any.whl

tomwer/core/process/reconstruction/axis/params.py

@@ -48,10 +48,8 @@ from .anglemode import CorAngleMode
 from .mode import AxisMode
 from .projectiontype import ProjectionType
 
-try:
-    from nabu.preproc.phase import PaganinPhaseRetrieval
-except ImportError:
-    from tomwer.third_party.nabu.preproc.phase import PaganinPhaseRetrieval
+from nabu.preproc.phase import PaganinPhaseRetrieval
+
 
 _logger = logging.getLogger(__name__)
 
@@ -734,10 +732,10 @@ class AxisRP:
         :return: center of rotation or none
         :rtype: Union[None, float]
         """
-        from .axis import AxisProcess  # avoid cyclic import
+        from .axis import AxisTask  # avoid cyclic import
 
         processes = Task.get_processes_frm_type(
-            process_file=file_path, process_type=AxisProcess, entry=entry
+            process_file=file_path, process_type=AxisTask, entry=entry
         )
         if len(processes) == 0:
             _logger.warning("failed to find any information on center " "of rotation")

tomwer/core/process/reconstruction/darkref/darkrefs.py

@@ -39,6 +39,8 @@ from processview.core.manager import DatasetState, ProcessManager
 from processview.core.superviseprocess import SuperviseProcess
 from silx.io.url import DataUrl
 from silx.io.utils import h5py_read_dataset
+from silx.utils.deprecation import deprecated_warning
+
 from tomoscan.framereducerbase import REDUCER_TARGET, ReduceMethod
 from tomoscan.io import HDF5File
 
@@ -56,8 +58,13 @@ from . import params as dkrf_reconsparams
 logger = logging.getLogger(__name__)
 
 
-class DarkRefs(
-    Task, SuperviseProcess, Queue, input_names=("data",), output_names=("data",)
+class DarkRefsTask(
+    Task,
+    SuperviseProcess,
+    Queue,  # TODO: fixme: this does not make much sense for a task
+    input_names=("data",),
+    output_names=("data",),
+    optional_input_names=("serialize_output_data",),
 ):
     """Compute median/mean dark and ref from originals (dark and ref files)"""
 
@@ -236,7 +243,9 @@ class DarkRefs(
         """
         if entry is None:
             with HDF5File(process_file, "r", swmr=True) as h5f:
-                entries = DarkRefs._get_process_nodes(root_node=h5f, process=DarkRefs)
+                entries = DarkRefsTask._get_process_nodes(
+                    root_node=h5f, process=DarkRefsTask
+                )
                 if len(entries) == 0:
                     logger.info("unable to find a DarkRef process in %s" % process_file)
                     return None
@@ -250,8 +259,8 @@ class DarkRefs(
             if entry not in h5f.keys():
                 logger.info(f"no dark found for {entry}")
                 return {}
-            dark_nodes = DarkRefs._get_process_nodes(
-                root_node=h5f[entry], process=DarkRefs
+            dark_nodes = DarkRefsTask._get_process_nodes(
+                root_node=h5f[entry], process=DarkRefsTask
             )
             index_to_path = {}
             for key, index in dark_nodes.items():
@@ -298,7 +307,9 @@ class DarkRefs(
         """
         if entry is None:
             with HDF5File(process_file, "r", swmr=True) as h5f:
-                entries = DarkRefs._get_process_nodes(root_node=h5f, process=DarkRefs)
+                entries = DarkRefsTask._get_process_nodes(
+                    root_node=h5f, process=DarkRefsTask
+                )
                 if len(entries) == 0:
                     logger.info("unable to find a DarkRef process in %s" % process_file)
                     return None
@@ -312,8 +323,8 @@ class DarkRefs(
             if entry not in h5f.keys():
                 logger.info(f"no flats found for {entry}")
                 return {}
-            dkref_nodes = DarkRefs._get_process_nodes(
-                root_node=h5f[entry], process=DarkRefs
+            dkref_nodes = DarkRefsTask._get_process_nodes(
+                root_node=h5f[entry], process=DarkRefsTask
             )
             index_to_path = {}
             for key, index in dkref_nodes.items():
@@ -407,7 +418,7 @@ class DarkRefs(
 
             self.outputs.data = None
             return
-        whats = (DarkRefs.WHAT_REF, DarkRefs.WHAT_DARK)
+        whats = (DarkRefsTask.WHAT_REF, DarkRefsTask.WHAT_DARK)
         overwrites = (
             self.recons_params.overwrite_flat,
             self.recons_params.overwrite_dark,
@@ -451,7 +462,7 @@ class DarkRefs(
                 interpretations["/".join(("darks", str(index)))] = "image"
             results["darks"] = f_darks
 
-            scan.save_reduced_darks(f_darks)
+            scan.save_reduced_darks(f_darks, overwrite=True)
         if (
             self.recons_params.flat_calc_method
             is not dkrf_reconsparams.ReduceMethod.NONE
@@ -465,7 +476,7 @@ class DarkRefs(
                 interpretations["/".join(("flats", str(index)))] = "image"
             results["flats"] = f_flats
 
-            scan.save_reduced_flats(f_flats)
+            scan.save_reduced_flats(f_flats, overwrite=True)
 
         if len(results) > 0:
             # if some processing to be registered
@@ -490,7 +501,7 @@ class DarkRefs(
             dataset=scan, state=DatasetState.SUCCEED, details=None
         )
 
-        if self._return_dict:
+        if self.get_input_value("serialize_output_data", True):
             self.outputs.data = scan.to_dict()
         else:
             self.outputs.data = scan
@@ -519,7 +530,9 @@ class DarkRefs(
             )
             scan.set_reduced_darks(darks=reduced_darks, darks_infos=metadata)
             try:
-                scan.save_reduced_darks(darks=reduced_darks, darks_infos=metadata)
+                scan.save_reduced_darks(
+                    darks=reduced_darks, darks_infos=metadata, overwrite=True
+                )
             except Exception as e:
                 logger.error(f"Fail to save reduced darks. Error is {e}")
         elif target is REDUCER_TARGET.FLATS:
@@ -530,7 +543,9 @@ class DarkRefs(
             )
             scan.set_reduced_flats(flats=reduced_flats, flats_infos=metadata)
             try:
-                scan.save_reduced_flats(flats=reduced_flats, flats_infos=metadata)
+                scan.save_reduced_flats(
+                    flats=reduced_flats, flats_infos=metadata, overwrite=True
+                )
             except Exception as e:
                 logger.error(f"Fail to save reduced flats. Error is {e}")
         else:
@@ -572,10 +587,11 @@ def requires_reduced_dark_and_flat(scan: TomwerScanBase, logger_=None) -> tuple:
         recons_params.dark_calc_method = dkrf_reconsparams.ReduceMethod.NONE
         recons_params.flat_calc_method = dkrf_reconsparams.ReduceMethod.FIRST
 
-        drp = DarkRefs(
+        drp = DarkRefsTask(
             inputs={
                 "data": scan,
                 "dark_ref_params": recons_params,
+                "serialize_output_data": False,
             }
         )
         if logger_ is not None:
@@ -593,10 +609,11 @@ def requires_reduced_dark_and_flat(scan: TomwerScanBase, logger_=None) -> tuple:
         recons_params.dark_calc_method = dkrf_reconsparams.ReduceMethod.FIRST
         recons_params.flat_calc_method = dkrf_reconsparams.ReduceMethod.NONE
 
-        drp = DarkRefs(
+        drp = DarkRefsTask(
             inputs={
                 "data": scan,
                 "dark_ref_params": recons_params,
+                "serialize_output_data": False,
             }
         )
         if logger_ is not None:
@@ -607,3 +624,23 @@ def requires_reduced_dark_and_flat(scan: TomwerScanBase, logger_=None) -> tuple:
         computed.append("dark")
 
     return tuple(computed)
+
+
+class DarkRefs(DarkRefsTask):
+    def __init__(
+        self,
+        process_id=None,
+        varinfo=None,
+        inputs=None,
+        node_id=None,
+        node_attrs=None,
+        execinfo=None,
+    ):
+        deprecated_warning(
+            name="tomwer.core.process.reconstruction.darkref.darkref.DarkRefs",
+            type_="class",
+            reason="improve readibility",
+            since_version="1.2",
+            replacement="DarkRefsTask",
+        )
+        super().__init__(process_id, varinfo, inputs, node_id, node_attrs, execinfo)

tomwer/core/process/reconstruction/darkref/darkrefscopy.py

@@ -47,14 +47,14 @@ from tomoscan.esrf.scan.utils import (
 from tomoscan.framereducerbase import REDUCER_TARGET
 from tomoscan.io import HDF5File
 
-from tomwer.core.process.reconstruction.darkref.darkrefs import DarkRefs
+from tomwer.core.process.reconstruction.darkref.darkrefs import DarkRefsTask
 from tomwer.core.scan.scanbase import TomwerScanBase
 from tomwer.core.utils.scanutils import data_identifier_to_scan
 
 logger = logging.getLogger(__name__)
 
 
-class DarkRefsCopy(DarkRefs):
+class DarkRefsCopy(DarkRefsTask):
     """
     Reimplement Dark ref to deal with copy when there is no median/mean files
     """
@@ -166,6 +166,10 @@ class DarkRefsCopy(DarkRefs):
         if data is None:
             return
         flat_url = DarkRefsCopy.get_flats_url(save_dir)
+        # remove group if already exists: else can end up by copying several dark / flat from different datasets
+        with HDF5File(flat_url.file_path(), mode="a") as h5f:
+            if flat_url.data_path() in h5f:
+                del h5f[flat_url.data_path()]
         dicttoh5(
             data,
             h5file=flat_url.file_path(),
@@ -181,6 +185,10 @@ class DarkRefsCopy(DarkRefs):
         if data is None:
             return
         dark_url = DarkRefsCopy.get_darks_url(save_dir)
+        # remove group if already exists: else can end up by copying several dark / flat from different datasets
+        with HDF5File(dark_url.file_path(), mode="a") as h5f:
+            if dark_url.data_path() in h5f:
+                del h5f[dark_url.data_path()]
         dicttoh5(
             data,
             h5file=dark_url.file_path(),
@@ -227,6 +235,7 @@ class DarkRefsCopy(DarkRefs):
                 darks_url=self._darks_url,
                 save=True,
                 raise_error_if_url_empty=True,
+                overwrite=True,
             )
         if self.has_flat_stored():
             copy_h5_dict_flats_to(
@@ -234,6 +243,7 @@ class DarkRefsCopy(DarkRefs):
                 flats_url=self._flats_url,
                 save=True,
                 raise_error_if_url_empty=True,
+                overwrite=True,
             )
 
     def clean_save_files(self):

tomwer/core/process/reconstruction/lamino/__init__.py

@@ -1 +1 @@
-from .tofu import LaminoReconstruction  # noqa F401
+from .tofu import LaminoReconstructionTask  # noqa F401

tomwer/core/process/reconstruction/lamino/tofu.py

@@ -49,6 +49,7 @@ from tomwer.core.scan.scanbase import TomwerScanBase
 from tomwer.core.scan.scanfactory import ScanFactory
 from tomwer.core.utils import getDim1Dim2
 from tomwer.core.utils.char import PSI_CHAR
+from silx.utils.deprecation import deprecated_warning
 
 logger = logging.getLogger(__name__)
 
@@ -462,7 +463,12 @@ def has_tofu():
         return False
 
 
-class LaminoReconstruction(Task, input_names=("data",), output_names=("data",)):
+class LaminoReconstructionTask(
+    Task,
+    input_names=("data",),
+    optional_input_names=("serialize_output_data",),
+    output_names=("data",),
+):
     """
     Process to launch the lamino reconstruction
 
@@ -879,7 +885,7 @@ class LaminoReconstruction(Task, input_names=("data",), output_names=("data",)):
 
         if res is False:
             logger.error(f"Reconstruction of {_scan.path} failed")
-        if self._return_dict:
+        if self.get_input_value("serialize_output_data", True):
             self.outputs.data = _scan.to_dict()
         else:
             self.outputs.data = _scan
@@ -978,3 +984,17 @@ def getFlats(scan):
     else:
         logger.warning("Found more than two refHST files")
         return (os.path.join(scan, refHSTFiles[0]), os.path.join(scan, refHSTFiles[-1]))
+
+
+class LaminoReconstruction(LaminoReconstructionTask):
+    def __init__(
+        self, varinfo=None, inputs=None, node_id=None, node_attrs=None, execinfo=None
+    ):
+        deprecated_warning(
+            name="tomwer.core.process.reconstruction.lamino.tofu.LaminoReconstruction",
+            type_="class",
+            reason="improve readibility",
+            since_version="1.2",
+            replacement="LaminoReconstructionTask",
+        )
+        super().__init__(varinfo, inputs, node_id, node_attrs, execinfo)

tomwer/core/process/reconstruction/nabu/nabucommon.py

@@ -37,6 +37,8 @@ import os
 import subprocess
 import uuid
 from typing import Iterable, Optional, Union
+import signal
+import psutil
 
 import numpy
 from silx.io.url import DataUrl
@@ -234,6 +236,9 @@ class _NabuBaseReconstructor:
         self._target = Target.from_value(target)
         self._dry_run = dry_run
         self._process_name = process_name
+        self._process = None
+        self._cancelled = False
+        # nabu subprocess if run locally
         if isinstance(cluster_config, SlurmClusterConfiguration):
             self._cluster_config = cluster_config
         elif isinstance(cluster_config, dict):
@@ -322,6 +327,18 @@ class _NabuBaseReconstructor:
         else:
             raise ValueError(f"{self.target} is not recognized as a valid target")
 
+    @staticmethod
+    def _get_gpu_and_cpu_mem_fraction(config_to_dump: dict):
+        gpu_mem_fraction = config_to_dump.get("resources", {}).get(
+            "gpu_mem_fraction", config_to_dump.get("gpu_mem_fraction", 0.9)
+        )
+        assert gpu_mem_fraction <= 1
+        cpu_mem_fraction = config_to_dump.get("resources", {}).get(
+            "cpu_mem_fraction", config_to_dump.get("cpu_mem_fraction", 0.9)
+        )
+        assert cpu_mem_fraction <= 1
+        return gpu_mem_fraction, cpu_mem_fraction
+
     def _run_nabu_locally(
         self,
         conf_file: str,
@@ -341,18 +358,37 @@ class _NabuBaseReconstructor:
         """
         if not has_nabu:
             raise ImportError("Fail to import nabu")
+        assert isinstance(config_to_dump, dict)
+        gpu_mem_fraction, cpu_mem_fraction = self._get_gpu_and_cpu_mem_fraction(
+            config_to_dump
+        )
+
         command = " ".join(
-            ("python3", "-m", settings.NABU_FULL_FIELD_APP_PATH, conf_file)
+            (
+                "python3",
+                "-m",
+                settings.NABU_FULL_FIELD_APP_PATH,
+                conf_file,
+                "--gpu_mem_fraction",
+                str(gpu_mem_fraction),
+                "--cpu_mem_fraction",
+                str(cpu_mem_fraction),
+            )
         )
         _logger.info(f'call nabu from "{command}"')
-        process = subprocess.Popen(
+
+        self._process = subprocess.Popen(
             command,
             shell=True,
             cwd=self.scan.path,
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
         )
-        outs, errs = process.communicate()
+        try:
+            outs, errs = self._process.communicate()
+        except (TimeoutError, KeyboardInterrupt):
+            self._process.kill()
+            outs, errs = self._process.communicate()
 
         recons_urls = utils.get_recons_volume_identifier(
             file_prefix=config_to_dump["output"]["file_prefix"],
@@ -412,13 +448,21 @@ class _NabuBaseReconstructor:
         project_name = project_name.replace(" ", "_")
         cluster_config["job_name"] = project_name
 
+        # extract gpu_mem_fraction and cpu_mem_fraction
+        assert isinstance(config_to_dump, dict)
+        gpu_mem_fraction, cpu_mem_fraction = self._get_gpu_and_cpu_mem_fraction(
+            config_to_dump
+        )
+
         # submit job
         script_name = get_slurm_script_name(prefix="nabu")
         # for now force job name
         cluster_config["job_name"] = f"tomwer-nabu {conf_file}"
         job = SBatchScriptJob(
             slurm_config=cluster_config,
-            script=(f"python3 -m {settings.NABU_FULL_FIELD_APP_PATH} {conf_file}",),
+            script=(
+                f"python3 -m {settings.NABU_FULL_FIELD_APP_PATH} {conf_file} --gpu_mem_fraction {gpu_mem_fraction} --cpu_mem_fraction {cpu_mem_fraction}",
+            ),
             script_path=os.path.join(self.scan.path, "slurm_scripts", script_name),
             clean_script=False,
             working_directory=self.scan.working_directory,
@@ -445,10 +489,12 @@ class _NabuBaseReconstructor:
         """Return a tuple a potential callback to be launch once the future is done"""
         return tuple()
 
-    def _treateOutputSliceConfig(self, config):
+    def _treateOutputSliceConfig(self, config) -> tuple:
         """
         - add or overwrite some parameters of the dictionary
         - create the output directory if does not exist
+
+        :return: (config: dict, nabu_cfg_folder: str)
         """
         # handle phase
         pag = False
@@ -482,16 +528,26 @@ class _NabuBaseReconstructor:
             location = config["output"].get("location", None)
             if location not in ("", None):
                 location = format_output_location(location, scan=self.scan)
-                # if user specify the location
-                if not os.path.isdir(config["output"]["location"]):
-                    os.makedirs(location)
+                location_cfg_files = location
             else:
                 # otherwise default location will be the data root level
                 location = self.scan.path
+                location_cfg_files = location
+                # TODO: if is a single file - append prefix
+                if config["output"].get("file_format") in (
+                    NabuOutputFileFormat.EDF.value,
+                    NabuOutputFileFormat.TIFF.value,
+                    NabuOutputFileFormat.JP2K.value,
+                ):  # if user specify the location
+                    location = "/".join([location, _file_name])
+
             # add reconstruction path to the list. scan `reconstruction_paths` register all the existing path where
             # reconstruction are saved in order to be able to browse them all
             self.scan.add_reconstruction_path(location)
             config["output"]["location"] = location
+        else:
+            # don't think this could ever happen
+            location_cfg_files = self.scan.path
         # handle preproc
         if "preproc" not in config:
             config["preproc"] = {}
@@ -504,14 +560,14 @@ class _NabuBaseReconstructor:
 
         extra_infos = self.scan.intensity_normalization.get_extra_infos()
 
-        nabu_cfg_folders = os.path.join(
-            config["output"]["location"], settings.NABU_CFG_FILE_FOLDER
+        nabu_cfg_folder = os.path.join(
+            location_cfg_files, settings.NABU_CFG_FILE_FOLDER
         )
-        os.makedirs(nabu_cfg_folders, exist_ok=True)
+        os.makedirs(nabu_cfg_folder, exist_ok=True)
 
         # configuration file and nabu_tomwer_serving_hatch must be in the same folder
         serving_hatch_file = os.path.join(
-            nabu_cfg_folders, settings.NABU_TOMWER_SERVING_HATCH
+            nabu_cfg_folder, settings.NABU_TOMWER_SERVING_HATCH
         )
 
         source = extra_infos.get("source", INormSource.NONE)
@@ -566,13 +622,36 @@ class _NabuBaseReconstructor:
                 )
         else:
             raise NotImplementedError(f"source type {source.value} is not handled")
-
-        return config
+        return config, nabu_cfg_folder
 
     def _get_file_basename_reconstruction(self, pag, db, ctf):
         """return created file base name"""
         raise NotImplementedError("Base class")
 
+    def cancel(self):
+        self._cancelled = True
+        if self._process:
+            # kill childs processes
+            try:
+                parent = psutil.Process(self._process.pid)
+            except psutil.NoSuchProcess:
+                pass
+            else:
+                # TODO: see with Pierre. But from my point of view this
+                # might be handled by nabu it self....
+                childrens = parent.children(recursive=True)
+                for child in childrens:
+                    try:
+                        child.send_signal(signal.SIGKILL)
+                    except psutil.NoSuchProcess:
+                        pass
+                    else:
+                        child.wait()
+                self._process.kill()
+                self._process.wait()
+            finally:
+                self._process = None
+
 
 def dump_normalization_array_for_nabu(
     scan: TomwerScanBase, output_file: str, array: Union[numpy.ndarray, float, int]