def refill_quotas(self, last_poll_time, processing_rate,
active_shard_states):
"""Refill quotas for all active shards.
Args:
last_poll_time: Datetime with the last time the job state was updated.
processing_rate: How many items to process per second overall.
active_shard_states: All active shard states, list of ShardState.
"""
if not active_shard_states:
return
quota_manager = quota.QuotaManager(memcache.Client())
current_time = int(self._time())
last_poll_time = time.mktime(last_poll_time.timetuple())
total_quota_refill = processing_rate * max(
0, current_time - last_poll_time)
quota_refill = int(
math.ceil(1.0 * total_quota_refill / len(active_shard_states)))
if not quota_refill:
return
for shard_state in active_shard_states:
quota_manager.put(shard_state.shard_id, quota_refill)
def _schedule_shards(cls,
spec,
input_readers,
output_writers,
queue_name,
base_path):
"""Prepares shard states and schedules their execution.
Args:
spec: mapreduce specification as MapreduceSpec.
input_readers: list of InputReaders describing shard splits.
queue_name: The queue to run this job on.
base_path: The base url path of mapreduce callbacks.
"""
assert len(input_readers) == len(output_writers)
shard_states = []
for shard_number, input_reader in enumerate(input_readers):
shard_state = model.ShardState.create_new(spec.mapreduce_id, shard_number)
shard_state.shard_description = str(input_reader)
shard_states.append(shard_state)
existing_shard_states = db.get(shard.key() for shard in shard_states)
existing_shard_keys = set(shard.key() for shard in existing_shard_states
if shard is not None)
db.put((shard for shard in shard_states
if shard.key() not in existing_shard_keys),
config=util.create_datastore_write_config(spec))
processing_rate = int(spec.mapper.params.get(
"processing_rate") or model._DEFAULT_PROCESSING_RATE_PER_SEC)
quota_refill = processing_rate / len(shard_states)
quota_manager = quota.QuotaManager(memcache.Client())
for shard_state in shard_states:
quota_manager.put(shard_state.shard_id, quota_refill)
for shard_number, (input_reader, output_writer) in enumerate(
zip(input_readers, output_writers)):
shard_id = model.ShardState.shard_id_from_number(
spec.mapreduce_id, shard_number)
MapperWorkerCallbackHandler._schedule_slice(
shard_states[shard_number],
model.TransientShardState(
base_path, spec, shard_id, 0, input_reader,
output_writer=output_writer),
queue_name=queue_name)
def handle(self):
"""Handle request."""
tstate = model.TransientShardState.from_request(self.request)
spec = tstate.mapreduce_spec
self._start_time = self._time()
shard_id = tstate.shard_id
shard_state, control = db.get([
model.ShardState.get_key_by_shard_id(shard_id),
model.MapreduceControl.get_key_by_job_id(spec.mapreduce_id),
])
if not shard_state:
logging.error("State not found for shard ID %r; shutting down",
shard_id)
return
ctx = context.Context(spec,
shard_state,
task_retry_count=self.task_retry_count())
if control and control.command == model.MapreduceControl.ABORT:
logging.info("Abort command received by shard %d of job '%s'",
shard_state.shard_number, shard_state.mapreduce_id)
if tstate.output_writer:
tstate.output_writer.finalize(ctx, shard_state.shard_number)
shard_state.active = False
shard_state.result_status = model.ShardState.RESULT_ABORTED
shard_state.put(config=util.create_datastore_write_config(spec))
model.MapreduceControl.abort(spec.mapreduce_id)
return
input_reader = tstate.input_reader
if spec.mapper.params.get("enable_quota", True):
quota_consumer = quota.QuotaConsumer(
quota.QuotaManager(memcache.Client()), shard_id,
_QUOTA_BATCH_SIZE)
else:
quota_consumer = None
context.Context._set(ctx)
try:
if not quota_consumer or quota_consumer.check():
scan_aborted = False
entity = None
if not quota_consumer or quota_consumer.consume():
for entity in input_reader:
if isinstance(entity, db.Model):
shard_state.last_work_item = repr(entity.key())
else:
shard_state.last_work_item = repr(entity)[:100]
scan_aborted = not self.process_data(
entity, input_reader, ctx, tstate)
if (quota_consumer and not scan_aborted
and not quota_consumer.consume()):
scan_aborted = True
if scan_aborted:
break
else:
scan_aborted = True
if not scan_aborted:
logging.info("Processing done for shard %d of job '%s'",
shard_state.shard_number,
shard_state.mapreduce_id)
if quota_consumer:
quota_consumer.put(1)
shard_state.active = False
shard_state.result_status = model.ShardState.RESULT_SUCCESS
operation.counters.Increment(
"mapper-walltime-msec",
int((time.time() - self._start_time) * 1000))(ctx)
ctx.flush()
if not shard_state.active:
if tstate.output_writer:
tstate.output_writer.finalize(ctx,
shard_state.shard_number)
shard_state.put(config=util.create_datastore_write_config(spec))
finally:
context.Context._set(None)
if quota_consumer:
quota_consumer.dispose()
if shard_state.active:
self.reschedule(shard_state, tstate)
gc.collect()
def handle(self):
"""Handle request."""
tstate = model.TransientShardState.from_request(self.request)
spec = tstate.mapreduce_spec
self._start_time = self._time()
shard_id = tstate.shard_id
shard_state, control = db.get([
model.ShardState.get_key_by_shard_id(shard_id),
model.MapreduceControl.get_key_by_job_id(spec.mapreduce_id),
])
if not shard_state:
logging.error("State not found for shard ID %r; shutting down",
shard_id)
return
if not shard_state.active:
logging.error(
"Shard is not active. Looks like spurious task execution.")
return
ctx = context.Context(spec,
shard_state,
task_retry_count=self.task_retry_count())
if control and control.command == model.MapreduceControl.ABORT:
logging.info("Abort command received by shard %d of job '%s'",
shard_state.shard_number, shard_state.mapreduce_id)
if tstate.output_writer:
tstate.output_writer.finalize(ctx, shard_state.shard_number)
shard_state.active = False
shard_state.result_status = model.ShardState.RESULT_ABORTED
shard_state.put(config=util.create_datastore_write_config(spec))
model.MapreduceControl.abort(spec.mapreduce_id)
return
input_reader = tstate.input_reader
if spec.mapper.params.get("enable_quota", True):
quota_consumer = quota.QuotaConsumer(
quota.QuotaManager(memcache.Client()), shard_id,
_QUOTA_BATCH_SIZE)
else:
quota_consumer = None
context.Context._set(ctx)
try:
if not quota_consumer or quota_consumer.check():
scan_aborted = False
entity = None
try:
if not quota_consumer or quota_consumer.consume():
for entity in input_reader:
if isinstance(entity, db.Model):
shard_state.last_work_item = repr(entity.key())
else:
shard_state.last_work_item = repr(entity)[:100]
scan_aborted = not self.process_data(
entity, input_reader, ctx, tstate)
if (quota_consumer and not scan_aborted
and not quota_consumer.consume()):
scan_aborted = True
if scan_aborted:
break
else:
scan_aborted = True
if not scan_aborted:
logging.info(
"Processing done for shard %d of job '%s'",
shard_state.shard_number, shard_state.mapreduce_id)
if quota_consumer:
quota_consumer.put(1)
shard_state.active = False
shard_state.result_status = model.ShardState.RESULT_SUCCESS
operation.counters.Increment(
context.COUNTER_MAPPER_WALLTIME_MS,
int((time.time() - self._start_time) * 1000))(ctx)
ctx.flush()
except errors.FailJobError, e:
logging.error("Job failed: %s", e)
scan_aborted = True
shard_state.active = False
shard_state.result_status = model.ShardState.RESULT_FAILED
if not shard_state.active:
if tstate.output_writer:
tstate.output_writer.finalize(ctx,
shard_state.shard_number)
config = util.create_datastore_write_config(spec)
@db.transactional(retries=5)
def tx():
fresh_shard_state = db.get(
model.ShardState.get_key_by_shard_id(shard_id))
if (not fresh_shard_state.active
or "worker_active_state_collision"
in _TEST_INJECTED_FAULTS):
shard_state.active = False
logging.error(
"Spurious task execution. Aborting the shard.")
return
fresh_shard_state.copy_from(shard_state)
fresh_shard_state.put(config=config)
tx()
def handle(self):
"""Handle request."""
tstate = model.TransientShardState.from_request(self.request)
spec = tstate.mapreduce_spec
self._start_time = self._time()
shard_id = tstate.shard_id
shard_state, control = db.get([
model.ShardState.get_key_by_shard_id(shard_id),
model.MapreduceControl.get_key_by_job_id(spec.mapreduce_id),
])
if not shard_state:
logging.error("State not found for shard ID %r; shutting down",
shard_id)
return
if not shard_state.active:
logging.error("Shard is not active. Looks like spurious task execution.")
return
ctx = context.Context(spec, shard_state,
task_retry_count=self.task_retry_count())
if control and control.command == model.MapreduceControl.ABORT:
logging.info("Abort command received by shard %d of job '%s'",
shard_state.shard_number, shard_state.mapreduce_id)
shard_state.active = False
shard_state.result_status = model.ShardState.RESULT_ABORTED
shard_state.put(config=util.create_datastore_write_config(spec))
model.MapreduceControl.abort(spec.mapreduce_id)
return
input_reader = tstate.input_reader
if spec.mapper.params.get("enable_quota", True):
quota_consumer = quota.QuotaConsumer(
quota.QuotaManager(memcache.Client()),
shard_id,
_QUOTA_BATCH_SIZE)
else:
quota_consumer = None
if ndb is not None:
ndb_ctx = ndb.get_context()
ndb_ctx.set_cache_policy(lambda key: False)
ndb_ctx.set_memcache_policy(lambda key: False)
context.Context._set(ctx)
try:
if not quota_consumer or quota_consumer.check():
scan_aborted = False
entity = None
try:
if not quota_consumer or quota_consumer.consume():
for entity in input_reader:
if isinstance(entity, db.Model):
shard_state.last_work_item = repr(entity.key())
else:
shard_state.last_work_item = repr(entity)[:100]
scan_aborted = not self.process_data(
entity, input_reader, ctx, tstate)
if (quota_consumer and not scan_aborted and
not quota_consumer.consume()):
scan_aborted = True
if scan_aborted:
break
else:
scan_aborted = True
if not scan_aborted:
logging.info("Processing done for shard %d of job '%s'",
shard_state.shard_number, shard_state.mapreduce_id)
if quota_consumer:
quota_consumer.put(1)
shard_state.active = False
shard_state.result_status = model.ShardState.RESULT_SUCCESS
operation.counters.Increment(
context.COUNTER_MAPPER_WALLTIME_MS,
int((time.time() - self._start_time)*1000))(ctx)
ctx.flush()
except errors.RetrySliceError, e:
logging.error("Slice error: %s", e)
retry_count = int(
os.environ.get("HTTP_X_APPENGINE_TASKRETRYCOUNT") or 0)
if retry_count <= _RETRY_SLICE_ERROR_MAX_RETRIES:
raise
logging.error("Too many retries: %d, failing the job", retry_count)
scan_aborted = True
shard_state.active = False
shard_state.result_status = model.ShardState.RESULT_FAILED
except errors.FailJobError, e:
logging.error("Job failed: %s", e)
scan_aborted = True
shard_state.active = False
shard_state.result_status = model.ShardState.RESULT_FAILED
def handle(self):
"""Handle request."""
spec = model.MapreduceSpec.from_json_str(
self.request.get("mapreduce_spec"))
self._start_time = self._time()
shard_id = self.shard_id()
logging.debug("post: shard=%s slice=%s headers=%s", shard_id,
self.slice_id(), self.request.headers)
shard_state, control = db.get([
model.ShardState.get_key_by_shard_id(shard_id),
model.MapreduceControl.get_key_by_job_id(spec.mapreduce_id),
])
if not shard_state:
logging.error("State not found for shard ID %r; shutting down",
shard_id)
return
if control and control.command == model.MapreduceControl.ABORT:
logging.info("Abort command received by shard %d of job '%s'",
shard_state.shard_number, shard_state.mapreduce_id)
shard_state.active = False
shard_state.result_status = model.ShardState.RESULT_ABORTED
shard_state.put(config=util.create_datastore_write_config(spec))
model.MapreduceControl.abort(spec.mapreduce_id)
return
input_reader = self.input_reader(spec.mapper)
if spec.mapper.params.get("enable_quota", True):
quota_consumer = quota.QuotaConsumer(
quota.QuotaManager(memcache.Client()), shard_id,
_QUOTA_BATCH_SIZE)
else:
quota_consumer = None
ctx = context.Context(spec,
shard_state,
task_retry_count=self.task_retry_count())
context.Context._set(ctx)
try:
if not quota_consumer or quota_consumer.check():
scan_aborted = False
entity = None
if not quota_consumer or quota_consumer.consume():
for entity in input_reader:
if isinstance(entity, db.Model):
shard_state.last_work_item = repr(entity.key())
else:
shard_state.last_work_item = repr(entity)[:100]
scan_aborted = not self.process_entity(entity, ctx)
if (quota_consumer and not scan_aborted
and not quota_consumer.consume()):
scan_aborted = True
if scan_aborted:
break
else:
scan_aborted = True
if not scan_aborted:
logging.info("Processing done for shard %d of job '%s'",
shard_state.shard_number,
shard_state.mapreduce_id)
if quota_consumer:
quota_consumer.put(1)
shard_state.active = False
shard_state.result_status = model.ShardState.RESULT_SUCCESS
ctx.flush()
finally:
context.Context._set(None)
if quota_consumer:
quota_consumer.dispose()
if shard_state.active:
self.reschedule(spec, input_reader)