def __init__(self, commit_file, bundle_service, worker_dir, max_cache_size_bytes): super(LocalFileSystemDependencyManager, self).__init__() self.add_transition(DependencyStage.DOWNLOADING, self._transition_from_DOWNLOADING) self.add_terminal(DependencyStage.READY) self.add_terminal(DependencyStage.FAILED) self._state_committer = JsonStateCommitter(commit_file) self._bundle_service = bundle_service self._max_cache_size_bytes = max_cache_size_bytes self.dependencies_dir = os.path.join( worker_dir, LocalFileSystemDependencyManager.DEPENDENCIES_DIR_NAME ) if not os.path.exists(self.dependencies_dir): logger.info('{} doesn\'t exist, creating.'.format(self.dependencies_dir)) os.makedirs(self.dependencies_dir, 0o770) # Locks for concurrency self._dependency_locks = dict() # (parent_uuid, parent_path) -> threading.RLock self._global_lock = threading.RLock() # Used for add/remove actions self._paths_lock = threading.RLock() # Used for path name computations # File paths that are currently being used to store dependencies. Used to prevent conflicts self._paths = set() # (parent_uuid, parent_path) -> DependencyState self._dependencies = dict() # (parent_uuid, parent_path) -> WorkerThread(thread, success, failure_message) self._downloading = ThreadDict(fields={'success': False, 'failure_message': None}) self._load_state() self._stop = False self._main_thread = None
def __init__( self, docker_image_manager, dependency_manager, worker_docker_network, docker_network_internal, docker_network_external, docker_runtime, upload_bundle_callback, assign_cpu_and_gpu_sets_fn, ): super(LocalRunStateMachine, self).__init__() self.add_transition(LocalRunStage.PREPARING, self._transition_from_PREPARING) self.add_transition(LocalRunStage.RUNNING, self._transition_from_RUNNING) self.add_transition(LocalRunStage.CLEANING_UP, self._transition_from_CLEANING_UP) self.add_transition( LocalRunStage.UPLOADING_RESULTS, self._transition_from_UPLOADING_RESULTS ) self.add_transition(LocalRunStage.FINALIZING, self._transition_from_FINALIZING) self.add_terminal(LocalRunStage.FINISHED) self.dependency_manager = dependency_manager self.docker_image_manager = docker_image_manager self.worker_docker_network = worker_docker_network self.docker_network_external = docker_network_external self.docker_network_internal = docker_network_internal self.docker_runtime = docker_runtime # bundle_uuid -> {'thread': Thread, 'run_status': str} self.uploading = ThreadDict(fields={'run_status': 'Upload started', 'success': False}) # bundle_uuid -> {'thread': Thread, 'disk_utilization': int, 'running': bool} self.disk_utilization = ThreadDict( fields={'disk_utilization': 0, 'running': True, 'lock': None} ) self.upload_bundle_callback = upload_bundle_callback self.assign_cpu_and_gpu_sets_fn = assign_cpu_and_gpu_sets_fn
def __init__(self, commit_file, bundle_service, worker_dir, max_cache_size_bytes): super(LocalFileSystemDependencyManager, self).__init__() self.add_transition(DependencyStage.DOWNLOADING, self._transition_from_DOWNLOADING) self.add_terminal(DependencyStage.READY) self.add_terminal(DependencyStage.FAILED) self._state_committer = JsonStateCommitter(commit_file) self._bundle_service = bundle_service self._max_cache_size_bytes = max_cache_size_bytes self.dependencies_dir = os.path.join( worker_dir, LocalFileSystemDependencyManager.DEPENDENCIES_DIR_NAME ) if not os.path.exists(self.dependencies_dir): logger.info('{} doesn\'t exist, creating.'.format(self.dependencies_dir)) os.makedirs(self.dependencies_dir, 0o770) # Locks for concurrency self._dependency_locks = dict() # (parent_uuid, parent_path) -> threading.RLock self._global_lock = threading.RLock() # Used for add/remove actions self._paths_lock = threading.RLock() # Used for path name computations # File paths that are currently being used to store dependencies. Used to prevent conflicts self._paths = set() # (parent_uuid, parent_path) -> DependencyState self._dependencies = dict() # (parent_uuid, parent_path) -> WorkerThread(thread, success, failure_message) self._downloading = ThreadDict(fields={'success': False, 'failure_message': None}) self._load_state() self._stop = False self._main_thread = None
def __init__( self, docker_image_manager, dependency_manager, worker_docker_network, docker_network_internal, docker_network_external, docker_runtime, upload_bundle_callback, assign_cpu_and_gpu_sets_fn, ): super(LocalRunStateMachine, self).__init__() self.add_transition(LocalRunStage.PREPARING, self._transition_from_PREPARING) self.add_transition(LocalRunStage.RUNNING, self._transition_from_RUNNING) self.add_transition(LocalRunStage.CLEANING_UP, self._transition_from_CLEANING_UP) self.add_transition(LocalRunStage.UPLOADING_RESULTS, self._transition_from_UPLOADING_RESULTS) self.add_transition(LocalRunStage.FINALIZING, self._transition_from_FINALIZING) self.add_terminal(LocalRunStage.FINISHED) self.dependency_manager = dependency_manager self.docker_image_manager = docker_image_manager self.worker_docker_network = worker_docker_network self.docker_network_external = docker_network_external self.docker_network_internal = docker_network_internal self.docker_runtime = docker_runtime # bundle_uuid -> {'thread': Thread, 'run_status': str} self.uploading = ThreadDict(fields={ 'run_status': 'Upload started', 'success': False }) # bundle_uuid -> {'thread': Thread, 'disk_utilization': int, 'running': bool} self.disk_utilization = ThreadDict(fields={ 'disk_utilization': 0, 'running': True, 'lock': None }) self.upload_bundle_callback = upload_bundle_callback self.assign_cpu_and_gpu_sets_fn = assign_cpu_and_gpu_sets_fn
def __init__(self, commit_file, max_image_cache_size): """ Initializes a DockerImageManager :param commit_file: String path to where the state file should be committed :param max_image_cache_size: Total size in bytes that the image cache can use """ self._state_committer = JsonStateCommitter(commit_file) # type: JsonStateCommitter self._docker = docker.from_env() # type: DockerClient self._image_cache = {} # type: Dict[str, ImageCacheEntry] self._downloading = ThreadDict( fields={'success': False, 'status': 'Download starting.'}, lock=True ) self._max_image_cache_size = max_image_cache_size self._lock = threading.RLock() self._stop = False self._sleep_secs = 10 self._cleanup_thread = None self._load_state()
def __init__(self, commit_file, max_image_cache_size): """ Initializes a DockerImageManager :param commit_file: String path to where the state file should be committed :param max_image_cache_size: Total size in bytes that the image cache can use """ self._state_committer = JsonStateCommitter(commit_file) # type: JsonStateCommitter self._docker = docker.from_env() # type: DockerClient self._image_cache = {} # type: Dict[str, ImageCacheEntry] self._downloading = ThreadDict( fields={'success': False, 'status': 'Download starting.'}, lock=True ) self._max_image_cache_size = max_image_cache_size self._lock = threading.RLock() self._stop = False self._sleep_secs = 10 self._cleanup_thread = None self._load_state()
class LocalFileSystemDependencyManager(StateTransitioner, BaseDependencyManager): """ This dependency manager downloads dependency bundles from Codalab server to the local filesystem. It caches all downloaded dependencies but cleans up the old ones if the disk use hits the given threshold For this class dependencies are uniquely identified by (parent_uuid, parent_path) """ DEPENDENCIES_DIR_NAME = 'dependencies' DEPENDENCY_FAILURE_COOLDOWN = 10 # TODO(bkgoksel): The server writes these to the worker_dependencies table, which stores the dependencies # json as a SqlAlchemy LargeBinary, which defaults to MySQL BLOB, which has a size limit of # 65K. For now we limit this value to about 58K to avoid any issues but we probably want to do # something better (either specify MEDIUMBLOB in the SqlAlchemy definition of the table or change # the data format of how we store this) MAX_SERIALIZED_LEN = 60000 def __init__(self, commit_file, bundle_service, worker_dir, max_cache_size_bytes): super(LocalFileSystemDependencyManager, self).__init__() self.add_transition(DependencyStage.DOWNLOADING, self._transition_from_DOWNLOADING) self.add_terminal(DependencyStage.READY) self.add_terminal(DependencyStage.FAILED) self._state_committer = JsonStateCommitter(commit_file) self._bundle_service = bundle_service self._max_cache_size_bytes = max_cache_size_bytes self.dependencies_dir = os.path.join( worker_dir, LocalFileSystemDependencyManager.DEPENDENCIES_DIR_NAME ) if not os.path.exists(self.dependencies_dir): logger.info('{} doesn\'t exist, creating.'.format(self.dependencies_dir)) os.makedirs(self.dependencies_dir, 0o770) # Locks for concurrency self._dependency_locks = dict() # (parent_uuid, parent_path) -> threading.RLock self._global_lock = threading.RLock() # Used for add/remove actions self._paths_lock = threading.RLock() # Used for path name computations # File paths that are currently being used to store dependencies. Used to prevent conflicts self._paths = set() # (parent_uuid, parent_path) -> DependencyState self._dependencies = dict() # (parent_uuid, parent_path) -> WorkerThread(thread, success, failure_message) self._downloading = ThreadDict(fields={'success': False, 'failure_message': None}) self._load_state() self._stop = False self._main_thread = None def _save_state(self): with self._global_lock, self._paths_lock: self._state_committer.commit({'dependencies': self._dependencies, 'paths': self._paths}) def _load_state(self): state = self._state_committer.load(default={'dependencies': {}, 'paths': set()}) dependencies = {} dependency_locks = {} paths = set() for dep, dep_state in state['dependencies'].items(): full_path = os.path.join(self.dependencies_dir, dep_state.path) if os.path.exists(full_path): dependencies[dep] = dep_state dependency_locks[dep] = threading.RLock() else: logger.info( "Dependency {} in loaded state but its path {} doesn't exist in the filesystem".format( dep, full_path ) ) if dep_state.path not in state['paths']: state['paths'].add(dep_state.path) logger.info( "Dependency {} in loaded state but its path {} is not in the loaded paths {}".format( dep, dep_state.path, state['paths'] ) ) for path in state['paths']: full_path = os.path.join(self.dependencies_dir, path) if os.path.exists(full_path): paths.add(path) else: logger.info( "Path {} in loaded state but doesn't exist in the filesystem".format(full_path) ) with self._global_lock, self._paths_lock: self._dependencies = dependencies self._dependency_locks = dependency_locks self._paths = paths logger.info( '{} dependencies, {} paths in cache.'.format(len(self._dependencies), len(self._paths)) ) def start(self): def loop(self): while not self._stop: try: self._process_dependencies() self._save_state() self._cleanup() self._save_state() except Exception: traceback.print_exc() self._main_thread = threading.Thread(target=loop, args=[self]) self._main_thread.start() def stop(self): logger.info("Stopping local dependency manager") self._stop = True self._downloading.stop() self._main_thread.join() logger.info("Stopped local dependency manager. Exiting.") def _process_dependencies(self): for entry, state in self._dependencies.items(): with self._dependency_locks[entry]: self._dependencies[entry] = self.transition(state) def _prune_failed_dependencies(self): """ Prune failed dependencies older than 10 seconds so that further runs get to retry the download. Without pruning any future run depending on a failed dependency would automatically fail indefinitely. """ with self._global_lock: self._acquire_all_locks() failed_deps = { dep: state for dep, state in self._dependencies.items() if state.stage == DependencyStage.FAILED and time.time() - state.last_used > LocalFileSystemDependencyManager.DEPENDENCY_FAILURE_COOLDOWN } for dependency, dependency_state in failed_deps.items(): self._delete_dependency(dependency) self._release_all_locks() def _cleanup(self): """ Prune failed dependencies older than 10 seconds Limit the disk usage of the dependencies (both the bundle files and the serialized state file size) Deletes oldest failed dependencies first and then oldest finished dependencies. Doesn't touch downloading dependencies. """ self._prune_failed_dependencies() # With all the locks (should be fast if no cleanup needed, otherwise make sure nothing is corrupted while True: with self._global_lock: self._acquire_all_locks() bytes_used = sum(dep.size_bytes for dep in self._dependencies.values()) serialized_length = len(codalabworker.pyjson.dumps(self._dependencies)) if ( bytes_used > self._max_cache_size_bytes or serialized_length > LocalFileSystemDependencyManager.MAX_SERIALIZED_LEN ): logger.debug( '%d dependencies in cache, disk usage: %s (max %s), serialized size: %s (max %s)', len(self._dependencies), size_str(bytes_used), size_str(self._max_cache_size_bytes), size_str(serialized_length), LocalFileSystemDependencyManager.MAX_SERIALIZED_LEN, ) ready_deps = { dep: state for dep, state in self._dependencies.items() if state.stage == DependencyStage.READY and not state.dependents } failed_deps = { dep: state for dep, state in self._dependencies.items() if state.stage == DependencyStage.FAILED } if failed_deps: dep_to_remove = min( failed_deps.iteritems(), key=lambda dep_state: dep_state[1].last_used )[0] elif ready_deps: dep_to_remove = min( ready_deps.iteritems(), key=lambda dep_state: dep_state[1].last_used )[0] else: logger.info( 'Dependency quota full but there are only downloading dependencies, not cleaning up until downloads are over' ) break if dep_to_remove: self._delete_dependency(dep_to_remove) self._release_all_locks() else: self._release_all_locks() break def _delete_dependency(self, dependency): """ Remove the given dependency from the manager's state Also delete any known files on the filesystem if any exist """ if self._acquire_if_exists(dependency): try: path_to_remove = self._dependencies[dependency].path self._paths.remove(path_to_remove) remove_path(path_to_remove) except Exception: pass finally: del self._dependencies[dependency] self._dependency_locks[dependency].release() def has(self, dependency): """ Takes a dependency = (parent_uuid, parent_path) Returns true if the manager has processed this dependency """ with self._global_lock: return dependency in self._dependencies def get(self, uuid, dependency): """ Request the dependency for the run with uuid, registering uuid as a dependent of this dependency """ now = time.time() if not self._acquire_if_exists(dependency): # add dependency state if it does not exist with self._global_lock: self._dependency_locks[dependency] = threading.RLock() self._dependency_locks[dependency].acquire() self._dependencies[dependency] = DependencyState( stage=DependencyStage.DOWNLOADING, dependency=dependency, path=self._assign_path(dependency), size_bytes=0, dependents=set([uuid]), last_used=now, message="Starting download", killed=False, ) # update last_used as long as it isn't in FAILED if self._dependencies[dependency].stage != DependencyStage.FAILED: self._dependencies[dependency].dependents.add(uuid) self._dependencies[dependency] = self._dependencies[dependency]._replace(last_used=now) self._dependency_locks[dependency].release() return self._dependencies[dependency] def release(self, uuid, dependency): """ Register that the run with uuid is no longer dependent on this dependency If no more runs are dependent on this dependency, kill it """ if self._acquire_if_exists(dependency): dep_state = self._dependencies[dependency] if uuid in dep_state.dependents: dep_state.dependents.remove(uuid) if not dep_state.dependents: dep_state = dep_state._replace(killed=True) self._dependencies[dependency] = dep_state self._dependency_locks[dependency].release() def _acquire_if_exists(self, dependency): """ Safely acquires a lock for the given dependency if it exists Returns True if depedendency exists, False otherwise Callers should remember to release the lock """ with self._global_lock: if dependency in self._dependencies: self._dependency_locks[dependency].acquire() return True else: return False def _acquire_all_locks(self): """ Acquires all dependency locks in the thread it's called from """ with self._global_lock: for dependency, lock in self._dependency_locks.iteritems(): lock.acquire() def _release_all_locks(self): """ Releases all dependency locks in the thread it's called from """ with self._global_lock: for dependency, lock in self._dependency_locks.iteritems(): lock.release() def _assign_path(self, dependency): """ Normalize the path for the dependency by replacing / with _, avoiding conflicts """ parent_uuid, parent_path = dependency if parent_path: path = os.path.join(parent_uuid, parent_path) else: path = parent_uuid path = path.replace(os.path.sep, '_') # You could have a conflict between, for example a/b_c and # a_b/c. We have to avoid those. with self._paths_lock: while path in self._paths: path = path + '_' self._paths.add(path) return path def _store_dependency(self, dependency_path, fileobj, target_type): """ Copy the dependency fileobj to its path in the local filesystem Overwrite existing files by the same name if found (may happen if filesystem modified outside the dependency manager, for example during an update if the state gets reset but filesystem doesn't get cleared) """ try: if os.path.exists(dependency_path): logger.info('Path %s already exists, overwriting', dependency_path) if os.path.isdir(dependency_path): shutil.rmtree(dependency_path) else: os.remove(dependency_path) if target_type == 'directory': un_tar_directory(fileobj, dependency_path, 'gz') else: with open(dependency_path, 'wb') as f: logger.debug('copying file to %s', dependency_path) shutil.copyfileobj(fileobj, f) except Exception: raise @property def all_dependencies(self): with self._global_lock: return list(self._dependencies.keys()) def _transition_from_DOWNLOADING(self, dependency_state): def download(): def update_state_and_check_killed(bytes_downloaded): """ Callback method for bundle service client updates dependency state and raises DownloadAbortedException if download is killed by dep. manager """ with self._dependency_locks[dependency]: state = self._dependencies[dependency] if state.killed: raise DownloadAbortedException("Aborted by user") self._dependencies[dependency] = state._replace( size_bytes=bytes_downloaded, message="Downloading dependency: %s downloaded" % size_str(bytes_downloaded), ) dependency_path = os.path.join(self.dependencies_dir, dependency_state.path) logger.debug('Downloading dependency %s/%s', parent_uuid, parent_path) try: # Start async download to the fileobj fileobj, target_type = self._bundle_service.get_bundle_contents( parent_uuid, parent_path ) with closing(fileobj): # "Bug" the fileobj's read function so that we can keep # track of the number of bytes downloaded so far. old_read_method = fileobj.read bytes_downloaded = [0] def interruptable_read(*args, **kwargs): data = old_read_method(*args, **kwargs) bytes_downloaded[0] += len(data) update_state_and_check_killed(bytes_downloaded[0]) return data fileobj.read = interruptable_read # Start copying the fileobj to filesystem dependency path self._store_dependency(dependency_path, fileobj, target_type) logger.debug( 'Finished downloading %s dependency %s/%s to %s', target_type, parent_uuid, parent_path, dependency_path, ) with self._dependency_locks[dependency]: self._downloading[dependency]['success'] = True except Exception as e: with self._dependency_locks[dependency]: self._downloading[dependency]['success'] = False self._downloading[dependency][ 'failure_message' ] = "Dependency download failed: %s " % str(e) dependency = dependency_state.dependency parent_uuid, parent_path = dependency self._downloading.add_if_new(dependency, threading.Thread(target=download, args=[])) if self._downloading[dependency].is_alive(): return dependency_state success = self._downloading[dependency]['success'] failure_message = self._downloading[dependency]['failure_message'] self._downloading.remove(dependency) if success: return dependency_state._replace( stage=DependencyStage.READY, message="Download complete" ) else: with self._paths_lock: self._paths.remove(dependency_state.path) return dependency_state._replace(stage=DependencyStage.FAILED, message=failure_message)
class LocalFileSystemDependencyManager(StateTransitioner, BaseDependencyManager): """ This dependency manager downloads dependency bundles from Codalab server to the local filesystem. It caches all downloaded dependencies but cleans up the old ones if the disk use hits the given threshold For this class dependencies are uniquely identified by (parent_uuid, parent_path) """ DEPENDENCIES_DIR_NAME = 'dependencies' DEPENDENCY_FAILURE_COOLDOWN = 10 # TODO(bkgoksel): The server writes these to the worker_dependencies table, which stores the dependencies # json as a SqlAlchemy LargeBinary, which defaults to MySQL BLOB, which has a size limit of # 65K. For now we limit this value to about 58K to avoid any issues but we probably want to do # something better (either specify MEDIUMBLOB in the SqlAlchemy definition of the table or change # the data format of how we store this) MAX_SERIALIZED_LEN = 60000 def __init__(self, commit_file, bundle_service, worker_dir, max_cache_size_bytes): super(LocalFileSystemDependencyManager, self).__init__() self.add_transition(DependencyStage.DOWNLOADING, self._transition_from_DOWNLOADING) self.add_terminal(DependencyStage.READY) self.add_terminal(DependencyStage.FAILED) self._state_committer = JsonStateCommitter(commit_file) self._bundle_service = bundle_service self._max_cache_size_bytes = max_cache_size_bytes self.dependencies_dir = os.path.join( worker_dir, LocalFileSystemDependencyManager.DEPENDENCIES_DIR_NAME ) if not os.path.exists(self.dependencies_dir): logger.info('{} doesn\'t exist, creating.'.format(self.dependencies_dir)) os.makedirs(self.dependencies_dir, 0o770) # Locks for concurrency self._dependency_locks = dict() # (parent_uuid, parent_path) -> threading.RLock self._global_lock = threading.RLock() # Used for add/remove actions self._paths_lock = threading.RLock() # Used for path name computations # File paths that are currently being used to store dependencies. Used to prevent conflicts self._paths = set() # (parent_uuid, parent_path) -> DependencyState self._dependencies = dict() # (parent_uuid, parent_path) -> WorkerThread(thread, success, failure_message) self._downloading = ThreadDict(fields={'success': False, 'failure_message': None}) self._load_state() self._stop = False self._main_thread = None def _save_state(self): with self._global_lock, self._paths_lock: self._state_committer.commit({'dependencies': self._dependencies, 'paths': self._paths}) def _load_state(self): state = self._state_committer.load(default={'dependencies': {}, 'paths': set()}) dependencies = {} dependency_locks = {} paths = set() for dep, dep_state in state['dependencies'].items(): full_path = os.path.join(self.dependencies_dir, dep_state.path) if os.path.exists(full_path): dependencies[dep] = dep_state dependency_locks[dep] = threading.RLock() else: logger.info( "Dependency {} in loaded state but its path {} doesn't exist in the filesystem".format( dep, full_path ) ) if dep_state.path not in state['paths']: state['paths'].add(dep_state.path) logger.info( "Dependency {} in loaded state but its path {} is not in the loaded paths {}".format( dep, dep_state.path, state['paths'] ) ) for path in state['paths']: full_path = os.path.join(self.dependencies_dir, path) if os.path.exists(full_path): paths.add(path) else: logger.info( "Path {} in loaded state but doesn't exist in the filesystem".format(full_path) ) with self._global_lock, self._paths_lock: self._dependencies = dependencies self._dependency_locks = dependency_locks self._paths = paths logger.info( '{} dependencies, {} paths in cache.'.format(len(self._dependencies), len(self._paths)) ) def start(self): logger.info('Starting local dependency manager') def loop(self): while not self._stop: try: self._process_dependencies() self._save_state() self._cleanup() self._save_state() except Exception: traceback.print_exc() time.sleep(1) self._main_thread = threading.Thread(target=loop, args=[self]) self._main_thread.start() def stop(self): logger.info('Stopping local dependency manager') self._stop = True self._downloading.stop() self._main_thread.join() logger.info('Stopped local dependency manager') def _process_dependencies(self): for entry, state in self._dependencies.items(): with self._dependency_locks[entry]: self._dependencies[entry] = self.transition(state) def _prune_failed_dependencies(self): """ Prune failed dependencies older than DEPENDENCY_FAILURE_COOLDOWN seconds so that further runs get to retry the download. Without pruning, any future run depending on a failed dependency would automatically fail indefinitely. """ with self._global_lock: self._acquire_all_locks() failed_deps = { dep: state for dep, state in self._dependencies.items() if state.stage == DependencyStage.FAILED and time.time() - state.last_used > LocalFileSystemDependencyManager.DEPENDENCY_FAILURE_COOLDOWN } for dependency, dependency_state in failed_deps.items(): self._delete_dependency(dependency) self._release_all_locks() def _cleanup(self): """ Prune failed dependencies older than DEPENDENCY_FAILURE_COOLDOWN seconds. Limit the disk usage of the dependencies (both the bundle files and the serialized state file size) Deletes oldest failed dependencies first and then oldest finished dependencies. Doesn't touch downloading dependencies. """ self._prune_failed_dependencies() # With all the locks (should be fast if no cleanup needed, otherwise make sure nothing is corrupted while True: with self._global_lock: self._acquire_all_locks() bytes_used = sum(dep.size_bytes for dep in self._dependencies.values()) serialized_length = len(codalabworker.pyjson.dumps(self._dependencies)) if ( bytes_used > self._max_cache_size_bytes or serialized_length > LocalFileSystemDependencyManager.MAX_SERIALIZED_LEN ): logger.debug( '%d dependencies in cache, disk usage: %s (max %s), serialized size: %s (max %s)', len(self._dependencies), size_str(bytes_used), size_str(self._max_cache_size_bytes), size_str(serialized_length), LocalFileSystemDependencyManager.MAX_SERIALIZED_LEN, ) ready_deps = { dep: state for dep, state in self._dependencies.items() if state.stage == DependencyStage.READY and not state.dependents } failed_deps = { dep: state for dep, state in self._dependencies.items() if state.stage == DependencyStage.FAILED } if failed_deps: dep_to_remove = min( failed_deps.iteritems(), key=lambda dep_state: dep_state[1].last_used )[0] elif ready_deps: dep_to_remove = min( ready_deps.iteritems(), key=lambda dep_state: dep_state[1].last_used )[0] else: logger.info( 'Dependency quota full but there are only downloading dependencies, not cleaning up until downloads are over' ) break if dep_to_remove: self._delete_dependency(dep_to_remove) self._release_all_locks() else: self._release_all_locks() break def _delete_dependency(self, dependency): """ Remove the given dependency from the manager's state Also delete any known files on the filesystem if any exist """ if self._acquire_if_exists(dependency): try: path_to_remove = self._dependencies[dependency].path self._paths.remove(path_to_remove) remove_path(path_to_remove) except Exception: pass finally: del self._dependencies[dependency] self._dependency_locks[dependency].release() def has(self, dependency): """ Takes a dependency = (parent_uuid, parent_path) Returns true if the manager has processed this dependency """ with self._global_lock: return dependency in self._dependencies def get(self, uuid, dependency): """ Request the dependency for the run with uuid, registering uuid as a dependent of this dependency """ now = time.time() if not self._acquire_if_exists(dependency): # add dependency state if it does not exist with self._global_lock: self._dependency_locks[dependency] = threading.RLock() self._dependency_locks[dependency].acquire() self._dependencies[dependency] = DependencyState( stage=DependencyStage.DOWNLOADING, dependency=dependency, path=self._assign_path(dependency), size_bytes=0, dependents=set([uuid]), last_used=now, message="Starting download", killed=False, ) # update last_used as long as it isn't in FAILED if self._dependencies[dependency].stage != DependencyStage.FAILED: self._dependencies[dependency].dependents.add(uuid) self._dependencies[dependency] = self._dependencies[dependency]._replace(last_used=now) self._dependency_locks[dependency].release() return self._dependencies[dependency] def release(self, uuid, dependency): """ Register that the run with uuid is no longer dependent on this dependency If no more runs are dependent on this dependency, kill it """ if self._acquire_if_exists(dependency): dep_state = self._dependencies[dependency] if uuid in dep_state.dependents: dep_state.dependents.remove(uuid) if not dep_state.dependents: dep_state = dep_state._replace(killed=True) self._dependencies[dependency] = dep_state self._dependency_locks[dependency].release() def _acquire_if_exists(self, dependency): """ Safely acquires a lock for the given dependency if it exists Returns True if depedendency exists, False otherwise Callers should remember to release the lock """ with self._global_lock: if dependency in self._dependencies: self._dependency_locks[dependency].acquire() return True else: return False def _acquire_all_locks(self): """ Acquires all dependency locks in the thread it's called from """ with self._global_lock: for dependency, lock in self._dependency_locks.iteritems(): lock.acquire() def _release_all_locks(self): """ Releases all dependency locks in the thread it's called from """ with self._global_lock: for dependency, lock in self._dependency_locks.iteritems(): lock.release() def _assign_path(self, dependency): """ Normalize the path for the dependency by replacing / with _, avoiding conflicts """ parent_uuid, parent_path = dependency if parent_path: path = os.path.join(parent_uuid, parent_path) else: path = parent_uuid path = path.replace(os.path.sep, '_') # You could have a conflict between, for example a/b_c and # a_b/c. We have to avoid those. with self._paths_lock: while path in self._paths: path = path + '_' self._paths.add(path) return path def _store_dependency(self, dependency_path, fileobj, target_type): """ Copy the dependency fileobj to its path in the local filesystem Overwrite existing files by the same name if found (may happen if filesystem modified outside the dependency manager, for example during an update if the state gets reset but filesystem doesn't get cleared) """ try: if os.path.exists(dependency_path): logger.info('Path %s already exists, overwriting', dependency_path) if os.path.isdir(dependency_path): shutil.rmtree(dependency_path) else: os.remove(dependency_path) if target_type == 'directory': un_tar_directory(fileobj, dependency_path, 'gz') else: with open(dependency_path, 'wb') as f: logger.debug('copying file to %s', dependency_path) shutil.copyfileobj(fileobj, f) except Exception: raise @property def all_dependencies(self): with self._global_lock: return list(self._dependencies.keys()) def _transition_from_DOWNLOADING(self, dependency_state): def download(): def update_state_and_check_killed(bytes_downloaded): """ Callback method for bundle service client updates dependency state and raises DownloadAbortedException if download is killed by dep. manager """ with self._dependency_locks[dependency]: state = self._dependencies[dependency] if state.killed: raise DownloadAbortedException("Aborted by user") self._dependencies[dependency] = state._replace( size_bytes=bytes_downloaded, message="Downloading dependency: %s downloaded" % size_str(bytes_downloaded), ) dependency_path = os.path.join(self.dependencies_dir, dependency_state.path) logger.debug('Downloading dependency %s/%s', parent_uuid, parent_path) try: # Start async download to the fileobj fileobj, target_type = self._bundle_service.get_bundle_contents( parent_uuid, parent_path ) with closing(fileobj): # "Bug" the fileobj's read function so that we can keep # track of the number of bytes downloaded so far. old_read_method = fileobj.read bytes_downloaded = [0] def interruptable_read(*args, **kwargs): data = old_read_method(*args, **kwargs) bytes_downloaded[0] += len(data) update_state_and_check_killed(bytes_downloaded[0]) return data fileobj.read = interruptable_read # Start copying the fileobj to filesystem dependency path self._store_dependency(dependency_path, fileobj, target_type) logger.debug( 'Finished downloading %s dependency %s/%s to %s', target_type, parent_uuid, parent_path, dependency_path, ) with self._dependency_locks[dependency]: self._downloading[dependency]['success'] = True except Exception as e: with self._dependency_locks[dependency]: self._downloading[dependency]['success'] = False self._downloading[dependency][ 'failure_message' ] = "Dependency download failed: %s " % str(e) dependency = dependency_state.dependency parent_uuid, parent_path = dependency self._downloading.add_if_new(dependency, threading.Thread(target=download, args=[])) if self._downloading[dependency].is_alive(): return dependency_state success = self._downloading[dependency]['success'] failure_message = self._downloading[dependency]['failure_message'] self._downloading.remove(dependency) if success: return dependency_state._replace( stage=DependencyStage.READY, message="Download complete" ) else: with self._paths_lock: self._paths.remove(dependency_state.path) return dependency_state._replace(stage=DependencyStage.FAILED, message=failure_message)
class LocalRunStateMachine(StateTransitioner): """ Manages the state machine of the runs running on the local machine """ def __init__( self, docker_image_manager, dependency_manager, docker_network_internal_name, docker_network_external_name, docker_runtime, upload_bundle_callback, assign_cpu_and_gpu_sets_fn, ): super(LocalRunStateMachine, self).__init__() self.add_transition(LocalRunStage.PREPARING, self._transition_from_PREPARING) self.add_transition(LocalRunStage.RUNNING, self._transition_from_RUNNING) self.add_transition(LocalRunStage.CLEANING_UP, self._transition_from_CLEANING_UP) self.add_transition(LocalRunStage.UPLOADING_RESULTS, self._transition_from_UPLOADING_RESULTS) self.add_transition(LocalRunStage.FINALIZING, self._transition_from_FINALIZING) self.add_terminal(LocalRunStage.FINISHED) self.dependency_manager = dependency_manager self.docker_image_manager = docker_image_manager self.docker_network_external_name = docker_network_external_name self.docker_network_internal_name = docker_network_internal_name self.docker_runtime = docker_runtime # bundle_uuid -> {'thread': Thread, 'run_status': str} self.uploading = ThreadDict(fields={ 'run_status': 'Upload started', 'success': False }) # bundle_uuid -> {'thread': Thread, 'disk_utilization': int, 'running': bool} self.disk_utilization = ThreadDict(fields={ 'disk_utilization': 0, 'running': True, 'lock': None }) self.upload_bundle_callback = upload_bundle_callback self.assign_cpu_and_gpu_sets_fn = assign_cpu_and_gpu_sets_fn def stop(self): for uuid in self.disk_utilization.keys(): self.disk_utilization[uuid]['running'] = False self.disk_utilization.stop() self.uploading.stop() def _transition_from_PREPARING(self, run_state): """ 1- Request the docker image from docker image manager - if image is failed, move to CLEANING_UP state 2- Request the dependencies from dependency manager - if any are failed, move to CLEANING_UP state 3- If all dependencies and docker image are ready: - Set up the local filesystem for the run - Create symlinks to dependencies - Allocate resources and prepare the docker container - Start the docker container 4- If all is successful, move to RUNNING state """ if run_state.is_killed: return run_state._replace(stage=LocalRunStage.CLEANING_UP, container_id=None) dependencies_ready = True status_messages = [] bundle_uuid = run_state.bundle['uuid'] # get dependencies for dep in run_state.bundle['dependencies']: dependency = (dep['parent_uuid'], dep['parent_path']) dependency_state = self.dependency_manager.get( bundle_uuid, dependency) if dependency_state.stage == DependencyStage.DOWNLOADING: status_messages.append( 'Downloading dependency %s: %s done (archived size)' % (dep['child_path'], size_str(dependency_state.size_bytes))) dependencies_ready = False elif dependency_state.stage == DependencyStage.FAILED: # Failed to download dependency; -> CLEANING_UP run_state.info[ 'failure_message'] = 'Failed to download dependency %s: %s' % ( dep['child_path'], '', ) return run_state._replace(stage=LocalRunStage.CLEANING_UP, info=run_state.info) # get the docker image docker_image = run_state.resources['docker_image'] image_state = self.docker_image_manager.get(docker_image) if image_state.stage == DependencyStage.DOWNLOADING: status_messages.append('Pulling docker image: ' + (image_state.message or docker_image or "")) dependencies_ready = False elif image_state.stage == DependencyStage.FAILED: # Failed to pull image; -> CLEANING_UP run_state.info['failure_message'] = image_state.message return run_state._replace(stage=LocalRunStage.CLEANING_UP, info=run_state.info) # stop proceeding if dependency and image downloads aren't all done if not dependencies_ready: status_message = status_messages.pop() if status_messages: status_message += "(and downloading %d other dependencies and docker images)" % len( status_messages) return run_state._replace(run_status=status_message) # All dependencies ready! Set up directories, symlinks and container. Start container. # 1) Set up a directory to store the bundle. remove_path(run_state.bundle_path) os.mkdir(run_state.bundle_path) # 2) Set up symlinks dependencies = [] docker_dependencies_path = '/' + bundle_uuid + '_dependencies' for dep in run_state.bundle['dependencies']: child_path = os.path.normpath( os.path.join(run_state.bundle_path, dep['child_path'])) if not child_path.startswith(run_state.bundle_path): raise Exception('Invalid key for dependency: %s' % (dep['child_path'])) dependency_path = self.dependency_manager.get( bundle_uuid, (dep['parent_uuid'], dep['parent_path'])).path dependency_path = os.path.join( self.dependency_manager.dependencies_dir, dependency_path) docker_dependency_path = os.path.join(docker_dependencies_path, dep['child_path']) os.symlink(docker_dependency_path, child_path) # These are turned into docker volume bindings like: # dependency_path:docker_dependency_path:ro dependencies.append((dependency_path, docker_dependency_path)) # 3) Set up container if run_state.resources['request_network']: docker_network = self.docker_network_external_name else: docker_network = self.docker_network_internal_name try: cpuset, gpuset = self.assign_cpu_and_gpu_sets_fn( run_state.resources['request_cpus'], run_state.resources['request_gpus']) except Exception: run_state.info[ 'failure_message'] = "Cannot assign enough resources" return run_state._replace(stage=LocalRunStage.CLEANING_UP, info=run_state.info) # 4) Start container try: container = docker_utils.start_bundle_container( run_state.bundle_path, bundle_uuid, dependencies, run_state.bundle['command'], run_state.resources['docker_image'], network=docker_network, cpuset=cpuset, gpuset=gpuset, memory_bytes=run_state.resources['request_memory'], runtime=self.docker_runtime, ) except docker_utils.DockerException as e: run_state.info[ 'failure_message'] = 'Cannot start Docker container: {}'.format( e) return run_state._replace(stage=LocalRunStage.CLEANING_UP, info=run_state.info) return run_state._replace( stage=LocalRunStage.RUNNING, start_time=time.time(), run_status='Running job in Docker container', container_id=container.id, container=container, docker_image=image_state.digest, has_contents=True, cpuset=cpuset, gpuset=gpuset, ) def _transition_from_RUNNING(self, run_state): """ 1- Check run status of the docker container 2- If run is killed, kill the container 3- If run is finished, move to CLEANING_UP state """ bundle_uuid = run_state.bundle['uuid'] def check_and_report_finished(run_state): try: finished, exitcode, failure_msg = docker_utils.check_finished( run_state.container) except docker_utils.DockerException: traceback.print_exc() finished, exitcode, failure_msg = False, None, None new_info = dict(finished=finished, exitcode=exitcode, failure_message=failure_msg) run_state.info.update(new_info) run_state = run_state._replace(info=run_state.info) return run_state def check_resource_utilization(run_state): kill_messages = [] run_stats = docker_utils.get_container_stats(run_state.container) time_used = time.time() - run_state.start_time run_state = run_state._replace(time_used=time_used) run_state = run_state._replace(max_memory=max( run_state.max_memory, run_stats.get('memory', 0))) run_state = run_state._replace( disk_utilization=self.disk_utilization[bundle_uuid] ['disk_utilization']) if (run_state.resources['request_time'] and run_state.time_used > run_state.resources['request_time']): kill_messages.append( 'Time limit %s exceeded.' % duration_str(run_state.resources['request_time'])) if run_state.max_memory > run_state.resources['request_memory']: kill_messages.append( 'Memory limit %s exceeded.' % duration_str(run_state.resources['request_memory'])) if (run_state.resources['request_disk'] and run_state.disk_utilization > run_state.resources['request_disk']): kill_messages.append( 'Disk limit %sb exceeded.' % size_str(run_state.resources['request_disk'])) if kill_messages: new_info = run_state.info new_info['kill_message'] = ' '.join(kill_messages) run_state = run_state._replace(info=new_info, is_killed=True) return run_state def check_disk_utilization(): running = True while running: start_time = time.time() try: disk_utilization = get_path_size(run_state.bundle_path) self.disk_utilization[bundle_uuid][ 'disk_utilization'] = disk_utilization running = self.disk_utilization[bundle_uuid]['running'] except Exception: traceback.print_exc() end_time = time.time() # To ensure that we don't hammer the disk for this computation when # there are lots of files, we run it at most 10% of the time. time.sleep(max((end_time - start_time) * 10, 1.0)) self.disk_utilization.add_if_new( bundle_uuid, threading.Thread(target=check_disk_utilization, args=[])) run_state = check_and_report_finished(run_state) run_state = check_resource_utilization(run_state) if run_state.is_killed and run_state.container_id is not None: try: run_state.container.kill() except docker.errors.APIError: finished, _, _ = docker_utils.check_finished( run_state.container) if not finished: # If we can't kill a Running container, something is wrong # Otherwise all well traceback.print_exc() self.disk_utilization[bundle_uuid]['running'] = False self.disk_utilization.remove(bundle_uuid) return run_state._replace(stage=LocalRunStage.CLEANING_UP, container_id=None) if run_state.info['finished']: logger.debug( 'Finished run with UUID %s, exitcode %s, failure_message %s', bundle_uuid, run_state.info['exitcode'], run_state.info['failure_message'], ) self.disk_utilization[bundle_uuid]['running'] = False self.disk_utilization.remove(bundle_uuid) return run_state._replace(stage=LocalRunStage.CLEANING_UP, run_status='Uploading results') else: return run_state def _transition_from_CLEANING_UP(self, run_state): """ 1- delete the container if still existent 2- clean up the dependencies from bundle folder 3- release the dependencies in dependency manager 4- If bundle has contents to upload (i.e. was RUNNING at some point), move to UPLOADING_RESULTS state Otherwise move to FINALIZING state """ bundle_uuid = run_state.bundle['uuid'] if run_state.container_id is not None: while True: try: finished, _, _ = docker_utils.check_finished( run_state.container) if finished: run_state.container.remove(force=True) break except docker.errors.APIError: traceback.print_exc() time.sleep(1) for dep in run_state.bundle['dependencies']: self.dependency_manager.release( bundle_uuid, (dep['parent_uuid'], dep['parent_path'])) child_path = os.path.join(run_state.bundle_path, dep['child_path']) try: remove_path(child_path) except Exception: traceback.print_exc() if run_state.has_contents: return run_state._replace( stage=LocalRunStage.UPLOADING_RESULTS, run_status='Uploading results', container=None, ) else: return self.finalize_run(run_state) def _transition_from_UPLOADING_RESULTS(self, run_state): """ If bundle not already uploading: Use the RunManager API to upload contents at bundle_path to the server Pass the callback to that API such that if the bundle is killed during the upload, the callback returns false, allowing killable uploads. If uploading and not finished: Update run_status with upload progress If uploading and finished: Move to FINALIZING state """ def upload_results(): try: # Upload results logger.debug('Uploading results for run with UUID %s', bundle_uuid) def progress_callback(bytes_uploaded): run_status = 'Uploading results: %s done (archived size)' % size_str( bytes_uploaded) self.uploading[bundle_uuid]['run_status'] = run_status return True self.upload_bundle_callback(bundle_uuid, run_state.bundle_path, progress_callback) self.uploading[bundle_uuid]['success'] = True except Exception as e: self.uploading[bundle_uuid][ 'run_status'] = "Error while uploading: %s" % e traceback.print_exc() bundle_uuid = run_state.bundle['uuid'] self.uploading.add_if_new( bundle_uuid, threading.Thread(target=upload_results, args=[])) if self.uploading[bundle_uuid].is_alive(): return run_state._replace( run_status=self.uploading[bundle_uuid]['run_status']) elif not self.uploading[bundle_uuid]['success']: # upload failed failure_message = run_state.info.get('failure_message', None) if failure_message: run_state.info['failure_message'] = ( failure_message + '. ' + self.uploading[bundle_uuid]['run_status']) else: run_state.info['failure_message'] = self.uploading[ bundle_uuid]['run_status'] self.uploading.remove(bundle_uuid) return self.finalize_run(run_state) def finalize_run(self, run_state): """ Prepare the finalize message to be sent with the next checkin """ failure_message = run_state.info.get('failure_message', None) if 'exitcode' not in run_state.info: run_state.info['exitcode'] = None if not failure_message and run_state.is_killed: run_state.info['failure_message'] = run_state.info['kill_message'] run_state.info['finalized'] = False return run_state._replace(stage=LocalRunStage.FINALIZING, info=run_state.info, run_status="Finalizing bundle") def _transition_from_FINALIZING(self, run_state): """ If a full worker cycle has passed since we got into FINALIZING we already reported to server so can move on to FINISHED. Can also remove bundle_path now """ if run_state.info['finalized']: remove_path(run_state.bundle_path) return run_state._replace(stage=LocalRunStage.FINISHED, run_status='Finished') else: return run_state
class DockerImageManager: def __init__(self, commit_file, max_image_cache_size): """ Initializes a DockerImageManager :param commit_file: String path to where the state file should be committed :param max_image_cache_size: Total size in bytes that the image cache can use """ self._state_committer = JsonStateCommitter(commit_file) # type: JsonStateCommitter self._docker = docker.from_env() # type: DockerClient self._image_cache = {} # type: Dict[str, ImageCacheEntry] self._downloading = ThreadDict( fields={'success': False, 'status': 'Download starting.'}, lock=True ) self._max_image_cache_size = max_image_cache_size self._lock = threading.RLock() self._stop = False self._sleep_secs = 10 self._cleanup_thread = None self._load_state() def _save_state(self): with self._lock: self._state_committer.commit(self._image_cache) def _load_state(self): with self._lock: self._image_cache = self._state_committer.load() def start(self): if self._max_image_cache_size: def cleanup_loop(self): while not self._stop: try: self._cleanup() self._save_state() except Exception: traceback.print_exc() time.sleep(self._sleep_secs) self._cleanup_thread = threading.Thread(target=cleanup_loop, args=[self]) self._cleanup_thread.start() def stop(self): logger.info("Stopping docker image manager") self._stop = True logger.debug("Stopping docker image manager: stop the downloads threads") self._downloading.stop() if self._cleanup_thread: logger.debug("Stopping docker image manager: stop the cleanup thread") self._cleanup_thread.join() logger.info("Stopped docker image manager.") def _cleanup(self): """ Prunes the image cache for runs 1. Only care about images we (this DockerImageManager) downloaded and know about 2. We use sum of VirtualSize's, which is an upper bound on the disk use of our images: in case no images share any intermediate layers, this will be the real disk use, however if images share layers, the virtual size will count that layer's size for each image that uses it, even though it's stored only once in the disk. The 'Size' field accounts for the marginal size each image adds on top of the shared layers, but summing those is not accurate either since the shared base layers need to be counted once to get the total size. (i.e. summing marginal sizes would give us a lower bound on the total disk use of images). Calling df gives us an accurate disk use of ALL the images on the machine but because of (1) we don't want to use that. """ while not self._stop: time.sleep(self._sleep_secs) deletable_entries = set(self._image_cache.values()) disk_use = sum(cache_entry.virtual_size for cache_entry in deletable_entries) while disk_use > self._max_image_cache_size: entry_to_remove = min(deletable_entries, key=lambda entry: entry.last_used) logger.info( 'Disk use (%s) > max cache size (%s), pruning image: %s', disk_use, self._max_image_cache_size, entry_to_remove.digest, ) try: image_to_delete = self._docker.images.get(entry_to_remove.id) tags_to_delete = image_to_delete.tags for tag in tags_to_delete: self._docker.images.remove(tag) # if we successfully removed the image also remove its cache entry del self._image_cache[entry_to_remove.digest] except docker.errors.NotFound: # image doesn't exist anymore for some reason, stop tracking it del self._image_cache[entry_to_remove.digest] except docker.errors.APIError as err: # Maybe we can't delete this image because its container is still running # (think a run that takes 4 days so this is the oldest image but still in use) # In that case we just continue with our lives, hoping it will get deleted once # it's no longer in use and the cache becomes full again logger.error( "Cannot remove image %s from cache: %s", entry_to_remove.digest, err ) deletable_entries.remove(entry_to_remove) disk_use = sum(entry.virtual_size for entry in deletable_entries) logger.debug("Stopping docker image manager cleanup") def get(self, image_spec): """ Request the docker image for the run with uuid, registering uuid as a dependent of this docker image :param image_spec: Repo image_spec of docker image being requested :returns: A DockerAvailabilityState object with the state of the docker image """ if ':' not in image_spec: # Both digests and repo:tag kind of specs include the : character. The only case without it is when # a repo is specified without a tag (like 'latest') # When this is the case, different images API methods act differently: # - pull pulls all tags of the image # - get tries to get `latest` by default # That means if someone requests a docker image without a tag, and the image does not have a latest # tag pushed to Dockerhub, pull will succeed since it will pull all other tags, but later get calls # will fail since the `latest` tag won't be found on the system. # We don't want to assume what tag the user wanted so we want the pull step to fail if no tag is specified # and there's no latest tag on dockerhub. # Hence, we append the latest tag to the image spec if there's no tag specified otherwise at the very beginning image_spec += ':latest' try: image = self._docker.images.get(image_spec) digest = image.attrs.get('RepoDigests', [image_spec])[0] with self._lock: self._image_cache[digest] = ImageCacheEntry( id=image.id, digest=digest, last_used=time.time(), virtual_size=image.attrs['VirtualSize'], marginal_size=image.attrs['Size'], ) # We can remove the download thread if it still exists if image_spec in self._downloading: self._downloading.remove(image_spec) return ImageAvailabilityState( digest=digest, stage=DependencyStage.READY, message='Image ready' ) except docker.errors.ImageNotFound: return self._pull_or_report(image_spec) # type: DockerAvailabilityState except Exception as ex: return ImageAvailabilityState( digest=None, stage=DependencyStage.FAILED, message=str(ex) ) def _pull_or_report(self, image_spec): if image_spec in self._downloading: with self._downloading[image_spec]['lock']: if self._downloading[image_spec].is_alive(): return ImageAvailabilityState( digest=None, stage=DependencyStage.DOWNLOADING, message=self._downloading[image_spec]['status'], ) else: if self._downloading[image_spec]['success']: digest = self._docker.images.get(image_spec).attrs.get( 'RepoDigests', [image_spec] )[0] status = ImageAvailabilityState( digest=digest, stage=DependencyStage.READY, message=self._downloading[image_spec]['message'], ) else: status = ImageAvailabilityState( digest=None, stage=DependencyStage.FAILED, message=self._downloading[image_spec]['message'], ) self._downloading.remove(image_spec) return status else: def download(): logger.debug('Downloading Docker image %s', image_spec) try: self._docker.images.pull(image_spec) logger.debug('Download for Docker image %s complete', image_spec) self._downloading[image_spec]['success'] = True self._downloading[image_spec]['message'] = "Downloading image" except (docker.errors.APIError, docker.errors.ImageNotFound) as ex: logger.debug('Download for Docker image %s failed: %s', image_spec, ex) self._downloading[image_spec]['success'] = False self._downloading[image_spec]['message'] = "Can't download image: {}".format(ex) self._downloading.add_if_new(image_spec, threading.Thread(target=download, args=[])) return ImageAvailabilityState( digest=None, stage=DependencyStage.DOWNLOADING, message=self._downloading[image_spec]['status'], )
class DockerImageManager: def __init__(self, commit_file, max_image_cache_size): """ Initializes a DockerImageManager :param commit_file: String path to where the state file should be committed :param max_image_cache_size: Total size in bytes that the image cache can use """ self._state_committer = JsonStateCommitter( commit_file) # type: JsonStateCommitter self._docker = docker.from_env() # type: DockerClient self._image_cache = {} # type: Dict[str, ImageCacheEntry] self._downloading = ThreadDict(fields={ 'success': False, 'status': 'Download starting.' }, lock=True) self._max_image_cache_size = max_image_cache_size self._lock = threading.RLock() self._stop = False self._sleep_secs = 10 self._cleanup_thread = None self._load_state() def _save_state(self): with self._lock: self._state_committer.commit(self._image_cache) def _load_state(self): with self._lock: self._image_cache = self._state_committer.load() def start(self): if self._max_image_cache_size: def cleanup_loop(self): while not self._stop: try: self._cleanup() self._save_state() except Exception: traceback.print_exc() time.sleep(self._sleep_secs) self._cleanup_thread = threading.Thread(target=cleanup_loop, args=[self]) self._cleanup_thread.start() def stop(self): logger.info("Stopping docker image manager") self._stop = True logger.debug( "Stopping docker image manager: stop the downloads threads") self._downloading.stop() if self._cleanup_thread: logger.debug( "Stopping docker image manager: stop the cleanup thread") self._cleanup_thread.join() logger.info("Stopped docker image manager.") def _cleanup(self): """ Prunes the image cache for runs 1. Only care about images we (this DockerImageManager) downloaded and know about 2. We use sum of VirtualSize's, which is an upper bound on the disk use of our images: in case no images share any intermediate layers, this will be the real disk use, however if images share layers, the virtual size will count that layer's size for each image that uses it, even though it's stored only once in the disk. The 'Size' field accounts for the marginal size each image adds on top of the shared layers, but summing those is not accurate either since the shared base layers need to be counted once to get the total size. (i.e. summing marginal sizes would give us a lower bound on the total disk use of images). Calling df gives us an accurate disk use of ALL the images on the machine but because of (1) we don't want to use that. """ while not self._stop: time.sleep(self._sleep_secs) deletable_entries = set(self._image_cache.values()) disk_use = sum(cache_entry.virtual_size for cache_entry in deletable_entries) while disk_use > self._max_image_cache_size: entry_to_remove = min(deletable_entries, key=lambda entry: entry.last_used) logger.info( 'Disk use (%s) > max cache size (%s), pruning image: %s', disk_use, self._max_image_cache_size, entry_to_remove.digest, ) try: image_to_delete = self._docker.images.get( entry_to_remove.id) tags_to_delete = image_to_delete.tags for tag in tags_to_delete: self._docker.images.remove(tag) # if we successfully removed the image also remove its cache entry del self._image_cache[entry_to_remove.digest] except docker.errors.NotFound: # image doesn't exist anymore for some reason, stop tracking it del self._image_cache[entry_to_remove.digest] except docker.errors.APIError as err: # Maybe we can't delete this image because its container is still running # (think a run that takes 4 days so this is the oldest image but still in use) # In that case we just continue with our lives, hoping it will get deleted once # it's no longer in use and the cache becomes full again logger.error("Cannot remove image %s from cache: %s", entry_to_remove.digest, err) deletable_entries.remove(entry_to_remove) disk_use = sum(entry.virtual_size for entry in deletable_entries) logger.debug("Stopping docker image manager cleanup") def get(self, image_spec): """ Request the docker image for the run with uuid, registering uuid as a dependent of this docker image :param image_spec: Repo image_spec of docker image being requested :returns: A DockerAvailabilityState object with the state of the docker image """ if ':' not in image_spec: # Both digests and repo:tag kind of specs include the : character. The only case without it is when # a repo is specified without a tag (like 'latest') # When this is the case, different images API methods act differently: # - pull pulls all tags of the image # - get tries to get `latest` by default # That means if someone requests a docker image without a tag, and the image does not have a latest # tag pushed to Dockerhub, pull will succeed since it will pull all other tags, but later get calls # will fail since the `latest` tag won't be found on the system. # We don't want to assume what tag the user wanted so we want the pull step to fail if no tag is specified # and there's no latest tag on dockerhub. # Hence, we append the latest tag to the image spec if there's no tag specified otherwise at the very beginning image_spec += ':latest' try: image = self._docker.images.get(image_spec) digest = image.attrs.get('RepoDigests', [image_spec])[0] with self._lock: self._image_cache[digest] = ImageCacheEntry( id=image.id, digest=digest, last_used=time.time(), virtual_size=image.attrs['VirtualSize'], marginal_size=image.attrs['Size'], ) # We can remove the download thread if it still exists if image_spec in self._downloading: self._downloading.remove(image_spec) return ImageAvailabilityState(digest=digest, stage=DependencyStage.READY, message='Image ready') except docker.errors.ImageNotFound: return self._pull_or_report( image_spec) # type: DockerAvailabilityState except Exception as ex: return ImageAvailabilityState(digest=None, stage=DependencyStage.FAILED, message=str(ex)) def _pull_or_report(self, image_spec): if image_spec in self._downloading: with self._downloading[image_spec]['lock']: if self._downloading[image_spec].is_alive(): return ImageAvailabilityState( digest=None, stage=DependencyStage.DOWNLOADING, message=self._downloading[image_spec]['status'], ) else: if self._downloading[image_spec]['success']: digest = self._docker.images.get(image_spec).attrs.get( 'RepoDigests', [image_spec])[0] status = ImageAvailabilityState( digest=digest, stage=DependencyStage.READY, message=self._downloading[image_spec]['message'], ) else: status = ImageAvailabilityState( digest=None, stage=DependencyStage.FAILED, message=self._downloading[image_spec]['message'], ) self._downloading.remove(image_spec) return status else: def download(): logger.debug('Downloading Docker image %s', image_spec) try: self._docker.images.pull(image_spec) logger.debug('Download for Docker image %s complete', image_spec) self._downloading[image_spec]['success'] = True self._downloading[image_spec][ 'message'] = "Downloading image" except (docker.errors.APIError, docker.errors.ImageNotFound) as ex: logger.debug('Download for Docker image %s failed: %s', image_spec, ex) self._downloading[image_spec]['success'] = False self._downloading[image_spec][ 'message'] = "Can't download image: {}".format(ex) self._downloading.add_if_new( image_spec, threading.Thread(target=download, args=[])) return ImageAvailabilityState( digest=None, stage=DependencyStage.DOWNLOADING, message=self._downloading[image_spec]['status'], )
class LocalRunStateMachine(StateTransitioner): """ Manages the state machine of the runs running on the local machine """ def __init__( self, docker_image_manager, dependency_manager, worker_docker_network, docker_network_internal, docker_network_external, docker_runtime, upload_bundle_callback, assign_cpu_and_gpu_sets_fn, ): super(LocalRunStateMachine, self).__init__() self.add_transition(LocalRunStage.PREPARING, self._transition_from_PREPARING) self.add_transition(LocalRunStage.RUNNING, self._transition_from_RUNNING) self.add_transition(LocalRunStage.CLEANING_UP, self._transition_from_CLEANING_UP) self.add_transition( LocalRunStage.UPLOADING_RESULTS, self._transition_from_UPLOADING_RESULTS ) self.add_transition(LocalRunStage.FINALIZING, self._transition_from_FINALIZING) self.add_terminal(LocalRunStage.FINISHED) self.dependency_manager = dependency_manager self.docker_image_manager = docker_image_manager self.worker_docker_network = worker_docker_network self.docker_network_external = docker_network_external self.docker_network_internal = docker_network_internal self.docker_runtime = docker_runtime # bundle_uuid -> {'thread': Thread, 'run_status': str} self.uploading = ThreadDict(fields={'run_status': 'Upload started', 'success': False}) # bundle_uuid -> {'thread': Thread, 'disk_utilization': int, 'running': bool} self.disk_utilization = ThreadDict( fields={'disk_utilization': 0, 'running': True, 'lock': None} ) self.upload_bundle_callback = upload_bundle_callback self.assign_cpu_and_gpu_sets_fn = assign_cpu_and_gpu_sets_fn def stop(self): for uuid in self.disk_utilization.keys(): self.disk_utilization[uuid]['running'] = False self.disk_utilization.stop() self.uploading.stop() def _transition_from_PREPARING(self, run_state): """ 1- Request the docker image from docker image manager - if image is failed, move to CLEANING_UP state 2- Request the dependencies from dependency manager - if any are failed, move to CLEANING_UP state 3- If all dependencies and docker image are ready: - Set up the local filesystem for the run - Create symlinks to dependencies - Allocate resources and prepare the docker container - Start the docker container 4- If all is successful, move to RUNNING state """ if run_state.is_killed: return run_state._replace(stage=LocalRunStage.CLEANING_UP, container_id=None) dependencies_ready = True status_messages = [] bundle_uuid = run_state.bundle['uuid'] # get dependencies for dep in run_state.bundle['dependencies']: dependency = (dep['parent_uuid'], dep['parent_path']) dependency_state = self.dependency_manager.get(bundle_uuid, dependency) if dependency_state.stage == DependencyStage.DOWNLOADING: status_messages.append( 'Downloading dependency %s: %s done (archived size)' % (dep['child_path'], size_str(dependency_state.size_bytes)) ) dependencies_ready = False elif dependency_state.stage == DependencyStage.FAILED: # Failed to download dependency; -> CLEANING_UP run_state.info['failure_message'] = 'Failed to download dependency %s: %s' % ( dep['child_path'], '', ) return run_state._replace(stage=LocalRunStage.CLEANING_UP, info=run_state.info) # get the docker image docker_image = run_state.resources['docker_image'] image_state = self.docker_image_manager.get(docker_image) if image_state.stage == DependencyStage.DOWNLOADING: status_messages.append( 'Pulling docker image: ' + (image_state.message or docker_image or "") ) dependencies_ready = False elif image_state.stage == DependencyStage.FAILED: # Failed to pull image; -> CLEANING_UP run_state.info['failure_message'] = image_state.message return run_state._replace(stage=LocalRunStage.CLEANING_UP, info=run_state.info) # stop proceeding if dependency and image downloads aren't all done if not dependencies_ready: status_message = status_messages.pop() if status_messages: status_message += "(and downloading %d other dependencies and docker images)" % len( status_messages ) return run_state._replace(run_status=status_message) # All dependencies ready! Set up directories, symlinks and container. Start container. # 1) Set up a directory to store the bundle. remove_path(run_state.bundle_path) os.mkdir(run_state.bundle_path) # 2) Set up symlinks dependencies = [] docker_dependencies_path = '/' + bundle_uuid + '_dependencies' for dep in run_state.bundle['dependencies']: child_path = os.path.normpath(os.path.join(run_state.bundle_path, dep['child_path'])) if not child_path.startswith(run_state.bundle_path): raise Exception('Invalid key for dependency: %s' % (dep['child_path'])) dependency_path = self.dependency_manager.get( bundle_uuid, (dep['parent_uuid'], dep['parent_path']) ).path dependency_path = os.path.join( self.dependency_manager.dependencies_dir, dependency_path ) docker_dependency_path = os.path.join(docker_dependencies_path, dep['child_path']) os.symlink(docker_dependency_path, child_path) # These are turned into docker volume bindings like: # dependency_path:docker_dependency_path:ro dependencies.append((dependency_path, docker_dependency_path)) # 3) Set up container if run_state.resources['request_network']: docker_network = self.docker_network_external.name else: docker_network = self.docker_network_internal.name try: cpuset, gpuset = self.assign_cpu_and_gpu_sets_fn( run_state.resources['request_cpus'], run_state.resources['request_gpus'] ) except Exception: run_state.info['failure_message'] = "Cannot assign enough resources" return run_state._replace(stage=LocalRunStage.CLEANING_UP, info=run_state.info) # 4) Start container try: container = docker_utils.start_bundle_container( run_state.bundle_path, bundle_uuid, dependencies, run_state.bundle['command'], run_state.resources['docker_image'], network=docker_network, cpuset=cpuset, gpuset=gpuset, memory_bytes=run_state.resources['request_memory'], runtime=self.docker_runtime, ) self.worker_docker_network.connect(container) except docker_utils.DockerException as e: run_state.info['failure_message'] = 'Cannot start Docker container: {}'.format(e) return run_state._replace(stage=LocalRunStage.CLEANING_UP, info=run_state.info) return run_state._replace( stage=LocalRunStage.RUNNING, start_time=time.time(), run_status='Running job in Docker container', container_id=container.id, container=container, docker_image=image_state.digest, has_contents=True, cpuset=cpuset, gpuset=gpuset, ) def _transition_from_RUNNING(self, run_state): """ 1- Check run status of the docker container 2- If run is killed, kill the container 3- If run is finished, move to CLEANING_UP state """ bundle_uuid = run_state.bundle['uuid'] def check_and_report_finished(run_state): try: finished, exitcode, failure_msg = docker_utils.check_finished(run_state.container) except docker_utils.DockerException: traceback.print_exc() finished, exitcode, failure_msg = False, None, None new_info = dict(finished=finished, exitcode=exitcode, failure_message=failure_msg) run_state.info.update(new_info) run_state = run_state._replace(info=run_state.info) return run_state def check_resource_utilization(run_state): kill_messages = [] run_stats = docker_utils.get_container_stats(run_state.container) time_used = time.time() - run_state.start_time run_state = run_state._replace(time_used=time_used) run_state = run_state._replace( max_memory=max(run_state.max_memory, run_stats.get('memory', 0)) ) run_state = run_state._replace( disk_utilization=self.disk_utilization[bundle_uuid]['disk_utilization'] ) if ( run_state.resources['request_time'] and run_state.time_used > run_state.resources['request_time'] ): kill_messages.append( 'Time limit %s exceeded.' % duration_str(run_state.resources['request_time']) ) if ( run_state.max_memory > run_state.resources['request_memory'] or run_state.info.get('exitcode', '0') == '137' ): kill_messages.append( 'Memory limit %s exceeded.' % size_str(run_state.resources['request_memory']) ) if ( run_state.resources['request_disk'] and run_state.disk_utilization > run_state.resources['request_disk'] ): kill_messages.append( 'Disk limit %sb exceeded.' % size_str(run_state.resources['request_disk']) ) if kill_messages: new_info = run_state.info new_info['kill_message'] = ' '.join(kill_messages) run_state = run_state._replace(info=new_info, is_killed=True) return run_state def check_disk_utilization(): running = True while running: start_time = time.time() try: disk_utilization = get_path_size(run_state.bundle_path) self.disk_utilization[bundle_uuid]['disk_utilization'] = disk_utilization running = self.disk_utilization[bundle_uuid]['running'] except Exception: traceback.print_exc() end_time = time.time() # To ensure that we don't hammer the disk for this computation when # there are lots of files, we run it at most 10% of the time. time.sleep(max((end_time - start_time) * 10, 1.0)) self.disk_utilization.add_if_new( bundle_uuid, threading.Thread(target=check_disk_utilization, args=[]) ) run_state = check_and_report_finished(run_state) run_state = check_resource_utilization(run_state) if run_state.is_killed: try: run_state.container.kill() except docker.errors.APIError: finished, _, _ = docker_utils.check_finished(run_state.container) if not finished: # If we can't kill a Running container, something is wrong # Otherwise all well traceback.print_exc() self.disk_utilization[bundle_uuid]['running'] = False self.disk_utilization.remove(bundle_uuid) return run_state._replace(stage=LocalRunStage.CLEANING_UP) if run_state.info['finished']: logger.debug( 'Finished run with UUID %s, exitcode %s, failure_message %s', bundle_uuid, run_state.info['exitcode'], run_state.info['failure_message'], ) self.disk_utilization[bundle_uuid]['running'] = False self.disk_utilization.remove(bundle_uuid) return run_state._replace( stage=LocalRunStage.CLEANING_UP, run_status='Uploading results' ) else: return run_state def _transition_from_CLEANING_UP(self, run_state): """ 1- delete the container if still existent 2- clean up the dependencies from bundle folder 3- release the dependencies in dependency manager 4- If bundle has contents to upload (i.e. was RUNNING at some point), move to UPLOADING_RESULTS state Otherwise move to FINALIZING state """ bundle_uuid = run_state.bundle['uuid'] if run_state.container_id is not None: while True: try: finished, _, _ = docker_utils.check_finished(run_state.container) if finished: run_state.container.remove(force=True) break except docker.errors.APIError: traceback.print_exc() time.sleep(1) for dep in run_state.bundle['dependencies']: self.dependency_manager.release(bundle_uuid, (dep['parent_uuid'], dep['parent_path'])) child_path = os.path.join(run_state.bundle_path, dep['child_path']) try: remove_path(child_path) except Exception: traceback.print_exc() if run_state.has_contents: return run_state._replace( stage=LocalRunStage.UPLOADING_RESULTS, run_status='Uploading results', container=None, ) else: return self.finalize_run(run_state) def _transition_from_UPLOADING_RESULTS(self, run_state): """ If bundle not already uploading: Use the RunManager API to upload contents at bundle_path to the server Pass the callback to that API such that if the bundle is killed during the upload, the callback returns false, allowing killable uploads. If uploading and not finished: Update run_status with upload progress If uploading and finished: Move to FINALIZING state """ def upload_results(): try: # Upload results logger.debug('Uploading results for run with UUID %s', bundle_uuid) def progress_callback(bytes_uploaded): run_status = 'Uploading results: %s done (archived size)' % size_str( bytes_uploaded ) self.uploading[bundle_uuid]['run_status'] = run_status return True self.upload_bundle_callback(bundle_uuid, run_state.bundle_path, progress_callback) self.uploading[bundle_uuid]['success'] = True except Exception as e: self.uploading[bundle_uuid]['run_status'] = "Error while uploading: %s" % e traceback.print_exc() bundle_uuid = run_state.bundle['uuid'] self.uploading.add_if_new(bundle_uuid, threading.Thread(target=upload_results, args=[])) if self.uploading[bundle_uuid].is_alive(): return run_state._replace(run_status=self.uploading[bundle_uuid]['run_status']) elif not self.uploading[bundle_uuid]['success']: # upload failed failure_message = run_state.info.get('failure_message', None) if failure_message: run_state.info['failure_message'] = ( failure_message + '. ' + self.uploading[bundle_uuid]['run_status'] ) else: run_state.info['failure_message'] = self.uploading[bundle_uuid]['run_status'] self.uploading.remove(bundle_uuid) return self.finalize_run(run_state) def finalize_run(self, run_state): """ Prepare the finalize message to be sent with the next checkin """ failure_message = run_state.info.get('failure_message', None) if 'exitcode' not in run_state.info: run_state.info['exitcode'] = None if not failure_message and run_state.is_killed: run_state.info['failure_message'] = run_state.info['kill_message'] run_state.info['finalized'] = False return run_state._replace( stage=LocalRunStage.FINALIZING, info=run_state.info, run_status="Finalizing bundle" ) def _transition_from_FINALIZING(self, run_state): """ If a full worker cycle has passed since we got into FINALIZING we already reported to server so can move on to FINISHED. Can also remove bundle_path now """ if run_state.info['finalized']: remove_path(run_state.bundle_path) return run_state._replace(stage=LocalRunStage.FINISHED, run_status='Finished') else: return run_state