Skip to content

Commit

Permalink
backend for scriptbuilder
Browse files Browse the repository at this point in the history
  • Loading branch information
saurabhsood91 committed Feb 8, 2017
1 parent bd82326 commit 113a337
Show file tree
Hide file tree
Showing 10 changed files with 251 additions and 0 deletions.
1 change: 1 addition & 0 deletions .travis.yml
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ script:
- npm install

- cd $PROJECT_ROOT
- pip install -r requirements.txt
- echo $(pwd) >> $PATH_FILE
- cd sandstone_slurm
- bower install
Expand Down
1 change: 1 addition & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
jsonschema==2.5.1
18 changes: 18 additions & 0 deletions sandstone_slurm/handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from sandstone_slurm.mixins.slurm_mixin import SlurmCmdMixin
from sandstone.lib.handlers.base import BaseHandler
from sandstone_slurm.config_utils import ConfigLoader
from sandstone_slurm.script_template_utils import TemplateFinder
from sandstone import settings


Expand Down Expand Up @@ -56,3 +57,20 @@ def delete(self,jobid):
# This method should dequeue the job specified
# by jobid, via scancel.
pass

class ScriptTemplateHandler(BaseHandler):

@sandstone.lib.decorators.authenticated
def get(self):
# get all the templates
templates = TemplateFinder.find_all_templates()

parsed_json = {
'templates': []
}

for template in templates:
# parse the template
parsed_json['templates'].append(template.parse())

self.write(parsed_json)
123 changes: 123 additions & 0 deletions sandstone_slurm/script_template_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,123 @@
import re
import json
import os
from sandstone import settings
from jsonschema import Draft4Validator
from jsonschema.exceptions import SchemaError

class TemplateFinder(object):
"""
Finds all script templates
"""

@staticmethod
def find_all_templates():
# find all the templates and return
all_templates = []
base_path = settings.SANDSTONE_TEMPLATES_DIR
files = os.listdir(base_path)

for filename in files:
filepath = os.path.join(base_path, filename)
if os.path.isfile(filepath):
ext = os.path.splitext(filepath)[1]
if ext == '.sh':
all_templates.append(BashTemplate(filepath))
elif ext == '.json':
json_template = JSONTemplate(filepath)
if json_template.is_valid_schema():
all_templates.append(json_template)
else:
# ignore the template
print 'Template is not valid'
return all_templates

class BaseTemplate(object):
"""
Base Script Template
"""

def __init__(self, filename):
self.filename = filename
with open(self.filename) as template_file:
self._raw_file_data = template_file.read()

def parse(self):
raise NotImplementedError('Method not implemented')


class JSONTemplate(BaseTemplate):
"""
Represents a JSON template
"""

def __init__(self, filename):
super(JSONTemplate, self).__init__(filename)

def is_valid_schema(self):

parsed_json = self.parse()

if not (parsed_json['title'] and parsed_json['description'] and parsed_json['command']):
return False

variables = parsed_json['variables']
try:
Draft4Validator.check_schema(variables)
return True
except SchemaError:
return False

def parse(self):
return json.loads(self._raw_file_data)


class BashTemplate(BaseTemplate):
"""
Represents a Bash template
"""

_template_name = ''

def __init__(self, filename):
super(BashTemplate, self).__init__(filename)
self._template_name = filename

def parse(self):
# parsing logic
variables = []
commands = []

description = ''


for line in self._raw_file_data.split('\n'):
# is sbatch directive or a blank line, continue
if line.startswith('#SBATCH') or line == '' or line.startswith('#'):
continue
elif line.startswith('#SANDSTONE_DESCRIPTION'):
# specifies the description
description = line.split('=')[1]
else:
# it is a command
# get line till a '#' is seen
command = line.split('#')[0].strip()
commands.append(command)
# get all variables
variables.extend([m[0] for m in re.findall(r'(\$SANDSTONE(_(\w){1,}){1,})', command)])
parsed_object = {
'template_name': self._template_name,
'description': description,
'command': '\n'.join(commands),
'variables': {}
}

for variable in variables:
current_variable = {
'type': 'string',
'description': '',
'title': variable
}
parsed_object['variables'][variable] = current_variable

return parsed_object
26 changes: 26 additions & 0 deletions sandstone_slurm/tests/python/test_script_handlers.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
import os
import mock
import pwd
import json
from sandstone.lib.handlers.base import BaseHandler
from sandstone.lib.test_utils import TestHandlerBase
from sandstone import settings

EXEC_USER = pwd.getpwuid(os.getuid())[0]
TEMPLATES_DIR = os.path.join(os.path.dirname(__file__), 'testfiles')

class ScriptTemplateHandlerTestCase(TestHandlerBase):

@mock.patch.object(BaseHandler, 'get_secure_cookie', return_value=EXEC_USER)
def test_get_templates(self, m):
settings.SANDSTONE_TEMPLATES_DIR = TEMPLATES_DIR
response = self.fetch(
'/slurm/a/scripttemplates',
method='GET',
follow_redirects=False
)
self.assertEqual(response.code, 200)

