cryoemservices 0.1.3__py3-none-any.whl

cryoemservices/__init__.py

@@ -0,0 +1,3 @@
+from __future__ import annotations
+
+__version__ = "0.1.3"

cryoemservices/cli/resubmit_wrapper.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import workflows.transport.pika_transport as pt
+from workflows.recipe import RecipeWrapper
+
+
+def run():
+    parser = argparse.ArgumentParser(
+        description="Resubmit a failed zocalo wrapper script using the .recipewrap file"
+    )
+    parser.add_argument(
+        "-w",
+        "--wrapper",
+        help="Location of the .recipewrap wrapper file to resubmit",
+        dest="wrapper",
+        required=True,
+    )
+    parser.add_argument(
+        "-c",
+        "--config",
+        help="Transport configuration file for connecting to the message broker",
+        dest="config",
+        required=True,
+    )
+    args = parser.parse_args()
+
+    if not Path(args.wrapper).is_file():
+        print(f"{args.wrapper} cannot be found")
+        return
+    if not Path(args.config).is_file():
+        print(f"{args.config} cannot be found")
+        return
+
+    # Connect to the message transport
+    transport = pt.PikaTransport()
+    transport.load_configuration_file(args.config)
+    transport.connect()
+
+    # Load and submit the wrapper part of the recipe
+    with open(args.wrapper, "r") as wrap:
+        recipe = json.load(wrap)
+    rw = RecipeWrapper(message=recipe, transport=transport)
+    rw._send_to_destination(rw.recipe_pointer, None, rw.payload, {})
+
+
+if __name__ == "__main__":
+    run()

cryoemservices/pipeliner_plugins/combine_star_files.py

