U
    Ph                     @  s  d dl mZ d dlZd dlZd dlZd dlmZ d dlmZ d dl	m
Z
mZ d dlZd dlmZ d dlmZ d dlmZ d d	lmZ d d
lmZmZ d dlmZ d dlmZ d dlmZ d dlmZ d dl m!Z!m"Z"m#Z#m$Z$ d dl%m&Z&m'Z' ee(dZ)e$d\Z*Z+G dd dZ,dS )    )annotationsN)deepcopy)sleep)Anycast)	BundleGen)DataAnalyzer)EnsembleRunner)NNIGen)export_bundle_algo_historyimport_bundle_algo_history)
get_logger)algo_to_pickle)ConfigParser)	SaveImage)AlgoKeys
has_optionlook_up_optionoptional_import) check_kwargs_exist_in_class_initrun_cmd)module_namennic                   @  sH  e Zd ZU dZded< dEd	d
dddddd	ddddddddddZdd Zdd Zd	ddddZdFddd dddZ	dGdd d d!d"Z
dHdd d#d$d%ZdId&d'ddd d(d)d*ZdJd	dd d,d-d.Zdd d/d0d1ZdKdd d#d2d3ZdLdd d#d4d5ZdMdd d#d6d7Zd8d d9d:d;Zd<d=d>d?d@Zd<d=d>dAdBZdCdD ZdS )N
AutoRunneru  
    An interface for handling Auto3Dseg with minimal inputs and understanding of the internal states in Auto3Dseg.
    The users can run the Auto3Dseg with default settings in one line of code. They can also customize the advanced
    features Auto3Dseg in a few additional lines. Examples of customization include

        - change cross-validation folds
        - change training/prediction parameters
        - change ensemble methods
        - automatic hyperparameter optimization.

    The output of the interface is a directory that contains

        - data statistics analysis report
        - algorithm definition files (scripts, configs, pickle objects) and training results (checkpoints, accuracies)
        - the predictions on the testing datasets from the final algorithm ensemble
        - a copy of the input arguments in form of YAML
        - cached intermediate results

    Args:
        work_dir: working directory to save the intermediate and final results.
        input: the configuration dictionary or the file path to the configuration in form of YAML.
            The configuration should contain datalist, dataroot, modality, multigpu, and class_names info.
        algos: optionally specify algorithms to use.  If a dictionary, must be in the form
            {"algname": dict(_target_="algname.scripts.algo.AlgnameAlgo", template_path="algname"), ...}
            If a list or a string, defines a subset of names of the algorithms to use, e.g. 'segresnet' or
            ['segresnet', 'dints'] out of the full set of algorithm templates provided by templates_path_or_url.
            Defaults to None, to use all available algorithms.
        analyze: on/off switch to run DataAnalyzer and generate a datastats report. Defaults to None, to automatically
            decide based on cache, and run data analysis only if we have not completed this step yet.
        algo_gen: on/off switch to run AlgoGen and generate templated BundleAlgos. Defaults to None, to automatically
            decide based on cache, and run algorithm folders generation only if we have not completed this step yet.
        train: on/off switch to run training and generate algorithm checkpoints. Defaults to None, to automatically
            decide based on cache, and run training only if we have not completed this step yet.
        hpo: use hyperparameter optimization (HPO) in the training phase. Users can provide a list of
            hyper-parameter and a search will be performed to investigate the algorithm performances.
        hpo_backend: a string that indicates the backend of the HPO. Currently, only NNI Grid-search mode
            is supported
        ensemble: on/off switch to run model ensemble and use the ensemble to predict outputs in testing
            datasets.
        not_use_cache: if the value is True, it will ignore all cached results in data analysis,
            algorithm generation, or training, and start the pipeline from scratch.
        templates_path_or_url: the folder with the algorithm templates or a url. If None provided, the default template
            zip url will be downloaded and extracted into the work_dir.
        allow_skip: a switch passed to BundleGen process which determines if some Algo in the default templates
            can be skipped based on the analysis on the dataset from Auto3DSeg DataAnalyzer.
        mlflow_tracking_uri: a tracking URI for MLflow server which could be local directory or address of the remote
            tracking Server; MLflow runs will be recorded locally in algorithms' model folder if the value is None.
        mlflow_experiment_name: the name of the experiment in MLflow server.
        kwargs: image writing parameters for the ensemble inference. The kwargs format follows the SaveImage
            transform. For more information, check https://docs.monai.io/en/stable/transforms.html#saveimage.


    Examples:
        - User can use the one-liner to start the Auto3Dseg workflow

        .. code-block:: bash

            python -m monai.apps.auto3dseg AutoRunner run --input             '{"modality": "ct", "datalist": "dl.json", "dataroot": "/dr", "multigpu": true, "class_names": ["A", "B"]}'

        - User can also save the input dictionary as a input YAML file and use the following one-liner

        .. code-block:: bash

            python -m monai.apps.auto3dseg AutoRunner run --input=./input.yaml

        - User can specify work_dir and data source config input and run AutoRunner:

        .. code-block:: python

            work_dir = "./work_dir"
            input = "path/to/input_yaml"
            runner = AutoRunner(work_dir=work_dir, input=input)
            runner.run()

        - User can specify a subset of algorithms to use and run AutoRunner:

        .. code-block:: python

            work_dir = "./work_dir"
            input = "path/to/input_yaml"
            algos = ["segresnet", "dints"]
            runner = AutoRunner(work_dir=work_dir, input=input, algos=algos)
            runner.run()

        - User can specify a local folder with algorithms templates and run AutoRunner:

        .. code-block:: python

            work_dir = "./work_dir"
            input = "path/to/input_yaml"
            algos = "segresnet"
            templates_path_or_url = "./local_path_to/algorithm_templates"
            runner = AutoRunner(work_dir=work_dir, input=input, algos=algos, templates_path_or_url=templates_path_or_url)
            runner.run()

        - User can specify training parameters by:

        .. code-block:: python

            input = "path/to/input_yaml"
            runner = AutoRunner(input=input)
            train_param = {
                "num_epochs_per_validation": 1,
                "num_images_per_batch": 2,
                "num_epochs": 2,
            }
            runner.set_training_params(params=train_param)  # 2 epochs
            runner.run()

        - User can specify the fold number of cross validation

        .. code-block:: python

            input = "path/to/input_yaml"
            runner = AutoRunner(input=input)
            runner.set_num_fold(n_fold = 2)
            runner.run()

        - User can specify the prediction parameters during algo ensemble inference:

        .. code-block:: python

            input = "path/to/input_yaml"
            pred_params = {
                'files_slices': slice(0,2),
                'mode': "vote",
                'sigmoid': True,
            }
            runner = AutoRunner(input=input)
            runner.set_prediction_params(params=pred_params)
            runner.run()

        - User can define a grid search space and use the HPO during training.

        .. code-block:: python

            input = "path/to/input_yaml"
            runner = AutoRunner(input=input, hpo=True)
            runner.set_nni_search_space({"learning_rate": {"_type": "choice", "_value": [0.0001, 0.001, 0.01, 0.1]}})
            runner.run()

    Notes:
        Expected results in the work_dir as below::

            work_dir/
            ├── algorithm_templates # bundle algo templates (scripts/configs)
            ├── cache.yaml          # Autorunner will automatically cache results to save time
            ├── datastats.yaml      # datastats of the dataset
            ├── dints_0             # network scripts/configs/checkpoints and pickle object of the algo
            ├── ensemble_output     # the prediction of testing datasets from the ensemble of the algos
            ├── input.yaml          # copy of the input data source configs
            ├── segresnet_0         # network scripts/configs/checkpoints and pickle object of the algo
            ├── segresnet2d_0       # network scripts/configs/checkpoints and pickle object of the algo
            └── swinunetr_0         # network scripts/configs/checkpoints and pickle object of the algo

    zdict | Noneanalyze_params
