import logging
from pathlib import Path
from typing import Dict, List, Union, Optional
from ewokstools.submit import save_and_execute, wait_to_finish_queue, Submitted
from ewoksutils.task_utils import task_inputs
from ._utils import (
SLURM_JOB_PARAMETERS_INTEGRATE,
EwoksWorkflow,
_compute_pyfai_table_widths,
_confirm_submission,
_get_destination_filename,
_get_pyfai_config_filename,
_make_nxprocess_name,
_normalize_external_output_filenames,
_normalize_lima_url_template_id13,
_normalize_pyfai_configurations,
_normalize_pyfai_tag,
_normalize_slurm_job_parameters,
_normalize_tag,
_print_checking_template_format,
_print_integration_plan,
_print_job_progress,
_print_pyfai_config_row,
_print_scan_table_header,
_print_scan_table_row,
_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 INTEGRATION_WORKFLOW, INTEGRATION_WORKFLOW_SAVECONFIG
from .resources.models import PyFAIIntegrationModel
logger = logging.getLogger(__name__)
[docs]
def main_pyfai_integration(args):
for file in args.FILES:
if file.endswith((".yaml", ".yml")):
_main_pyfai_integration_from_template(filename=file)
elif file.endswith(".h5"):
kwargs = vars(args)
for key in ("FILES", "command"):
kwargs.pop(key, None)
_main_pyfai_integration_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_pyfai_integration_from_template(filename: Union[str, Path]):
_print_checking_template_format()
is_valid, result = _validate_yaml_template_model(filename, PyFAIIntegrationModel)
if not is_valid:
_print_template_invalid(filename, result)
return
_print_template_valid()
_main_pyfai_integration(**result)
def _main_pyfai_integration_from_cli(filename: Union[str, Path], **kwargs):
if kwargs.get("pyfai_configurations") is None:
logger.error("PyFAI configuration file is required for integration.")
return
integration_params = {
"bliss_filenames": [filename],
"detectors": kwargs.pop("detectors"),
"scans": kwargs.pop("scans"),
"save_separate_files": kwargs.pop("save_separate_files"),
"output_root_folder": kwargs.pop("output_root_folder"),
"tag": kwargs.pop("tag"),
}
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),
}
pyfai_configurations = [
{"filename": file} for file in kwargs.pop("pyfai_configurations")
]
for pyfai_config in pyfai_configurations:
pyfai_config.update({k: v for k, v in kwargs.items() if v is not None})
_main_pyfai_integration(
pyfai_configurations=pyfai_configurations,
submit_parameters=submit_parameters,
**integration_params,
**kwargs,
)
def _main_pyfai_integration(
bliss_filenames: Union[str, List[str]],
pyfai_configurations: Union[Dict, List[Dict], List[str]],
detectors: Union[str, List[str]] = None,
scans: Union[int, List[int]] = None,
subscans: Union[int, List[int]] = None,
save_separate_files: bool = False,
save_pyfai_config: bool = True,
output_root_folder: Optional[str] = None,
tag: Optional[str] = "",
submit_parameters: Optional[Dict] = None,
wait_for_completion: Optional[bool] = True,
**kwargs,
) -> List[Submitted]:
"""
Main function to submit a simple pyFAI integration workflow.
param: bliss_filenames: str or list of str
One or several bliss files containing the scans to reprocess. Can be provided as a list of filenames or as a single string.
param: pyfai_configurations:
One or several pyFAI configurations to apply for the reprocessing.
Can be provided as a list of dictionaries or as a single dictionary.
Each configuration should contain:
- "filename": the path to the .poni/.json file with the pyFAI configuration.
- "tag" (optional): a tag to differentiate the output files if several pyFAI configurations are provided.
- **additional_parameters for pyFAI worker
param: detectors: str or list of str, optional
The detector name(s) to reprocess. If not provided, all 2D detectors will be reprocessed.
param: scans: int or list of int, optional
The scan number(s) to reprocess. If not provided, all scans will be reprocessed.
param: subscans: int or list of int, optional
The subscan number(s) to reprocess. If not provided, all subscans will be reprocessed.
param: save_separate_files: bool, optional
Whether to save integrated data in the dataset file or linked to separated files.
param: save_pyfai_config: bool, optional
Whether to save the pyFAI configuration file used for the integration in the same folder as the integrated files.
param: output_root_folder: str, optional
The root output folder where the reprocessed files will be saved. If not provided, the files will be saved in PROCESSED_DATA or in the same directory as the input files.
param: tag: str, optional
A tag to add to the output files and job names to differentiate them from other reprocessings.
param: submit_parameters: dict, optional
Additional parameters for job submission, such as whether to actually submit the jobs or just do a dry run, and parameters for slurm job submission.
**kwargs:
All the kwargs should be dictionaries whose keys are the id of the nodes in the workflow.
Hence, you can pass any input for any task through this dictionary, for example:
"integrate" : {
"prioritize_nonnative_h5items": False,
}
"""
workflow_path = (
INTEGRATION_WORKFLOW_SAVECONFIG if save_pyfai_config else INTEGRATION_WORKFLOW
)
workflow_integration = EwoksWorkflow(workflow_path)
submit_parameters = submit_parameters or {}
dry_run = not submit_parameters.get("submit", True)
list_scans_info = get_list_scan_info_from_bliss_filenames_id13(
bliss_filenames=bliss_filenames,
detectors=detectors,
scans=scans,
subscans=subscans,
output_root_folder=output_root_folder,
**kwargs,
)
pyfai_configurations = _normalize_pyfai_configurations(pyfai_configurations)
nb_pyfai_configs = len(pyfai_configurations)
nb_scans = len(list_scans_info)
nb_total_jobs = nb_pyfai_configs * nb_scans
_print_integration_plan(nb_scans, nb_pyfai_configs, nb_total_jobs)
_confirm_submission()
if dry_run:
_warning_dry_run_mode()
else:
_warning_if_too_many_jobs(nb_total_jobs)
submitted_jobs = []
index_job = 0
tag = _normalize_tag(tag, dry_run)
scan_table_widths = _print_scan_table_header(list_scans_info)
pyfai_table_widths = _compute_pyfai_table_widths(pyfai_configurations)
for scan_info in list_scans_info:
_print_scan_table_row(scan_info, scan_table_widths)
output_filename, external_output_filename = (
_normalize_external_output_filenames(
scan_info=scan_info,
save_separate_files=save_separate_files,
)
)
for index_pyfai, pyfai_config in enumerate(pyfai_configurations):
_print_pyfai_config_row(pyfai_config, pyfai_table_widths)
index_job += 1
inputs_ewoks_reprocess = []
suffix = f"{tag}{_normalize_pyfai_tag(pyfai_config, index_pyfai)}"
# Inputs for pyFAI config generation
inputs_ewoks_reprocess += task_inputs(
id=workflow_integration.pyfai_config.id,
task_identifier=workflow_integration.pyfai_config.task_identifier,
inputs={
"filename": pyfai_config.get("filename"),
"integration_options": pyfai_config,
},
)
# Inputs for Saving pyFAI config file
if save_pyfai_config:
inputs_ewoks_reprocess += task_inputs(
id=workflow_integration.save_pyfai_config.id,
task_identifier=workflow_integration.save_pyfai_config.task_identifier,
inputs={
"output_filename": _get_pyfai_config_filename(
output_filename,
scan_info.scan_nb,
scan_info.detector_name,
suffix,
),
},
)
# Inputs for IntegrateBlissScan task
nxprocess_name = _make_nxprocess_name(
scan_info.detector_name, workflow_integration.integrate.id, suffix
)
inputs_ewoks_reprocess += task_inputs(
id=workflow_integration.integrate.id,
task_identifier=workflow_integration.integrate.task_identifier,
inputs={
"filename": scan_info.filename_raw_dataset,
"detector_name": scan_info.detector_name,
"scan": scan_info.scan_nb,
"subscan": scan_info.subscan,
"lima_url_template": _normalize_lima_url_template_id13(
lima_url_template=scan_info.lima_url_template,
**kwargs.get("integrate", {}),
),
"lima_url_template_args": scan_info.lima_url_template_args,
"nxprocess_name": nxprocess_name,
"output_filename": output_filename,
"external_output_filename": external_output_filename,
**kwargs.get("integrate", {}),
},
)
submit_parameters_ = submit_parameters.copy()
submitted_jobs.append(
save_and_execute(
dry_run=dry_run,
workflow=workflow_path,
inputs=inputs_ewoks_reprocess,
destination_filename=_get_destination_filename(
output_filename, scan_info.scan_nb, nxprocess_name
),
slurm_job_parameters=_normalize_slurm_job_parameters(
SLURM_JOB_PARAMETERS_INTEGRATE,
submit_parameters_.pop("slurm_job_parameters", {}),
),
**submit_parameters_,
)
)
_print_job_progress(index_job, nb_total_jobs, dry_run)
if wait_for_completion:
wait_to_finish_queue(submitted_jobs)
return submitted_jobs