Select Git revision
Code owners
Assign users and groups as approvers for specific file changes. Learn more.
Launch.py 10.31 KiB
# ===================================================================================== #
# author: Natalie Danezi <anatoli.danezi@surfsara.nl> -- SURFsara #
# helpdesk: Grid Services <grid.support@surfsara.nl> -- SURFsara #
# #
# usage: python pilot.py [picas_db_name] [picas_username] [picas_pwd] #
# description: #
# Connect to PiCaS server with [picas_username] [picas_pwd] #
# Get the next token in todo View #
# Fetch the token parameters #
# Run the main job (master_step23_v3.sh) with the proper input arguments #
# Get sterr and stdout in the output field of the token #
# ===================================================================================== #
import sys
# python imports
import sys, os
import time
from GRID_PiCaS_Launcher import couchdb
from GRID_PiCaS_Launcher.get_picas_credentials import PicasCred
import subprocess
import shutil
import glob
import warnings
import traceback
import logging
# picas imports
from GRID_PiCaS_Launcher.picas.actors import RunActor
from GRID_PiCaS_Launcher.picas.clients import CouchClient
from GRID_PiCaS_Launcher.picas.iterators import BasicViewIterator
from GRID_PiCaS_Launcher.picas.modifiers import BasicTokenModifier
from GRID_PiCaS_Launcher.picas.executers import execute
from GRID_PiCaS_Launcher.launcher_logging import logger
# token imports
from GRID_PiCaS_Launcher.update_token_status import update_status
from GRID_PiCaS_Launcher.set_token_field import set_token_field
from GRID_PiCaS_Launcher.upload_attachment import upload_attachment
from GRID_PiCaS_Launcher.tok_to_bash import export_dict_to_env
from GRID_PiCaS_Launcher.tok_to_bash import export_variable
from GRID_PiCaS_Launcher.singularity import parse_singularity_link
from GRID_PiCaS_Launcher.singularity import parse_json_payload
from GRID_PiCaS_Launcher.upload_results import GSIUploader
from GRID_PiCaS_Launcher.upload_results import uploader
# from tok_to_bash import export_tok_keys
from GRID_PiCaS_Launcher import sandbox
import pdb
from multiprocessing import Process
from GRID_PiCaS_Launcher import __file__
class ExampleActor(RunActor):
def __init__(self, iterator, modifier):
self.RUNDIR = os.getcwd()
self.iterator = iterator
self.modifier = modifier
self.client = iterator.client
def run(self):
self.p_creds = PicasCred()
self.p_creds.user = self.user
self.p_creds.password = self.password
self.p_creds.database = self.database
self.p_creds.put_picas_creds_in_env()
super(ExampleActor, self).run()
def create_sandbox(self, json_payload=None):
if not json_payload:
json_payload = self.config
if "sandbox" not in json_payload.keys():
logging.warn("No sandbox configuration")
logging.warn("json_payload keys are {0}".format(json_payload.keys()))
return
sbx = sandbox.Sandbox(config_json=json_payload["sandbox"])
sbx.build_sandbox(True) # Not needed
def get_image(self, config=None):
"""get_image: Downloads the image in the cwd
and sets the environment variable $SIMG to point to
the local downloaded image.
:param config:Optional config JSON, if None, we use
the ExampleActor.config member
"""
if not config:
config = self.config
logging.info("getting image from {0}".format(config))
simg_config = parse_json_payload(config)
for key in simg_config.keys():
export_variable(key, simg_config[key])
image_location = parse_singularity_link(
simg_config["SIMG"], simg_config.get("SIMG_COMMIT"))
if os.path.isfile(image_location):
os.environ['SIMG'] = image_location
else:
raise IOError("SIMG does not exist!")
@staticmethod
def get_variables_from_config(config, variables=None):
if not variables:
variables = {}
if "variables" in config.keys():
logging.info("Getting variables from config file")
_vars = config["variables"]
for var in _vars:
logging.debug(
"Setting Environment variable {0} to {1}".format(var, _vars[var])
)
variables[var] = _vars[var]
else:
logging.warn(
"No Variables found in the token config. Nothing is put in the environment!"
)
return {}
return variables
def download_sandbox(self, token):
if "SBXloc" in token.keys():
location = token["SBXloc"]
else:
return None
sbx = sandbox.Sandbox(location=location)
sbx.download_sandbox()
def upload_logs(self, logs_dir):
os.chdir(logs_dir)
logsout = "logs_out"
logserr = "logs_.err"
if os.path.isfile(logsout):
upload_attachment(
token_id=self.token_id,
attachment=logsout,
picas_credentials=self.p_creds,
)
if os.path.isfile(logserr):
upload_attachment(
token_id=self.token_id,
attachment=logsout,
picas_credentials=self.p_creds,
)
def process_token(self, key, token):
# Print token information
os.environ["TOKEN"] = token["_id"]
self.token_id = token["_id"]
logging.info(
"Working on token {0} from databse {1} as user {2}".format(
self.token_id, self.database, self.user
)
)
self.config = token["config.json"]
self.pc = PicasCred(usr=self.user, pwd=self.password, dbn=self.database)
variables = self.get_variables_from_config(self.config)
if "container" in self.config.keys() or "singularity" in self.config.keys():
set_token_field(token["_id"], "status", "pulling_container", self.pc)
self.get_image()
set_token_field(
token["_id"],
"status",
"building_sandbox",
self.pc
)
p = Process(target=self.create_sandbox)
p.start()
logging.info("Creating Sandbox from config: {0}".format(self.config["sandbox"]))
p.join()
with open(os.devnull, "w") as FNULL:
subprocess.call(
["chmod", "a+x", "master.sh"], stdout=FNULL, stderr=subprocess.STDOUT
)
export_dict_to_env(
self.client.db, variables, self.token_id, db_name=self.database
)
logging.info("Working on token: " + token["_id"])
## Read tokvar values from token and write to bash variables if not already exist! Save attachments and export abs filename to variable
set_token_field(token["_id"], "status", "launched", self.pc)
# The launched script is simply master.sh with token and picas authen stored in env vars
# master.sh takes the variables straight from the token.
command = "/usr/bin/time -v ./master.sh 2> logs_.err 1> logs_out"
logging.info("executing " + command)
out = execute(command, shell=True)
logging.info("master.sh exit status is " + str(out))
set_token_field(token["_id"], "output", out[0], self.pc)
if out[0] == 0:
set_token_field(token["_id"], "status", "done", self.pc)
logging.info("Job exited OK")
else:
set_token_field(token["_id"], "status", "error", self.pc)
logging.error("Job exited with status {0}".format(out[0]))
self.upload_logs(self.RUNDIR)
# Just attaches all png files in the working directory to the token
self.find_and_upload_files()
self.find_and_upload_files("*.fits")
self.client.modify_token(self.modifier.close(self.client.db[self.token_id]))
def find_and_upload_files(self, filepattern="*.png"):
sols_search = subprocess.Popen(
["find", ".", "-name", filepattern], stdout=subprocess.PIPE
)
result = sols_search.communicate()[0]
for filename in result.split():
if isinstance(filename, bytes):
filename = filename.decode()
upload_attachment(
token_id=self.token_id,
attachment=filename,
picas_credentials=self.p_creds,
)
os.remove(filename)
def main(
url="https://picas-lofar.grid.surfsara.nl:6984",
db=None,
username=None,
password=None,
view="todo",
):
"""The main function locks a token from a CouchDB view and executes
process_token on the token which builds sandbox and runs scripts
"""
# setup connection to db
client = CouchClient(url=url, db=db, username=username, password=password)
# Create token modifier
modifier = BasicTokenModifier()
# Create iterator, point to the right todo view
iterator = BasicViewIterator(client, sys.argv[-1] + "/" + view, modifier)
# Create actor, takes one token from todo view
actor = ExampleActor(iterator, modifier)
actor.user = username
actor.database = db
actor.password = password
# Start work!
try:
actor.run()
except Exception as e:
exc_info = sys.exc_info()
traceback.print_exception(*exc_info)
del exc_info
print("Exception occured")
print(str(e.args))
set_token_field(actor.token_id, "launcher_status", str(e.args), actor.pc)
finally:
with open(
"{0}/GRID_PiCaS_Launcher.log".format(__file__.split("__init__")[0])
) as f:
print(f.read())
if __name__ == "__main__":
"""Entry point of the Launcher.
The options are either 1: database, 2:username 3: password and 4: picas_token_type
or 1: picas_token_type. In this case, we get the picas credentials from ~/.picasrc
or from the environment variables"""
if len(sys.argv) > 4:
db = str(sys.argv[1])
username = str(sys.argv[2])
password = str(sys.argv[3])
else:
pc = PicasCred()
db = pc.database
username = pc.user
password = pc.password
main(db=db, username=username, password=password)