Repository: codalab/codalab-worksheets

Issue:
Slurm: different workers on the same VM don't share the same cache
For the slurm worker manager

Discussion:

Patch:
diff --git a/codalab/worker/dependency_manager.py b/codalab/worker/dependency_manager.py
index d9bd496e..bfee1a96 100644
--- a/codalab/worker/dependency_manager.py
+++ b/codalab/worker/dependency_manager.py
@@ -4,9 +4,11 @@ import threading
 import traceback
 import time
 import shutil
-from contextlib import closing
+import uuid
 from collections import namedtuple
-from typing import Dict, Set
+from contextlib import closing
+from datetime import timedelta
+from typing import Dict, Set, Union, List
 
 import codalab.worker.pyjson
 from .bundle_service_client import BundleServiceClient
@@ -15,19 +17,27 @@ from codalab.worker.file_util import remove_path
 from codalab.worker.un_tar_directory import un_tar_directory
 from codalab.worker.fsm import BaseDependencyManager, DependencyStage, StateTransitioner
 from codalab.worker.worker_thread import ThreadDict
-from codalab.worker.state_committer import JsonStateCommitter
 from codalab.worker.bundle_state import DependencyKey
+from codalab.worker.state_committer import JsonStateCommitter
+
+from flufl.lock import Lock, AlreadyLockedError, NotLockedError  # noqa: E402
+
+logging.getLogger('flufl.lock').setLevel(logging.WARNING)
+
 
 logger = logging.getLogger(__name__)
 
 DependencyState = namedtuple(
-    'DependencyState', 'stage dependency_key path size_bytes dependents last_used message killed'
+    'DependencyState',
+    # downloading_by - worker id of which worker is downloading / has downloaded the dependency
+    'stage downloading_by dependency_key path size_bytes dependents last_used last_downloading message killed',
 )
 
+DependencyManagerState = Dict[str, Union[Dict[DependencyKey, DependencyState], Set[str]]]
+
 
 class DownloadAbortedException(Exception):
     """
-
     Exception raised by the download if a download is killed before it is complete
     """
 
@@ -35,13 +45,55 @@ class DownloadAbortedException(Exception):
         super(DownloadAbortedException, self).__init__(message)
 
 
+class NFSLock:
+    def __init__(self, path):
+        # Specify the path to a file that will be used to synchronize the lock.
+        # Per the flufl.lock documentation, use a file that does not exist.
+        self._lock = Lock(path)
+
+        # Locks have a lifetime (default 15 seconds) which is the period of time that the process expects
+        # to keep the lock once it has been acquired. We set the lifetime to be 5 minutes as we expect
+        # all operations that require locks to be completed within that time.
+        self._lock.lifetime = timedelta(minutes=5)
+
+        # Ensure multiple threads within a process run NFSLock operations one at a time.
+        # We must acquire the reentrant lock before acquiring the flufl lock and only release after
+        # the flufl lock is released.
+        self._r_lock = threading.RLock()
+
+    def acquire(self):
+        self._r_lock.acquire()
+        try:
+            self._lock.lock()
+        except AlreadyLockedError:
+            # Safe to re-attempt to acquire a lock
+            pass
+
+    def release(self):
+        try:
+            self._lock.unlock()
+        except NotLockedError:
+            # Safe to re-attempt to release a lock
+            pass
+        self._r_lock.release()
+
+    def __enter__(self):
+        self.acquire()
+
+    def __exit__(self, t, v, tb):
+        self.release()
+
+    @property
+    def is_locked(self):
+        return self._lock.is_locked
+
+
 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
+    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'
@@ -53,6 +105,9 @@ class DependencyManager(StateTransitioner, BaseDependencyManager):
     # 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,
@@ -66,6 +121,7 @@ class DependencyManager(StateTransitioner, BaseDependencyManager):
         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
@@ -75,110 +131,128 @@ class DependencyManager(StateTransitioner, BaseDependencyManager):
             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
