Source code for ewoksid13.tasks.slurm

import logging

from ewokscore import Task
from ewokstools.submit import save_and_execute_remote

logger = logging.getLogger(__name__)




class SlurmTask(
    Task,
    optional_input_names=[
        "slurm_user",
        "slurm_url",
        "slurm_token",
        "slurm_job_parameters",
        "slurm_celery_queue",
        "worker_module",
        "submit_to_slurm",
        "wait_to_finish",
        "destination_file",
    ],
    output_names=["job_id"],
):
    """
    A base class for tasks that can be executed on a SLURM cluster.
    It provides methods to handle SLURM job submission and execution.
    """



    def execute_remote(self, **kwargs):
        # The job will be submitted:
        #   - To SLURM directly using ewokstools-Pyslurmutils, this is to be used by esrfusers (offline execution).
        #    (slurm_client=True, slurm_celery_queue=None)
        #   - To a Celery slurm queue, for example, plugged to a Beacon server (executed by opid during the experiment)
        #    (slurm_client=False, slurm_celery_queue=lid13gpu3_slurm (for example))

        submitted = save_and_execute_remote(
            workflow=kwargs.get("workflow"),
            inputs=kwargs.get("inputs"),
            destination_filename=self.get_input_value("destination_file", None),
            queue=self.get_input_value("slurm_celery_queue", None),
            slurm_user=self.get_input_value("slurm_user", None),
            slurm_url=self.get_input_value("slurm_url", None),
            slurm_token=self.get_input_value("slurm_token", None),
            worker_module=self.get_input_value("worker_module", None),
            slurm_job_parameters=self.get_input_value("slurm_job_parameters", {}),
            wait_to_finish=self.get_input_value("wait_to_finish", False),
            python_package="ewoksid13",
        )
        self.outputs.job_id = submitted.job_id