def test_cleanup_removes_correct_files(self): job_bundler = JobBundler('fake_name', {}, None, None) job_bundler.cleanup() remove_archive = call(self.temp_directory + '/fake_name.tgz') remove_bin = call('fake_name.bin') remove_config = call(self.temp_directory + '/fake_name.config.yaml') self.mock_os_remove.assert_has_calls([remove_archive, remove_bin, remove_config])
class BucketJobDeployment(object): def __init__(self, job_name, job, job_source_bundle, code_bucket, results_bucket): from foundations_contrib.global_state import config_manager from foundations_contrib.job_bundler import JobBundler self._config = {} self._config.update(config_manager.config()) self._config['_is_deployment'] = True self._code_bucket = code_bucket self._result_bucket = results_bucket self._job_name = job_name self._job = job self._job_bundler = JobBundler( self._job_name, self._config, self._job, job_source_bundle) self._job_results = None def config(self): return self._config def job_name(self): return self._job_name def deploy(self): self._job_bundler.bundle() try: self._run() finally: self._job_bundler.cleanup() def is_job_complete(self): return self._result_bucket.exists(self._job_archive_name()) def fetch_job_results(self): import os import tarfile from foundations_internal.serializer import deserialize_from_file from foundations_contrib.simple_tempfile import SimpleTempfile if self._job_results is None: with SimpleTempfile('w+b') as temp_file: self._result_bucket.download_to_file( self._job_archive_name(), temp_file) with tarfile.open(temp_file.name, "r:gz") as tar: for tarinfo in tar: if os.path.splitext(tarinfo.name)[1] == ".pkl": file = tar.extractfile(tarinfo) self._job_results = deserialize_from_file(file) file.close() return self._job_results def _job_archive_name(self): return self._job_bundler.job_archive_name() def upload_to_result_bucket(self): self._bucket_upload_from_file(self._result_bucket) def _run(self): self._bucket_upload_from_file(self._code_bucket) def _bucket_upload_from_file(self, bucket): with open(self._job_archive(), 'rb') as file: bucket.upload_from_file(self._job_archive_name(), file) def _job_archive_name(self): return self._job_bundler.job_archive_name() def _job_archive(self): return self._job_bundler.job_archive()