+        # 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()
-        self._dependencies: Dict[DependencyKey, DependencyState] = 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)
-            )
-        )
+        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. self._dependencies, self._dependency_locks, and self._paths: populated from dependencies-state.json
-            in function _load_state()
+        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)
         """
-        # 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
-        )
+        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}.'
+                )
 
-        # 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)
+            # 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 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)
+                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._save_state()
+            # 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')
+        logger.info('Starting local dependency manager...')
 
         def loop(self):
             while not self._stop:
                 try:
-                    self._process_dependencies()
-                    self._save_state()
+                    self._transition_dependencies()
                     self._cleanup()
-                    self._save_state()
                 except Exception:
                     traceback.print_exc()
                 time.sleep(1)
@@ -187,16 +261,27 @@ class DependencyManager(StateTransitioner, BaseDependencyManager):
         self._main_thread.start()
 
     def stop(self):
-        logger.info('Stopping local dependency manager')
+        logger.info('Stopping local dependency manager...')
         self._stop = True
         self._downloading.stop()
         self._main_thread.join()
-        logger.info('Stopped local dependency manager')
+        self._state_lock.release()
+        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 _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):
         """
@@ -204,18 +289,28 @@ class DependencyManager(StateTransitioner, BaseDependencyManager):
         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()
+        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):
         """
@@ -225,19 +320,27 @@ class DependencyManager(StateTransitioner, BaseDependencyManager):
         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))
+            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 in cache, disk usage: %s (max %s), serialized size: %s (max %s)',
-                        len(self._dependencies),
+                        '%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),
@@ -245,14 +348,15 @@ class DependencyManager(StateTransitioner, BaseDependencyManager):
                     )
                     ready_deps = {
                         dep_key: dep_state
-                        for dep_key, dep_state in self._dependencies.items()
+                        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 self._dependencies.items()
+                        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
@@ -263,129 +367,112 @@ class DependencyManager(StateTransitioner, BaseDependencyManager):
                         )[0]
                     else:
                         logger.info(
-                            'Dependency quota full but there are only downloading dependencies, not cleaning up until downloads are over'
+                            '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()
+                        self._delete_dependency(dep_key_to_remove, dependencies, paths)
+                        self._commit_state(dependencies, paths)
                 else:
-                    self._release_all_locks()
                     break
 
-    def _delete_dependency(self, dependency_key):
+    def _delete_dependency(self, dep_key, dependencies, paths):
         """
         Remove the given dependency from the manager's state
-        Also delete any known files on the filesystem if any exist
+        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.
         """
-        if self._acquire_if_exists(dependency_key):
+        assert self._state_lock.is_locked
+
+        if dep_key in dependencies:
             try:
-                path_to_remove = self._dependencies[dependency_key].path
-                self._paths.remove(path_to_remove)
+                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 self._dependencies[dependency_key]
-                self._dependency_locks[dependency_key].release()
+                del dependencies[dep_key]
+                logger.info(f"Deleted dependency {dep_key}.")
 
     def has(self, dependency_key):
         """
-        Takes a DependencyKey
-        Returns true if the manager has processed this dependency
+        Takes a DependencyKey and returns true if the manager has processed this dependency
         """
-        with self._global_lock:
-            return dependency_key in self._dependencies
+        with self._state_lock:
+            dependencies: Dict[DependencyKey, DependencyState] = self._fetch_dependencies()
+            return dependency_key in dependencies
 
-    def get(self, uuid, dependency_key):
+    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
         """
-        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(
+        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(dependency_key),
+                    path=self._assign_path(paths, dependency_key),
                     size_bytes=0,
-                    dependents=set([uuid]),
+                    dependents={uuid},
                     last_used=now,
+                    last_downloading=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]
+            # 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
+        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):
+        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:
         """
-        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
+        Checks the current path against `paths`.
+        Normalize the path for the dependency by replacing / with _, avoiding conflicts.
+        Adds the new path to `paths`.
         """
-        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: 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. We have to avoid those.
-        with self._paths_lock:
-            while path in self._paths:
-                path = path + '_'
-            self._paths.add(path)
+        # 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):
@@ -413,28 +500,44 @@ class DependencyManager(StateTransitioner, BaseDependencyManager):
             raise
 
     @property