./work_dirNFr   Tstrzdict[str, Any] | str | Nonezdict | list | str | Nonezbool | Noneboolz
str | Noner   )work_dirinputalgosanalyzealgo_gentrainhpohpo_backendensemblenot_use_cachetemplates_path_or_url
allow_skipmlflow_tracking_urimlflow_experiment_namekwargsc                 K  s  |d krLt jt jt j|drLt jt j|d}td|  t | _t	|trf|| _nBt	|t
rt j|rt|| _td|  nt| dd| jkr| jd }t j|| _td| j  t j| jdd t j| jd| _|| _|| _|| _|
| _t j| jd	| _|  | _|   |d krP| jd
  n|| _|d krl| jd  n|| _|| _|	| _|ot| _|| _|| _ || _!t"|| _#dD ]6}|| jkrt	| j| t$rt%| || j|  qdD ]$}|| jkrt%| || j|  qdddh&| j' }t(|dkr>td| t j)| jd sftd| jd  t j| jt j*| jd }|| jd krz&t+,| jd | td|  W n t+j-k
r   Y nX | j.|d| _/d| jkrt0| jd }td| d n| j/}td| d| d || jd< tj1| j| jddd | jd | _2t j| jd| _3|| _4| 5  | 6  | 7  | 8  | 9  | j:|d d| _;i | _<| j= d krt>d!| jot| _| ?  i | _@d| _Ad"| j#krd"| jkr| jd" | j#d"< d S )#Nz
input.yamlz0Input config is not provided, using the default zLoading input config z is not a valid file or dictr   z AutoRunner using work directory T)exist_okz
cache.yamlr!   r"   )r!   r"   r#   r$   r&   r'   r)   )r    r%   r(   r*   r+   datarootdatalistmodalityr   zConfig keys are missing zDatalist file is not found z!Datalist was copied to work_dir: )datalist_filenamenum_foldzSetting num_fold z based on the input config. based on the input datalist .yamlF)configfilepathfmt	sort_keyszdatastats.yaml)r2   r   z HPOGen backend only supports NNIsigmoid)BospathisfilejoinabspathloggerinfodictZdata_src_cfg
isinstancer   r   load_config_file
ValueErrorr   makedirsdata_src_cfg_namer    r(   r)   r'   cache_filename
read_cachecacheexport_cacher!   r"   r#   r&   has_nnir$   r%   r*   r+   r   r,   r   setattr
differencekeyslenexistsbasenameshutilcopyfileSameFileErrorinspect_datalist_foldsmax_foldintexport_config_filer.   datastats_filenamer1   set_training_paramsset_device_infoset_prediction_paramsset_analyze_paramsset_ensemble_methodset_num_foldgpu_customizationgpu_customization_specslowerNotImplementedErrorset_hpo_paramssearch_space	hpo_tasks)selfr   r   r    r!   r"   r#   r$   r%   r&   r'   r(   r)   r*   r+   r,   parammissing_keysr1   r2    rk   U/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/apps/auto3dseg/auto_runner.py__init__   s    &




	
   zAutoRunner.__init__c                 C  s   ddddd}| j s"tj| js&|S t| j}| D ]\}}||| q:|d rt	|d t
