def __init__(self, commit_file, bundle_service, worker_dir, max_cache_size_bytes):
super(DependencyManager, 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, DependencyManager.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() # type: Dict[DependencyKey, 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()
# DependencyKey -> DependencyState
self._dependencies = dict()
# DependencyKey -> WorkerThread(thread, success, failure_message)
self._downloading = ThreadDict(fields={'success': False, 'failure_message': None})
self._load_state()
# Sync states between dependency-state.json and dependency directories on the local file system.
self._sync_state()
self._stop = False
self._main_thread = None
def __init__(
self,
commit_file: str,
bundle_service: BundleServiceClient,
worker_dir: str,
max_cache_size_bytes: int,
download_dependencies_max_retries: int,
):
super(DependencyManager, self).__init__()
self.add_transition(DependencyStage.DOWNLOADING, self._transition_from_DOWNLOADING)
self.add_terminal(DependencyStage.READY)
self.add_terminal(DependencyStage.FAILED)
self._id: str = "worker-dependency-manager-{}".format(uuid.uuid4().hex[:8])
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, DependencyManager.DEPENDENCIES_DIR_NAME)
self._download_dependencies_max_retries = download_dependencies_max_retries
if not os.path.exists(self.dependencies_dir):
logger.info('{} doesn\'t exist, creating.'.format(self.dependencies_dir))
os.makedirs(self.dependencies_dir, 0o770)
# Create a lock for concurrency over NFS
# Create a separate locks directory to hold the lock files.
# Each lock file is created when a process tries to claim the main lock.
locks_claims_dir: str = os.path.join(worker_dir, 'locks_claims')
try:
os.makedirs(locks_claims_dir)
except FileExistsError:
logger.info(f"A locks directory at {locks_claims_dir} already exists.")
self._state_lock = NFSLock(os.path.join(locks_claims_dir, 'state.lock'))
# File paths that are currently being used to store dependencies. Used to prevent conflicts
self._paths: Set[str] = set()
# DependencyKey -> WorkerThread(thread, success, failure_message)
self._downloading = ThreadDict(fields={'success': False, 'failure_message': None})
# Sync states between dependency-state.json and dependency directories on the local file system.
self._sync_state()
self._stop = False
self._main_thread = None
logger.info(f"Initialized Dependency Manager with ID: {self._id}")
class JsonStateCommitterTest(unittest.TestCase):
def setUp(self):
self.test_dir = tempfile.mkdtemp()
self.state_file = 'test-state.json'
self.state_path = os.path.join(self.test_dir, self.state_file)
self.committer = JsonStateCommitter(self.state_path)
def tearDown(self):
try:
os.remove(self.state_path)
except OSError:
pass
os.rmdir(self.test_dir)
def test_path_parsing(self):
""" Simple test to ensure we don't mess up the state file path"""
self.assertEqual(self.committer._state_file, self.state_path)
def test_commit(self):
"""Make sure state is committed correctly"""
test_state = {'state': 'value'}
test_state_json_str = '{\"state\": \"value\"}'
self.committer.commit(test_state)
with open(self.state_path) as f:
self.assertEqual(test_state_json_str, f.read())
self.assertFalse(os.path.exists(self.committer.temp_file))
def test_load(self):
""" Make sure load loads the state file if it exists """
test_state = {'state': 'value'}
test_state_json_str = '{\"state\": \"value\"}'
with open(self.state_path, 'w') as f:
f.write(test_state_json_str)
loaded_state = self.committer.load()
self.assertDictEqual(test_state, loaded_state)
def test_default(self):
""" Make sure load with a default works if state file doesn't exist """
default_state = {'state': 'value'}
loaded_state = self.committer.load(default=default_state)
self.assertDictEqual(default_state, loaded_state)
def __init__(
self,
worker, # type: Worker
image_manager, # type: DockerImageManager
dependency_manager, # type: LocalFileSystemDependencyManager
commit_file, # type: str
cpuset, # type: Set[str]
gpuset, # type: Set[str]
work_dir, # type: str
docker_runtime=docker_utils.DEFAULT_RUNTIME, # type: str
docker_network_prefix='codalab_worker_network', # type: str
):
self._worker = worker
self._state_committer = JsonStateCommitter(commit_file)
self._reader = LocalReader()
self._docker = docker.from_env()
self._bundles_dir = os.path.join(work_dir, LocalRunManager.BUNDLES_DIR_NAME)
if not os.path.exists(self._bundles_dir):
logger.info('{} doesn\'t exist, creating.'.format(self._bundles_dir))
os.makedirs(self._bundles_dir, 0o770)
self._image_manager = image_manager
self._dependency_manager = dependency_manager
self._cpuset = cpuset
self._gpuset = gpuset
self._stop = False
self._work_dir = work_dir
self._runs = {} # bundle_uuid -> LocalRunState
self._lock = threading.RLock()
self._init_docker_networks(docker_network_prefix)
self._run_state_manager = LocalRunStateMachine(
docker_image_manager=self._image_manager,
dependency_manager=self._dependency_manager,
worker_docker_network=self.worker_docker_network,
docker_network_internal=self.docker_network_internal,
docker_network_external=self.docker_network_external,
docker_runtime=docker_runtime,
upload_bundle_callback=self._worker.upload_bundle_contents,
assign_cpu_and_gpu_sets_fn=self.assign_cpu_and_gpu_sets,
)
def __init__(self, commit_file: str, max_image_cache_size: int,
max_image_size: int):
"""
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
:param max_image_size: Total size in bytes that the image can have
"""
super().__init__(max_image_size, max_image_cache_size)
self._state_committer = JsonStateCommitter(
commit_file) # type: JsonStateCommitter
self._docker = docker.from_env(
timeout=DEFAULT_DOCKER_TIMEOUT) # type: DockerClient
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, max_image_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
:param max_image_size: Total size in bytes that the image can have
"""
self._state_committer = JsonStateCommitter(commit_file) # type: JsonStateCommitter
self._docker = docker.from_env(timeout=DEFAULT_DOCKER_TIMEOUT) # type: DockerClient
self._downloading = ThreadDict(
fields={'success': False, 'status': 'Download starting'}, lock=True
)
self._max_image_cache_size = max_image_cache_size
self._max_image_size = max_image_size
self._stop = False
self._sleep_secs = 10
self._cleanup_thread = None
def setUp(self):
self.test_dir = tempfile.mkdtemp()
self.state_file = 'test-state.json'
self.state_path = os.path.join(self.test_dir, self.state_file)
self.committer = JsonStateCommitter(self.state_path)
class DependencyManager(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 DependencyKey
"""
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(DependencyManager, 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, DependencyManager.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(
) # type: Dict[DependencyKey, 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()
# DependencyKey -> DependencyState
self._dependencies = dict()
# DependencyKey -> WorkerThread(thread, success, failure_message)
self._downloading = ThreadDict(fields={
'success': False,
'failure_message': None
})
self._load_state()
# Sync states between dependency-state.json and dependency directories on the local file system.
self._sync_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):
"""
Load states from dependencies-state.json, which contains information about bundles (e.g., state, dependencies,
last used, etc.) and populates values for self._dependencies, self._dependency_locks, and self._paths
"""
state = self._state_committer.load(default={
'dependencies': {},
'paths': set()
})
dependencies = {}
dependency_locks = {}
for dep, dep_state in state['dependencies'].items():
dependencies[dep] = dep_state
dependency_locks[dep] = threading.RLock()
with self._global_lock, self._paths_lock:
self._dependencies = dependencies
self._dependency_locks = dependency_locks
self._paths = state['paths']
logger.info('Loaded {} dependencies, {} paths from cache.'.format(
len(self._dependencies), len(self._paths)))
def _sync_state(self):
"""
Synchronize dependency states between dependencies-state.json and the local file system as follows:
1. self._dependencies, self._dependency_locks, and self._paths: populated from dependencies-state.json
in function _load_state()
2. directories on the local file system: the bundle contents
This function forces the 1 and 2 to be in sync by taking the intersection (e.g., deleting bundles from the
local file system that don't appear in the dependencies-state.json and vice-versa)
"""
# Get the paths that exist in dependency state, loaded path and
# the local file system (the dependency directories under self.dependencies_dir)
local_directories = set(os.listdir(self.dependencies_dir))
paths_in_loaded_state = [
dep_state.path for dep_state in self._dependencies.values()
]
self._paths = self._paths.intersection(
paths_in_loaded_state).intersection(local_directories)
# Remove the orphaned dependencies from self._dependencies and
# self._dependency_locks if they don't exist in self._paths (intersection of paths in dependency state,
# loaded paths and the paths on the local file system)
dependencies_to_remove = [
dep for dep, dep_state in self._dependencies.items()
if dep_state.path not in self._paths
]
for dep in dependencies_to_remove:
logger.info(
"Dependency {} in dependency state but its path {} doesn't exist on the local file system. "
"Removing it from dependency state.".format(
dep,
os.path.join(self.dependencies_dir,
self._dependencies[dep].path)))
del self._dependencies[dep]
del self._dependency_locks[dep]
# Remove the orphaned directories from the local file system
directories_to_remove = local_directories - self._paths
for dir in directories_to_remove:
full_path = os.path.join(self.dependencies_dir, dir)
logger.info(
"Remove orphaned directory {} from the local file system.".
format(full_path))
remove_path(full_path)
# Save the current synced state back to the state file: dependency-state.json as
# the current state might have been changed during the state syncing phase
self._save_state()
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 dep_key, dep_state in self._dependencies.items():
with self._dependency_locks[dep_key]:
self._dependencies[dep_key] = self.transition(dep_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_key: dep_state
for dep_key, dep_state in self._dependencies.items()
if dep_state.stage == DependencyStage.FAILED and time.time() -
dep_state.last_used >
DependencyManager.DEPENDENCY_FAILURE_COOLDOWN
}
for dep_key, dep_state in failed_deps.items():
self._delete_dependency(dep_key)
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_state.size_bytes
for dep_state in self._dependencies.values())
serialized_length = len(
codalab.worker.pyjson.dumps(self._dependencies))
if (bytes_used > self._max_cache_size_bytes
or serialized_length >
DependencyManager.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),
DependencyManager.MAX_SERIALIZED_LEN,
)
ready_deps = {
dep_key: dep_state
for dep_key, dep_state in self._dependencies.items()
if dep_state.stage == DependencyStage.READY
and not dep_state.dependents
}
failed_deps = {
dep_key: dep_state
for dep_key, dep_state in self._dependencies.items()
if dep_state.stage == DependencyStage.FAILED
}
if failed_deps:
dep_key_to_remove = min(
failed_deps.items(),
key=lambda dep: dep[1].last_used)[0]
elif ready_deps:
dep_key_to_remove = min(
ready_deps.items(),
key=lambda dep: dep[1].last_used)[0]
else:
logger.info(
'Dependency quota full but there are only downloading dependencies, not cleaning up until downloads are over'
)
self._release_all_locks()
break
if dep_key_to_remove:
self._delete_dependency(dep_key_to_remove)
self._release_all_locks()
else:
self._release_all_locks()
break
def _delete_dependency(self, dependency_key):
"""
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_key):
try:
path_to_remove = self._dependencies[dependency_key].path
self._paths.remove(path_to_remove)
remove_path(path_to_remove)
except Exception:
pass
finally:
del self._dependencies[dependency_key]
self._dependency_locks[dependency_key].release()
def has(self, dependency_key):
"""
Takes a DependencyKey
Returns true if the manager has processed this dependency
"""
with self._global_lock:
return dependency_key in self._dependencies
def get(self, uuid, dependency_key):
"""
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_key): # add dependency state if it does not exist
with self._global_lock:
self._dependency_locks[dependency_key] = threading.RLock()
self._dependency_locks[dependency_key].acquire()
self._dependencies[dependency_key] = DependencyState(
stage=DependencyStage.DOWNLOADING,
dependency_key=dependency_key,
path=self._assign_path(dependency_key),
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_key].stage != DependencyStage.FAILED:
self._dependencies[dependency_key].dependents.add(uuid)
self._dependencies[dependency_key] = self._dependencies[
dependency_key]._replace(last_used=now)
self._dependency_locks[dependency_key].release()
return self._dependencies[dependency_key]
def release(self, uuid, dependency_key):
"""
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_key):
dep_state = self._dependencies[dependency_key]
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_key] = dep_state
self._dependency_locks[dependency_key].release()
def _acquire_if_exists(self, dependency_key):
"""
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_key in self._dependencies:
self._dependency_locks[dependency_key].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.items():
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.items():
lock.release()
def _assign_path(self, dependency_key):
"""
Normalize the path for the dependency by replacing / with _, avoiding conflicts
"""
if dependency_key.parent_path:
path = os.path.join(dependency_key.parent_uuid,
dependency_key.parent_path)
else:
path = dependency_key.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 on 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.dependency_key]:
state = self._dependencies[dependency_state.dependency_key]
if state.killed:
raise DownloadAbortedException("Aborted by user")
self._dependencies[
dependency_state.dependency_key] = 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',
dependency_state.dependency_key)
try:
# Start async download to the fileobj
fileobj, target_type = self._bundle_service.get_bundle_contents(
dependency_state.dependency_key.parent_uuid,
dependency_state.dependency_key.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 to %s',
target_type,
dependency_state.dependency_key,
dependency_path,
)
with self._dependency_locks[dependency_state.dependency_key]:
self._downloading[
dependency_state.dependency_key]['success'] = True
except Exception as e:
with self._dependency_locks[dependency_state.dependency_key]:
self._downloading[
dependency_state.dependency_key]['success'] = False
self._downloading[dependency_state.dependency_key][
'failure_message'] = "Dependency download failed: %s " % str(
e)
self._downloading.add_if_new(
dependency_state.dependency_key,
threading.Thread(target=download, args=[]))
if self._downloading[dependency_state.dependency_key].is_alive():
return dependency_state
success = self._downloading[dependency_state.dependency_key]['success']
failure_message = self._downloading[
dependency_state.dependency_key]['failure_message']
self._downloading.remove(dependency_state.dependency_key)
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 DependencyManager(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. It's also NFS-safe.
In this class, dependencies are uniquely identified by DependencyKey.
"""
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
# If it has been this long since a worker has downloaded anything, another worker will take over downloading.
DEPENDENCY_DOWNLOAD_TIMEOUT_SECONDS = 5 * 60
def __init__(
self,
commit_file: str,
bundle_service: BundleServiceClient,
worker_dir: str,
max_cache_size_bytes: int,
download_dependencies_max_retries: int,
):
super(DependencyManager, self).__init__()
self.add_transition(DependencyStage.DOWNLOADING, self._transition_from_DOWNLOADING)
self.add_terminal(DependencyStage.READY)
self.add_terminal(DependencyStage.FAILED)
self._id: str = "worker-dependency-manager-{}".format(uuid.uuid4().hex[:8])
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, DependencyManager.DEPENDENCIES_DIR_NAME)
self._download_dependencies_max_retries = download_dependencies_max_retries
if not os.path.exists(self.dependencies_dir):
logger.info('{} doesn\'t exist, creating.'.format(self.dependencies_dir))
os.makedirs(self.dependencies_dir, 0o770)
# Create a lock for concurrency over NFS
# Create a separate locks directory to hold the lock files.
# Each lock file is created when a process tries to claim the main lock.
locks_claims_dir: str = os.path.join(worker_dir, 'locks_claims')
try:
os.makedirs(locks_claims_dir)
except FileExistsError:
logger.info(f"A locks directory at {locks_claims_dir} already exists.")
self._state_lock = NFSLock(os.path.join(locks_claims_dir, 'state.lock'))
# File paths that are currently being used to store dependencies. Used to prevent conflicts
self._paths: Set[str] = set()
# DependencyKey -> WorkerThread(thread, success, failure_message)
self._downloading = ThreadDict(fields={'success': False, 'failure_message': None})
# Sync states between dependency-state.json and dependency directories on the local file system.
self._sync_state()
self._stop = False
self._main_thread = None
logger.info(f"Initialized Dependency Manager with ID: {self._id}")
def _sync_state(self):
"""
Synchronize dependency states between dependencies-state.json and the local file system as follows:
1. dependencies and paths: populated from dependencies-state.json
2. directories on the local file system: the bundle contents
This function forces the 1 and 2 to be in sync by taking the intersection (e.g., deleting bundles from the
local file system that don't appear in the dependencies-state.json and vice-versa)
"""
with self._state_lock:
# Load states from dependencies-state.json, which contains information about bundles (e.g., state,
# dependencies, last used, etc.).
if self._state_committer.state_file_exists:
# If the state file exists, do not pass in a default. It's critical that we read the contents
# of the state file, as this method prunes dependencies. If we can't read the contents of the
# state file, fail immediately.
dependencies, paths = self._fetch_state()
logger.info(
'Found {} dependencies, {} paths from cache.'.format(
len(dependencies), len(paths)
)
)
else:
dependencies: Dict[DependencyKey, DependencyState] = dict()
paths: Set[str] = set()
logger.info(
f'State file did not exist. Will create one at path {self._state_committer.path}.'
)
# Get the paths that exist in dependency state, loaded path and
# the local file system (the dependency directories under self.dependencies_dir)
local_directories = set(os.listdir(self.dependencies_dir))
paths_in_loaded_state = [dep_state.path for dep_state in dependencies.values()]
paths = paths.intersection(paths_in_loaded_state).intersection(local_directories)
# Remove the orphaned dependencies if they don't exist in paths
# (intersection of paths in dependency state, loaded paths and the paths on the local file system)
dependencies_to_remove = [
dep for dep, dep_state in dependencies.items() if dep_state.path not in paths
]
for dep in dependencies_to_remove:
logger.info(
"Dependency {} in dependency state but its path {} doesn't exist on the local file system. "
"Removing it from dependency state.".format(
dep, os.path.join(self.dependencies_dir, dependencies[dep].path)
)
)
del dependencies[dep]
# Remove the orphaned directories from the local file system
directories_to_remove = local_directories - paths
for directory in directories_to_remove:
full_path = os.path.join(self.dependencies_dir, directory)
if os.path.exists(full_path):
logger.info(
"Remove orphaned directory {} from the local file system.".format(full_path)
)
remove_path(full_path)
# Save the current synced state back to the state file: dependency-state.json as
# the current state might have been changed during the state syncing phase
self._commit_state(dependencies, paths)
def _fetch_state(self, default=None):
"""
Fetch state from JSON file stored on disk.
NOT NFS-SAFE - Caller should acquire self._state_lock before calling this method.
WARNING: If a value for `default` is specified, errors will be silently handled.
"""
assert self._state_lock.is_locked
state: DependencyManagerState = self._state_committer.load(default)
dependencies: Dict[DependencyKey, DependencyState] = state['dependencies']
paths: Set[str] = state['paths']
return dependencies, paths
def _fetch_dependencies(self, default=None) -> Dict[DependencyKey, DependencyState]:
"""
Fetch dependencies from JSON file stored on disk.
NOT NFS-SAFE - Caller should acquire self._state_lock before calling this method.
WARNING: If a value for `default` is specified, errors will be silently handled.
"""
assert self._state_lock.is_locked
dependencies, _ = self._fetch_state(default)
return dependencies
def _commit_state(self, dependencies: Dict[DependencyKey, DependencyState], paths: Set[str]):
"""
Update state in dependencies JSON file stored on disk.
NOT NFS-SAFE - Caller should acquire self._state_lock before calling this method.
"""
assert self._state_lock.is_locked
state: DependencyManagerState = {'dependencies': dependencies, 'paths': paths}
self._state_committer.commit(state)
def start(self):
logger.info('Starting local dependency manager...')
def loop(self):
while not self._stop:
try:
self._transition_dependencies()
self._cleanup()
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()
self._state_lock.release()
logger.info('Stopped local dependency manager.')
def _transition_dependencies(self):
with self._state_lock:
try:
dependencies, paths = self._fetch_state()
# Update the class variable _paths as the transition function may update it
self._paths = paths
for dep_key, dep_state in dependencies.items():
dependencies[dep_key] = self.transition(dep_state)
self._commit_state(dependencies, self._paths)
except (ValueError, EnvironmentError):
# Do nothing if an error is thrown while reading from the state file
logging.exception("Error reading from state file while transitioning dependencies")
pass
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._state_lock:
try:
dependencies, paths = self._fetch_state()
failed_deps: Dict[DependencyKey, DependencyState] = {
dep_key: dep_state
for dep_key, dep_state in dependencies.items()
if dep_state.stage == DependencyStage.FAILED
and time.time() - dep_state.last_used
> DependencyManager.DEPENDENCY_FAILURE_COOLDOWN
}
if len(failed_deps) == 0:
return
for dep_key, dep_state in failed_deps.items():
self._delete_dependency(dep_key, dependencies, paths)
self._commit_state(dependencies, paths)
except (ValueError, EnvironmentError):
# Do nothing if an error is thrown while reading from the state file
logging.exception(
"Error reading from state file while pruning failed dependencies."
)
pass
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()
while True:
with self._state_lock:
try:
dependencies, paths = self._fetch_state()
except (ValueError, EnvironmentError):
# Do nothing if an error is thrown while reading from the state file
logging.exception(
"Error reading from state file when cleaning up dependencies. Trying again..."
)
continue
bytes_used = sum(dep_state.size_bytes for dep_state in dependencies.values())
serialized_length = len(codalab.worker.pyjson.dumps(dependencies))
if (
bytes_used > self._max_cache_size_bytes
or serialized_length > DependencyManager.MAX_SERIALIZED_LEN
):
logger.debug(
'%d dependencies, disk usage: %s (max %s), serialized size: %s (max %s)',
len(dependencies),
size_str(bytes_used),
size_str(self._max_cache_size_bytes),
size_str(serialized_length),
DependencyManager.MAX_SERIALIZED_LEN,
)
ready_deps = {
dep_key: dep_state
for dep_key, dep_state in dependencies.items()
if dep_state.stage == DependencyStage.READY and not dep_state.dependents
}
failed_deps = {
dep_key: dep_state
for dep_key, dep_state in dependencies.items()
if dep_state.stage == DependencyStage.FAILED
}
if failed_deps:
dep_key_to_remove = min(
failed_deps.items(), key=lambda dep: dep[1].last_used
)[0]
elif ready_deps:
dep_key_to_remove = min(
ready_deps.items(), key=lambda dep: dep[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_key_to_remove:
self._delete_dependency(dep_key_to_remove, dependencies, paths)
self._commit_state(dependencies, paths)
else:
break
def _delete_dependency(self, dep_key, dependencies, paths):
"""
Remove the given dependency from the manager's state
Modifies `dependencies` and `paths` that are passed in.
Also deletes any known files on the filesystem if any exist.
NOT NFS-SAFE - Caller should acquire self._state_lock before calling this method.
"""
assert self._state_lock.is_locked
if dep_key in dependencies:
try:
path_to_remove = dependencies[dep_key].path
paths.remove(path_to_remove)
# Deletes dependency content from disk
remove_path(path_to_remove)
except Exception:
pass
finally:
del dependencies[dep_key]
logger.info(f"Deleted dependency {dep_key}.")
def has(self, dependency_key):
"""
Takes a DependencyKey and returns true if the manager has processed this dependency
"""
with self._state_lock:
dependencies: Dict[DependencyKey, DependencyState] = self._fetch_dependencies()
return dependency_key in dependencies
def get(self, uuid: str, dependency_key: DependencyKey) -> DependencyState:
"""
Request the dependency for the run with uuid, registering uuid as a dependent of this dependency
"""
with self._state_lock:
dependencies, paths = self._fetch_state()
now = time.time()
# Add dependency state if it does not exist
if dependency_key not in dependencies:
dependencies[dependency_key] = DependencyState(
stage=DependencyStage.DOWNLOADING,
downloading_by=None,
dependency_key=dependency_key,
path=self._assign_path(paths, dependency_key),
size_bytes=0,
dependents={uuid},
last_used=now,
last_downloading=now,
message="Starting download",
killed=False,
)
# Update last_used as long as it isn't in a FAILED stage
if dependencies[dependency_key].stage != DependencyStage.FAILED:
dependencies[dependency_key].dependents.add(uuid)
dependencies[dependency_key] = dependencies[dependency_key]._replace(last_used=now)
self._commit_state(dependencies, paths)
return dependencies[dependency_key]
def release(self, uuid, dependency_key):
"""
Register that the run with uuid is no longer dependent on this dependency
If no more runs are dependent on this dependency, kill it.
"""
with self._state_lock:
dependencies, paths = self._fetch_state()
if dependency_key in dependencies:
dep_state = dependencies[dependency_key]
if uuid in dep_state.dependents:
dep_state.dependents.remove(uuid)
if not dep_state.dependents:
dep_state = dep_state._replace(killed=True)
dependencies[dependency_key] = dep_state
self._commit_state(dependencies, paths)
def _assign_path(self, paths: Set[str], dependency_key: DependencyKey) -> str:
"""
Checks the current path against `paths`.
Normalize the path for the dependency by replacing / with _, avoiding conflicts.
Adds the new path to `paths`.
"""
path: str = (
os.path.join(dependency_key.parent_uuid, dependency_key.parent_path)
if dependency_key.parent_path
else dependency_key.parent_uuid
)
path = path.replace(os.path.sep, '_')
# You could have a conflict between, for example a/b_c and a_b/c
while path in paths:
path = path + '_'
paths.add(path)
return path
def _store_dependency(self, dependency_path, fileobj, target_type):
"""
Copy the dependency fileobj to its path on 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) -> List[DependencyKey]:
with self._state_lock:
dependencies: Dict[DependencyKey, DependencyState] = self._fetch_dependencies(
default={'dependencies': {}, 'paths': set()}
)
return list(dependencies.keys())
def _transition_from_DOWNLOADING(self, dependency_state: DependencyState):
"""
Checks if the dependency is downloading or not.
NOT NFS-SAFE - Caller should acquire self._state_lock before calling this method.
"""
assert self._state_lock.is_locked
def download():
"""
Runs in a separate thread. Only one worker should be running this in a thread at a time.
"""
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
Note: This function needs to be fast, since it's called every time fileobj.read is called.
Therefore, we keep a copy of the state in memory (self._downloading) and copy over
non-critical fields (last_downloading, size_bytes and message) when the download transition
function is executed.
"""
state = self._downloading[dependency_state.dependency_key]['state']
if state.killed:
raise DownloadAbortedException("Aborted by user")
self._downloading[dependency_state.dependency_key]['state'] = state._replace(
last_downloading=time.time(),
size_bytes=bytes_downloaded,
message=f"Downloading dependency: {str(bytes_downloaded)} downloaded",
)
dependency_path = os.path.join(self.dependencies_dir, dependency_state.path)
logger.debug('Downloading dependency %s', dependency_state.dependency_key)
attempt = 0
while attempt < self._download_dependencies_max_retries:
try:
# Start async download to the fileobj
target_type = self._bundle_service.get_bundle_info(
dependency_state.dependency_key.parent_uuid,
dependency_state.dependency_key.parent_path,
)["type"]
fileobj = self._bundle_service.get_bundle_contents(
dependency_state.dependency_key.parent_uuid,
dependency_state.dependency_key.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.
original_read_method = fileobj.read
bytes_downloaded = [0]
def interruptable_read(*args, **kwargs):
data = original_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
# Note: Overwrites if something already exists at dependency_path, such as when
# another worker partially downloads a dependency and then goes offline.
self._store_dependency(dependency_path, fileobj, target_type)
logger.debug(
'Finished downloading %s dependency %s to %s',
target_type,
dependency_state.dependency_key,
dependency_path,
)
self._downloading[dependency_state.dependency_key]['success'] = True
except Exception as e:
attempt += 1
if attempt >= self._download_dependencies_max_retries:
self._downloading[dependency_state.dependency_key]['success'] = False
self._downloading[dependency_state.dependency_key][
'failure_message'
] = f"Dependency download failed: {e} "
else:
logger.warning(
f'Failed to download {dependency_state.dependency_key} after {attempt} attempt(s) '
f'due to {e}. Retrying up to {self._download_dependencies_max_retries} times...',
exc_info=True,
)
else:
# Break out of the retry loop if no exceptions were thrown
break
# Start downloading if either:
# 1. No other dependency manager is downloading the dependency
# 2. There was a dependency manager downloading a dependency, but it has been longer than
# DEPENDENCY_DOWNLOAD_TIMEOUT_SECONDS since it last downloaded anything for the particular dependency.
now = time.time()
if not dependency_state.downloading_by or (
dependency_state.downloading_by
and now - dependency_state.last_downloading
>= DependencyManager.DEPENDENCY_DOWNLOAD_TIMEOUT_SECONDS
):
if not dependency_state.downloading_by:
logger.info(
f"{self._id} will start downloading dependency: {dependency_state.dependency_key}."
)
else:
logger.info(
f"{dependency_state.downloading_by} stopped downloading "
f"dependency: {dependency_state.dependency_key}. {self._id} will restart downloading."
)
self._downloading.add_if_new(
dependency_state.dependency_key, threading.Thread(target=download, args=[])
)
self._downloading[dependency_state.dependency_key]['state'] = dependency_state
dependency_state = dependency_state._replace(downloading_by=self._id)
# If there is already another worker downloading the dependency,
# just return the dependency state as downloading is in progress.
if dependency_state.downloading_by != self._id:
logger.debug(
f"Waiting for {dependency_state.downloading_by} "
f"to download dependency: {dependency_state.dependency_key}"
)
return dependency_state
if (
dependency_state.dependency_key in self._downloading
and self._downloading[dependency_state.dependency_key].is_alive()
):
logger.debug(
f"This dependency manager ({dependency_state.downloading_by}) "
f"is downloading dependency: {dependency_state.dependency_key}"
)
state = self._downloading[dependency_state.dependency_key]['state']
# Copy over the values of the non-critical fields of the state in memory
# that is being updated by the download thread.
return dependency_state._replace(
last_downloading=state.last_downloading,
size_bytes=state.size_bytes,
message=state.message,
)
# At this point, no thread is downloading the dependency, but the dependency is still
# assigned to the current worker. Check if the download finished.
success: bool = self._downloading[dependency_state.dependency_key]['success']
failure_message: str = self._downloading[dependency_state.dependency_key]['failure_message']
dependency_state = dependency_state._replace(downloading_by=None)
self._downloading.remove(dependency_state.dependency_key)
logger.info(
f"Download complete. Removing downloading thread for {dependency_state.dependency_key}."
)
if success:
return dependency_state._replace(
stage=DependencyStage.READY, message="Download complete"
)
else:
self._paths.remove(dependency_state.path)
logger.error(
f"Dependency {dependency_state.dependency_key} download failed: {failure_message}"
)
return dependency_state._replace(stage=DependencyStage.FAILED, message=failure_message)
class LocalRunManager(BaseRunManager):
"""
LocalRunManager executes the runs locally, each one in its own Docker
container. It manages its cache of local Docker images and its own local
Docker network.
"""
# Network buffer size to use while proxying with netcat
NETCAT_BUFFER_SIZE = 4096
# Number of seconds to wait for bundle kills to propagate before forcing kill
KILL_TIMEOUT = 100
# Directory name to store running bundles in worker filesystem
BUNDLES_DIR_NAME = 'runs'
def __init__(
self,
worker, # type: Worker
image_manager, # type: DockerImageManager
dependency_manager, # type: LocalFileSystemDependencyManager
commit_file, # type: str
cpuset, # type: Set[str]
gpuset, # type: Set[str]
work_dir, # type: str
docker_runtime=docker_utils.DEFAULT_RUNTIME, # type: str
docker_network_prefix='codalab_worker_network', # type: str
):
self._worker = worker
self._state_committer = JsonStateCommitter(commit_file)
self._reader = LocalReader()
self._docker = docker.from_env()
self._bundles_dir = os.path.join(work_dir, LocalRunManager.BUNDLES_DIR_NAME)
if not os.path.exists(self._bundles_dir):
logger.info('{} doesn\'t exist, creating.'.format(self._bundles_dir))
os.makedirs(self._bundles_dir, 0o770)
self._image_manager = image_manager
self._dependency_manager = dependency_manager
self._cpuset = cpuset
self._gpuset = gpuset
self._stop = False
self._work_dir = work_dir
self._runs = {} # bundle_uuid -> LocalRunState
self._lock = threading.RLock()
self._init_docker_networks(docker_network_prefix)
self._run_state_manager = LocalRunStateMachine(
docker_image_manager=self._image_manager,
dependency_manager=self._dependency_manager,
worker_docker_network=self.worker_docker_network,
docker_network_internal=self.docker_network_internal,
docker_network_external=self.docker_network_external,
docker_runtime=docker_runtime,
upload_bundle_callback=self._worker.upload_bundle_contents,
assign_cpu_and_gpu_sets_fn=self.assign_cpu_and_gpu_sets,
)
def _init_docker_networks(self, docker_network_prefix):
"""
Set up docker networks for runs: one with external network access and one without
"""
def create_or_get_network(name, internal):
try:
logger.debug('Creating docker network %s', name)
return self._docker.networks.create(name, internal=internal, check_duplicate=True)
except docker.errors.APIError:
logger.debug('Network %s already exists, reusing', name)
return self._docker.networks.list(names=[name])[0]
self.worker_docker_network = create_or_get_network(docker_network_prefix, True)
self.docker_network_external = create_or_get_network(docker_network_prefix + "_ext", False)
self.docker_network_internal = create_or_get_network(docker_network_prefix + "_int", True)
def save_state(self):
# Remove complex container objects from state before serializing, these can be retrieved
simple_runs = {uuid: state._replace(container=None) for uuid, state in self._runs.items()}
self._state_committer.commit(simple_runs)
def load_state(self):
runs = self._state_committer.load()
# Retrieve the complex container objects from the Docker API
for uuid, run_state in runs.items():
if run_state.container_id:
try:
run_state = run_state._replace(
container=self._docker.containers.get(run_state.container_id)
)
except docker.errors.NotFound as ex:
logger.debug('Error getting the container for the run: %s', ex)
run_state = run_state._replace(container_id=None)
finally:
self._runs[uuid] = run_state
def start(self):
"""
Load your state from disk, and start your sub-managers
"""
self.load_state()
self._image_manager.start()
self._dependency_manager.start()
def stop(self):
"""
Starts any necessary cleanup and propagates to its other managers
Blocks until cleanup is complete and it is safe to quit
"""
logger.info("Stopping Local Run Manager")
self._stop = True
self._image_manager.stop()
self._dependency_manager.stop()
self._run_state_manager.stop()
self.save_state()
try:
self.docker_network_internal.remove()
self.docker_network_external.remove()
except docker.errors.APIError as e:
logger.error("Cannot clear docker networks: {}".format(str(e)))
logger.info("Stopped Local Run Manager. Exiting")
def kill_all(self):
"""
Kills all runs
"""
logger.debug("Killing all bundles")
# Set all bundle statuses to killed
with self._lock:
for uuid in self._runs.keys():
run_state = self._runs[uuid]
run_state.info['kill_message'] = 'Worker stopped'
run_state = run_state._replace(info=run_state.info, is_killed=True)
self._runs[uuid] = run_state
# Wait until all runs finished or KILL_TIMEOUT seconds pas
for attempt in range(LocalRunManager.KILL_TIMEOUT):
with self._lock:
self._runs = {
k: v for k, v in self._runs.items() if v.stage != LocalRunStage.FINISHED
}
if len(self._runs) > 0:
logger.debug(
"Waiting for {} more bundles. {} seconds until force quit.".format(
len(self._runs), LocalRunManager.KILL_TIMEOUT - attempt
)
)
time.sleep(1)
def process_runs(self):
""" Transition each run then filter out finished runs """
with self._lock:
# transition all runs
for bundle_uuid in self._runs.keys():
run_state = self._runs[bundle_uuid]
self._runs[bundle_uuid] = self._run_state_manager.transition(run_state)
# filter out finished runs
finished_container_ids = [
run.container
for run in self._runs.values()
if (run.stage == LocalRunStage.FINISHED or run.stage == LocalRunStage.FINALIZING)
and run.container_id is not None
]
for container_id in finished_container_ids:
try:
container = self._docker.containers.get(container_id)
container.remove(force=True)
except (docker.errors.NotFound, docker.errors.NullResource):
pass
self._runs = {k: v for k, v in self._runs.items() if v.stage != LocalRunStage.FINISHED}
def create_run(self, bundle, resources):
"""
Creates and starts processing a new run with the given bundle and
resources
"""
if self._stop:
# Run Manager stopped, refuse more runs
return
bundle_uuid = bundle['uuid']
bundle_path = os.path.join(self._bundles_dir, bundle_uuid)
now = time.time()
run_state = LocalRunState(
stage=LocalRunStage.PREPARING,
run_status='',
bundle=bundle,
bundle_path=os.path.realpath(bundle_path),
resources=resources,
start_time=now,
container_id=None,
container=None,
docker_image=None,
is_killed=False,
has_contents=False,
cpuset=None,
gpuset=None,
time_used=0,
max_memory=0,
disk_utilization=0,
info={},
)
with self._lock:
self._runs[bundle_uuid] = run_state
def assign_cpu_and_gpu_sets(self, request_cpus, request_gpus):
"""
Propose a cpuset and gpuset to a bundle based on given requested resources.
Note: no side effects (this is important: we don't want to maintain more state than necessary)
Arguments:
request_cpus: integer
request_gpus: integer
Returns a 2-tuple:
cpuset: assigned cpuset (str indices).
gpuset: assigned gpuset (str indices).
Throws an exception if unsuccessful.
"""
cpuset, gpuset = set(self._cpuset), set(self._gpuset)
with self._lock:
for run_state in self._runs.values():
if run_state.stage == LocalRunStage.RUNNING:
cpuset -= run_state.cpuset
gpuset -= run_state.gpuset
if len(cpuset) < request_cpus or len(gpuset) < request_gpus:
raise Exception("Not enough cpus or gpus to assign!")
def propose_set(resource_set, request_count):
return set(str(el) for el in list(resource_set)[:request_count])
return propose_set(cpuset, request_cpus), propose_set(gpuset, request_gpus)
def get_run(self, uuid):
"""
Returns the state of the run with the given UUID if it is managed
by this RunManager, returns None otherwise
"""
with self._lock:
return self._runs.get(uuid, None)
def mark_finalized(self, uuid):
"""
Marks the run as finalized server-side so it can be discarded
"""
if uuid in self._runs:
with self._lock:
self._runs[uuid].info['finalized'] = True
def read(self, run_state, path, dep_paths, args, reply):
"""
Use your Reader helper to invoke the given read command
"""
self._reader.read(run_state, path, dep_paths, args, reply)
def write(self, run_state, path, dep_paths, string):
"""
Write `string` (string) to path in bundle with uuid.
"""
if os.path.normpath(path) in dep_paths:
return
with open(os.path.join(run_state.bundle_path, path), 'w') as f:
f.write(string)
def netcat(self, run_state, port, message, reply):
"""
Write `message` (string) to port of bundle with uuid and read the response.
Returns a stream with the response contents (bytes).
"""
# TODO: handle this in a thread since this could take a while
container_ip = docker_utils.get_container_ip(
self.worker_docker_network.name, run_state.container
)
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.connect((container_ip, port))
s.sendall(message.encode())
total_data = []
while True:
data = s.recv(LocalRunManager.NETCAT_BUFFER_SIZE)
if not data:
break
total_data.append(data)
s.close()
reply(None, {}, b''.join(total_data))
def kill(self, run_state):
"""
Kill bundle with uuid
"""
with self._lock:
run_state.info['kill_message'] = 'Kill requested'
run_state = run_state._replace(info=run_state.info, is_killed=True)
self._runs[run_state.bundle['uuid']] = run_state
@property
def all_runs(self):
"""
Returns a list of all the runs managed by this RunManager
"""
with self._lock:
result = {
bundle_uuid: {
'run_status': run_state.run_status,
'start_time': run_state.start_time,
'docker_image': run_state.docker_image,
'info': run_state.info,
'state': LocalRunStage.WORKER_STATE_TO_SERVER_STATE[run_state.stage],
'remote': self._worker.id,
}
for bundle_uuid, run_state in self._runs.items()
}
return result
@property
def all_dependencies(self):
"""
Returns a list of all dependencies available in this RunManager
"""
return self._dependency_manager.all_dependencies
@property
def cpus(self):
"""
Total number of CPUs this RunManager has
"""
return len(self._cpuset)
@property
def gpus(self):
"""
Total number of GPUs this RunManager has
"""
return len(self._gpuset)
@property
def memory_bytes(self):
"""
Total installed memory of this RunManager
"""
try:
return os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES')
except ValueError:
# Fallback to sysctl when os.sysconf('SC_PHYS_PAGES') fails on OS X
return int(check_output(['sysctl', '-n', 'hw.memsize']).strip())
@property
def free_disk_bytes(self):
"""
Available disk space by bytes of this RunManager.
"""
error_msg = "Failed to run command {}".format("df " + self._work_dir)
try:
p = Popen(["df", self._work_dir], stdout=PIPE)
output, error = p.communicate()
# Return None when there is an error.
if error:
logger.error(error.strip())
return None
if output:
lines = output.decode().split("\n")
index = lines[0].split().index("Available")
# We convert the original result from df command in unit of 1KB blocks into bytes.
return int(lines[1].split()[index]) * 1024
except Exception as e:
logger.error("{}: {}".format(error_msg, str(e)))
return None
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):
logger.info("Starting docker image manager")
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:
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)
digests = image.attrs.get('RepoDigests', [image_spec])
if len(digests) == 0:
return ImageAvailabilityState(
digest=None,
stage=DependencyStage.FAILED,
message=
'No digest available for {}, probably because it was built locally; delete the Docker image on the worker and try again'
.format(image_spec),
)
digest = digests[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'],
)