From 2fa4e69d67078f9ea9b658594c910cb5c726c4aa Mon Sep 17 00:00:00 2001 From: Jacob Freck Date: Fri, 22 Dec 2017 13:35:52 -0500 Subject: [PATCH] Bug: fix logic for worker custom scripts (#295) --- custom-scripts/simple.sh | 4 ++++ node_scripts/install/install.py | 5 +++-- node_scripts/install/scripts.py | 25 +++++++++++++++---------- 3 files changed, 22 insertions(+), 12 deletions(-) diff --git a/custom-scripts/simple.sh b/custom-scripts/simple.sh index dcc2a1b0..a408ab36 100644 --- a/custom-scripts/simple.sh +++ b/custom-scripts/simple.sh @@ -4,5 +4,9 @@ if [ "$IS_MASTER" = "1" ]; then echo "This is a custom script running on just the master!" fi +if [ "$IS_WORKER" = "1" ]; then + echo "This is a custom script running on just the workers!" +fi + echo "This is a custom script running all workers and the master!" diff --git a/node_scripts/install/install.py b/node_scripts/install/install.py index 445ccfab..82a3b75a 100644 --- a/node_scripts/install/install.py +++ b/node_scripts/install/install.py @@ -2,6 +2,7 @@ import os from core import config from install import pick_master, spark, scripts + def setup_node(): client = config.batch_client @@ -16,11 +17,11 @@ def setup_node(): if is_master: setup_as_master() - scripts.run_custom_scripts(is_master = True) + scripts.run_custom_scripts(is_master=True, is_worker=True) else: setup_as_worker() - scripts.run_custom_scripts(is_master = False) + scripts.run_custom_scripts(is_master=False, is_worker=True) def setup_as_master(): diff --git a/node_scripts/install/scripts.py b/node_scripts/install/scripts.py index abf54153..6bb942ea 100644 --- a/node_scripts/install/scripts.py +++ b/node_scripts/install/scripts.py @@ -18,10 +18,10 @@ def _read_yaml_file(path=None): return custom_scripts -def _run_on_this_node(script_obj=None, is_master=False): +def _run_on_this_node(script_obj=None, is_master=False, is_worker=False): if script_obj['runOn'] == 'master' and is_master is True: return True - if script_obj['runOn'] == 'worker' and is_master is False: + if script_obj['runOn'] == 'worker' and is_worker is True: return True if script_obj['runOn'] == 'all-nodes': return True @@ -29,16 +29,16 @@ def _run_on_this_node(script_obj=None, is_master=False): return False -def _run_script_or_scripts_dir(scripts=None, is_master=False, custom_scripts_dir=None): +def _run_script_or_scripts_dir(scripts=None, is_master=False, is_worker=False, custom_scripts_dir=None): for script_obj in scripts: - if _run_on_this_node(script_obj, is_master): + if _run_on_this_node(script_obj, is_master, is_worker): path = Path(os.path.join(custom_scripts_dir, script_obj['script'])) if path.is_dir(): _run_scripts_dir(str(path)) else: _run_script(str(path)) - + def _run_script(script_path: str = None): if not os.path.isfile(script_path): print("Cannot run script: {0} file does not exist".format(script_path)) @@ -47,7 +47,7 @@ def _run_script(script_path: str = None): os.chmod(script_path, file_stat.st_mode | 0o777) print("Running custom script:", script_path) try: - subprocess.call([script_path], shell = True) + subprocess.call([script_path], shell = True) except Exception as e: print(e) @@ -56,22 +56,27 @@ def _run_scripts_dir(root: str = None): for path, subdirs, files in os.walk(root): for name in files: script_path = os.path.join(path, name) - _run_script(script_path) + _run_script(script_path) except FileNotFoundError as e: print(e) except IOError as e: print(e) -def run_custom_scripts(is_master: bool = False): +def run_custom_scripts(is_master: bool = False, is_worker: bool = False): if is_master: os.environ["IS_MASTER"] = "1" else: os.environ["IS_MASTER"] = "0" - + + if is_worker: + os.environ["IS_WORKER"] = "1" + else: + os.environ["IS_WORKER"] = "0" + custom_scripts_dir = os.path.join(os.environ['DOCKER_WORKING_DIR'], 'custom-scripts') custom_scripts = _read_yaml_file(os.path.join(custom_scripts_dir, 'custom-scripts.yaml')) if custom_scripts is not None: - _run_script_or_scripts_dir(custom_scripts, is_master, custom_scripts_dir) + _run_script_or_scripts_dir(custom_scripts, is_master, is_worker, custom_scripts_dir)