diff --git a/.travis.yml b/.travis.yml index fc915f9..3909bc7 100644 --- a/.travis.yml +++ b/.travis.yml @@ -36,6 +36,7 @@ script: - npm install - cd $PROJECT_ROOT + - pip install -r requirements.txt - echo $(pwd) >> $PATH_FILE - cd sandstone_slurm - bower install diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..88a2bfd --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +jsonschema==2.5.1 diff --git a/sandstone_slurm/handlers.py b/sandstone_slurm/handlers.py index aa70ee2..804ded8 100644 --- a/sandstone_slurm/handlers.py +++ b/sandstone_slurm/handlers.py @@ -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 @@ -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) diff --git a/sandstone_slurm/script_template_utils.py b/sandstone_slurm/script_template_utils.py new file mode 100644 index 0000000..55fc625 --- /dev/null +++ b/sandstone_slurm/script_template_utils.py @@ -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 diff --git a/sandstone_slurm/tests/python/test_script_handlers.py b/sandstone_slurm/tests/python/test_script_handlers.py new file mode 100644 index 0000000..e66ce64 --- /dev/null +++ b/sandstone_slurm/tests/python/test_script_handlers.py @@ -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) diff --git a/sandstone_slurm/tests/python/test_script_templates.py b/sandstone_slurm/tests/python/test_script_templates.py new file mode 100644 index 0000000..f6ca241 --- /dev/null +++ b/sandstone_slurm/tests/python/test_script_templates.py @@ -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']) diff --git a/sandstone_slurm/tests/python/testfiles/test_complex_bash.sh b/sandstone_slurm/tests/python/testfiles/test_complex_bash.sh new file mode 100644 index 0000000..8c6f77e --- /dev/null +++ b/sandstone_slurm/tests/python/testfiles/test_complex_bash.sh @@ -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) diff --git a/sandstone_slurm/tests/python/testfiles/test_json.json b/sandstone_slurm/tests/python/testfiles/test_json.json new file mode 100644 index 0000000..add5e23 --- /dev/null +++ b/sandstone_slurm/tests/python/testfiles/test_json.json @@ -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" + } + } + } +} diff --git a/sandstone_slurm/tests/python/testfiles/test_script.sh b/sandstone_slurm/tests/python/testfiles/test_script.sh new file mode 100644 index 0000000..2fe4654 --- /dev/null +++ b/sandstone_slurm/tests/python/testfiles/test_script.sh @@ -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 diff --git a/sandstone_slurm/urls.py b/sandstone_slurm/urls.py index a484776..3bc04d8 100644 --- a/sandstone_slurm/urls.py +++ b/sandstone_slurm/urls.py @@ -1,6 +1,7 @@ from sandstone_slurm.handlers import FormConfigHandler from sandstone_slurm.handlers import JobListHandler from sandstone_slurm.handlers import JobHandler +from sandstone_slurm.handlers import ScriptTemplateHandler @@ -8,4 +9,5 @@ (r"/slurm/a/config", FormConfigHandler), (r"/slurm/a/jobs", JobListHandler), (r"/slurm/a/jobs/(?P[0-9]+)?", JobHandler), + (r"/slurm/a/scripttemplates", ScriptTemplateHandler), ]