Bug: fix logic for worker custom scripts (#295)

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!"
 

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():

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)