def tcp_serve_task(self, conn, addr, coro=None):
        def job_request_task(msg):
            try:
                _job = unserialize(msg)
            except:
                _dispy_logger.debug('Ignoring job request from %s', addr[0])
                # _dispy_logger.debug(traceback.format_exc())
                raise StopIteration

            compute = self.computations.get(_job.compute_id, None)
            if compute is not None:
                if compute.scheduler_ip_addr != self.scheduler['ip_addr'] or \
                   compute.scheduler_port != self.scheduler['port'] or \
                   compute.auth not in self.scheduler['auth']:
                    _dispy_logger.debug('Invalid scheduler IP address: scheduler %s:%s != %s:%s',
                                        compute.scheduler_ip_addr, compute.scheduler_port,
                                        self.scheduler['ip_addr'], self.scheduler['port'])
                    compute = None
            if self.avail_cpus == 0:
                try:
                    yield conn.send_msg('NAK (all cpus busy)'.encode())
                except:
                    pass
                raise StopIteration
            elif compute is None:
                _dispy_logger.warning('Invalid computation %s', _job.compute_id)
                try:
                    yield conn.send_msg(('NAK (invalid computation %s)' %
                                         _job.compute_id).encode())
                except:
                    pass
                raise StopIteration

            for xf in _job.xfer_files:
                if MaxFileSize and xf.stat_buf.st_size > MaxFileSize:
                    try:
                        yield conn.send_msg('NAK'.encode())
                    except:
                        pass
                    raise StopIteration

            reply_addr = (compute.scheduler_ip_addr, compute.job_result_port)
            _dispy_logger.debug('New job id %s from %s/%s',
                                _job.uid, addr[0], compute.scheduler_ip_addr)

            if compute.type == _Compute.func_type:
                reply = _JobReply(_job, self.ext_ip_addr)
                reply.start_time = time.time()
                job_info = _DispyJobInfo(reply, reply_addr, compute, _job.xfer_files)

                args = (job_info, self.certfile, self.keyfile, compute.name,
                        _job.args, _job.kwargs, (compute.code, _job.code),
                        compute.globals, compute.dest_path, self.reply_Q)
                try:
                    yield conn.send_msg('ACK')
                except:
                    _dispy_logger.warning('Failed to send response for new job to %s', str(addr))
                    raise StopIteration
                proc = multiprocessing.Process(target=_dispy_job_func, args=args)
                self.avail_cpus -= 1
                compute.pending_jobs += 1
                self.thread_lock.acquire()
                self.job_infos[_job.uid] = job_info
                self.thread_lock.release()
                try:
                    proc.start()
                except:
                    job_info.job_reply.status = DispyJob.Terminated
                    job_info.job_reply.exception = traceback.format_exc()
                    job_info.job_reply.end_time = time.time()
                    job_info.proc = None
                    self.reply_Q.put(job_info.job_reply)
                else:
                    job_info.proc = proc
                    job_info.job_reply.status = DispyJob.Running
                raise StopIteration
            elif compute.type == _Compute.prog_type:
                try:
                    yield conn.send_msg('ACK')
                except:
                    _dispy_logger.warning('Failed to send response for new job to %s', str(addr))
                    raise StopIteration
                reply = _JobReply(_job, self.ext_ip_addr)
                reply.start_time = time.time()
                job_info = _DispyJobInfo(reply, reply_addr, compute, _job.xfer_files)
                job_info.job_reply.status = DispyJob.Running
                self.thread_lock.acquire()
                self.job_infos[_job.uid] = job_info
                self.thread_lock.release()
                self.avail_cpus -= 1
                compute.pending_jobs += 1
                prog_thread = threading.Thread(target=self.__job_program, args=(_job, job_info))
                prog_thread.start()
                raise StopIteration
            else:
                try:
                    yield conn.send_msg(('NAK (invalid computation type "%s")' %
                                         compute.type).encode())
                except:
                    _dispy_logger.warning('Failed to send response for new job to %s', str(addr))

        def add_computation_task(msg):
            try:
                compute = unserialize(msg)
            except:
                try:
                    yield conn.send_msg(('Invalid computation request ignored').encode())
                except:
                    pass
                raise StopIteration
            if not ((self.scheduler['ip_addr'] is None and not self.scheduler['auth']) or
                    (self.scheduler['ip_addr'] == compute.scheduler_ip_addr and
                     self.scheduler['port'] == compute.scheduler_port and
                     self.service_available())):
                _dispy_logger.debug('Ignoring computation request from %s: %s, %s, %s',
                                    compute.scheduler_ip_addr, self.scheduler['ip_addr'],
                                    self.avail_cpus, self.num_cpus)
                try:
                    yield conn.send_msg(('Node busy').encode())
                except:
                    pass
                raise StopIteration

            if MaxFileSize:
                for xf in compute.xfer_files:
                    if xf.stat_buf.st_size > MaxFileSize:
                        try:
                            yield conn.send_msg(('File "%s" is too big; limit is %s' %
                                                 (xf.name, MaxFileSize)).encode())
                        except:
                            pass
                        raise StopIteration
            compute.xfer_files = set()
            dest = os.path.join(self.dest_path_prefix, compute.scheduler_ip_addr)
            if not os.path.isdir(dest):
                try:
                    os.mkdir(dest)
                except:
                    yield conn.send_msg(('Could not create destination path').encode())
                    raise StopIteration
            if compute.dest_path and isinstance(compute.dest_path, str):
                # TODO: get os.sep from client and convert (in case of mixed environments)?
                if not compute.dest_path.startswith(os.sep):
                    compute.dest_path = os.path.join(dest, compute.dest_path)
                if not os.path.isdir(compute.dest_path):
                    try:
                        os.makedirs(compute.dest_path)
                    except:
                        try:
                            yield conn.send_msg(('Could not create destination path').encode())
                        except:
                            pass
                        raise StopIteration
            else:
                compute.dest_path = tempfile.mkdtemp(prefix=compute.name + '_', dir=dest)
            os.chmod(compute.dest_path, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR)

            if compute.id in self.computations:
                _dispy_logger.warning('Computation "%s" (%s) is being replaced',
                                      compute.name, compute.id)
            setattr(compute, 'last_pulse', time.time())
            setattr(compute, 'pending_jobs', 0)
            setattr(compute, 'pending_results', 0)
            setattr(compute, 'zombie', False)
            setattr(compute, 'globals', {})
            setattr(compute, 'ante_modules', set(sys.modules.iterkeys()))
            setattr(compute, 'file_uses', {})

            if compute.code:
                try:
                    code = compute.code
                    code += self.__init_code
                    code = compile(code, '<string>', 'exec')
                except:
                    if os.path.isdir(compute.dest_path):
                        os.rmdir(compute.dest_path)
                    try:
                        yield conn.send_msg(('%s: Computation "%s" could not be compiled' %
                                             (self.ext_ip_addr, compute.name)).encode())
                    except:
                        pass
                    raise StopIteration
                compute.code = marshal.dumps(code)

            if compute.type == _Compute.prog_type:
                compute.name = os.path.join(compute.dest_path, os.path.basename(compute.name))

            if not ((self.scheduler['ip_addr'] is None) or
                    (self.scheduler['ip_addr'] == compute.scheduler_ip_addr and
                     self.scheduler['port'] == compute.scheduler_port)):
                if os.path.isdir(compute.dest_path):
                    try:
                        os.rmdir(compute.dest_path)
                        yield conn.send_msg(serialize(-1))
                    except:
                        pass
                raise StopIteration

            self.computations[compute.id] = compute
            self.scheduler['ip_addr'] = compute.scheduler_ip_addr
            self.scheduler['port'] = compute.scheduler_port
            self.scheduler['auth'].add(compute.auth)
            compute_save = os.path.join(self.dest_path_prefix, '%s_%s' % (compute.id, compute.auth))
            fd = open(compute_save, 'wb')
            pickle.dump(compute, fd)
            fd.close()

            # add variables needed for 'dispy_provisional_result' and
            # 'dispy_send_file' to compute.globals; but in Windows
            # compute.globals can't be passed via multiprocessing.Process
            if os.name == 'nt':
                compute.globals = {}
            else:
                for var in ('AsyncSocket', 'DispyJob', 'serialize', 'unserialize', '_XferFile',
                            'MaxFileSize', 'MsgTimeout'):
                    compute.globals[var] = globals()[var]
                compute.globals.update(self.__init_modules)
            compute.globals['_DispyNode'] = None

            try:
                yield conn.send_msg(serialize(self.avail_cpus))
            except:
                del self.computations[compute.id]
                compute.globals = {}
                self.scheduler['ip_addr'] = None
                self.scheduler['auth'].discard(compute.auth)
                os.remove(compute_save)
                if os.path.isdir(compute.dest_path):
                    try:
                        os.rmdir(compute.dest_path)
                    except:
                        pass
            else:
                self.pulse_interval = num_min(self.pulse_interval, compute.pulse_interval)
                if not self.pulse_interval:
                    self.pulse_interval = 10 * 60
                if self.zombie_interval:
                    self.pulse_interval = num_min(self.pulse_interval, self.zombie_interval / 5.0)
                self.timer_coro.resume(True)

        def xfer_file_task(msg):
            try:
                xf = unserialize(msg)
            except:
                _dispy_logger.debug('Ignoring file trasnfer request from %s', addr[0])
                raise StopIteration

            compute = self.computations.get(xf.compute_id, None)
            if not compute or (MaxFileSize and xf.stat_buf.st_size > MaxFileSize):
                _dispy_logger.error('Invalid file transfer for "%s"', xf.name)
                yield conn.send_msg(serialize(-1))
                raise StopIteration
            tgt = os.path.join(compute.dest_path, os.path.basename(xf.name))
            if os.path.isfile(tgt) and _same_file(tgt, xf):
                if tgt in compute.file_uses:
                    compute.file_uses[tgt] += 1
                else:
                    compute.file_uses[tgt] = 2
                yield conn.send_msg(serialize(xf.stat_buf.st_size))
            else:
                try:
                    fd = open(tgt, 'wb')
                    recvd = 0
                    _dispy_logger.debug('Copying file %s to %s (%s)',
                                        xf.name, tgt, xf.stat_buf.st_size)
                    while recvd < xf.stat_buf.st_size:
                        yield conn.send_msg(serialize(recvd))
                        data = yield conn.recvall(min(xf.stat_buf.st_size-recvd, 1024000))
                        if not data:
                            break
                        fd.write(data)
                        recvd += len(data)
                    yield conn.send_msg(serialize(recvd))
                    fd.close()
                    _dispy_logger.debug('Copied file %s, %s / %s',
                                        tgt, recvd, xf.stat_buf.st_size)
                    assert recvd == xf.stat_buf.st_size
                    os.utime(tgt, (xf.stat_buf.st_atime, xf.stat_buf.st_mtime))
                    os.chmod(tgt, stat.S_IMODE(xf.stat_buf.st_mode))
                except:
                    _dispy_logger.warning('Copying file "%s" failed with "%s"',
                                          xf.name, traceback.format_exc())
                    os.remove(tgt)
                else:
                    if tgt in compute.file_uses:
                        compute.file_uses[tgt] += 1
                    else:
                        compute.file_uses[tgt] = 1
            raise StopIteration  # xfer_file_task

        def setup_computation(msg):
            try:
                compute_id = unserialize(msg)
                compute = self.computations[compute_id]
                assert isinstance(compute.setup, _Function)
                os.chdir(compute.dest_path)
                localvars = {'_dispy_setup_args': compute.setup.args,
                             '_dispy_setup_kwargs': compute.setup.kwargs}
                if os.name == 'nt':
                    globalvars = globals()
                else:
                    globalvars = compute.globals
                exec(marshal.loads(compute.code)) in globalvars, localvars
                exec('assert %s(*_dispy_setup_args, **_dispy_setup_kwargs) == 0' %
                     compute.setup.name) in globalvars, localvars
                if os.name == 'nt':
                    compute.globals.update({var: globals()[var] for var in globals()
                                            if var not in self.__init_globals})
            except:
                _dispy_logger.debug('Setup failed')
                resp = traceback.format_exc().encode()
            else:
                resp = 'ACK'
            if resp != 'ACK':
                if not compute.cleanup:
                    compute.cleanup = True
                compute.zombie = True
                self.cleanup_computation(compute)
            yield conn.send_msg(resp)

        def terminate_job_task(compute, job_info):
            if not job_info.proc:
                raise StopIteration
            _dispy_logger.debug('Terminating job %s of "%s"',
                                job_info.job_reply.uid, compute.name)
            job_info.proc.terminate()
            if isinstance(job_info.proc, multiprocessing.Process):
                for x in range(20):
                    if job_info.proc.is_alive():
                        yield coro.sleep(0.1)
                    else:
                        _dispy_logger.debug('Process "%s" for job %s terminated',
                                            compute.name, job_info.job_reply.uid)
                        break
                else:
                    _dispy_logger.warning('Could not kill process %s', compute.name)
                    raise StopIteration
            else:
                assert isinstance(job_info.proc, subprocess.Popen)
                for x in range(20):
                    rc = job_info.proc.poll()
                    _dispy_logger.debug('Program "%s" for job %s terminated with %s',
                                        compute.name, job_info.job_reply.uid, rc)
                    if rc is not None:
                        break
                    if x == 10:
                        _dispy_logger.debug('Killing job %s', job_info.job_reply.uid)
                        job_info.proc.kill()
                    yield coro.sleep(0.1)
                else:
                    _dispy_logger.warning('Could not kill process %s', compute.name)
                    raise StopIteration
            job_info.job_reply.end_time = time.time()
            job_info.proc = None
            self.thread_lock.acquire()
            if self.job_infos.get(job_info.job_reply.uid, None) == job_info:
                job_info.job_reply.status = DispyJob.Terminated
                self.reply_Q.put(job_info.job_reply)
            self.thread_lock.release()

        def retrieve_job_task(msg):
            # generator

            def send_reply(reply):
                try:
                    yield conn.send_msg(serialize(reply))
                except:
                    raise StopIteration(-1)
                raise StopIteration(0)

            try:
                req = unserialize(msg)
                uid = req['uid']
                compute_id = req['compute_id']
                auth = req['auth']
                job_hash = req['hash']
            except:
                yield send_reply(None)
                raise StopIteration

            pkl_path = os.path.join(self.dest_path_prefix, '%s_%s' % (compute_id, auth))
            compute = self.computations.get(compute_id, None)
            if not compute:
                fd = open(pkl_path, 'rb')
                compute = pickle.load(fd)
                fd.close()
            if not compute or compute.auth != auth:
                yield send_reply(None)
                raise StopIteration

            info_file = os.path.join(compute.dest_path, '_dispy_job_reply_%s' % uid)
            if not os.path.isfile(info_file):
                yield send_reply(None)
                raise StopIteration
            try:
                fd = open(info_file, 'rb')
                job_reply = pickle.load(fd)
                fd.close()
                assert job_reply.hash == job_hash
            except:
                yield send_reply(None)
                raise StopIteration

            try:
                yield conn.send_msg(serialize(job_reply))
                ack = yield conn.recv_msg()
                assert ack == 'ACK'
                compute.pending_results -= 1
                fd = open(pkl_path, 'wb')
                pickle.dump(compute, fd)
                fd.close()
            except:
                pass
            else:
                try:
                    os.remove(info_file)
                except:
                    pass
                if compute.pending_results == 0:
                    self.cleanup_computation(compute)

        # tcp_serve_task starts
        try:
            req = yield conn.recvall(len(self.auth))
        except:
            _dispy_logger.warning('Ignoring request; invalid client authentication?')
            conn.close()
            raise StopIteration
        msg = yield conn.recv_msg()
        if req != self.auth:
            if msg.startswith('PING:'):
                pass
            else:
                _dispy_logger.warning('Ignoring request; invalid client authentication?')
                conn.close()
                raise StopIteration
        if not msg:
            conn.close()
            raise StopIteration
        if msg.startswith('JOB:'):
            msg = msg[len('JOB:'):]
            yield job_request_task(msg)
            conn.close()
        elif msg.startswith('COMPUTE:'):
            msg = msg[len('COMPUTE:'):]
            yield add_computation_task(msg)
            conn.close()
        elif msg.startswith('FILEXFER:'):
            msg = msg[len('FILEXFER:'):]
            yield xfer_file_task(msg)
            conn.close()
        elif msg.startswith('SETUP:'):
            msg = msg[len('SETUP:'):]
            yield setup_computation(msg)
            conn.close()
        elif msg.startswith('CLOSE:'):
            msg = msg[len('CLOSE:'):]
            try:
                info = unserialize(msg)
                compute_id = info['compute_id']
                auth = info['auth']
                terminate_pending = info.get('terminate_pending', False)
            except:
                _dispy_logger.debug('Deleting computation failed with %s', traceback.format_exc())
            else:
                compute = self.computations.get(compute_id, None)
                if compute is None or compute.auth != auth:
                    _dispy_logger.warning('Computation "%s" is not valid', compute_id)
                else:
                    compute.zombie = True
                    if terminate_pending:
                        self.thread_lock.acquire()
                        job_infos = [job_info for job_info in self.job_infos.values()
                                     if job_info.compute_id == compute_id]
                        self.thread_lock.release()
                        for job_info in job_infos:
                            yield terminate_job_task(compute, job_info)
                    self.cleanup_computation(compute)
            yield conn.send_msg('ACK')
            conn.close()
        elif msg.startswith('TERMINATE_JOB:'):
            msg = msg[len('TERMINATE_JOB:'):]
            try:
                _job = unserialize(msg)
                compute = self.computations[_job.compute_id]
                # assert addr[0] == compute.scheduler_ip_addr
                self.thread_lock.acquire()
                job_info = self.job_infos.get(_job.uid, None)
                self.thread_lock.release()
                assert job_info is not None
            except:
                _dispy_logger.debug('Invalid terminate job request from %s, %s',
                                    addr[0], compute.scheduler_ip_addr)
            else:
                yield terminate_job_task(compute, job_info)
            conn.close()
        elif msg.startswith('RESEND_JOB_RESULTS:'):
            msg = msg[len('RESEND_JOB_RESULTS:'):]
            try:
                info = unserialize(msg)
                compute_id = info['compute_id']
                auth = info['auth']
            except:
                reply = 0
            else:
                compute = self.computations.get(compute_id, None)
                if compute is None or compute.auth != auth:
                    try:
                        fd = open(os.path.join(self.dest_path_prefix,
                                               '%s_%s' % (compute_id, auth)), 'rb')
                        compute = pickle.load(fd)
                        fd.close()
                    except:
                        pass
                if compute is None:
                    reply = 0
                else:
                    reply = compute.pending_results + compute.pending_jobs
            yield conn.send_msg(serialize(reply))
            conn.close()
            if reply > 0:
                yield self.resend_job_results(compute, coro=coro)
        elif msg.startswith('PING:'):
            try:
                info = unserialize(msg[len('PING:'):])
                if (info['version'] == _dispy_version and
                   not self.scheduler['ip_addr'] and not self.job_infos):
                    Coro(self.send_pong_msg, info, addr)
            except:
                _dispy_logger.debug(traceback.format_exc())
            conn.close()
        elif msg.startswith('PENDING_JOBS:'):
            msg = msg[len('PENDING_JOBS:'):]
            reply = {'done': [], 'pending': 0}
            try:
                info = unserialize(msg)
                compute_id = info['compute_id']
                auth = info['auth']
            except:
                pass
            else:
                compute = self.computations.get(compute_id, None)
                if compute is None or compute.auth != auth:
                    fd = open(os.path.join(self.dest_path_prefix,
                                           '%s_%s' % (compute_id, auth)), 'rb')
                    compute = pickle.load(fd)
                    fd.close()
                if compute is not None:
                    done = []
                    if compute.pending_results:
                        for result_file in glob.glob(os.path.join(compute.dest_path,
                                                                  '_dispy_job_reply_*')):
                            result_file = os.path.basename(result_file)
                            try:
                                uid = int(result_file[len('_dispy_job_reply_'):])
                            except:
                                pass
                            else:
                                done.append(uid)
                                # limit so as not to take up too much time
                                if len(done) > 50:
                                    break
                    reply['done'] = done
                    reply['pending'] = compute.pending_jobs
            yield conn.send_msg(serialize(reply))
            conn.close()
        elif msg.startswith('RETRIEVE_JOB:'):
            msg = msg[len('RETRIEVE_JOB:'):]
            yield retrieve_job_task(msg)
            conn.close()
        else:
            _dispy_logger.warning('Invalid request "%s" from %s',
                                  msg[:min(10, len(msg))], addr[0])
            resp = ('NAK (invalid command: %s)' % (msg[:min(10, len(msg))])).encode()
            try:
                yield conn.send_msg(resp)
            except:
                _dispy_logger.warning('Failed to send reply to %s', str(addr))
            conn.close()
dataset/ETHPy150Open pgiri/dispy/py2/dispy/dispynode.py/_DispyNode.tcp_serve_task/original