-    def all_dependencies(self):
-        with self._global_lock:
-            return list(self._dependencies.keys())
+    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 _transition_from_DOWNLOADING(self, dependency_state):
         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.
                 """
-                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),
-                    )
+                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",
+                )
 
-            # TODO(Ashwin): make this not fs-specific.
             dependency_path = os.path.join(self.dependencies_dir, dependency_state.path)
             logger.debug('Downloading dependency %s', dependency_state.dependency_key)
 
@@ -453,11 +556,11 @@ class DependencyManager(StateTransitioner, BaseDependencyManager):
                     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
+                        original_read_method = fileobj.read
                         bytes_downloaded = [0]
 
                         def interruptable_read(*args, **kwargs):
-                            data = old_read_method(*args, **kwargs)
+                            data = original_read_method(*args, **kwargs)
                             bytes_downloaded[0] += len(data)
                             update_state_and_check_killed(bytes_downloaded[0])
                             return data
@@ -465,6 +568,8 @@ class DependencyManager(StateTransitioner, BaseDependencyManager):
                         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(
@@ -473,44 +578,95 @@ class DependencyManager(StateTransitioner, BaseDependencyManager):
                         dependency_state.dependency_key,
                         dependency_path,
                     )
-                    with self._dependency_locks[dependency_state.dependency_key]:
-                        self._downloading[dependency_state.dependency_key]['success'] = True
+                    self._downloading[dependency_state.dependency_key]['success'] = True
 
                 except Exception as e:
                     attempt += 1
                     if attempt >= self._download_dependencies_max_retries:
-                        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[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 {str(e)}. Retrying up to {self._download_dependencies_max_retries} times...'
+                            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
 
-        self._downloading.add_if_new(
-            dependency_state.dependency_key, threading.Thread(target=download, args=[])
-        )
+        # 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."
+                )
 
-        if self._downloading[dependency_state.dependency_key].is_alive():
+            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
 
-        success = self._downloading[dependency_state.dependency_key]['success']
-        failure_message = self._downloading[dependency_state.dependency_key]['failure_message']
+        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:
-            with self._paths_lock:
-                self._paths.remove(dependency_state.path)
-
-            logger.error(failure_message)
+            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)
diff --git a/codalab/worker/main.py b/codalab/worker/main.py
index f73db3e4..eb9fc3b0 100644
--- a/codalab/worker/main.py
+++ b/codalab/worker/main.py
@@ -304,7 +304,9 @@ def main():
     worker = Worker(
         image_manager,
         dependency_manager,
-        os.path.join(args.work_dir, 'worker-state.json'),
+        # Include the worker ID in the worker state JSON path, so multiple workers
+        # sharing the same work directory maintain their own state.
+        os.path.join(args.work_dir, f'worker-state-{args.id}.json'),
         args.cpuset,
         args.gpuset,
         args.max_memory,
diff --git a/codalab/worker/state_committer.py b/codalab/worker/state_committer.py
index 05f1d16d..f8e9e190 100644
--- a/codalab/worker/state_committer.py
+++ b/codalab/worker/state_committer.py
@@ -2,6 +2,7 @@ import logging
 import os
 import tempfile
 import shutil
+
 from . import pyjson
 
 
@@ -20,30 +21,44 @@ class BaseStateCommitter(object):
 
 class JsonStateCommitter(BaseStateCommitter):
     def __init__(self, json_path):
-        self.temp_file = None
         self._state_file = json_path
 
+    @property
+    def path(self):
+        return self._state_file
+
+    @property
+    def state_file_exists(self) -> bool:
+        return os.path.isfile(self._state_file)
+
     def load(self, default=None):
+        """
+        Loads and reads from state file. If an error occurs, `default` will be returned, if it exists.
+        """
         try:
             with open(self._state_file) as json_data:
                 return pyjson.load(json_data)
-        except (ValueError, EnvironmentError):
-            return dict() if default is None else default
+        except (ValueError, EnvironmentError) as e:
+            if default is not None:
+                logger.warning(
+                    f"Failed to load state from {self.path} due to {e}. Returning default: {default}.",
+                    exc_info=True,
+                )
+                return default
+            logger.error(f"Failed to load state from {self.path}: {e}", exc_info=True)
+            raise e
 
     def commit(self, state):
         """ Write out the state in JSON format to a temporary file and rename it into place """
         with tempfile.NamedTemporaryFile(delete=False) as f:
             try:
-                self.temp_file = f.name
                 f.write(pyjson.dumps(state).encode())
                 f.flush()
-                shutil.copyfile(self.temp_file, self._state_file)
+                shutil.copyfile(f.name, self._state_file)
             finally:
                 try:
-                    os.unlink(self.temp_file)
+                    os.unlink(f.name)
                 except FileNotFoundError:
                     logger.error(
-                        "Problem occurred in deleting temp file {} via os.unlink".format(
-                            self.temp_file
-                        )
+                        "Problem occurred in deleting temp file {} via os.unlink".format(f.name)
                     )
diff --git a/codalab/worker/worker.py b/codalab/worker/worker.py
index fdd4baca..e564c27f 100644
--- a/codalab/worker/worker.py
+++ b/codalab/worker/worker.py
@@ -194,7 +194,8 @@ class Worker:
         self.state_committer.commit(runs)
 
     def load_state(self):
-        runs = self.state_committer.load()
+        # If the state file doesn't exist yet, have the state committer return an empty state.
+        runs = self.state_committer.load(default=dict())
         # Retrieve the complex container objects from the Docker API
         for uuid, run_state in runs.items():
             if run_state.container_id:
diff --git a/codalab/worker/worker_run_state.py b/codalab/worker/worker_run_state.py
index 3cd45aed..008c8e2c 100644
--- a/codalab/worker/worker_run_state.py
+++ b/codalab/worker/worker_run_state.py
@@ -5,6 +5,7 @@ import os
 import threading
 import time
 import traceback
+from typing import Dict
 
 import codalab.worker.docker_utils as docker_utils
 
@@ -251,13 +252,26 @@ class RunStateMachine(StateTransitioner):
 
         dependencies_ready = True
         status_messages = []
+        dependency_keys_to_paths: Dict[DependencyKey, str] = dict()
 
         if not self.shared_file_system:
             # No need to download dependencies if we're in the shared FS,
             # since they're already in our FS
             for dep in run_state.bundle.dependencies:
                 dep_key = DependencyKey(dep.parent_uuid, dep.parent_path)
-                dependency_state = self.dependency_manager.get(run_state.bundle.uuid, dep_key)
+
+                try:
+                    # Fetching dependencies from the Dependency Manager can fail.
+                    # Just update the download status on the next iteration of this transition function.
+                    dependency_state = self.dependency_manager.get(run_state.bundle.uuid, dep_key)
+                    dependency_keys_to_paths[dep_key] = os.path.join(
+                        self.dependency_manager.dependencies_dir, dependency_state.path
+                    )
+                except Exception:
+                    status_messages.append(f'Downloading dependency {dep.child_path} failed')
+                    dependencies_ready = False
+                    continue
+
                 if dependency_state.stage == DependencyStage.DOWNLOADING:
                     status_messages.append(
                         'Downloading dependency %s: %s done (archived size)'
@@ -311,7 +325,7 @@ class RunStateMachine(StateTransitioner):
                 if run_state.bundle_dir_wait_num_tries == 0:
                     message = (
                         "Bundle directory cannot be found on the shared filesystem. "
-                        "Please ensure the shared fileystem between the server and "
+                        "Please ensure the shared filesystem between the server and "
                         "your worker is mounted properly or contact your administrators."
                     )
                     log_bundle_transition(
@@ -344,7 +358,13 @@ class RunStateMachine(StateTransitioner):
         for dep in run_state.bundle.dependencies:
             full_child_path = os.path.normpath(os.path.join(run_state.bundle_path, dep.child_path))
             to_mount = []
-            dependency_path = self._get_dependency_path(run_state, dep)
+            if self.shared_file_system:
+                # TODO(Ashwin): make this not fs-specific.
+                # On a shared FS, we know where the dependency is stored and can get the contents directly
+                dependency_path = os.path.realpath(os.path.join(dep.location, dep.parent_path))
+            else:
+                dep_key = DependencyKey(dep.parent_uuid, dep.parent_path)
+                dependency_path = dependency_keys_to_paths[dep_key]
 
             if dep.child_path == RunStateMachine._CURRENT_DIRECTORY:
                 # Mount all the content of the dependency_path to the top-level of the bundle
@@ -441,19 +461,6 @@ class RunStateMachine(StateTransitioner):
             gpuset=gpuset,
         )
 
-    def _get_dependency_path(self, run_state, dependency):
-        if self.shared_file_system:
-            # TODO(Ashwin): make this not fs-specific.
-            # On a shared FS, we know where the dependency is stored and can get the contents directly
-            return os.path.realpath(os.path.join(dependency.location, dependency.parent_path))
-        else:
-            # On a dependency_manager setup, ask the manager where the dependency is
-            dep_key = DependencyKey(dependency.parent_uuid, dependency.parent_path)
-            return os.path.join(
-                self.dependency_manager.dependencies_dir,
-                self.dependency_manager.get(run_state.bundle.uuid, dep_key).path,
-            )
-
     def _transition_from_RUNNING(self, run_state):
         """
         1- Check run status of the docker container
