NFS-safe Dependency Manager

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,

codalab/worker/state_committer.py

@@ -2,6 +2,7 @@
 import os
 import tempfile
 import shutil
+
 from . import pyjson
 
 
@@ -20,30 +21,44 @@ def commit(self, state):
 
 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)
                     )

codalab/worker/worker.py

@@ -194,7 +194,8 @@ def save_state(self):
         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:

codalab/worker/worker_run_state.py

@@ -5,6 +5,7 @@
 import threading
 import time
 import traceback
+from typing import Dict
 
 import codalab.worker.docker_utils as docker_utils
 
@@ -251,13 +252,26 @@ def mount_dependency(dependency, shared_file_system):
 
         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}')
+                    dependencies_ready = False
+                    continue
+
                 if dependency_state.stage == DependencyStage.DOWNLOADING:
                     status_messages.append(
                         'Downloading dependency %s: %s done (archived size)'
@@ -344,7 +358,13 @@ def mount_dependency(dependency, shared_file_system):
         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 @@ def mount_dependency(dependency, shared_file_system):
             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
@@ -610,10 +617,15 @@ def remove_path_no_fail(path):
                     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 Exception:
+            return run_state
 
         # Clean up dependencies paths
         for path in run_state.paths_to_remove or []:

codalab/worker_manager/slurm_batch_worker_manager.py

@@ -267,8 +267,9 @@ def setup_codalab_worker(self, worker_id):
         else:
             work_dir_prefix = Path()
 
+        work_dir_name = f"{self.username}-codalab-slurm-worker-shared"
         worker_work_dir = work_dir_prefix.joinpath(
-            Path('{}-codalab-SlurmBatchWorkerManager-scratch'.format(self.username), worker_id)
+            Path('{}-codalab-SlurmBatchWorkerManager-scratch'.format(self.username), work_dir_name)
         )
         command = self.build_command(worker_id, str(worker_work_dir))
 

codalab_service.py

@@ -40,7 +40,7 @@
     'monitor',
     'worker-manager-cpu',
     'worker-manager-gpu',
-    'worker-singularity',
+    'worker2',
 ]
 
 ALL_NO_SERVICES = [
@@ -58,7 +58,7 @@
     'worker-manager-gpu': 'server',
     'monitor': 'server',
     'worker': 'worker',
-    'worker-singularity': 'worker',
+    'worker2': 'worker',
 }
 
 # Max timeout in seconds to wait for request to a service to get through
@@ -952,7 +952,7 @@ def start_services(self):
             self.bring_up_service('worker-shared-file-system')
         else:
             self.bring_up_service('worker')
-        self.bring_up_service('worker-singularity')
+        self.bring_up_service('worker2')
 
         self.bring_up_service('monitor')
 

docker_config/compose_files/docker-compose.yml

@@ -381,24 +381,22 @@ services:
       - rest-server
     shm_size: '500mb'
 
-  worker-singularity:
+  worker2:
     image: codalab/worker:${CODALAB_VERSION}
-    user: ${CODALAB_UID}
     command: |
       cl-worker
       --server http://rest-server:${CODALAB_REST_PORT}
       --work-dir ${CODALAB_WORKER_DIR}
       --network-prefix ${CODALAB_WORKER_NETWORK_PREFIX}
-      --id ${HOSTNAME}
-      --container-runtime singularity
+      --id ${HOSTNAME}-worker2
       --verbose
     <<: *codalab-base
+    <<: *codalab-root  # Not ideal since worker files saved as root, but without it, can't use docker
     depends_on:
       - rest-server
     volumes:
+      - /var/run/docker.sock:/var/run/docker.sock
       - ${CODALAB_WORKER_DIR}:${CODALAB_WORKER_DIR}
-      - ${CODALAB_HOME}:${CODALAB_HOME}
-      - ${CODALAB_BUNDLE_MOUNT}:${CODALAB_BUNDLE_MOUNT}
     networks:
       - service
       - worker

docs/Server-Setup.md

@@ -118,6 +118,13 @@ If you want to delete all the data associated with this, then do:
 
     ./codalab_service.py delete
 
+## Run two workers locally
+
+Sometimes you may want to run two workers locally. In that case, you should run:
+
+    ./codalab_service.py start -bds default worker2
+
+
 ## Azure Blob Storage
 
 To start the server in dev mode with Azurite (an Azure Blob Storage emulator) enabled, run:

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