Python MesosSchedulerDriver.sendFrameworkMessage方法代码示例

# 需要导入模块: from pymesos import MesosSchedulerDriver [as 别名]
# 或者: from pymesos.MesosSchedulerDriver import sendFrameworkMessage [as 别名]
def test_send_framework_message(mocker):
    ID = str(uuid.uuid4())
    sched = mocker.Mock()
    framework = {'id': {'value': ID}}
    master = mocker.Mock()
    driver = MesosSchedulerDriver(sched, framework, master)
    driver._send = mocker.Mock()
    executor_id = {'value': str(uuid.uuid4())}
    agent_id = {'value': str(uuid.uuid4())}
    message = ''.join(random.choice(string.printable)
                      for _ in range(random.randint(1, 100)))
    message = encode_data(message.encode('utf-8'))
    driver.sendFrameworkMessage(executor_id, agent_id, message)
    driver._send.assert_not_called()
    driver._stream_id = 'a-stream-id'
    driver.sendFrameworkMessage(executor_id, agent_id, message)
    driver._send.assert_called_once_with({
        'type': 'MESSAGE',
        'framework_id': {
            'value': ID
        },
        'message': {
            'agent_id': agent_id,
            'executor_id': executor_id,
            'data': message,
        }
    })

# 需要导入模块: from pymesos import MesosSchedulerDriver [as 别名]
# 或者: from pymesos.MesosSchedulerDriver import sendFrameworkMessage [as 别名]

#.........这里部分代码省略.........
            if slave_id is not None:
                if slave_id in self.slave_to_proc:
                    self.slave_to_proc[slave_id].remove(proc_id)
            else:
                for slave_id, procs in self.slave_to_proc.iteritems():
                    if proc_id in procs:
                        procs.remove(proc_id)

            proc._finished(success, message, data)

    def statusUpdate(self, driver, update):
        with self._lock:
            proc_id = int(update.task_id.value)
            logger.info("Status update for proc, id=%s, state=%s" % (proc_id, update.state))
            if update.state == mesos_pb2.TASK_RUNNING:
                if update.slave_id.value in self.slave_to_proc:
                    self.slave_to_proc[update.slave_id.value].add(proc_id)
                else:
                    self.slave_to_proc[update.slave_id.value] = set([proc_id])

                proc = self.procs_launched[proc_id]
                proc._started()

            elif update.state >= mesos_pb2.TASK_FINISHED:
                slave_id = update.slave_id.value
                success = update.state == mesos_pb2.TASK_FINISHED
                message = update.message
                data = update.data and pickle.loads(update.data)
                self._call_finished(proc_id, success, message, data, slave_id)
                driver.reviveOffers()

    def offerRescinded(self, driver, offer_id):
        with self._lock:
            if self.procs_pending:
                logger.info("Revive offers for pending procs")
                driver.reviveOffers()

    def slaveLost(self, driver, slave_id):
        with self._lock:
            for proc_id in self.slave_to_proc.pop(slave_id, []):
                self._call_finished(proc_id, False, "Slave lost", None, slave_id)

    def error(self, driver, message):
        with self._lock:
            for proc in self.procs_pending.values():
                self._call_finished(proc.id, False, "Stopped", None)

            for proc in self.procs_launched.values():
                self._call_finished(proc.id, False, "Stopped", None)

        self.stop()

    def start(self):
        self.driver.start()

    def stop(self):
        assert not self.driver.aborted
        self.driver.stop()

    def submit(self, proc):
        if self.driver.aborted:
            raise RuntimeError("driver already aborted")

        with self._lock:
            if proc.id not in self.procs_pending:
                logger.info("Try submit proc, id=%s", (proc.id,))
                self.procs_pending[proc.id] = proc
                if len(self.procs_pending) == 1:
                    logger.info("Revive offers for pending procs")
                    self.driver.reviveOffers()
            else:
                raise ValueError("Proc with same id already submitted")

    def cancel(self, proc):
        if self.driver.aborted:
            raise RuntimeError("driver already aborted")

        with self._lock:
            if proc.id in self.procs_pending:
                del self.procs_pending[proc.id]
            elif proc.id in self.procs_launched:
                del self.procs_launched[proc.id]
                self.driver.killTask(mesos_pb2.TaskID(value=str(proc.id)))

            for slave_id, procs in self.slave_to_proc.items():
                procs.pop(proc.id)
                if not procs:
                    del self.slave_to_proc[slave_id]

    def send_data(self, pid, type, data):
        if self.driver.aborted:
            raise RuntimeError("driver already aborted")

        msg = pickle.dumps((pid, type, data))
        for slave_id, procs in self.slave_to_proc.iteritems():
            if pid in procs:
                self.driver.sendFrameworkMessage(self.executor.executor_id, mesos_pb2.SlaveID(value=slave_id), msg)
                return

        raise RuntimeError("Cannot find slave for pid %s" % (pid,))