@@ -608,10 +615,19 @@ class RunStateMachine(StateTransitioner):
                     logger.error(traceback.format_exc())
                     time.sleep(1)
 
-        for dep in run_state.bundle.dependencies:
-            if not self.shared_file_system:  # No dependencies if shared fs worker
-                dep_key = DependencyKey(dep.parent_uuid, dep.parent_path)
-                self.dependency_manager.release(run_state.bundle.uuid, dep_key)
+        try:
+            # Fetching dependencies from the Dependency Manager can fail.
+            # Finish cleaning up on the next iteration of this transition function.
+            for dep in run_state.bundle.dependencies:
+                if not self.shared_file_system:  # No dependencies if shared fs worker
+                    dep_key = DependencyKey(dep.parent_uuid, dep.parent_path)
+                    self.dependency_manager.release(run_state.bundle.uuid, dep_key)
+        except (ValueError, EnvironmentError):
+            # Do nothing if an error is thrown while reading from the state file
+            logging.exception(
+                f"Error reading from dependencies state file while releasing a dependency from {run_state.bundle.uuid}"
+            )
+            return run_state
 
         # Clean up dependencies paths
         for path in run_state.paths_to_remove or []:
diff --git a/codalab/worker_manager/slurm_batch_worker_manager.py b/codalab/worker_manager/slurm_batch_worker_manager.py
index c8ca6b57..f6f18f69 100644
--- a/codalab/worker_manager/slurm_batch_worker_manager.py
+++ b/codalab/worker_manager/slurm_batch_worker_manager.py
@@ -268,7 +268,7 @@ class SlurmBatchWorkerManager(WorkerManager):
             work_dir_prefix = Path()
 
         worker_work_dir = work_dir_prefix.joinpath(
-            Path('{}-codalab-SlurmBatchWorkerManager-scratch'.format(self.username), worker_id)
+            Path('{}-codalab-SlurmBatchWorkerManager-scratch'.format(self.username), "workdir")
         )
         command = self.build_command(worker_id, str(worker_work_dir))
 