@@ -0,0 +1,251 @@
+from __future__ import annotations
+
+import argparse
+from math import ceil
+from pathlib import Path
+from typing import List
+
+import starfile
+
+
+def write_empty_particles_file(file_to_write, optics_dataframe, particles_dataframe):
+    """Write a particles star file with no particles, ready for appending to"""
+    with open(file_to_write, "w") as optics_file:
+        optics_file.write("data_optics\n\nloop_\n")
+        for optics_loop_tag in optics_dataframe.keys():
+            optics_file.write(f"_{optics_loop_tag}\n")
+        optics_file.write(" ".join(optics_dataframe.to_numpy(dtype=str)[0]))
+        optics_file.write("\n\n\n")
+
+        optics_file.write("data_particles\n\nloop_\n")
+        for particles_loop_tag in particles_dataframe.keys():
+            optics_file.write(f"_{particles_loop_tag}\n")
+
+
+def combine_star_files(files_to_process: List[Path], output_dir: Path):
+    """Combines any number of particle star files into a single file.
+
+    Parameters:
+    files_to_process: A list of the particle star files to combine
+    output_dir: The directory in which to save the combined "particles_all.star" file
+    """
+    total_particles = 0
+
+    # Never read particles_all.star first as it will be big
+    if files_to_process[0].name == "particles_all.star":
+        files_to_process.append(files_to_process[0])
+        files_to_process = files_to_process[1:]
+
+    # Make a temporary star file to get the table headings from
+    reference_optics = None
+    with open(files_to_process[0], "r") as full_starfile, open(
+        output_dir / ".particles_tmp.star", "w"
+    ) as tmp_starfile:
+        for line_counter in range(50):
+            line = full_starfile.readline()
+            if line.startswith("opticsGroup"):
+                reference_optics = line.split()
+            if not line:
+                break
+            tmp_starfile.write(line)
+
+    star_dictionary = starfile.read(output_dir / ".particles_tmp.star")
+    (output_dir / ".particles_tmp.star").unlink()
+
+    if not reference_optics:
+        raise IndexError(f"Cannot find optics group in {files_to_process[0]}")
+
+    write_empty_particles_file(
+        output_dir / ".particles_all_tmp.star",
+        star_dictionary["optics"],
+        star_dictionary["particles"],
+    )
+
+    number_of_star_files = 0
+    # Add the remaining files using append mode for speed and memory efficiency
+    for split_file in files_to_process:
+        # Check that the files have the same optics tables
+        with open(split_file, "r") as added_starfile:
+            while True:
+                optics_line = added_starfile.readline()
+                if not optics_line:
+                    raise IndexError(f"Cannot find optics group in {split_file}")
+                if optics_line.startswith("opticsGroup"):
+                    new_optics = optics_line.split()
+                    break
+
+        if len(new_optics) != len(reference_optics):
+            raise IndexError(
+                "Cannot combine star files with different length optics tables."
+            )
+        for optics_label in range(len(reference_optics)):
+            ref_value = reference_optics[optics_label]
+            new_value = new_optics[optics_label]
+            if ref_value[0].isdigit() and new_value[0].isdigit():
+                ref_value = float(ref_value)
+                new_value = float(new_value)
+            if ref_value != new_value:
+                print(ref_value, new_value)
+                raise IndexError(
+                    "Cannot combine star files with different values in optics tables."
+                )
+
+        # Add the particles lines to the final star file
+        file_particles_count = 0
+        with open(split_file, "r") as added_starfile, open(
+            output_dir / ".particles_all_tmp.star", "a"
+        ) as particles_file:
+            while True:
+                particle_line = added_starfile.readline()
+                if not particle_line:
+                    break
+                particle_split_line = particle_line.split()
+                if len(particle_split_line) > 0 and particle_split_line[0][0].isdigit():
+                    file_particles_count += 1
+                    total_particles += 1
+                    particles_file.write(particle_line)
+
+        print(f"Adding {split_file} with {file_particles_count} particles")
+        number_of_star_files += 1
+
+    (output_dir / ".particles_all_tmp.star").rename(output_dir / "particles_all.star")
+    print(
+        f"Combined {number_of_star_files} files into particles_all.star "
+        f"with {total_particles} particles"
+    )
+
+
+def split_star_file(
+    file_to_process: Path,
+    output_dir: Path,
+    number_of_splits: int or None = None,
+    split_size: int or None = None,
+):
+    """Splits a star file into subfiles.
+
+    The number of subfiles can be given with number_of_splits
+    or is determined by split_size, the number of particles for in each file.
+    """
+
+    # Make a temporary star file to get the table headings from
+    with open(file_to_process, "r") as full_starfile, open(
+        output_dir / ".particles_tmp.star", "w"
+    ) as tmp_starfile:
+        for line_counter in range(50):
+            line = full_starfile.readline()
+            if not line:
+                break
+            tmp_starfile.write(line)
+
+    star_dictionary = starfile.read(output_dir / ".particles_tmp.star")
+    (output_dir / ".particles_tmp.star").unlink()
+
+    # Find the number of lines in the full file
+    starfile_starter_lines = line_counter - star_dictionary["particles"].shape[0]
+    count = 0
+    with open(file_to_process, "r") as full_starfile:
+        for count, line in enumerate(full_starfile):
+            pass
+    number_of_particles = count + 1 - starfile_starter_lines
+
+    # Determine the number of files and size of the splits
+    if number_of_splits:
+        if split_size:
+            print(
+                "Warning: "
+                "Both number_of_splits and split_size have been given, "
+                "using number_of_splits.",
+            )
+        split_size = ceil(number_of_particles / number_of_splits)
+    elif split_size:
+        number_of_splits = ceil(number_of_particles / split_size)
+    else:
+        raise KeyError("Either number_of_splits or split_size must be given.")
+
+    with open(file_to_process, "r") as full_starfile:
+        # Read in the full file line by line, removing the header lines first
+        for start_line in range(starfile_starter_lines):
+            full_starfile.readline()
+
+        for split in range(number_of_splits):
+            # Give each new file the header information
+            write_empty_particles_file(
+                output_dir / f".particles_split{split+1}_tmp.star",
+                star_dictionary["optics"],
+                star_dictionary["particles"],
+            )
+
+            # Write particles to the split files by reading in lines from the full file
+            with open(
+                output_dir / f".particles_split{split+1}_tmp.star", "a"
+            ) as split_file:
+                for count in range(split_size):
+                    particle_line = full_starfile.readline()
+                    if not particle_line:
+                        break
+                    split_file.write(particle_line)
+
+            (output_dir / f".particles_split{split+1}_tmp.star").rename(
+                output_dir / f"particles_split{split+1}.star"
+            )
+
+    print(
+        f"Split {number_of_particles} particles into "
+        f"{number_of_splits} files with {split_size} particles in each."
+    )
+
+
+def create_parser():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "files_to_process",
+        type=Path,
+        nargs="+",
+        help="The star files from which to combine particles.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        dest="output_dir",
+        type=Path,
+        help="Folder in which to save the new star files.",
+    )
+
+    parser.add_argument(
+        "--split",
+        dest="do_split",
+        action="store_true",
+        default=False,
+        help="Whether to split the particles again into new star files.",
+    )
+    parser.add_argument(
+        "--n_files",
+        dest="n_files",
+        type=int,
+        default=None,
+        help="Number of files to split the particles into.",
+    )
+    parser.add_argument(
+        "--split_size",
+        dest="split_size",
+        type=int,
+        default=None,
+        help="Number of particles to write per file.",
+    )
+
+    return parser
+
+
+def main():
+    arg_parser = create_parser()
+    run_args = vars(arg_parser.parse_args())
+
+    combine_star_files(run_args["files_to_process"], run_args["output_dir"])
+
+    if run_args["do_split"]:
+        split_star_file(
+            run_args["output_dir"] / "particles_all.star",
+            run_args["output_dir"],
+            run_args["n_files"],
+            run_args["split_size"],
+        )

