import logging
from pathlib import Path
from typing import Dict, List, Union
from ewokstools.submit import save_and_execute, wait_to_finish_queue
from ewoksutils.task_utils import task_inputs
from ._utils import (
SLURM_JOB_PARAMETERS_BACKGROUND,
EwoksWorkflow,
_confirm_submission,
_get_destination_filename,
_normalize_slurm_job_parameters,
_print_checking_template_format,
_print_job_progress,
_print_number_of_jobs_to_submit,
_print_template_invalid,
_print_template_valid,
_validate_yaml_template_model,
_warning_dry_run_mode,
_warning_if_too_many_jobs,
get_list_scan_info_from_bliss_filenames_id13,
)
from .resources import BACKGROUND_WORKFLOW
from .resources.models import BackgroundModel
logger = logging.getLogger(__name__)
[docs]
def main_background(args):
for file in args.FILES:
if file.endswith((".yaml", ".yml")):
_main_background_from_template(filename=file)
elif file.endswith(".h5"):
kwargs = vars(args)
for key in ("FILES", "command"):
kwargs.pop(key, None)
_main_background_from_cli(filename=file, **kwargs)
else:
logger.warning(
f"File {file} has an unsupported extension. Skipping it. Supported extensions are .yaml, .yml and .h5."
)
def _main_background_from_template(filename: Union[str, Path]):
_print_checking_template_format()
is_valid, result = _validate_yaml_template_model(filename, BackgroundModel)
if not is_valid:
_print_template_invalid(filename, result)
return
_print_template_valid()
_main_background(**result)
def _main_background_from_cli(filename: Union[str, Path], **kwargs):
integration_params = {
"processed_filenames": [filename],
"nxprocess_path_diffmap": kwargs.pop("nxprocess_path_diffmap"),
"nxprocess_path_filtered": kwargs.pop("nxprocess_path_filtered"),
}
submit_parameters = {
"submit": not kwargs.pop("dry_run"),
"slurm_job_parameters": kwargs.pop("slurm_job_parameters") or {},
"queue": kwargs.pop("queue", None),
"local_execution": kwargs.pop("local_execution", False),
}
_main_background(
submit_parameters=submit_parameters,
**integration_params,
**kwargs,
)
def _main_background(
processed_filenames: Union[str, List[str]],
nxprocess_path_diffmap: str,
nxprocess_path_filtered: str,
submit_parameters: Dict = None,
**kwargs,
):
workflow_background = EwoksWorkflow(BACKGROUND_WORKFLOW)
submit_parameters = submit_parameters or {}
dry_run = not submit_parameters.get("submit")
list_scans_info = get_list_scan_info_from_bliss_filenames_id13(
bliss_filenames=processed_filenames,
**kwargs,
)
nb_total_jobs = len(list_scans_info)
_print_number_of_jobs_to_submit(nb_total_jobs)
_confirm_submission()
if dry_run:
_warning_dry_run_mode()
else:
_warning_if_too_many_jobs(nb_total_jobs)
submitted_jobs = []
for index_job, scan_info in enumerate(list_scans_info):
output_filename = scan_info.filename_raw_dataset
external_output_filename = output_filename
nxprocess_name = nxprocess_path_filtered.split("/")[-1]
inputs_ewoks = task_inputs(
id=workflow_background.background.id,
task_identifier=workflow_background.background.task_identifier,
inputs={
"nxdata_url": f"{scan_info.filename_raw_dataset}::{nxprocess_path_diffmap}/diffmap",
"nxprocess_name": nxprocess_name,
"use_neuralnetwork": False,
"destination_file": _get_destination_filename(
output_filename, scan_info.scan_nb, nxprocess_name
),
"worker_module": "scattering",
"wait_to_finish": False,
"output_filename": output_filename,
"external_output_filename": external_output_filename,
"submit_to_slurm": False,
"do_background_removal": True,
"force_training": False,
**kwargs.get("background", {}),
},
)
submit_parameters_ = submit_parameters.copy()
submitted = save_and_execute(
python_package="ewoksid13",
dry_run=dry_run,
workflow=BACKGROUND_WORKFLOW,
inputs=inputs_ewoks,
destination_filename=_get_destination_filename(
output_filename, scan_info.scan_nb, nxprocess_name
),
slurm_job_parameters=_normalize_slurm_job_parameters(
SLURM_JOB_PARAMETERS_BACKGROUND,
submit_parameters_.pop("slurm_job_parameters", {}),
),
**submit_parameters_,
)
submitted_jobs.append(submitted)
_print_job_progress(index_job + 1, nb_total_jobs, dry_run)
wait_to_finish_queue(submitted_jobs)