Python Worker.reservation方法代码示例

# 需要导入模块: from worker import Worker [as 别名]
# 或者: from worker.Worker import reservation [as 别名]
    def test_request_to_booted(self):
        self.aws_mock.my_booted_machine = mock.MagicMock()
        self.aws_mock.my_booted_machine.return_value = 'instance', 'ip'

        from machine_midwife import MachineMidwife
        Apprentice = MachineMidwife.Apprentice
        from job import Job
        from worker import Worker

        apprentice = Apprentice()
        apprentice.settings = mock.MagicMock()
        apprentice.settings.aws_req_max_wait = 10
        apprentice.client = mock.MagicMock()
        apprentice.client.exists.return_value = True
        job = Job('requested', 'batch-')
        worker = Worker(None, None)
        worker.reservation = 'some'
        worker.request_time = datetime.now() - timedelta(minutes=5)

        apprentice.client.keys.return_value = ['jm-']
        apprentice.client.get.side_effect = [pickle.dumps(worker), pickle.dumps(job)]
        apprentice.client.set = mock.MagicMock()
        apprentice.client.publish = mock.MagicMock()

        apprentice.check_newborn()

        assert apprentice.client.keys.call_count == 1
        assert apprentice.client.get.call_count == 2
        assert apprentice.client.set.call_count == 2
        assert apprentice.client.publish.call_count == 1
        assert pickle.loads(apprentice.client.set.call_args_list[0][0][1]).instance == 'instance'
        assert pickle.loads(apprentice.client.set.call_args_list[1][0][1]).state == 'booted'

# 需要导入模块: from worker import Worker [as 别名]
# 或者: from worker.Worker import reservation [as 别名]
    def test_delayed_machine_state_flow_for_requested_with_recycle(self):
        from machine_midwife import MachineMidwife
        from job import Job
        from worker import Worker

        midwife = MachineMidwife()
        midwife.apprentice = mock.MagicMock()
        midwife.settings = mock.MagicMock()
        midwife.client = mock.MagicMock()
        midwife.job_pub_sub = mock.MagicMock()
        midwife.job_pub_sub.listen.return_value = [{'data': 'test'}]
        midwife.client.exists.return_value = True
        job = Job('delayed', 'batch-')
        worker = Worker(None, 'batch-')
        worker.reservation = 'reservation'
        worker.request_time = datetime.now()
        midwife.client.keys.return_value = ['jm-']
        midwife.client.get.side_effect = [pickle.dumps(job), pickle.dumps(worker)]
        midwife.client.set = mock.MagicMock()
        midwife.client.publish = mock.MagicMock()

        midwife.run()

        assert midwife.client.exists.call_count == 2
        assert len(midwife.client.set.call_args_list) == 2
        assert pickle.loads(midwife.client.set.call_args_list[0][0][1]).job_id == 'test'
        assert pickle.loads(midwife.client.set.call_args_list[1][0][1]).state == 'booted'

# 需要导入模块: from worker import Worker [as 别名]
# 或者: from worker.Worker import reservation [as 别名]
 def run(self):
     self.apprentice.start()
     for item in self.job_pub_sub.listen():
         job_id = item['data']
         if job_id == 'KILL':
             self.apprentice.halt()
             self.job_pub_sub.unsubscribe()
             logging.info('MachineMidwife: Stopping.')
             return
         #
         queue_full = self.choke_full()
         #
         logging.debug('MachineMidwife: Redis signals for job: %s' % job_id)
         if not self.client.exists(job_id):
             logging.warning('MachineMidwife: Redis signaled for non-existing job: %s' % job_id)
             continue
         job = pickle.loads(self.client.get(job_id))
         logging.debug('MachineMidwife: Job %s has state %s' % (job_id, job.state))
         if job.state != 'received' and job.state != 'delayed':
             continue
         # Recycle
         recycled = False
         for worker_id in [worker_key for worker_key in self.client.keys() if worker_key.startswith('jm-')]:  # Redis keys(pattern='*') does not filter at all.
             if not self.client.exists(worker_id):
                 continue
             existing_worker = pickle.loads(self.client.get(worker_id))
             if existing_worker.batch_id == job.batch_id and existing_worker.job_id is None:
                 logging.info('MachineMidwife: Recycling worker %s for job %s of batch %s.' % (worker_id, job_id, job.batch_id))
                 job.state = 'booted'
                 existing_worker.job_id = job_id
                 self.client.set(worker_id, pickle.dumps(existing_worker))
                 self.client.set(job_id, pickle.dumps(job))
                 self.client.publish('jobs', job_id)
                 recycled = True
                 break
         if recycled:
             continue
         # New machine
         if not queue_full:
             worker_id, reservation = aws.start_machine(job.ami, job.instance_type)
             if worker_id is not None:
                 logging.info('MachineMidwife: Started new worker for job %s of batch %s.' % (job_id, job.batch_id))
                 job.state = 'requested'
                 worker = Worker(job_id, job.batch_id)
                 worker.request_time = datetime.now()
                 worker.reservation = reservation
                 self.client.set(worker_id, pickle.dumps(worker))
                 self.client.set(job_id, pickle.dumps(job))
                 self.client.publish('jobs', job_id)
             else:
                 logging.warning('MachineMidwife: AWS failed to start a new machine.')
                 job.state = 'delayed'
                 self.client.set(job_id, pickle.dumps(job))
         else:
             job.state = 'delayed'
             self.client.set(job_id, pickle.dumps(job))