diff --git a/docs/Worker-Managers.md b/docs/Worker-Managers.md
index 09342464..d7df0932 100644
--- a/docs/Worker-Managers.md
+++ b/docs/Worker-Managers.md
@@ -12,6 +12,17 @@ We support the following Worker Managers:
 | slurm-batch      | Worker manager for submitting jobs using Slurm Batch.        | 
 | kubernetes       | Worker manager for submitting jobs to a Kubernetes cluster.  |
 
+## Setting a shared cache
+
+To use a shared cache among workers, have all the workers use the same working directory by specifying
+the same path for `--work-dir`. The working directory is set by `-worker-work-dir-prefix` when starting 
+a worker manager. The dependency managers can be used over NFS, so a working directory can be on a network disk.
+
+```commandline
+cl-worker-manager --worker-work-dir-prefix /juice  slurm-batch --cpus 4 --gpus 1 --memory-mb 16000 
+```
+
+In the example worker manager command above, `juice` is a directory on a network disk.
 
 ## AWS Batch Worker Manager
 
diff --git a/requirements.txt b/requirements.txt
index 8f0acfaf..aaed0d8d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -28,3 +28,4 @@ wheel==0.35.1
 urllib3==1.26.5
 retry==0.9.2
 spython==0.1.14
+flufl.lock==6.0
\ No newline at end of file