rvtj|d sd|d< d|d< |d rt| jdd}t|dkrd|d< |d	 rt| jd
d}t|dkrd|d	< |S )af  
        Check if the intermediate result is cached after each step in the current working directory

        Returns:
            a dict of cache results. If not_use_cache is set to True, or there is no cache file in the
            working directory, the result will be ``empty_cache`` in which all ``has_cache`` keys are
            set to False.
        FN)r!   	datastatsr"   r#   r!   rn   r"   Zonly_trainedr   r#   T)r'   r;   r<   r=   rH   r   rD   items
setdefaultrC   r   r   r   rP   )rh   empty_cacherJ   kvhistoryZtrained_historyrk   rk   rl   rI   S  s&    
zAutoRunner.read_cachec                 K  s(   | j | tj| j | jdddd dS )zQ
        Save the cache state as ``cache.yaml`` in the working directory
        r5   NF)r8   default_flow_styler9   )rJ   updater   rY   rH   )rh   r,   rk   rk   rl   rK   x  s        zAutoRunner.export_cacherX   )r1   returnc                 C  s  t |}d|kr"tdt| dd |d D }t|dkrt|d }td| d| d	 tt||krtd
|d  nbd|krft|d dkrftd |d D ]}d|d< q|d D ]}d|d< qdd |d D }tdt| d| |d D ]&}|d |krd|d< ||d = q|d t	|
  |d< t j||ddd d}nd}td| d| d ddlm} ||ddd}t||d D ]*\}	\}
}|D ]}|	|d | d< qqt j||ddd |S )a  
        Returns number of folds in the datalist file, and assigns fold numbers if not provided.

        Args:
            datalist_filename: path to the datalist file.

        Notes:
            If the fold key is not provided, it auto generates 5 folds assignments in the training key list.
            If validation key list is available, then it assumes a single fold validation.
        trainingz#Datalist files has no training key:c                 S  s    g | ]}d |krt |d  qS )fold)rX   .0drk   rk   rl   
