async def test_start_monitoring_for_existing_job(
self, real_sagemaker_job_monitor: SageMakerJobMonitor):
job_name = "pytorch-rnn-2019-01-04-11-20-03" # Job we have run in our AWS account
job = real_sagemaker_job_monitor.create_job(job_name=job_name)
monitoring_task = real_sagemaker_job_monitor.start(job)
await monitoring_task
assert job.status == JobStatus.FINISHED
async def test_start_monitor_for_finished_job(
self, sagemaker_job_monitor: SageMakerJobMonitor):
job_name = "foobar"
sagemaker_job_monitor.sagemaker_helper.get_job_status.return_value = JobStatus.FINISHED
job = sagemaker_job_monitor.create_job(job_name=job_name,
poll_interval=0.5)
monitoring_task = sagemaker_job_monitor.start(job)
await asyncio.wait_for(monitoring_task, timeout=1) # Should finish
sagemaker_job_monitor.sagemaker_helper.wait_for_job_finish.assert_called_once(
)
sagemaker_job_monitor.notify_finish.assert_called_with(job)
def test_create_running_job(self,
sagemaker_job_monitor: SageMakerJobMonitor):
job_name = "foobar"
sagemaker_job_monitor.sagemaker_helper.get_job_status.return_value = JobStatus.RUNNING
job = sagemaker_job_monitor.create_job(job_name=job_name)
assert job.name == job_name
assert job.status == JobStatus.RUNNING
def _build_api(service, cloud_client): # pylint: disable=too-many-locals
# Cyclic import check is cancelled here; pylint complains about it even though it's only imported in the daemon
# process, which is a clean one and therefore there are no cyclic imports.
from meeshkan.core.api import Api # pylint: disable=cyclic-import
from meeshkan.notifications.notifiers import CloudNotifier, LoggingNotifier, NotifierCollection
from meeshkan.core.tasks import TaskPoller
from meeshkan.core.scheduler import Scheduler, QueueProcessor
from meeshkan.core.config import ensure_base_dirs as ensure_base_dirs_
from meeshkan.core.logger import setup_logging as setup_logging_
from meeshkan.core.sagemaker_monitor import SageMakerJobMonitor
ensure_base_dirs_()
setup_logging_(silent=True)
cloud_notifier = CloudNotifier(
name="Cloud Service",
post_payload=cloud_client.post_payload,
upload_file=cloud_client.post_payload_with_file)
logging_notifier = LoggingNotifier(name="Local Service")
task_poller = TaskPoller(cloud_client.pop_tasks)
queue_processor = QueueProcessor()
notifier_collection = NotifierCollection(
*[cloud_notifier, logging_notifier])
scheduler = Scheduler(queue_processor=queue_processor,
notifier=notifier_collection)
sagemaker_job_monitor = SageMakerJobMonitor(
notify_start=notifier_collection.notify_job_start,
notify_update=notifier_collection.notify,
notify_finish=notifier_collection.notify_job_end)
api = Api(scheduler=scheduler,
service=service,
task_poller=task_poller,
notifier=notifier_collection,
sagemaker_job_monitor=sagemaker_job_monitor)
api.add_stop_callback(cloud_client.close)
return api
def test_start_monitoring_for_non_existing_job(
self, real_sagemaker_job_monitor: SageMakerJobMonitor):
job_name = "foobar"
with pytest.raises(exceptions.JobNotFoundException):
real_sagemaker_job_monitor.create_job(job_name=job_name)
def real_sagemaker_job_monitor(event_loop):
return SageMakerJobMonitor(event_loop=event_loop)
def sagemaker_job_monitor(event_loop, mock_sagemaker_helper):
return SageMakerJobMonitor(event_loop=event_loop,
sagemaker_helper=mock_sagemaker_helper,
notify_finish=MagicMock())