Test patch:
diff --git a/tests/unit/worker/dependency_manager_test.py b/tests/unit/worker/dependency_manager_test.py
new file mode 100644
index 00000000..0153c92c
--- /dev/null
+++ b/tests/unit/worker/dependency_manager_test.py
@@ -0,0 +1,147 @@
+import os
+import time
+import unittest
+import shutil
+import tempfile
+from concurrent.futures import ProcessPoolExecutor
+from unittest.mock import MagicMock
+
+from codalab.worker.bundle_state import DependencyKey
+
+try:
+    from codalab.worker.dependency_manager import DependencyManager
+
+    module_failed = False
+except ImportError:
+    module_failed = True
+
+
+class DependencyManagerTest(unittest.TestCase):
+    def setUp(self):
+        if module_failed:
+            self.skipTest('Issue with ratarmountcore.')
+
+        self.work_dir = tempfile.mkdtemp()
+        self.state_path = os.path.join(self.work_dir, "dependencies-state.json")
+        self.dependency_manager = DependencyManager(
+            commit_file=self.state_path,
+            bundle_service=None,
+            worker_dir=self.work_dir,
+            max_cache_size_bytes=1024,
+            download_dependencies_max_retries=1,
+        )
+
+    def tearDown(self):
+        shutil.rmtree(self.work_dir)
+
+    def test_get_has(self):
+        dependent_uuid = "0x2"
+        dependency_key = DependencyKey(parent_uuid="0x1", parent_path="parent")
+        state = self.dependency_manager.get(dependent_uuid, dependency_key)
+        self.assertTrue(self.dependency_manager.has(dependency_key))
+        self.assertEqual(state.stage, "DOWNLOADING")
+        self.assertEqual(state.path, "0x1_parent")
+        self.assertEqual(state.dependents, {dependent_uuid})
+
+    def test_release(self):
+        dependency_key = DependencyKey(parent_uuid="0x1", parent_path="parent")
+        self.dependency_manager.get("0x2", dependency_key)
+        state = self.dependency_manager.get("0x3", dependency_key)
+        # Passing in the same dependency key with a different dependent, will just add the dependent
+        self.assertEqual(state.dependents, {"0x2", "0x3"})
+
+        # Release 0x2 as a dependent
+        self.dependency_manager.release("0x2", dependency_key)
+        with self.dependency_manager._state_lock:
+            dependencies = self.dependency_manager._fetch_dependencies()
+        state = dependencies[dependency_key]
+        self.assertEqual(state.dependents, {"0x3"})
+
+        # Release 0x3 as a dependent - should be left with no dependents
+        self.dependency_manager.release("0x3", dependency_key)
+        with self.dependency_manager._state_lock:
+            dependencies = self.dependency_manager._fetch_dependencies()
+        state = dependencies[dependency_key]
+        self.assertEqual(len(state.dependents), 0)
+
+    def test_all_dependencies(self):
+        dependency_key = DependencyKey(parent_uuid="0x1", parent_path="parent")
+        self.dependency_manager.get("0x2", dependency_key)
+        dependency_key = DependencyKey(parent_uuid="0x3", parent_path="parent2")
+        self.dependency_manager.get("0x4", dependency_key)
+        dependency_keys = self.dependency_manager.all_dependencies
+        self.assertEqual(len(dependency_keys), 2)
+
+    @unittest.skip(
+        "Flufl.lock doesn't seem to work on GHA for some reason, "
+        "even though this test passes on other machines."
+    )
+    def test_concurrency(self):
+        num_of_dependency_managers = 10
+        executor = ProcessPoolExecutor(max_workers=num_of_dependency_managers)
+
+        random_file_path = os.path.join(self.work_dir, "random_file")
+        with open(random_file_path, "wb") as f:
+            f.seek((1024 * 1024 * 1024) - 1)  # 1 GB
+            f.write(b"\0")
+
+        futures = [
+            executor.submit(task, self.work_dir, self.state_path, random_file_path)
+            for _ in range(num_of_dependency_managers)
+        ]
+        for future in futures:
+            print(future.result())
+            self.assertIsNone(future.exception())
+        executor.shutdown()
+
+
+def task(work_dir, state_path, random_file_path):
+    """
+    Runs the end-to-end workflow of the Dependency Manager.
+    Note: ProcessPoolExecutor must serialize everything before sending it to the worker,
+          so this function needs to be defined at the top-level.
+    # """
+    # Mock Bundle Service to return a random file object
+    mock_bundle_service = MagicMock()
+    mock_bundle_service.get_bundle_info = MagicMock(return_value={'type': "file"})
+    file_obj = open(random_file_path, "rb")
+    mock_bundle_service.get_bundle_contents = MagicMock(return_value=file_obj)
+
+    # Create and start a dependency manager
+    process_id = os.getpid()
+    print(f"{process_id}: Starting a DependencyManager...")
+    dependency_manager = DependencyManager(
+        commit_file=state_path,
+        bundle_service=mock_bundle_service,
+        worker_dir=work_dir,
+        max_cache_size_bytes=2048,
+        download_dependencies_max_retries=1,
+    )
+    dependency_manager.start()
+    print(f"{process_id}: Started with work directory: {work_dir}.")
+
+    # Register a run's UUID as a dependent of a parent bundle with UUID 0x1
+    dependency_key = DependencyKey(parent_uuid="0x1", parent_path="parent")
+    run_uuid = f"0x{process_id}"
+    state = dependency_manager.get(run_uuid, dependency_key)
+    assert (
+        run_uuid in state.dependents
+    ), f"{process_id}: Expected {run_uuid} as one of the dependents."
+
+    # Release the run bundle as a dependent
+    dependency_manager.release(run_uuid, dependency_key)
+    dependencies = dependency_manager._fetch_dependencies()
+    if dependency_key in dependencies:
+        state = dependencies[dependency_key]
+        print(f"{process_id}: Checking {run_uuid} in {state.dependents}")
+        assert (
+            run_uuid not in state.dependents
+        ), f"{process_id}: Dependent should not be in the list of dependents after unregistering."
+
+    # Keep the dependency manager running for some time to test the loop
+    time.sleep(30)
+
+    # Stop the Dependency Manager
+    print(f"{process_id}: Stopping DependencyManger...")
+    dependency_manager.stop()
+    print(f"{process_id}: Done.")
diff --git a/tests/unit/worker/state_committer_test.py b/tests/unit/worker/state_committer_test.py
index 4b0de32d..b4b481bb 100644
--- a/tests/unit/worker/state_committer_test.py
+++ b/tests/unit/worker/state_committer_test.py
@@ -21,7 +21,12 @@ class JsonStateCommitterTest(unittest.TestCase):
 
     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)
