add running molecule sims with fdmnes (#17)

Author: jacobfilik

database/tables.sql

@@ -23,7 +23,7 @@ INSERT INTO "simulation_type" VALUES(2,'FDMNES');
 INSERT INTO "simulation_type" VALUES(3,'Quantum ESPRESSO');
 
 
-CREATE TYPE simulation_status_enum AS ENUM('requested', 'submitted', 'running', 'completed', 'failed', 'error');
+CREATE TYPE simulation_status_enum AS ENUM('requested', 'submitted', 'running', 'completed', 'failed', 'error', 'request_cancel', 'cancelled');
 
 CREATE TABLE simulation (
     id INTEGER PRIMARY KEY GENERATED ALWAYS AS IDENTITY,
@@ -68,7 +68,7 @@ CREATE TABLE molecular_structure (
 COMMENT ON TABLE molecular_structure IS 'Table to hold molecular structures';
 
 
-CREATE TYPE ibrav_enum AS ENUM('1', '2', '3','-3', '4', '5','-5', '6', '7', '8', '9','-9','91', '10', '11', '12','-12', '13','-13', '14');
+-- CREATE TYPE ibrav_enum AS ENUM('1', '2', '3','-3', '4', '5','-5', '6', '7', '8', '9','-9','91', '10', '11', '12','-12', '13','-13', '14');
 
 
 CREATE TABLE crystal_structure (
@@ -81,7 +81,6 @@ CREATE TABLE crystal_structure (
     alpha NUMERIC NOT NULL,
     beta NUMERIC NOT NULL,
     gamma NUMERIC NOT NULL,
-    ibrav ibrav_enum NOT NULL,
     person_id INTEGER NOT NULL
 );
 
@@ -129,17 +128,21 @@ CREATE TABLE orca_simulation_spectrum(
 
 CREATE TYPE edge_enum AS ENUM('k', 'l1', 'l2', 'l3', 'm1', 'm2', 'm3', 'm4', 'm5');
 
-CREATE TYPE structure_enum AS ENUM('crystal', 'molecule');
+-- MAKE JUST ONE FDMNES ENTRY
+
 
 CREATE TABLE fdmnes_simulation (
     simulation_id INTEGER PRIMARY KEY,
     simulation_type_id INTEGER GENERATED ALWAYS AS (2) STORED,
-    crystal_structure_id INTEGER NOT NULL,
+    crystal_structure_id INTEGER,
+    molecular_structure_id INTEGER,
     element INTEGER NOT NULL,
     edge edge_enum NOT NULL,
     greens_approach BOOLEAN NOT NULL,
-    structure_type structure_enum NOT NULL,
 
+    CONSTRAINT chk_only_one_is_not_null CHECK (num_nonnulls(crystal_structure_id, molecular_structure_id) = 1),
+
+    FOREIGN KEY(molecular_structure_id) REFERENCES molecular_structure (id),
     FOREIGN KEY(crystal_structure_id) REFERENCES crystal_structure (id),
     FOREIGN KEY(simulation_id, simulation_type_id) REFERENCES simulation (id,simulation_type_id)
 );
@@ -169,11 +172,11 @@ O    0.0000000    0.0000000   -0.1653507
 H   -0.7493682    0.0000000    0.4424329');
 
 
-INSERT INTO crystal_structure(label, person_id,a,b,c,alpha,beta,gamma,ibrav,structure) VALUES('Silver',1,4.1043564,4.1043564,4.1043564,90,90,90,'2','Ag 0.0 0.0 0.0
+INSERT INTO crystal_structure(label, person_id,a,b,c,alpha,beta,gamma,structure) VALUES('Silver',1,4.1043564,4.1043564,4.1043564,90,90,90,'Ag 0.0 0.0 0.0
 Ag 0.5 0.5 0.0
 Ag 0.5 0.0 0.5
 Ag 0.0 0.5 0.5');
 
 
-INSERT INTO crystal_structure(label, person_id,a,b,c,alpha,beta,gamma,ibrav,structure) VALUES('KCl',1,6.28,6.28,6.28,90,90,90,'2','K 0.0 0.0 0.0
-Cl 0.5 0.0 0.0');
+INSERT INTO crystal_structure(label, person_id,a,b,c,alpha,beta,gamma,structure) VALUES('KCl',1,4.44,4.44,4.44,60,60,60,'K 0.0 0.0 0.0
+Cl 0.5 0.5 0.5');

web-conexs-api/pyproject.toml

@@ -19,6 +19,7 @@ dependencies = [
     "psycopg2",
     "numpy",
     "requests",
+    "pymatgen"
 ] # Add project dependencies here, e.g. ["click", "numpy"]
 dynamic = ["version"]
 license.file = "LICENSE"

web-conexs-api/src/slurm_submission_service/crud.py

@@ -6,6 +6,7 @@
     build_fdmnes_inputfile,
     build_orca_input_file,
     build_qe_inputfile,
+    fdmnes_molecule_to_crystal,
 )
 from web_conexs_api.models.models import (
     CrystalStructure,
@@ -79,9 +80,22 @@ def get_crystal_structure(session, id):
 
 def get_fdmnes_jobfile(session, id):
     fdmnes_simulation = get_fdmnes_simulation(session, id)
-    structure = get_crystal_structure(session, fdmnes_simulation.crystal_structure_id)
 
-    return build_fdmnes_inputfile(fdmnes_simulation, structure)
+    if fdmnes_simulation.crystal_structure_id is not None:
+        structure = get_crystal_structure(
+            session, fdmnes_simulation.crystal_structure_id
+        )
+        crystalIsMolecule = False
+    else:
+        molecule = get_molecular_structure(
+            session, fdmnes_simulation.molecular_structure_id
+        )
+        structure = fdmnes_molecule_to_crystal(molecule)
+        crystalIsMolecule = True
+
+    print(structure.structure)
+
+    return build_fdmnes_inputfile(fdmnes_simulation, structure, crystalIsMolecule)
 
 
 def get_submitted_simulations(session) -> List[Simulation]:
@@ -91,6 +105,13 @@ def get_submitted_simulations(session) -> List[Simulation]:
     return session.exec(statement).all()
 
 
+def get_request_cancelled_simulations(session) -> List[Simulation]:
+    statement = select(Simulation).where(
+        Simulation.status == SimulationStatus.request_cancel
+    )
+    return session.exec(statement).all()
+
+
 def update_simulation(session, simulation: Simulation):
     session.add(simulation)
     session.commit()

web-conexs-api/src/slurm_submission_service/submitter.py

@@ -21,6 +21,7 @@
     get_fdmnes_jobfile,
     get_orca_jobfile_with_technique,
     get_qe_jobfile,
+    get_request_cancelled_simulations,
     get_submitted_simulations,
     update_simulation,
 )
@@ -62,15 +63,13 @@ def get_token():
         return fh.read()
 
 
-def build_job_and_run(
-    script, user_id, working_dir, job_name, cpus, memory, cluster_dir
-):
+def build_job_and_run(script, job_name, cpus, memory, cluster_dir, as_tasks):
     job_request = {
         "job": {
             "partition": SLURM_PARTITION,
             "name": job_name,
             "nodes": 1,
-            "tasks": cpus,
+            "tasks": 1 if as_tasks else cpus,
             "memory_per_node": int(memory * 1000),
             "time_limit": 30,
             "current_working_directory": str(cluster_dir),
@@ -83,6 +82,13 @@ def build_job_and_run(
         "script": script,
     }
 
+    if as_tasks:
+        job_request["job"]["ntasks_per_node"] = 1
+        job_request["job"]["cpus_per_task"] = cpus
+        job_request["job"]["environment"]["OMP_NUM_THREADS"] = cpus
+
+    logger.debug(job_request)
+
     url_submit = SLURM_API + "/job/submit"
 
     slurm_token = get_token()
@@ -160,7 +166,7 @@ def submit_orca(session, sim: Simulation):
 
     try:
         job_id = build_job_and_run(
-            script, user, working_dir, job_name, sim.n_cores, sim.memory, cluster_dir
+            script, job_name, sim.n_cores, sim.memory, cluster_dir, False
         )
         sim.job_id = job_id
         sim.working_directory = working_dir
@@ -171,6 +177,7 @@ def submit_orca(session, sim: Simulation):
         logger.exception("Error submitting orca job")
         sim.status = SimulationStatus.failed
         update_simulation(session, sim)
+        logger.exception("Error submitting ORCA job")
 
 
 def submit_fdmnes(session, sim: Simulation):
@@ -200,12 +207,12 @@ def submit_fdmnes(session, sim: Simulation):
 
     script = (
         "#!/bin/bash\n"
-        + f"singularity exec {fdmnes_sif} mpirun -np 4 fdmnes_mpi > fdmnes_result.txt"
+        + f"singularity exec {fdmnes_sif} mpirun -np 1 fdmnes_mpi > fdmnes_result.txt"
     )
 
     try:
         job_id = build_job_and_run(
-            script, user, working_dir, job_name, sim.n_cores, sim.memory, cluster_dir
+            script, job_name, sim.n_cores, sim.memory, cluster_dir, True
         )
         sim.job_id = job_id
         sim.working_directory = working_dir
@@ -215,10 +222,41 @@ def submit_fdmnes(session, sim: Simulation):
     except Exception:
         sim.status = SimulationStatus.failed
         update_simulation(session, sim)
+        logger.exception("Error submitting FDMNES job")
+
+
+def clean_request_cancelled(session):
+    request_cancel = get_request_cancelled_simulations(session)
+
+    for sim in request_cancel:
+        if sim.job_id is None:
+            # No job id, hasn't been submitted, so just flag as cancelled
+            sim.status = SimulationStatus.cancelled
+            update_simulation(session, sim)
+        else:
+            url = SLURM_API + "/job/" + str(sim.job_id)
+            slurm_token = get_token()
+
+            headers = {
+                "X-SLURM-USER-NAME": SLURM_USER,
+                "X-SLURM-USER-TOKEN": slurm_token,
+                "Content-Type": "application/json",
+            }
+
+            r = requests.delete(url, headers=headers)
+
+            if r.status_code != 200:
+                logger.error(f"Job delete response not successful {r.status_code}")
+                continue
+
+            sim.status = SimulationStatus.cancelled
+            update_simulation(session, sim)
 
 
 def run_update():
     with contextmanager(get_session)() as session:
+        clean_request_cancelled(session)
+
         # sessions = get_session()
         # session = next(sessions)
         sims = get_submitted_simulations(session)
@@ -342,7 +380,7 @@ def submit_qe(session, sim: Simulation):
 
     try:
         job_id = build_job_and_run(
-            script, user, working_dir, job_name, sim.n_cores, sim.memory, cluster_dir
+            script, job_name, sim.n_cores, sim.memory, cluster_dir, False
         )
         sim.job_id = job_id
         sim.working_directory = working_dir
@@ -352,14 +390,18 @@ def submit_qe(session, sim: Simulation):
     except Exception:
         sim.status = SimulationStatus.failed
         update_simulation(session, sim)
+        logger.exception("Error submitting QE job")
 
 
-def test_read():
-    with contextmanager(get_session)() as session:
-        sims = get_submitted_simulations(session)
-        for sim in sims:
-            if sim.simulation_type_id == 1:
-                jf, calc = get_orca_jobfile_with_technique(session, sim.id)
+# def test_read():
+#     with contextmanager(get_session)() as session:
+#         sims = get_submitted_simulations(session)
+#         for sim in sims:
+#             if sim.simulation_type_id == 1:
+#                 jf, calc = get_orca_jobfile_with_technique(session, sim.id)
+#             elif sim.simulation_type_id == 2:
+#                 jobfile = get_fdmnes_jobfile(session, sim.id)
+#                 logger.error(jobfile)
 
 
 def main():
@@ -377,6 +419,8 @@ def main():
 
     logger.info("Running main loop")
 
+    # test_read()
+
     while True:
         try:
             run_update()

web-conexs-api/src/web_conexs_api/crud.py

@@ -27,6 +27,7 @@
     QESimulation,
     QESimulationInput,
     Simulation,
+    SimulationStatus,
 )
 
 
@@ -139,8 +140,12 @@ def get_simulations_page(session, user_id) -> CursorPage[Simulation]:
     return paginate(session, statement.order_by(Simulation.id.desc()))
 
 
-def get_simulation(session, id, user_id):
-    statement = select(Simulation).join(Person).where(Simulation.id == id)
+def get_simulation(session, id, user_id) -> Simulation:
+    statement = (
+        select(Simulation)
+        .join(Person)
+        .where(and_(Person.identifier == user_id, Simulation.id == id))
+    )
 
     simulation = session.exec(statement).first()
 
@@ -424,3 +429,19 @@ def get_qe_xas(session, id, user_id):
     output = {"energy": out[:, 0].tolist(), "xas": out[:, 1].tolist()}
 
     return output
+
+
+def request_cancel_simulation(session, id, user_id):
+    sim = get_simulation(session, id, user_id)
+
+    if (
+        sim.status == SimulationStatus.running
+        or sim.status == SimulationStatus.requested
+        or sim.status == SimulationStatus.submitted
+    ):
+        sim.status = SimulationStatus.request_cancel
+        session.add(sim)
+        session.commit()
+        session.refresh(sim)
+
+    return sim