cryoemservices/pipeliner_plugins/combine_star_job.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from pipeliner.data_structure import SELECT_DIR
+from pipeliner.display_tools import mini_montage_from_starfile
+from pipeliner.job_options import BooleanJobOption, IntJobOption, StringJobOption
+from pipeliner.nodes import NODE_PARTICLESDATA, Node
+from pipeliner.pipeliner_job import ExternalProgram, PipelinerJob
+
+try:
+    from pipeliner.pipeliner_job import PipelinerCommand
+except ImportError:
+    PipelinerCommand = None
+
+COMBINE_STAR_NAME = "combine_star_files"
+
+
+class ProcessStarFiles(PipelinerJob):
+    PROCESS_NAME = "combine_star_files_job"
+    OUT_DIR = SELECT_DIR
+
+    def __init__(self):
+        super().__init__()
+
+        self.jobinfo.display_name = "Particle star file merging"
+        self.jobinfo.short_desc = "Combine and split star files of particles"
+        self.jobinfo.long_desc = (
+            "Combine star files of particles, then optionally split them again."
+        )
+
+        self.jobinfo.programs = [ExternalProgram(command=COMBINE_STAR_NAME)]
+
+        self.joboptions["files_to_process"] = StringJobOption(
+            label="The star files from which to combine particles",
+            is_required=True,
+            help_text="The names of the star files, separated by spaces.",
+        )
+        self.joboptions["do_split"] = BooleanJobOption(
+            label="Whether to split the combined star file", default_value=False
+        )
+        self.joboptions["n_files"] = IntJobOption(
+            label="Number of files to split the combined file into",
+            default_value=-1,
+            help_text=(
+                "Provide either the number of files to split into,"
+                " or the number of particles per file."
+            ),
+            deactivate_if=[("do_split", "is", "False")],
+        )
+        self.joboptions["split_size"] = IntJobOption(
+            label="Number of particles to put in each split",
+            default_value=-1,
+            help_text=(
+                "Provide either the number of files to split into,"
+                " or the number of particles per file."
+            ),
+            deactivate_if=[("do_split", "is", "False")],
+        )
+
+        self.set_joboption_order(
+            ["files_to_process", "do_split", "n_files", "split_size"]
+        )
+
+    def get_commands(self):
+        """Construct the command for combining and splitting star files"""
+        command = [COMBINE_STAR_NAME]
+        file_list = self.joboptions["files_to_process"].get_string().split(" ")
+        command.extend(file_list)
+        command.extend(["--output_dir", str(self.output_dir)])
+
+        if self.joboptions["do_split"].get_boolean():
+            command.extend(["--split"])
+
+            if (
+                self.joboptions["n_files"].get_number() <= 0
+                and self.joboptions["split_size"].get_number() <= 0
+            ):
+                raise ValueError(
+                    "ERROR: When splitting the combined STAR file into subsets,"
+                    " set n_files or split_size to a positive value"
+                )
+
+            if self.joboptions["n_files"].get_number() > 0:
+                command.extend(["--n_files", self.joboptions["n_files"].get_string()])
+
+            if self.joboptions["split_size"].get_number() > 0:
+                command.extend(
+                    ["--split_size", self.joboptions["split_size"].get_string()]
+                )
+
+        # Add files as input nodes, as long as they are not also the output node
+        for particle_file in file_list:
+            if particle_file != str(Path(self.output_dir) / "particles_all.star"):
+                self.input_nodes.append(Node(particle_file, NODE_PARTICLESDATA))
+        self.output_nodes.append(
+            Node(str(Path(self.output_dir) / "particles_all.star"), NODE_PARTICLESDATA)
+        )
+
+        if PipelinerCommand is None:
+            return [command]
+        pipeliner_commands = [PipelinerCommand([command], relion_control=False)]
+        return pipeliner_commands
+
+    def create_output_nodes(self):
+        self.add_output_node("particles_all.star", NODE_PARTICLESDATA, ["relion"])
+
+    def post_run_actions(self):
+        """Find any output files produced by the splitting"""
+        output_files = Path(self.output_dir).glob("particles_split*.star")
+        for split in output_files:
+            self.output_nodes.append(Node(str(split), NODE_PARTICLESDATA))
+
+    def create_results_display(self):
+        with open(
+            Path(self.output_dir) / "class_averages.star", "r"
+        ) as all_classes, open(
+            Path(self.output_dir) / ".class_display_tmp.star", "w"
+        ) as display_classes:
+            for line in range(200):
+                class_line = all_classes.readline()
+                if not class_line:
+                    break
+                display_classes.write(class_line)
+        with open(
+            Path(self.output_dir) / "particles_all.star", "r"
+        ) as all_particles, open(
+            Path(self.output_dir) / ".particles_display_tmp.star", "w"
+        ) as display_particles:
+            for line in range(200):
+                particles_line = all_particles.readline()
+                if not particles_line:
+                    break
+                display_particles.write(particles_line)
+        output_dobs = [
+            mini_montage_from_starfile(
+                starfile=str(Path(self.output_dir) / ".class_display_tmp.star"),
+                block="",
+                column="_rlnReferenceImage",
+                outputdir=self.output_dir,
+                nimg=100,
+                title="Selected 2D classes",
+            ),
+            mini_montage_from_starfile(
+                starfile=str(Path(self.output_dir) / ".particles_display_tmp.star"),
+                block="particles",
+                column="_rlnImageName",
+                outputdir=self.output_dir,
+                nimg=20,
+                title="Examples of selected particles",
+            ),
+        ]
+        (Path(self.output_dir) / ".class_display_tmp.star").unlink()
+        (Path(self.output_dir) / ".particles_display_tmp.star").unlink()
+        return output_dobs

