def map(message):
"""Map Step: Create COGs."""
if isinstance(message, str):
message = json.loads(message)
jobid = message["jobid"]
partid = message["partid"]
mosaicid = message["metadata"]["mosaicid"]
logger.info(f"[MAP] Processing {jobid} - {partid}")
with Part(jobid, partid, progress=progress):
src_path = message["src_path"]
# with rasterio.open(src_path) as src_dst:
# witdh, height, dtype = src_dst.width, src_dst.height, src_dst.dtypes[0]
# size = numpy.array([0], dtype=dtype).nbytes * witdh * height
# logger.info(f"Size of image {size}")
bname = os.path.splitext(os.path.basename(src_path))[0]
out_key = os.path.join("cogs", mosaicid, f"{bname}_cog.tif")
translator.process(
src_path,
os.environ["MOSAIC_BUCKET"],
out_key,
message["profile_name"],
message["profile_options"],
**message["options"],
)
return True
def test_Part_job_not_done(aws_send_message, monkeypatch):
monkeypatch.setenv('WorkTopic', 'abc123')
monkeypatch.setenv('ProgressTable', 'arn::table/foo')
with Part(jobid='123', partid=1, progress=MockProgress()):
pass
aws_send_message.assert_not_called()
def test_Part_job_done(aws_send_message, monkeypatch):
monkeypatch.setenv('WorkTopic', 'abc123')
monkeypatch.setenv('ProgressTable', 'arn::table/foo')
class CustomProgress(MockProgress):
def complete_part(self, jobid, partid):
return True
with Part(jobid='123', partid=1, progress=CustomProgress()):
pass
aws_send_message.assert_called_once()
assert aws_send_message.call_args[1].get('subject') == 'reduce'
def test_Part_fail_job_on(status, fail_job, aws_send_message, monkeypatch):
monkeypatch.setenv('WorkTopic', 'abc123')
monkeypatch.setenv('ProgressTable', 'arn::table/foo')
class CustomException(Exception):
pass
# the job WILL be marked as failed, the exeption is still re-raised
with pytest.raises(CustomException):
with Part(jobid=2, partid=2, fail_job_on=[CustomException]):
raise CustomException()
fail_job.assert_called_once_with(2, 2)
aws_send_message.assert_not_called()
def test_progress_reduce_once(aws_send_message, sns_worker, monkeypatch):
"""Ensure that reduce message is only sent once,
even if parts are completed multiple times
"""
monkeypatch.setenv('WorkTopic', 'abc123')
progress = RedisProgress()
jobid = create_job(parts, progress=progress)
for i in range(3):
with Part(jobid, i, progress=progress):
pass
aws_send_message.assert_called_once()
assert aws_send_message.call_args[1].get('subject') == 'reduce'
aws_send_message.reset_mock()
# Finishing already completed parts should not trigger reduce message, only warn
with pytest.warns(UserWarning) as record:
for i in range(3):
with Part(jobid, i, progress=progress):
pass
assert 'skip' in record[0].message.args[0]
aws_send_message.assert_not_called()
def test_progress_all_done(aws_send_message, sns_worker, monkeypatch):
monkeypatch.setenv('WorkTopic', 'abc123')
progress = RedisProgress()
jobid = create_job(parts, progress=progress)
for i in range(3):
with Part(jobid, i, progress=progress):
pass
assert progress.status(jobid)['remaining'] == 0
# 3 map messages and 1 reduce message sent
sns_worker.assert_called_once()
assert len(sns_worker.call_args[0][0]) == 3
assert sns_worker.call_args[1].get('subject') == 'map'
aws_send_message.assert_called_once()
assert aws_send_message.call_args[1].get('subject') == 'reduce'
def test_progress_one_done(aws_send_message, sns_worker, monkeypatch):
monkeypatch.setenv('WorkTopic', 'abc123')
sns_worker.side_effect = [True]
progress = RedisProgress()
jobid = create_job(parts, progress=progress)
with Part(jobid, 0, progress=progress):
pass
assert progress.status(jobid)['remaining'] == 2
# 3 map messages, No reduce message sent
sns_worker.assert_called_once()
assert len(sns_worker.call_args[0][0]) == 3
assert sns_worker.call_args[1].get('subject') == 'map'
aws_send_message.assert_not_called()
def test_Part_job_done_on_reduce_callback(monkeypatch):
monkeypatch.setenv('WorkTopic', 'abc123')
monkeypatch.setenv('ProgressTable', 'arn::table/foo')
class CustomProgress(MockProgress):
def complete_part(self, jobid, partid):
return True
on_reduce = Mock()
with Part(jobid='123',
partid=1,
progress=CustomProgress(),
on_reduce=on_reduce):
pass
on_reduce.assert_called_once()
assert on_reduce.call_args[1].get('subject') == 'reduce'
def test_Part_already_failed(aws_send_message, monkeypatch):
monkeypatch.setenv('WorkTopic', 'abc123')
monkeypatch.setenv('ProgressTable', 'arn::table/foo')
class CustomProgress(MockProgress):
def status(self, x):
return {'progress': 0.3, 'failed': True}
class CustomException(Exception):
pass
msg = {'source': 'a.tif', 'partid': 1, 'jobid': '123'}
with pytest.raises(JobFailed) as e:
with Part(fail_job_on=[CustomException],
progress=CustomProgress(),
**msg):
# sees that the overall job failed and will not execute this code
raise NotImplementedError("You'll never get here")
assert 'already failed' in str(e)
aws_send_message.assert_not_called()
def doit(msg, jobid=jobid, progress=progress):
partid = msg['partid']
with Part(partid=partid, jobid=jobid, progress=progress):
pass
from watchbot_progress import Part, create_job
parts = [{'source': 'a.tif'}, {'source': 'b.tif'}, {'source': 'c.tif'}]
def annotate_parts(parts):
return [
dict(**part, partid=partid, jobid=jobid)
for partid, part in enumerate(parts)
]
if __name__ == "__main__":
# ------------- Subject: start-job
jobid = create_job(parts)
# ------------- Subject: map
for msg in annotate_parts(parts):
with Part(msg):
partid = msg['partid']
print(f'PROCCESS {partid}')
print('PROCCESS complete')
print("complete")
def test_Part_bad_progress(aws_send_message):
with pytest.raises(ProgressTypeError):
with Part(partid=1, jobid='1', progress=Exception()):
pass
aws_send_message.assert_not_called()