+        self.assertEqual(self.committer.path, self.state_path)
+
+    def test_state_file_exists(self):
+        self.assertFalse(self.committer.state_file_exists)
+        self.committer.commit({'state': 'value'})
+        self.assertTrue(self.committer.state_file_exists)
 
     def test_commit(self):
         """Make sure state is committed correctly"""
@@ -30,7 +35,6 @@ class JsonStateCommitterTest(unittest.TestCase):
         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 """
diff --git a/tests/unit/worker_manager/slurm_batch_worker_manager_test.py b/tests/unit/worker_manager/slurm_batch_worker_manager_test.py
index 863f63ca..7b4d40ff 100644
--- a/tests/unit/worker_manager/slurm_batch_worker_manager_test.py
+++ b/tests/unit/worker_manager/slurm_batch_worker_manager_test.py
@@ -40,7 +40,7 @@ class SlurmBatchWorkerManagerTest(unittest.TestCase):
 
         expected_command_str = (
             "cl-worker --server some_server --verbose --exit-when-idle --idle-seconds 888 "
-            "--work-dir /some/path/some_user-codalab-SlurmBatchWorkerManager-scratch/some_worker_id "
+            "--work-dir /some/path/some_user-codalab-SlurmBatchWorkerManager-scratch/workdir "
             "--id $(hostname -s)-some_worker_id --network-prefix cl_worker_some_worker_id_network --tag some_tag "
             "--group some_group --exit-after-num-runs 8 --download-dependencies-max-retries 5 "
             "--max-work-dir-size 88g --checkin-frequency-seconds 30 --shared-memory-size-gb 10 "