cryoemservices/services/cluster_submission.py

@@ -0,0 +1,229 @@
+from __future__ import annotations
+
+import datetime
+import json
+import logging
+import math
+import os
+from pathlib import Path
+from typing import Optional
+
+import requests
+import workflows.recipe
+import zocalo.configuration
+from importlib_metadata import entry_points
+from pydantic import BaseModel, Field
+from workflows.services.common_service import CommonService
+from zocalo.util import slurm
+
+
+class JobSubmissionParameters(BaseModel):
+    scheduler: str = "slurm"
+    cluster: Optional[str]
+    partition: Optional[str]
+    prefer: Optional[str]
+    job_name: Optional[str]
+    environment: Optional[dict[str, str]] = None
+    cpus_per_task: Optional[int] = None
+    tasks: Optional[int] = None
+    nodes: Optional[int]
+    memory_per_node: Optional[int] = None
+    gpus_per_node: Optional[str] = None
+    min_memory_per_cpu: Optional[int] = Field(
+        None, description="Minimum real memory per cpu (MB)"
+    )
+    time_limit: Optional[datetime.timedelta] = None
+    gpus: Optional[int] = None
+    exclusive: bool = False
+    account: Optional[str]
+    commands: str | list[str]
+    qos: Optional[str]
+
+
+def submit_to_slurm(
+    params: JobSubmissionParameters,
+    working_directory: Path,
+    logger: logging.Logger,
+    zc: zocalo.configuration,
+) -> int | None:
+    api = slurm.SlurmRestApi.from_zocalo_configuration(zc)
+
+    script = params.commands
+    if not isinstance(script, str):
+        script = "\n".join(script)
+    script = f"#!/bin/bash\n. /etc/profile.d/modules.sh\n{script}"
+
+    logger.debug(f"Submitting script to Slurm:\n{script}")
+    if params.time_limit:
+        time_limit_minutes = math.ceil(params.time_limit.total_seconds() / 60)
+    else:
+        time_limit_minutes = None
+    job_submission = slurm.models.JobSubmission(
+        script=script,
+        job=slurm.models.JobProperties(
+            partition=params.partition,
+            prefer=params.prefer,
+            name=params.job_name,
+            cpus_per_task=params.cpus_per_task,
+            tasks=params.tasks,
+            nodes=[params.nodes, params.nodes] if params.nodes else params.nodes,
+            gpus_per_node=params.gpus_per_node,
+            memory_per_node=params.memory_per_node,
+            environment=os.environ
+            if params.environment is None
+            else params.environment,
+            memory_per_cpu=params.min_memory_per_cpu,
+            time_limit=time_limit_minutes,
+            gpus=params.gpus,
+            # exclusive=params.exclusive,
+            account=params.account,
+            current_working_directory=os.fspath(working_directory),
+            qos=params.qos,
+        ),
+    )
+    try:
+        response = api.submit_job(job_submission)
+    except requests.HTTPError as e:
+        logger.error(f"Failed Slurm job submission: {e}\n" f"{e.response.text}")
+        return None
+    if response.errors:
+        error_message = "\n".join(f"{e.errno}: {e.error}" for e in response.errors)
+        logger.error(f"Failed Slurm job submission: {error_message}")
+        return None
+    return response.job_id
+
+
+class ClusterSubmission(CommonService):
+    """A service to start new jobs on a slurm cluster."""
+
+    # Human readable service name
+    _service_name = "EMCluster"
+
+    # Logger name
+    _logger_name = "cryoemservices.services.cluster"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.schedulers = {}
+
+    def initializing(self):
+        """Subscribe to the cluster submission queue.
+        Received messages must be acknowledged."""
+        self.log.info("Cluster submission service starting")
+
+        self.schedulers = {
+            f.name: f.load()
+            for f in entry_points(group="cryoemservices.services.cluster.schedulers")
+        }
+        self.log.debug(f"Supported schedulers: {', '.join(self.schedulers.keys())}")
+        workflows.recipe.wrap_subscribe(
+            self._transport,
+            "cluster.submission",
+            self.run_submit_job,
+            acknowledgement=True,
+            log_extender=self.extend_log,
+        )
+
+    def run_submit_job(self, rw, header, message):
+        """Submit cluster job according to message."""
+
+        parameters = rw.recipe_step["parameters"]
+        cluster_params = JobSubmissionParameters(**parameters.get("cluster", {}))
+
+        if not isinstance(cluster_params.commands, str):
+            cluster_params.commands = "\n".join(cluster_params.commands)
+
+        if "recipefile" in parameters:
+            recipefile = parameters["recipefile"]
+            try:
+                Path(recipefile).parent.mkdir(parents=True, exist_ok=True)
+            except OSError:
+                self.log.error(f"Cannot make directory for {recipefile}")
+                self._transport.nack(header)
+                return
+            self.log.debug("Writing recipe to %s", recipefile)
+            cluster_params.commands = cluster_params.commands.replace(
+                "$RECIPEFILE", recipefile
+            )
+            with open(recipefile, "w") as fh:
+                fh.write(rw.recipe.pretty())
+        if "recipeenvironment" in parameters:
+            recipeenvironment = parameters["recipeenvironment"]
+            try:
+                Path(recipeenvironment).parent.mkdir(parents=True, exist_ok=True)
+            except OSError:
+                self.log.error(f"Cannot make directory for {recipeenvironment}")
+                self._transport.nack(header)
+                return
+            self.log.debug("Writing recipe environment to %s", recipeenvironment)
+            cluster_params.commands = cluster_params.commands.replace(
+                "$RECIPEENV", recipeenvironment
+            )
+            with open(recipeenvironment, "w") as fh:
+                json.dump(
+                    rw.environment, fh, sort_keys=True, indent=2, separators=(",", ": ")
+                )
+        if "recipewrapper" in parameters:
+            recipewrapper = parameters["recipewrapper"]
+            try:
+                Path(recipewrapper).parent.mkdir(parents=True, exist_ok=True)
+            except OSError:
+                self.log.error(f"Cannot make directory for {recipewrapper}")
+                self._transport.nack(header)
+                return
+            self.log.debug("Storing serialized recipe wrapper in %s", recipewrapper)
+            cluster_params.commands = cluster_params.commands.replace(
+                "$RECIPEWRAP", recipewrapper
+            )
+            with open(recipewrapper, "w") as fh:
+                json.dump(
+                    {
+                        "recipe": rw.recipe.recipe,
+                        "recipe-pointer": rw.recipe_pointer,
+                        "environment": rw.environment,
+                        "recipe-path": rw.recipe_path,
+                        "payload": rw.payload,
+                    },
+                    fh,
+                    indent=2,
+                    separators=(",", ": "),
+                )
+
+        if "workingdir" not in parameters or not parameters["workingdir"].startswith(
+            "/"
+        ):
+            self.log.error(
+                "No absolute working directory specified. Will not run cluster job"
+            )
+            self._transport.nack(header)
+            return
+        working_directory = Path(parameters["workingdir"])
+        try:
+            working_directory.mkdir(parents=True, exist_ok=True)
+        except OSError as e:
+            self.log.error(
+                "Could not create working directory: %s", str(e), exc_info=True
+            )
+            self._transport.nack(header)
+            return
+
+        submit_to_scheduler = self.schedulers.get(cluster_params.scheduler)
+
+        jobnumber = submit_to_scheduler(
+            cluster_params, working_directory, self.log, zc=self.config
+        )
+        if not jobnumber:
+            self._transport.nack(header)
+            return
+
+        # Conditionally acknowledge receipt of the message
+        txn = self._transport.transaction_begin(subscription_id=header["subscription"])
+        self._transport.ack(header, transaction=txn)
+
+        # Send results onwards
+        rw.set_default_channel("job_submitted")
+        rw.send({"jobid": jobnumber}, transaction=txn)
+
+        # Commit transaction
+        self._transport.transaction_commit(txn)
+        self.log.info(f"Submitted job {jobnumber} to {cluster_params.cluster}")