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