Initial staging commit

Author: hategan

src/psij/exceptions.py

@@ -62,3 +62,21 @@ def __init__(self, message: str, exception: Optional[Exception] = None,
         conditions such an error would persist across subsequent re-tries until correct credentials
         are used.
         """
+
+class CompositeException(Exception):
+    def __init__(self, ex: Exception) -> None:
+        self.exceptions = [ex]
+
+    def add_exception(self, ex: Exception) -> None:
+        self.exceptions.append(ex)
+
+
+class LauncherException(Exception):
+    def __init__(self, message: str) -> None:
+        super().__init__('Launcher failure: %s' % message)
+
+
+class JobException(Exception):
+    def __init__(self, exit_code: int) -> None:
+        super().__init__('Job exited with exit code %s' % exit_code)
+        self.exit_code = exit_code

src/psij/executors/batch/batch_scheduler_executor.py

@@ -1,5 +1,12 @@
+import atexit
+import fcntl
+import io
 import logging
 import os
+import tempfile
+
+import psutil
+import socket
 import subprocess
 import time
 import traceback
@@ -14,6 +21,7 @@
 from psij import JobExecutor, JobExecutorConfig, Launcher, Job, SubmitException, \
     JobStatus, JobState
 from psij.executors.batch.template_function_library import ALL as FUNCTION_LIBRARY
+from psij.utils import SingletonThread
 
 UNKNOWN_ERROR = 'PSIJ: Unknown error'
 
@@ -265,6 +273,8 @@ def cancel(self, job: Job) -> None:
             except SubmitException:
                 # re-raise
                 raise
+        finally:
+            self._status_update_thread.unregister_job(job)
 
     def attach(self, job: Job, native_id: str) -> None:
         """Attaches a job to a native job.
@@ -531,9 +541,11 @@ def _set_job_status(self, job: Job, status: JobStatus) -> None:
             # is_greater_than returns T/F if the states are comparable and None if not, so
             # we have to check explicitly for the boolean value rather than truthiness
             return
-        if status.state.final and job.native_id:
-            self._clean_submit_script(job)
-            self._read_aux_files(job, status)
+        if status.state.final:
+            self._status_update_thread.unregister_job(job)
+            if job.native_id:
+                self._clean_submit_script(job)
+                self._read_aux_files(job, status)
         super()._set_job_status(job, status)
 
     def _clean_submit_script(self, job: Job) -> None:
@@ -638,13 +650,19 @@ def __init__(self, name: str, config: BatchSchedulerExecutorConfig,
         # counts consecutive errors while invoking qstat or equivalent
         self._poll_error_count = 0
         self._jobs_lock = RLock()
+        self._status_updater = _StatusUpdater(config, executor)
 
     def run(self) -> None:
         logger.debug('Executor %s: queue poll thread started', self.executor)
         time.sleep(self.config.initial_queue_polling_delay)
         while True:
             self._poll()
-            time.sleep(self.config.queue_polling_interval)
+            start = time.time()
+            now = start
+            while now - start < self.config.queue_polling_interval:
+                self._status_updater.step()
+                time.sleep(1)
+                now = time.time()
 
     def _poll(self) -> None:
         with self._jobs_lock:
@@ -686,6 +704,8 @@ def _poll(self) -> None:
                 if status.state.final:
                     with self._jobs_lock:
                         del self._jobs[native_id]
+                    for job in job_list:
+                        self._status_updater.unregister_job(job)
         except Exception as ex:
             msg = traceback.format_exc()
             self._handle_poll_error(True, ex, 'Error updating job statuses {}'.format(msg))
@@ -713,9 +733,11 @@ def _handle_poll_error(self, immediate: bool, ex: Exception, msg: str) -> None:
                 self._jobs.clear()
             for job_list in jobs_copy.values():
                 for job in job_list:
+                    self._status_updater.unregister_job(job)
                     self.executor._set_job_status(job, JobStatus(JobState.FAILED, message=msg))
 
     def register_job(self, job: Job) -> None:
+        self._status_updater.register_job(job)
         assert job.native_id
         logger.info('Job %s: registering', job.id)
         with self._jobs_lock:
@@ -724,3 +746,91 @@ def register_job(self, job: Job) -> None:
                 self._jobs[native_id] = [job]
             else:
                 self._jobs[job.native_id].append(job)
+
+class _StatusUpdater:
+    # we are expecting short messages in the form <jobid> <status>
+    RECV_BUFSZ = 2048
+
+    def __init__(self, config: BatchSchedulerExecutorConfig,
+                 executor: BatchSchedulerExecutor) -> None:
+        self.config = config
+        self.executor = executor
+        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self.socket.setblocking(False)
+        self.socket.bind(('', 0))
+        self.port = self.socket.getsockname()[1]
+        self.ips = self._get_ips()
+        print('IPS: %s' % self.ips)
+        print('Port: %s' % self.port)
+        self._create_update_file()
+        print('Update file: %s' % self.update_file.name)
+        self.partial_file_data = ''
+        self.partial_net_data = ''
+        self._jobs = {}
+        self._jobs_lock = RLock()
+
+    def _get_ips(self) -> List[str]:
+        addrs = psutil.net_if_addrs()
+        r = []
+        for name, l in addrs.items():
+            if name == 'lo':
+                continue
+            for a in l:
+                if a.family == socket.AddressFamily.AF_INET:
+                    r.append(a.address)
+        return r
+
+    def _create_update_file(self) -> None:
+        f = tempfile.NamedTemporaryFile(dir=self.config.work_directory, prefix='supd_',
+                                        delete=False)
+        name = f.name
+        atexit.register(os.remove, name)
+        f.close()
+        self.update_file = open(name, 'r+b')
+        self.update_file.seek(0, io.SEEK_END)
+        self.update_file_pos = self.update_file.tell()
+
+    def register_job(self, job: Job) -> None:
+        with self._jobs_lock:
+            self._jobs[job.id] = job
+
+    def unregister_job(self, job: Job) -> None:
+        with self._jobs_lock:
+            del self._jobs[job.id]
+
+    def step(self) -> None:
+        self.update_file.seek(0, io.SEEK_END)
+        pos = self.update_file.tell()
+        if pos > self.update_file_pos:
+            self.update_file.seek(self.update_file_pos, io.SEEK_SET)
+            n = pos - self.update_file_pos
+            self._process_update_data(self.update_file.read(n))
+            self.update_file_pos = pos
+        else:
+            try:
+                data = self.socket.recv(_StatusUpdater.RECV_BUFSZ)
+                self._process_update_data(data)
+            except socket.error as e:
+                pass
+
+    def _process_update_data(self, data: bytes) -> None:
+        sdata = data.decode('utf-8')
+        lines = sdata.splitlines()
+        for line in lines:
+            print('Status update line: %s' % line)
+            els = line.split()
+            if len(els) != 2:
+                logger.warning('Invalid status update message received: %s' % line)
+                continue
+            job_id = els[0]
+            state = JobState.from_name(els[1])
+            job = None
+            with self._jobs_lock:
+                try:
+                    job = self._jobs[job_id]
+                except KeyError:
+                    logger.warning('Received status updated for inexistent job with id %s' % job_id)
+            if job:
+                self.executor._set_job_status(job, JobStatus(state))
+
+

src/psij/executors/batch/common/batch_lib.mustache

@@ -0,0 +1,80 @@
+_UPDATE_MODE="none"
+
+update_status() {
+    STATUS="$1"
+    if [ "$_UPDATE_MODE" == "none" ]; then
+        if which nc >/dev/null; then
+            $_UPDATE_MODE="nc"
+        else
+            $_UPDATE_MODE="file"
+        fi
+    fi
+
+    if [ "$_UPDATE_MODE" == "nc" ]; then
+        for ADDR in "{{psij.us_addrs}}"; do
+            echo "{{psij.job_id}} $STATUS" | nc -q0 -w0 -4 -u $ADDR {{psij.us_port}}
+        done
+    else
+        echo "{{psij.job_id}} $STATUS" >> {{psij.us_file}}
+    fi
+}
+
+do_stagein() {
+    do_stage "$@" 0
+}
+
+do_stage() {
+    SOURCE="$1"
+    TARGET="$2"
+    MODE="$3"
+    MISSING_OK="$4"
+
+    if [ ! -a "$SOURCE" ] && [ "$MISSING_OK" == "0" ]; then
+        echo "Missing source file: $SOURCE"
+        exit 1
+    fi
+
+    if [ "$MODE" == "1" ]; then
+        # copy
+        cp -r -T "$SOURCE" "$TARGET"
+    elif [ "$MODE" == "2" ]; tben
+        # link
+        ln -s "$SOURCE" "$TARGET"
+    elif [ "$MODE" == "3" ]; then
+        # move
+        mv -T -f "$SOURCE" "$TARGET"
+    fi
+}
+
+_FLAG_IF_PRESENT=1
+_FLAG_ON_SUCCESS=2
+_FLAG_ON_ERROR=4
+_FLAG_ON_CANCEL=8
+
+do_stageout() {
+    SOURCE="$1"
+    TARGET="$2"
+    MODE="$3"
+    FLAGS="$4"
+    FAILED="$5"
+
+    if [ "$FAILED" == "0" ] && $((FLAGS & _FLAG_ON_SUCCESS)) ; then
+        do_stage "$SOURCE" "$TARGET" "$MODE" $((FLAGS & _FLAG_IF_PRESENT))
+    elif [ "$FAILED" != "0" ] && $((FLAGS & _FLAG_ON_ERROR)) ; then
+        do_stage "$SOURCE" "$TARGET" "$MODE" $((FLAGS & _FLAG_IF_PRESENT))
+    fi
+}
+
+do_cleanup() {
+    TARGET="$1"
+
+    case "$TARGET" in
+        "{{job.spec.directory}}"*)
+            echo "rm -rf $TARGET" >>~/cleanup.log
+            ;;
+        *)
+            echo "Cannot clean $TARGET outside of job directory {{job.spec.directory}}"
+            exit 1
+            ;;
+    esac
+}

src/psij/executors/batch/script_generator.py

@@ -72,7 +72,8 @@ def __init__(self, config: JobExecutorConfig, template_path: pathlib.Path,
         super().__init__(config)
         with template_path.open('r') as template_file:
             self.template = pystache.parse(template_file.read())
-        self.renderer = pystache.Renderer(escape=escape)
+        common_dir = pathlib.Path(__file__).parent / 'common'
+        self.renderer = pystache.Renderer(escape=escape, search_dirs=[str(common_dir)])
 
     def generate_submit_script(self, job: Job, context: Dict[str, object], out: IO[str]) -> None:
         """See :func:`~SubmitScriptGenerator.generate_submit_script`.