Merge pull request #51 from tilde-lab/aiida

Generalize AiiDA integration w.r.t. supported engines

Author: ansobolev

yascheduler/aiida_plugin.py

@@ -1,9 +1,12 @@
 """
-Aiida plugin for yascheduler
+Aiida plugin for yascheduler,
+with respect to the supported yascheduler engines
 """
 
 import aiida.schedulers
 from aiida.schedulers.datastructures import JobState, JobInfo, NodeNumberJobResource
+from aiida.orm import load_node
+
 
 _MAP_STATUS_YASCHEDULER = {
     "QUEUED": JobState.QUEUED,
@@ -67,12 +70,15 @@ def _get_submit_script_header(self, job_tmpl):
         Return the submit script header, using the parameters from the
         job_tmpl.
         """
-        lines = []
+        aiida_code = load_node(job_tmpl.codes_info[0]['code_uuid'])
+
+        # We map the lowercase code labels onto yascheduler engines,
+        # so that the required input file(s) can be deduced
+        lines = ["ENGINE={}".format(aiida_code.label.lower())]
+
         if job_tmpl.job_name:
             lines.append("LABEL={}".format(job_tmpl.job_name))
 
-        # TODO too specific for engine.pcrystal
-        lines += ["INPUT=INPUT", "STRUCT=fort.34"]
         return "\n".join(lines)
 
     def _get_submit_command(self, submit_script):
@@ -88,7 +94,8 @@ def _parse_submit_output(self, retval, stdout, stderr):
         """
         if stderr.strip():
             self.logger.warning("Stderr when submitting: {}".format(stderr.strip()))
-        return stdout.split(":")[1].strip()
+
+        return stdout.strip()
 
     def _parse_joblist_output(self, retval, stdout, stderr):
         """

yascheduler/utils.py

@@ -17,33 +17,43 @@
 
 
 def submit():
-    parser = argparse.ArgumentParser(description="Submit task to yascheduler daemon")
+    parser = argparse.ArgumentParser(description="Submit task to yascheduler via AiiDA script")
     parser.add_argument("script")
 
     args = parser.parse_args()
     if not os.path.isfile(args.script):
         raise ValueError("Script parameter is not a file name")
 
-    inputs = {}
+    script_params = {}
     with open(args.script) as f:
         for l in f.readlines():
             try:
                 k, v = l.split("=")
-                inputs[k.strip()] = v.strip()
+                script_params[k.strip()] = v.strip()
             except ValueError:
                 pass
+
+    label = script_params.get('LABEL', 'AiiDA job')
+    options = {
+        "local_folder": os.getcwd()
+    }
+    if not script_params.get('ENGINE'):
+        raise ValueError("Script has not defined an engine")
+
     config = ConfigParser()
     config.read(CONFIG_FILE)
     yac = Yascheduler(config)
-    task_id = yac.queue_submit_task(
-        inputs["LABEL"],
-        {
-            "structure": open(inputs["STRUCT"]).read(),
-            "input": open(inputs["INPUT"]).read(),
-            "local_folder": os.getcwd(),
-        },
-    )  # TODO
 
+    if script_params['ENGINE'] not in yac.engines:
+        raise ValueError("Script refers to unknown engine")
+
+    for input_file in yac.engines[script_params['ENGINE']].input_files:
+        if not os.path.exists(os.path.join(options["local_folder"], input_file)):
+            raise ValueError("Script was not supplied with the required input")
+
+        options[input_file] = open(input_file).read()
+
+    task_id = yac.queue_submit_task(label, options, script_params['ENGINE'])
     print("Successfully submitted task: {}".format(task_id))
     yac.connection.close()