resourceManager.py

import os
import time
import json
import subprocess
from pyworkflow.object import Boolean
from pyworkflow.protocol import getProtocolFromDb
import configparser

config = configparser.ConfigParser()
config.read(os.path.join(os.path.dirname(__file__), 'config.yaml'))
standard_queue = os.getenv('QUEUE_STANDARD_QUEUE_NAME') or config['QUEUE'].get('STANDARD_QUEUE_NAME')
priority_queue = os.getenv('QUEUE_PRIORITY_QUEUE_NAME') or config['QUEUE'].get('PRIORITY_QUEUE_NAME')
if os.getenv('QUEUE_NODE_LIST'):
  node_list = os.getenv('QUEUE_NODE_LIST')
elif config['QUEUE'].getboolean('USE_SPECIFIC_NODES'):
  node_list = config['QUEUE'].get('NODE_LIST')
else:
  node_list = None

#TODO: add function to set whether the user wants to use slurm or not

def waitOutput(project, prot, outputAttributeName, sleepTime=10, timeOut=432000):
  """ Wait until the output is being generated by the protocol. """

  def _loadProt():
    # Load the last version of the protocol from its own database
    loadedProt = getProtocolFromDb(prot.getProject().path,
                                   prot.getDbPath(),
                                   prot.getObjId())
    # Close DB connections
    loadedProt.getProject().closeMapper()
    loadedProt.closeMappers()
    return loadedProt

  counter = 1
  prot2 = _loadProt()

  numberOfSleeps = timeOut / sleepTime

  while (not prot2.hasAttribute(outputAttributeName)) and prot2.isActive():
    time.sleep(sleepTime)
    prot2 = _loadProt()
    if counter > numberOfSleeps:
      print("Timeout (%s) reached waiting for %s at %s" % (timeOut, outputAttributeName, prot))
      break
    counter += 1

  # Update the protocol instance to get latest changes
  project._updateProtocol(prot)

def waitOutputFile(project, prot, outputFileName, sleepTime=10, timeOut=432000):
  """ Wait until the output file is being generated by the protocol. """
  def _loadProt():
    # Load the last version of the protocol from its own database
    loadedProt = getProtocolFromDb(prot.getProject().path,
                                   prot.getDbPath(),
                                   prot.getObjId())
    # Close DB connections
    loadedProt.getProject().closeMapper()
    loadedProt.closeMappers()
    return loadedProt

  counter = 1
  prot2 = _loadProt()

  numberOfSleeps = timeOut / sleepTime
  while (not os.path.exists(os.path.join(project.getPath(), prot2._getExtraPath(outputFileName)))) and prot2.isActive():
    time.sleep(sleepTime)
    prot2 = _loadProt()
    if counter > numberOfSleeps:
      print("Timeout (%s) reached waiting for %s at %s" % (timeOut, outputFileName, prot))
      break
    counter += 1

def waitUntilFinishes(project, prot, sleepTime=10, timeOut=432000):
  """ Wait until the protocol finishes. """

  def _loadProt():
    # Load the last version of the protocol from its own database
    loadedProt = getProtocolFromDb(prot.getProject().path,
                                   prot.getDbPath(),
                                   prot.getObjId())
    # Close DB connections
    loadedProt.getProject().closeMapper()
    loadedProt.closeMappers()
    return loadedProt

  counter = 1
  prot2 = _loadProt()

  numberOfSleeps = timeOut / sleepTime

  while not prot2.isFinished() and not prot2.isFailed() and not prot2.isAborted():
    time.sleep(sleepTime)
    prot2 = _loadProt()
    if counter > numberOfSleeps:
      print("Timeout (%s) reached waiting for %s to finish" % (timeOut, prot))
      break
    counter += 1

  # Update the protocol instance to get latest changes
  project._updateProtocol(prot)

#TODO: allow the user to add parameters (e.g. prot.gpuList.set(7))

def sendToSlurm(prot, memory=8192, hours=48, GPU=False, nGPUs=1, nMPIs=None, nThreads=None, priority=False):
    prot._useQueue.set(Boolean(True))
    QUEUE_PARAMS = (priority_queue if priority else standard_queue, {'JOB_MEMORY': memory, 'JOB_TIME': hours, 'GPU_COUNT': nGPUs if GPU else 0, 'JOB_NODES': nMPIs if nMPIs else int(prot.numberOfMpi), 'JOB_THREADS': nThreads if nThreads else int(prot.numberOfThreads)})
    prot._queueParams.set(json.dumps(QUEUE_PARAMS))

def skipSlurm(prot, GPUId):
    prot._useQueue.set(Boolean(False))
    prot.gpuList.set(GPUId)

def createScriptForSlurm(jobname, path, command, nTasks=1, cpusPerTask=1, memory=8192, nGPUs=0, hours=48, priority=False):
  script = """#!/bin/bash
#SBATCH --export=ALL
#SBATCH -J %s
#SBATCH -o %s
#SBATCH -e %s
#SBATCH --open-mode=append
%s
#SBATCH -p %s
#SBATCH --time=%d:00:00 --ntasks=%d --cpus-per-task=%d --mem=%d --gres=gpu:%d
  
%s""" % (jobname,
         os.path.join(path, jobname + '.job.out'),
         os.path.join(path, jobname + '.job.err'),
         f'#SBATCH --nodelist={node_list}' if node_list else '',
         priority_queue if priority else standard_queue,
         hours, nTasks, cpusPerTask, memory, nGPUs, command)
  with open(os.path.join(path, jobname + '.sh'), 'w') as archivo:
    archivo.write(script)
  return os.path.join(path, jobname + '.sh')

def checkIfJobFinished(jobname):
  command = 'squeue -o "%.500j"'
  process = subprocess.Popen(command.split(), stdout=subprocess.PIPE)
  output, _ = process.communicate()
  output = output.decode('utf-8')

  return jobname not in output