json_response = json.loads(response.body)
self.assertIsNotNone(json_response['templates'])
self.assertEqual(len(json_response['templates']), 3)
47 changes: 47 additions & 0 deletions sandstone_slurm/tests/python/test_script_templates.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
import unittest
import mock
import json
import os
from sandstone import settings
from sandstone_slurm.script_template_utils import BashTemplate, JSONTemplate, TemplateFinder

TEST_JSON_FILE = os.path.join(os.path.dirname(__file__), 'testfiles/test_json.json')
TEST_BASH_FILE = os.path.join(os.path.dirname(__file__), 'testfiles/test_script.sh')
TEST_COMPLEX_BASH_FILE = os.path.join(os.path.dirname(__file__), 'testfiles/test_complex_bash.sh')
TEMPLATES_DIR = os.path.join(os.path.dirname(__file__), 'testfiles')

class TemplateFinderTestCase(unittest.TestCase):
def setUp(self):
settings.SANDSTONE_TEMPLATES_DIR = TEMPLATES_DIR

def test_template_finder(self):
all_templates = TemplateFinder.find_all_templates()
self.assertEqual(len(all_templates), 3)


class JSONTemplateParseTestCase(unittest.TestCase):
def setUp(self):
self.json_template = JSONTemplate(TEST_JSON_FILE)

def test_valid_json(self):
self.assertTrue(self.json_template.is_valid_schema())

class BashTemplateParseTestCase(unittest.TestCase):

def test_template_properties(self):
self.bash_template = BashTemplate(TEST_BASH_FILE)
parsed_object = self.bash_template.parse()
self.assertEqual(parsed_object['template_name'], TEST_BASH_FILE)
self.assertEqual(parsed_object['description'], "")
self.assertEqual(parsed_object['command'], "python $SANDSTONE_FILENAME")
self.assertEqual(len(parsed_object['variables'].keys()), 1)
self.assertItemsEqual(parsed_object['variables'].keys(), ['$SANDSTONE_FILENAME'])

def test_complex_template_properties(self):
self.bash_template = BashTemplate(TEST_COMPLEX_BASH_FILE)
parsed_object = self.bash_template.parse()
self.assertEqual(parsed_object['template_name'], TEST_COMPLEX_BASH_FILE)
self.assertEqual(parsed_object['description'], "")
self.assertEqual(parsed_object['command'], "export OMP_NUM_THREADS=$SANDSTONE_NUM_THREADS\nsrun -n 20 -c 4 $SANDSTONE_SCRIPT_NAME")
self.assertEqual(len(parsed_object['variables'].keys()), 2)
self.assertItemsEqual(parsed_object['variables'].keys(), ['$SANDSTONE_NUM_THREADS', '$SANDSTONE_SCRIPT_NAME'])
11 changes: 11 additions & 0 deletions sandstone_slurm/tests/python/testfiles/test_complex_bash.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
#!/bin/bash -l

#SBATCH -p regular
#SBATCH -N 64
#SBATCH -t 12:00:00
#SBATCH -L project
#SBATCH -C haswell

#to run with pure MPI
export OMP_NUM_THREADS=$SANDSTONE_NUM_THREADS
srun -n 20 -c 4 $SANDSTONE_SCRIPT_NAME # -c is optional since this example is fully packed pure MPI (32 MPI tasks per node)
14 changes: 14 additions & 0 deletions sandstone_slurm/tests/python/testfiles/test_json.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
{
"title": "Some JSON template",
"description": "Some Description",
"command": "ls | grep $SANDSTONE_SEARCH_EXPRESSION",
"variables": {
"type": "object",
"properties": {
"$SANDSTONE_SEARCH_EXPRESSION": {
"type": "string",
"description": "Some Sample description"
}
}
}
}
8 changes: 8 additions & 0 deletions sandstone_slurm/tests/python/testfiles/test_script.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
#!/bin/bash
#SBATCH --job-name test-job
#SBATCH --nodes 2
#SBATCH --account=crcsupport
#SBATCH --output output.out
#python main.py

python $SANDSTONE_FILENAME
2 changes: 2 additions & 0 deletions sandstone_slurm/urls.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
from sandstone_slurm.handlers import FormConfigHandler
from sandstone_slurm.handlers import JobListHandler
from sandstone_slurm.handlers import JobHandler
from sandstone_slurm.handlers import ScriptTemplateHandler



URL_SCHEMA = [
(r"/slurm/a/config", FormConfigHandler),
(r"/slurm/a/jobs", JobListHandler),
(r"/slurm/a/jobs/(?P<jobid>[0-9]+)?", JobHandler),
(r"/slurm/a/scripttemplates", ScriptTemplateHandler),
]

0 comments on commit 113a337

Please sign in to comment.