<listcomp>  s      z5AutoRunner.inspect_datalist_folds.<locals>.<listcomp>r      zFound num_fold r3   r4   z*Fold numbers are not continuous from 0 to 
validationzUNo fold numbers provided, attempting to use a single fold based on the validation keyrz   c                 S  s   i | ]}d |kr|d  |qS )labelrk   r{   rk   rk   rl   
<dictcomp>  s       z5AutoRunner.inspect_datalist_folds.<locals>.<dictcomp>zFound z8 items in the validation key, saving updated datalist tor   json   )r8   indent   z Datalist has no folds specified z...Generating z[ folds randomly.Please consider presaving fold numbers beforehand for repeated experiments.)KFoldT)Zn_splitsshufflerandom_state)r   rD   rE   r   rP   maxr@   rA   setlistvaluesrY   warningswarnZsklearn.model_selectionr   	enumeratesplit)rh   r1   r/   Z	fold_listr2   r}   Z
val_labelsr   kfi_Z	valid_idxvirk   rk   rl   rV     sL    



 z!AutoRunner.inspect_datalist_foldszdict[str, Any] | None)ra   rb   rx   c                 C  s   || _ |dk	r|| _| S )a  
        Set options for GPU-based parameter customization/optimization.

        Args:
            gpu_customization: the switch to determine automatically customize/optimize bundle script/config
                parameters for each bundleAlgo based on gpus. Custom parameters are obtained through dummy
                training to simulate the actual model training process and hyperparameter optimization (HPO)
                experiments.
            gpu_customization_specs (optional): the dictionary to enable users overwrite the HPO settings. user can
                overwrite part of variables as follows or all of them. The structure is as follows.

                .. code-block:: python

                    gpu_customization_specs = {
                        'ALGO': {
                            'num_trials': 6,
                            'range_num_images_per_batch': [1, 20],
                            'range_num_sw_batch_size': [1, 20]
                        }
                    }

            ALGO: the name of algorithm. It could be one of algorithm names (e.g., 'dints') or 'universal' which
                would apply changes to all algorithms. Possible options are

                - {``"universal"``, ``"dints"``, ``"segresnet"``, ``"segresnet2d"``, ``"swinunetr"``}.

            num_trials: the number of HPO trials/experiments to run.
            range_num_images_per_batch: the range of number of images per mini-batch.
            range_num_sw_batch_size: the range of batch size in sliding-window inferer.
        N)ra   rb   )rh   ra   rb   rk   rk   rl   set_gpu_customization  s    !z AutoRunner.set_gpu_customizationr   )r2   rx   c                 C  sD   |dkrt d| || jkr:t d| j d| j d|| _| S )z
        Set the number of cross validation folds for all algos.

        Args:
            num_fold: a positive integer to define the number of folds.
        r   zEnum_fold is expected to be an integer greater than zero. Now it gets z1num_fold is greater than the maximum fold number z in r4   )rE   rW   r1   r2   )rh   r2   rk   rk   rl   r`     s    
zAutoRunner.set_num_fold)paramsrx   c                 C  s0   |dk	rt |ni | _d| jkr,tdt | S )a  
        Set the training params for all algos.

        Args:
            params: a dict that defines the overriding key-value pairs during training. The overriding method
                is defined by the algo class.

        Examples:
            For BundleAlgo objects, the training parameter to shorten the training time to a few epochs can be
                {"num_epochs": 2, "num_epochs_per_validation": 1}

        NCUDA_VISIBLE_DEVICESz]CUDA_VISIBLE_DEVICES is deprecated from 'set_training_params'. Use 'set_device_info' instead.)r   train_paramsr   r   DeprecationWarningrh   r   rk   rk   rl   r[     s    
zAutoRunner.set_training_paramszlist[int] | str | Nonez
int | None)cuda_visible_devices	num_nodesmn_start_method
cmd_prefixrx   c                 C  sT  i | _ |dkrtjd}|dkrXddd ttj D | j d< tj | j d< ntt	|t
r|| j d< t|d| j d< nJt	|ttfrddd |D | j d< t|| j d< ntd| d	 |dkrttjd
d}|| j d
< |dkrtjdd}|| j d< |dkr*tjdd}|| j d< |dk	rPtd| d | S )a  
        Set the device related info

        Args:
            cuda_visible_devices: define GPU ids for data analyzer, training, and ensembling.
                List of GPU ids [0,1,2,3] or a string "0,1,2,3".
                Default using env "CUDA_VISIBLE_DEVICES" or all devices available.
            num_nodes: number of nodes for training and ensembling.
                Default using env "NUM_NODES" or 1 if "NUM_NODES" is unset.
            mn_start_method: multi-node start method. Autorunner will use the method to start multi-node processes.
                Default using env "MN_START_METHOD" or 'bcprun' if "MN_START_METHOD" is unset.
            cmd_prefix: command line prefix for subprocess running in BundleAlgo and EnsembleRunner.
                Default using env "CMD_PREFIX" or None, examples are:

                    - single GPU/CPU or multinode bcprun: "python " or "/opt/conda/bin/python3.8 ",
                    - single node multi-GPU running "torchrun --nnodes=1 --nproc_per_node=2 "

                If user define this prefix, please make sure --nproc_per_node matches cuda_visible_device or
                os.env['CUDA_VISIBLE_DEVICES']. Also always set --nnodes=1. Set num_nodes for multi-node.
        Nr   ,c                 S  s   g | ]}t |qS rk   r   r|   xrk   rk   rl   r~   4  s     z.AutoRunner.set_device_info.<locals>.<listcomp>	n_devicesc                 S  s   g | ]}t |qS rk   r   r   rk   rk   rl   r~   :  s     z%Wrong format of cuda_visible_devices z, devices not setZ	NUM_NODESr   ZMN_START_METHODZbcprunZ
CMD_PREFIX z*Using user defined command running prefix z, will override other settings)device_settingr;   environgetr>   rangetorchcudadevice_countrC   r   rP   r   r   tupler@   r   rX   rA   )rh   r   r   r   r   rk   rk   rl   r\     s2    $







zAutoRunner.set_device_infoAlgoEnsembleBestByFold)ensemble_method_namer,   rx   c                 K  s"   t |ddgd| _| j| | S )a  
        Set the bundle ensemble method name and parameters for save image transform parameters.

        Args:
            ensemble_method_name: the name of the ensemble method. Only two methods are supported "AlgoEnsembleBestN"
                and "AlgoEnsembleBestByFold".
            kwargs: the keyword arguments used to define the ensemble method. Currently only ``n_best`` for
                ``AlgoEnsembleBestN`` is supported.
        AlgoEnsembleBestNr   )	supported)r   r   r,   rw   )rh   r   r,   rk   rk   rl   r_   P  s    
 zAutoRunner.set_ensemble_method)r,   rx   c                 K  s2   t t|\}}|r | j| nt| d| S )a#  
        Set the ensemble output transform.

        Args:
            kwargs: image writing parameters for the ensemble inference. The kwargs format follows SaveImage
                transform. For more information, check https://docs.monai.io/en/stable/transforms.html#saveimage.

        z are not supported in monai.transforms.SaveImage,Check https://docs.monai.io/en/stable/transforms.html#saveimage for more information.)r   r   r,   rw   rE   )rh   r,   Zare_all_args_present
extra_argsrk   rk   rl   set_image_save_transforma  s    
z#AutoRunner.set_image_save_transformc                 C  s   |dk	rt |ni | _| S )a  
        Set the prediction params for all algos.

        Args:
            params: a dict that defines the overriding key-value pairs during prediction. The overriding method
                is defined by the algo class.

        Examples:

            For BundleAlgo objects, this set of param will specify the algo ensemble to only inference the first
                two files in the testing datalist {"file_slices": slice(0, 2)}

        N)r   pred_paramsr   rk   rk   rl   r]   v  s    z AutoRunner.set_prediction_paramsc                 C  s$   |dkrddd| _ n
t|| _ | S )z
        Set the data analysis extra params.

        Args:
            params: a dict that defines the overriding key-value pairs during training. The overriding method
                is defined by the algo class.

        NFr   )Zdo_ccpdevice)r   r   r   rk   rk   rl   r^     s    	
zAutoRunner.set_analyze_paramsc                 C  s   |dkr| j n|| _| S )a  
        Set parameters for the HPO module and the algos before the training. It will attempt to (1) override bundle
        templates with the key-value pairs in ``params`` (2) change the config of the HPO module (e.g. NNI) if the
        key is found to be one of:

            - "trialCodeDirectory"
            - "trialGpuNumber"
            - "trialConcurrency"
            - "maxTrialNumber"
            - "maxExperimentDuration"
            - "tuner"
            - "trainingService"

        and (3) enable the dry-run mode if the user would generate the NNI configs without starting the NNI service.

        Args:
            params: a dict that defines the overriding key-value pairs during instantiation of the algo. For
                BundleAlgo, it will override the template config filling.

        Notes:
            Users can set ``nni_dry_run`` to ``True`` in the ``params`` to enable the dry-run mode for the NNI backend.

        N)r   
hpo_paramsr   rk   rk   rl   re     s    zAutoRunner.set_hpo_paramszdict[str, Any])rf   rx   c                 C  sX   d}|  D ]:\}}d|kr6t| d| d| d|t|d 9 }q|| _|| _| S )a$  
        Set the search space for NNI parameter search.

        Args:
            search_space: hyper parameter search space in the form of dict. For more information, please check
                NNI documentation: https://nni.readthedocs.io/en/v2.2/Tutorial/SearchSpaceSpec.html .
        r   _valuez key z value z has not _value)rp   rE   rP   rf   rg   )rh   rf   Zvalue_combinationsrs   rt   rk   rk   rl   set_nni_search_space  s    zAutoRunner.set_nni_search_spacezlist[dict[str, Any]]None)ru   rx   c                 C  sn   |D ]d}|t j }t|jdr0|| j| j n|| j | }tt j|i}t	|fd|j
i| qdS )a  
        Train the Algos in a sequential scheme. The order of training is randomized.

        Args:
            history: the history of generated Algos. It is a list of dicts. Each element has the task name
                (e.g. "dints_0" for dints network in fold 0) as the key and the algo object as the value.
                After the training, the algo object with the ``best_metric`` will be saved as a pickle file.

        Note:
            The final results of the model training will be written to all the generated algorithm's output
            folders under the working directory. The results include the model checkpoints, a
            progress.yaml, accuracies in CSV and a pickle file of the Algo object.
        r   template_pathN)r   ALGOr   r#   r   r   	get_scorer   SCOREr   r   )rh   ru   	algo_dictalgoaccZalgo_meta_datark   rk   rl   _train_algo_in_sequence  s    
z"AutoRunner._train_algo_in_sequencec              	   C  s  dt j dddddiddd	d
}tt| jdd}| jdd}|D ]v}|tj	 }|tj
 }t|| jd}| }	t|}
| jD ]}||
kr| j| |
|< q|
d|i |
d| ji d|	 d | j }|
d|i tjtj| j| d}tj|
|ddd t| jtt|d }d| d }|rJtd|  qHt| dd tt| jdd}|| |k rtd tt| jdd}qld}t| dd td|  |}qHdS )a!  
        Train the Algos using HPO.

        Args:
            history: the history of generated Algos. It is a list of dicts. Each element has the task name
                (e.g. "dints_0" for dints network in fold 0) as the key and the algo object as the value.
                After the training, the algo object with the ``best_metric`` will be saved as a pickle file.

        Note:
            The final results of the model training will not be written to all the previously generated
            algorithm's output folders. Instead, HPO will generate a new algo during the searching, and
            the new algo will be saved under the working directory with a different format of the name.
            For example, if the searching space has "learning_rate", the result of HPO will be written to
            a folder name with original task name and the param (e.g. "dints_0_learning_rate_0.001").
            The results include the model checkpoints, a progress.yaml, accuracies in CSV and a pickle
            file of the Algo object.

        r4   r   
   1hnameZ
GridSearchlocalT)platformZuseActiveGpu)ZtrialCodeDirectoryZtrialGpuNumberZtrialConcurrencymaxTrialNumberZmaxExperimentDurationZtunerZtrainingServicero   Znni_dry_runF)r   r   ZexperimentNamerf   z/python -m monai.apps.auto3dseg NNIGen run_algo  ZtrialCommandz_nni_config.yamlr5   N)r8   rv   r   znnictl create --config z --port 8088z;AutoRunner HPO is in dry-run mode. Please manually launch: )checkznnictl stop --allzNNI completes HPO on )r   r   r   rP   r   r   r   popr   IDr   r
   Zget_obj_filenamer   rw   rf   r;   r<   r?   r>   r   rY   minrg   r   rX   r@   rA   r   r   r   )rh   ru   Zdefault_nni_configZlast_total_tasksZmode_dry_runr   r   r   Znni_genZobj_filenameZ
nni_configkeyZ	trial_cmdZnni_config_filenameZ	max_trialcmdZn_trainingsrk   rk   rl   _train_algo_in_nni  sN    




zAutoRunner._train_algo_in_nnic              	   C  sB  | j r`| jdk	r`td t| j| jfd| ji| j}|  d}t	j
  | jd| jd n
td | jrtj| jstd| j dt| j| j| j| j| j| j| jd	}| jr|j| j| j| j| j| jd
 n|j| j| j| jd | }t| | jdd n
td | j dk}| j s>|r| j!d st"| jdd}t#|dkrltd| j d|rdd |D }|rtd| d dd |D }t#|dkr| j$s| %| n
| &| | jdd n
td | j'r4t(f | j| j| j| j)t*| j+d dkd| j,| j-}|.| j+ td dS )z-
        Run the AutoRunner pipeline
        NzRunning data analysis...Zoutput_pathT)r!   rn   zSkipping data analysis...z"Could not find the datastats file z=. Possibly the required data analysis step was not completed.)r    Z	algo_pathr(   Zdata_stats_filenamerG   r*   r+   )r2   ra   rb   r)   )r2   r)   )r"   z Skipping algorithm generation...r#   Fro   r   z#Could not find training scripts in zE. Possibly the required algorithms generation step was not completed.c                 S  s    g | ]}|t j r|t j qS rk   )r   
IS_TRAINEDr   r|   hrk   rk   rl   r~   d  s     
 z"AutoRunner.run.<locals>.<listcomp>zSkipping already trained algos z3.Set option train=True to always retrain all algos.c                 S  s   g | ]}|t j s|qS rk   )r   r   r   rk   rk   rl   r~   j  s     
 )r#   zSkipping algorithm training...r   r   )rG   r   r2   r   Zmgpuz-Auto3Dseg pipeline is completed successfully.)/r!   r   r@   rA   r   r1   r.   rZ   Zget_all_case_statsr   r   rr   rK   r"   r;   r<   r=   rE   r   r    r   r(   rG   r*   r+   ra   generater2   rb   r)   Zget_historyr   r#   rJ   r   rP   r$   r   r   r&   r	   r   rX   r   r,   r   run)rh   daZbundle_generatorru   Zauto_train_choiceZ
skip_algosZensemble_runnerrk   rk   rl   r   #  s    
 







	zAutoRunner.run)r   NNNNNFr   TFNTNN)FN)r   )N)NNNN)r   )N)N)N)__name__
__module____qualname____doc____annotations__rm   rI   rK   rV   r   r`   r[   r\   r_   r   r]   r^   re   r   r   r   r   rk   rk   rk   rl   r   (   sR   
               , %	E   '    ;Cr   )-
__future__r   r;   rS   r   copyr   timer   typingr   r   r   Zmonai.apps.auto3dseg.bundle_genr   Z"monai.apps.auto3dseg.data_analyzerr   Z%monai.apps.auto3dseg.ensemble_builderr	   Zmonai.apps.auto3dseg.hpo_genr
   Zmonai.apps.auto3dseg.utilsr   r   monai.apps.utilsr   Zmonai.auto3dseg.utilsr   Zmonai.bundler   monai.transformsr   monai.utilsr   r   r   r   monai.utils.miscr   r   r   r@   r   rL   r   rk   rk   rk   rl   <module>   s*   
