class RunStateMachine(StateTransitioner):
"""
Manages the state machine of the runs running on the local machine
Note that in general there are two types of errors:
- User errors (fault of bundle) - we fail the bundle (move to CLEANING_UP state).
- System errors (fault of worker) - we freeze this worker (Exception is thrown up).
It's not always clear where the line is.
"""
_ROOT = '/'
_CURRENT_DIRECTORY = '.'
RESTAGED_REASON = 'The bundle is not in terminal states {READY, FAILED, KILLED} when the worker checks termination'
def __init__(
self,
image_manager, # Component to request docker images from
dependency_manager, # Component to request dependency downloads from
worker_docker_network, # Docker network to add all bundles to
docker_network_internal, # Docker network to add non-net connected bundles to
docker_network_external, # Docker network to add internet connected bundles to
docker_runtime, # Docker runtime to use for containers (nvidia or runc)
upload_bundle_callback, # Function to call to upload bundle results to the server
assign_cpu_and_gpu_sets_fn, # Function to call to assign CPU and GPU resources to each run
shared_file_system, # If True, bundle mount is shared with server
):
super(RunStateMachine, self).__init__()
self.add_transition(RunStage.PREPARING,
self._transition_from_PREPARING)
self.add_transition(RunStage.RUNNING, self._transition_from_RUNNING)
self.add_transition(RunStage.CLEANING_UP,
self._transition_from_CLEANING_UP)
self.add_transition(RunStage.UPLOADING_RESULTS,
self._transition_from_UPLOADING_RESULTS)
self.add_transition(RunStage.FINALIZING,
self._transition_from_FINALIZING)
self.add_terminal(RunStage.FINISHED)
self.add_terminal(RunStage.RESTAGED)
self.dependency_manager = dependency_manager
self.image_manager = image_manager
self.worker_docker_network = worker_docker_network
self.docker_network_external = docker_network_external
self.docker_network_internal = docker_network_internal
# todo aditya: docker_runtime will be None if the worker is a singularity worker. handle this.
self.docker_runtime = docker_runtime
# bundle.uuid -> {'thread': Thread, 'run_status': str}
self.uploading = ThreadDict(fields={
'run_status': 'Upload started',
'success': False
})
# bundle.uuid -> {'thread': Thread, 'disk_utilization': int, 'running': bool}
self.disk_utilization = ThreadDict(fields={
'disk_utilization': 0,
'running': True,
'lock': None
})
self.upload_bundle_callback = upload_bundle_callback
self.assign_cpu_and_gpu_sets_fn = assign_cpu_and_gpu_sets_fn
self.shared_file_system = shared_file_system
def stop(self):
for uuid in self.disk_utilization.keys():
self.disk_utilization[uuid]['running'] = False
self.disk_utilization.stop()
self.uploading.stop()
def _transition_from_PREPARING(self, run_state):
"""
1- Request the docker image from docker image manager
- if image is failed, move to CLEANING_UP state
2- Request the dependencies from dependency manager
- if any are failed, move to CLEANING_UP state
3- If all dependencies and docker image are ready:
- Set up the local filesystem for the run
- Create symlinks to dependencies
- Allocate resources and prepare the docker container
- Start the docker container
4- If all is successful, move to RUNNING state
"""
def mount_dependency(dependency, shared_file_system):
if not shared_file_system:
# Set up symlinks for the content at dependency path
Path(dependency.child_path).parent.mkdir(parents=True,
exist_ok=True)
os.symlink(dependency.docker_path, dependency.child_path)
# The following will be converted into a Docker volume binding like:
# dependency_path:docker_dependency_path:ro
docker_dependencies.append(
(dependency.parent_path, dependency.docker_path))
if run_state.is_killed or run_state.is_restaged:
log_bundle_transition(
bundle_uuid=run_state.bundle.uuid,
previous_stage=run_state.stage,
next_stage=RunStage.CLEANING_UP,
reason=
f'the bundle was {"killed" if run_state.is_killed else "restaged"}',
)
return run_state._replace(stage=RunStage.CLEANING_UP)
# Check CPU and GPU availability
try:
cpuset, gpuset = self.assign_cpu_and_gpu_sets_fn(
run_state.resources.cpus, run_state.resources.gpus)
except Exception as e:
message = "Unexpectedly unable to assign enough resources to bundle {}: {}".format(
run_state.bundle.uuid, str(e))
logger.error(message)
logger.error(traceback.format_exc())
return run_state._replace(run_status=message)
dependencies_ready = True
status_messages = []
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)
if dependency_state.stage == DependencyStage.DOWNLOADING:
status_messages.append(
'Downloading dependency %s: %s done (archived size)' %
(dep.child_path, size_str(
dependency_state.size_bytes)))
dependencies_ready = False
elif dependency_state.stage == DependencyStage.FAILED:
# Failed to download dependency; -> CLEANING_UP
log_bundle_transition(
bundle_uuid=run_state.bundle.uuid,
previous_stage=run_state.stage,
next_stage=RunStage.CLEANING_UP,
reason=
f'Dependency has failed for this bundle. Dependency child uuid: {dep.child_uuid}. Dependency child path: {dep.child_path}',
)
return run_state._replace(
stage=RunStage.CLEANING_UP,
failure_message='Failed to download dependency %s: %s'
% (dep.child_path, dependency_state.message),
)
# get the docker image
docker_image = run_state.resources.docker_image
image_state = self.image_manager.get(docker_image)
if image_state.stage == DependencyStage.DOWNLOADING:
status_messages.append('Pulling docker image %s %s' %
(docker_image, image_state.message))
dependencies_ready = False
elif image_state.stage == DependencyStage.FAILED:
# Failed to pull image; -> CLEANING_UP
message = 'Failed to download Docker image: %s' % image_state.message
logger.error(message)
return run_state._replace(stage=RunStage.CLEANING_UP,
failure_message=message)
# stop proceeding if dependency and image downloads aren't all done
if not dependencies_ready:
status_message = status_messages.pop()
if status_messages:
status_message += "(and downloading %d other dependencies and docker images)" % len(
status_messages)
logger.info(
f'bundle is not ready yet. uuid: {run_state.bundle.uuid}. status message: {status_message}'
)
return run_state._replace(run_status=status_message)
# All dependencies ready! Set up directories, symlinks and container. Start container.
# 1) Set up a directory to store the bundle.
if self.shared_file_system:
if not os.path.exists(run_state.bundle_path):
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 "
"your worker is mounted properly or contact your administrators."
)
log_bundle_transition(
bundle_uuid=run_state.bundle.uuid,
previous_stage=run_state.stage,
next_stage=RunStage.CLEANING_UP,
reason=
"Bundle directory cannot be found on the shared filesystem.",
)
return run_state._replace(stage=RunStage.CLEANING_UP,
failure_message=message)
next_bundle_dir_wait_num_tries = run_state.bundle_dir_wait_num_tries - 1
logger.info(
f'Waiting for bundle directory to be created by the server, uuid: {run_state.bundle.uuid}, bundle_dir_wait_num_tries: {next_bundle_dir_wait_num_tries}'
)
return run_state._replace(
run_status=
"Waiting for bundle directory to be created by the server",
bundle_dir_wait_num_tries=next_bundle_dir_wait_num_tries,
)
else:
remove_path(run_state.bundle_path)
os.makedirs(run_state.bundle_path)
# 2) Set up symlinks
docker_dependencies = []
docker_dependencies_path = (
RunStateMachine._ROOT + run_state.bundle.uuid +
('_dependencies' if not self.shared_file_system else ''))
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 dep.child_path == RunStateMachine._CURRENT_DIRECTORY:
# Mount all the content of the dependency_path to the top-level of the bundle
for child in os.listdir(dependency_path):
child_path = os.path.normpath(
os.path.join(run_state.bundle_path, child))
to_mount.append(
DependencyToMount(
docker_path=os.path.join(docker_dependencies_path,
child),
child_path=child_path,
parent_path=os.path.join(dependency_path, child),
))
run_state = run_state._replace(
paths_to_remove=(run_state.paths_to_remove or []) +
[child_path])
else:
to_mount.append(
DependencyToMount(
docker_path=os.path.join(docker_dependencies_path,
dep.child_path),
child_path=full_child_path,
parent_path=dependency_path,
))
first_element_of_path = Path(dep.child_path).parts[0]
if first_element_of_path == RunStateMachine._ROOT:
run_state = run_state._replace(
paths_to_remove=(run_state.paths_to_remove or []) +
[full_child_path])
else:
# child_path can be a nested path, so later remove everything from the first element of the path
path_to_remove = os.path.join(run_state.bundle_path,
first_element_of_path)
run_state = run_state._replace(
paths_to_remove=(run_state.paths_to_remove or []) +
[path_to_remove])
for dependency in to_mount:
try:
mount_dependency(dependency, self.shared_file_system)
except OSError as e:
log_bundle_transition(
bundle_uuid=run_state.bundle.uuid,
previous_stage=run_state.stage,
next_stage=RunStage.CLEANING_UP,
reason=str(e.__class__),
level=logging.ERROR,
)
return run_state._replace(stage=RunStage.CLEANING_UP,
failure_message=str(e))
if run_state.resources.network:
docker_network = self.docker_network_external.name
else:
docker_network = self.docker_network_internal.name
# 3) Start container
try:
container = docker_utils.start_bundle_container(
run_state.bundle_path,
run_state.bundle.uuid,
docker_dependencies,
run_state.bundle.command,
run_state.resources.docker_image,
network=docker_network,
cpuset=cpuset,
gpuset=gpuset,
memory_bytes=run_state.resources.memory,
runtime=self.docker_runtime,
)
self.worker_docker_network.connect(container)
except docker_utils.DockerUserErrorException as e:
message = 'Cannot start Docker container: {}'.format(e)
log_bundle_transition(
bundle_uuid=run_state.bundle.uuid,
previous_stage=run_state.stage,
next_stage=RunStage.CLEANING_UP,
reason='Cannot start Docker container.',
level=logging.ERROR,
)
return run_state._replace(stage=RunStage.CLEANING_UP,
failure_message=message)
except Exception as e:
message = 'Cannot start container: {}'.format(e)
logger.error(message)
logger.error(traceback.format_exc())
raise
return run_state._replace(
stage=RunStage.RUNNING,
run_status='Running job in container',
container_id=container.id,
container=container,
docker_image=image_state.digest,
has_contents=True,
cpuset=cpuset,
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
2- If run is killed, kill the container
3- If run is finished, move to CLEANING_UP state
"""
def check_and_report_finished(run_state):
try:
finished, exitcode, failure_msg = docker_utils.check_finished(
run_state.container)
except docker_utils.DockerException:
logger.error(traceback.format_exc())
finished, exitcode, failure_msg = False, None, None
return run_state._replace(finished=finished,
exitcode=exitcode,
failure_message=failure_msg)
def check_resource_utilization(run_state: RunState):
logger.info(
f'Checking resource utilization for bundle. uuid: {run_state.bundle.uuid}'
)
cpu_usage, memory_usage = docker_utils.get_container_stats_with_docker_stats(
run_state.container)
run_state = run_state._replace(cpu_usage=cpu_usage,
memory_usage=memory_usage)
run_state = run_state._replace(memory_usage=memory_usage)
kill_messages = []
run_stats = docker_utils.get_container_stats(run_state.container)
run_state = run_state._replace(max_memory=max(
run_state.max_memory, run_stats.get('memory', 0)))
run_state = run_state._replace(
disk_utilization=self.disk_utilization[
run_state.bundle.uuid]['disk_utilization'])
container_time_total = docker_utils.get_container_running_time(
run_state.container)
run_state = run_state._replace(
container_time_total=container_time_total,
container_time_user=run_stats.get(
'container_time_user', run_state.container_time_user),
container_time_system=run_stats.get(
'container_time_system', run_state.container_time_system),
)
if run_state.resources.time and container_time_total > run_state.resources.time:
kill_messages.append(
'Time limit exceeded. (Container uptime %s > time limit %s)'
% (duration_str(container_time_total),
duration_str(run_state.resources.time)))
if run_state.max_memory > run_state.resources.memory or run_state.exitcode == 137:
kill_messages.append('Memory limit %s exceeded.' %
size_str(run_state.resources.memory))
if run_state.resources.disk and run_state.disk_utilization > run_state.resources.disk:
kill_messages.append('Disk limit %sb exceeded.' %
size_str(run_state.resources.disk))
if kill_messages:
run_state = run_state._replace(
kill_message=' '.join(kill_messages), is_killed=True)
return run_state
def check_disk_utilization():
logger.info(
f'Checking disk utilization for bundle. uuid: {run_state.bundle.uuid}'
)
running = True
while running:
start_time = time.time()
try:
disk_utilization = get_path_size(run_state.bundle_path)
self.disk_utilization[run_state.bundle.uuid][
'disk_utilization'] = disk_utilization
running = self.disk_utilization[
run_state.bundle.uuid]['running']
except Exception:
logger.error(traceback.format_exc())
end_time = time.time()
# To ensure that we don't hammer the disk for this computation when
# there are lots of files, we run it at most 10% of the time.
time.sleep(max((end_time - start_time) * 10, 1.0))
self.disk_utilization.add_if_new(
run_state.bundle.uuid,
threading.Thread(target=check_disk_utilization, args=[]))
run_state = check_and_report_finished(run_state)
run_state = check_resource_utilization(run_state)
if run_state.is_killed or run_state.is_restaged:
log_bundle_transition(
bundle_uuid=run_state.bundle.uuid,
previous_stage=run_state.stage,
next_stage=RunStage.CLEANING_UP,
reason=
f'the bundle was {"killed" if run_state.is_killed else "restaged"}',
)
if docker_utils.container_exists(run_state.container):
try:
run_state.container.kill()
except docker.errors.APIError:
finished, _, _ = docker_utils.check_finished(
run_state.container)
if not finished:
logger.error(traceback.format_exc())
self.disk_utilization[run_state.bundle.uuid]['running'] = False
self.disk_utilization.remove(run_state.bundle.uuid)
return run_state._replace(stage=RunStage.CLEANING_UP)
if run_state.finished:
logger.debug(
'Finished run with UUID %s, exitcode %s, failure_message %s',
run_state.bundle.uuid,
run_state.exitcode,
run_state.failure_message,
)
self.disk_utilization[run_state.bundle.uuid]['running'] = False
self.disk_utilization.remove(run_state.bundle.uuid)
return run_state._replace(stage=RunStage.CLEANING_UP,
run_status='Uploading results')
else:
return run_state
def _transition_from_CLEANING_UP(self, run_state):
"""
1- delete the container if still existent
2- clean up the dependencies from bundle directory
3- release the dependencies in dependency manager
4- If bundle has contents to upload (i.e. was RUNNING at some point),
move to UPLOADING_RESULTS state
Otherwise move to FINALIZING state
"""
def remove_path_no_fail(path):
try:
remove_path(path)
except Exception:
logger.error(traceback.format_exc())
if run_state.container_id is not None:
while docker_utils.container_exists(run_state.container):
try:
finished, _, _ = docker_utils.check_finished(
run_state.container)
if finished:
run_state.container.remove(force=True)
run_state = run_state._replace(container=None,
container_id=None)
break
else:
try:
run_state.container.kill()
except docker.errors.APIError:
logger.error(traceback.format_exc())
time.sleep(1)
except docker.errors.APIError:
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)
# Clean up dependencies paths
for path in run_state.paths_to_remove or []:
remove_path_no_fail(path)
run_state = run_state._replace(paths_to_remove=[])
if run_state.is_restaged:
log_bundle_transition(
bundle_uuid=run_state.bundle.uuid,
previous_stage=run_state.stage,
next_stage=RunStage.RESTAGED,
reason=self.RESTAGED_REASON,
)
return run_state._replace(stage=RunStage.RESTAGED)
if not self.shared_file_system and run_state.has_contents:
log_bundle_transition(
bundle_uuid=run_state.bundle.uuid,
previous_stage=run_state.stage,
next_stage=RunStage.UPLOADING_RESULTS,
)
return run_state._replace(stage=RunStage.UPLOADING_RESULTS,
run_status='Uploading results',
container=None)
else:
# No need to upload results since results are directly written to bundle store
# Delete any files that match the exclude_patterns .
for exclude_pattern in run_state.bundle.metadata[
"exclude_patterns"]:
full_pattern = os.path.join(run_state.bundle_path,
exclude_pattern)
for file_path in glob.glob(full_pattern, recursive=True):
# Only remove files that are subpaths of run_state.bundle_path, in case
# that exclude_pattern is something like "../../../".
if path_is_parent(parent_path=run_state.bundle_path,
child_path=file_path):
remove_path(file_path)
return self.finalize_run(run_state)
def _transition_from_UPLOADING_RESULTS(self, run_state):
"""
If bundle not already uploading:
Use the RunManager API to upload contents at bundle_path to the server
Pass the callback to that API such that if the bundle is killed during the upload,
the callback returns false, allowing killable uploads.
If uploading and not finished:
Update run_status with upload progress
If uploading and finished:
Move to FINALIZING state
"""
if run_state.is_restaged:
log_bundle_transition(
bundle_uuid=run_state.bundle.uuid,
previous_stage=run_state.stage,
next_stage=RunStage.RESTAGED,
reason=self.RESTAGED_REASON,
)
return run_state._replace(stage=RunStage.RESTAGED)
def upload_results():
try:
# Upload results
logger.debug('Uploading results for run with UUID %s',
run_state.bundle.uuid)
def progress_callback(bytes_uploaded):
run_status = 'Uploading results: %s done (archived size)' % size_str(
bytes_uploaded)
self.uploading[
run_state.bundle.uuid]['run_status'] = run_status
return True
self.upload_bundle_callback(
run_state.bundle.uuid,
run_state.bundle_path,
run_state.bundle.metadata["exclude_patterns"],
progress_callback,
)
self.uploading[run_state.bundle.uuid]['success'] = True
except Exception as e:
self.uploading[run_state.bundle.uuid]['run_status'] = (
"Error while uploading: %s" % e)
logger.error(traceback.format_exc())
self.uploading.add_if_new(
run_state.bundle.uuid,
threading.Thread(target=upload_results, args=[]))
if self.uploading[run_state.bundle.uuid].is_alive():
return run_state._replace(
run_status=self.uploading[run_state.bundle.uuid]['run_status'])
elif not self.uploading[run_state.bundle.uuid]['success']:
# upload failed
failure_message = run_state.failure_message
if failure_message:
failure_message = (
f'{failure_message}. {self.uploading[run_state.bundle.uuid]["run_status"]}'
)
else:
failure_message = self.uploading[
run_state.bundle.uuid]['run_status']
logger.info(
f'Upload failed. uuid: {run_state.bundle.uuid}. failure message: {failure_message}'
)
run_state = run_state._replace(failure_message=failure_message)
self.uploading.remove(run_state.bundle.uuid)
return self.finalize_run(run_state)
def finalize_run(self, run_state):
"""
Prepare the finalize message to be sent with the next checkin
"""
if run_state.is_killed:
# Append kill_message, which contains more useful info on why a run was killed, to the failure message.
failure_message = ("{}. {}".format(run_state.failure_message,
run_state.kill_message)
if run_state.failure_message else
run_state.kill_message)
log_bundle_transition(
bundle_uuid=run_state.bundle.uuid,
previous_stage=run_state.stage,
next_stage=RunStage.FINALIZING,
reason=
f'Bundle is killed. uuid: {run_state.bundle.uuid}. failure message: {failure_message}',
)
run_state = run_state._replace(failure_message=failure_message)
else:
log_bundle_transition(
bundle_uuid=run_state.bundle.uuid,
previous_stage=run_state.stage,
next_stage=RunStage.FINALIZING,
)
return run_state._replace(stage=RunStage.FINALIZING,
run_status="Finalizing bundle")
def _transition_from_FINALIZING(self, run_state):
"""
If a full worker cycle has passed since we got into the FINALIZING state we already reported to
server, if bundle is going be sent back to the server, move on to the RESTAGED state. Otherwise,
move on to the FINISHED state. Can also remove bundle_path now.
"""
if run_state.is_restaged:
log_bundle_transition(
bundle_uuid=run_state.bundle.uuid,
previous_stage=run_state.stage,
next_stage=RunStage.RESTAGED,
reason=
'the bundle is restaged, as `pass-down-termination` is specified for worker',
)
return run_state._replace(stage=RunStage.RESTAGED)
elif run_state.finalized:
if not self.shared_file_system:
remove_path(
run_state.bundle_path) # don't remove bundle if shared FS
return run_state._replace(stage=RunStage.FINISHED,
run_status='Finished')
else:
return run_state
class RunStateMachine(StateTransitioner):
"""
Manages the state machine of the runs running on the local machine
Note that in general there are two types of errors:
- User errors (fault of bundle) - we fail the bundle (move to CLEANING_UP state).
- System errors (fault of worker) - we freeze this worker (Exception is thrown up).
It's not always clear where the line is.
"""
def __init__(
self,
docker_image_manager, # Component to request docker images from
dependency_manager, # Component to request dependency downloads from
worker_docker_network, # Docker network to add all bundles to
docker_network_internal, # Docker network to add non-net connected bundles to
docker_network_external, # Docker network to add internet connected bundles to
docker_runtime, # Docker runtime to use for containers (nvidia or runc)
upload_bundle_callback, # Function to call to upload bundle results to the server
assign_cpu_and_gpu_sets_fn, # Function to call to assign CPU and GPU resources to each run
shared_file_system, # If True, bundle mount is shared with server
):
super(RunStateMachine, self).__init__()
self.add_transition(RunStage.PREPARING,
self._transition_from_PREPARING)
self.add_transition(RunStage.RUNNING, self._transition_from_RUNNING)
self.add_transition(RunStage.CLEANING_UP,
self._transition_from_CLEANING_UP)
self.add_transition(RunStage.UPLOADING_RESULTS,
self._transition_from_UPLOADING_RESULTS)
self.add_transition(RunStage.FINALIZING,
self._transition_from_FINALIZING)
self.add_terminal(RunStage.FINISHED)
self.dependency_manager = dependency_manager
self.docker_image_manager = docker_image_manager
self.worker_docker_network = worker_docker_network
self.docker_network_external = docker_network_external
self.docker_network_internal = docker_network_internal
self.docker_runtime = docker_runtime
# bundle.uuid -> {'thread': Thread, 'run_status': str}
self.uploading = ThreadDict(fields={
'run_status': 'Upload started',
'success': False
})
# bundle.uuid -> {'thread': Thread, 'disk_utilization': int, 'running': bool}
self.disk_utilization = ThreadDict(fields={
'disk_utilization': 0,
'running': True,
'lock': None
})
self.upload_bundle_callback = upload_bundle_callback
self.assign_cpu_and_gpu_sets_fn = assign_cpu_and_gpu_sets_fn
self.shared_file_system = shared_file_system
def stop(self):
for uuid in self.disk_utilization.keys():
self.disk_utilization[uuid]['running'] = False
self.disk_utilization.stop()
self.uploading.stop()
def _transition_from_PREPARING(self, run_state):
"""
1- Request the docker image from docker image manager
- if image is failed, move to CLEANING_UP state
2- Request the dependencies from dependency manager
- if any are failed, move to CLEANING_UP state
3- If all dependencies and docker image are ready:
- Set up the local filesystem for the run
- Create symlinks to dependencies
- Allocate resources and prepare the docker container
- Start the docker container
4- If all is successful, move to RUNNING state
"""
if run_state.is_killed:
return run_state._replace(stage=RunStage.CLEANING_UP)
dependencies_ready = True
status_messages = []
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_key, dep in run_state.bundle.dependencies.items():
dependency_state = self.dependency_manager.get(
run_state.bundle.uuid, dep_key)
if dependency_state.stage == DependencyStage.DOWNLOADING:
status_messages.append(
'Downloading dependency %s: %s done (archived size)' %
(dep.child_path, size_str(
dependency_state.size_bytes)))
dependencies_ready = False
elif dependency_state.stage == DependencyStage.FAILED:
# Failed to download dependency; -> CLEANING_UP
return run_state._replace(
stage=RunStage.CLEANING_UP,
failure_message='Failed to download dependency %s: %s'
% (dep.child_path, dependency_state.message),
)
# get the docker image
docker_image = run_state.resources.docker_image
image_state = self.docker_image_manager.get(docker_image)
if image_state.stage == DependencyStage.DOWNLOADING:
status_messages.append('Pulling docker image: ' +
(image_state.message or docker_image or ""))
dependencies_ready = False
elif image_state.stage == DependencyStage.FAILED:
# Failed to pull image; -> CLEANING_UP
message = 'Failed to download Docker image: %s' % image_state.message
logger.error(message)
return run_state._replace(stage=RunStage.CLEANING_UP,
failure_message=message)
# stop proceeding if dependency and image downloads aren't all done
if not dependencies_ready:
status_message = status_messages.pop()
if status_messages:
status_message += "(and downloading %d other dependencies and docker images)" % len(
status_messages)
return run_state._replace(run_status=status_message)
# All dependencies ready! Set up directories, symlinks and container. Start container.
# 1) Set up a directory to store the bundle.
if self.shared_file_system:
if not os.path.exists(run_state.bundle_path):
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 "
"your worker is mounted properly or contact your administrators."
)
logger.error(message)
return run_state._replace(stage=RunStage.CLEANING_UP,
failure_message=message)
return run_state._replace(
run_status=
"Waiting for bundle directory to be created by the server",
bundle_dir_wait_num_tries=run_state.
bundle_dir_wait_num_tries - 1,
)
else:
remove_path(run_state.bundle_path)
os.mkdir(run_state.bundle_path)
# 2) Set up symlinks
docker_dependencies = []
docker_dependencies_path = (
'/' + run_state.bundle.uuid +
('_dependencies' if not self.shared_file_system else ''))
for dep_key, dep in run_state.bundle.dependencies.items():
full_child_path = os.path.normpath(
os.path.join(run_state.bundle_path, dep.child_path))
if not full_child_path.startswith(run_state.bundle_path):
# Dependencies should end up in their bundles (ie prevent using relative paths like ..
# to get out of their parent bundles
message = 'Invalid key for dependency: %s' % (dep.child_path)
logger.error(message)
return run_state._replace(stage=RunStage.CLEANING_UP,
failure_message=message)
docker_dependency_path = os.path.join(docker_dependencies_path,
dep.child_path)
if self.shared_file_system:
# On a shared FS, we know where the dep is stored and can get the contents directly
dependency_path = os.path.realpath(
os.path.join(dep.location, dep.parent_path))
else:
# On a dependency_manager setup ask the manager where the dependency is
dependency_path = os.path.join(
self.dependency_manager.dependencies_dir,
self.dependency_manager.get(run_state.bundle.uuid,
dep_key).path,
)
os.symlink(docker_dependency_path, full_child_path)
# These are turned into docker volume bindings like:
# dependency_path:docker_dependency_path:ro
docker_dependencies.append(
(dependency_path, docker_dependency_path))
# 3) Set up container
if run_state.resources.network:
docker_network = self.docker_network_external.name
else:
docker_network = self.docker_network_internal.name
try:
cpuset, gpuset = self.assign_cpu_and_gpu_sets_fn(
run_state.resources.cpus, run_state.resources.gpus)
except Exception as e:
message = "Cannot assign enough resources: %s" % str(e)
logger.error(message)
logger.error(traceback.format_exc())
return run_state._replace(run_status=message)
# 4) Start container
try:
container = docker_utils.start_bundle_container(
run_state.bundle_path,
run_state.bundle.uuid,
docker_dependencies,
run_state.bundle.command,
run_state.resources.docker_image,
network=docker_network,
cpuset=cpuset,
gpuset=gpuset,
memory_bytes=run_state.resources.memory,
runtime=self.docker_runtime,
)
self.worker_docker_network.connect(container)
except Exception as e:
message = 'Cannot start Docker container: {}'.format(e)
logger.error(message)
logger.error(traceback.format_exc())
raise
return run_state._replace(
stage=RunStage.RUNNING,
run_status='Running job in Docker container',
container_id=container.id,
container=container,
docker_image=image_state.digest,
has_contents=True,
cpuset=cpuset,
gpuset=gpuset,
)
def _transition_from_RUNNING(self, run_state):
"""
1- Check run status of the docker container
2- If run is killed, kill the container
3- If run is finished, move to CLEANING_UP state
"""
def check_and_report_finished(run_state):
try:
finished, exitcode, failure_msg = docker_utils.check_finished(
run_state.container)
except docker_utils.DockerException:
logger.error(traceback.format_exc())
finished, exitcode, failure_msg = False, None, None
return run_state._replace(finished=finished,
exitcode=exitcode,
failure_message=failure_msg)
def check_resource_utilization(run_state):
kill_messages = []
run_stats = docker_utils.get_container_stats(run_state.container)
run_state = run_state._replace(max_memory=max(
run_state.max_memory, run_stats.get('memory', 0)))
run_state = run_state._replace(
disk_utilization=self.disk_utilization[
run_state.bundle.uuid]['disk_utilization'])
container_time_total = docker_utils.get_container_running_time(
run_state.container)
run_state = run_state._replace(
container_time_total=container_time_total,
container_time_user=run_stats.get(
'container_time_user', run_state.container_time_user),
container_time_system=run_stats.get(
'container_time_system', run_state.container_time_system),
)
if run_state.resources.time and container_time_total > run_state.resources.time:
kill_messages.append(
'Time limit exceeded. (Container uptime %s > time limit %s)'
% (duration_str(container_time_total),
duration_str(run_state.resources.time)))
if run_state.max_memory > run_state.resources.memory or run_state.exitcode == 137:
kill_messages.append('Memory limit %s exceeded.' %
size_str(run_state.resources.memory))
if run_state.resources.disk and run_state.disk_utilization > run_state.resources.disk:
kill_messages.append('Disk limit %sb exceeded.' %
size_str(run_state.resources.disk))
if kill_messages:
run_state = run_state._replace(
kill_message=' '.join(kill_messages), is_killed=True)
return run_state
def check_disk_utilization():
running = True
while running:
start_time = time.time()
try:
disk_utilization = get_path_size(run_state.bundle_path)
self.disk_utilization[run_state.bundle.uuid][
'disk_utilization'] = disk_utilization
running = self.disk_utilization[
run_state.bundle.uuid]['running']
except Exception:
logger.error(traceback.format_exc())
end_time = time.time()
# To ensure that we don't hammer the disk for this computation when
# there are lots of files, we run it at most 10% of the time.
time.sleep(max((end_time - start_time) * 10, 1.0))
self.disk_utilization.add_if_new(
run_state.bundle.uuid,
threading.Thread(target=check_disk_utilization, args=[]))
run_state = check_and_report_finished(run_state)
run_state = check_resource_utilization(run_state)
if run_state.is_killed:
if docker_utils.container_exists(run_state.container):
try:
run_state.container.kill()
except docker.errors.APIError:
finished, _, _ = docker_utils.check_finished(
run_state.container)
if not finished:
logger.error(traceback.format_exc())
self.disk_utilization[run_state.bundle.uuid]['running'] = False
self.disk_utilization.remove(run_state.bundle.uuid)
return run_state._replace(stage=RunStage.CLEANING_UP)
if run_state.finished:
logger.debug(
'Finished run with UUID %s, exitcode %s, failure_message %s',
run_state.bundle.uuid,
run_state.exitcode,
run_state.failure_message,
)
self.disk_utilization[run_state.bundle.uuid]['running'] = False
self.disk_utilization.remove(run_state.bundle.uuid)
return run_state._replace(stage=RunStage.CLEANING_UP,
run_status='Uploading results')
else:
return run_state
def _transition_from_CLEANING_UP(self, run_state):
"""
1- delete the container if still existent
2- clean up the dependencies from bundle directory
3- release the dependencies in dependency manager
4- If bundle has contents to upload (i.e. was RUNNING at some point),
move to UPLOADING_RESULTS state
Otherwise move to FINALIZING state
"""
if run_state.container_id is not None:
while docker_utils.container_exists(run_state.container):
try:
finished, _, _ = docker_utils.check_finished(
run_state.container)
if finished:
run_state.container.remove(force=True)
run_state = run_state._replace(container=None,
container_id=None)
break
else:
try:
run_state.container.kill()
except docker.errors.APIError:
logger.error(traceback.format_exc())
time.sleep(1)
except docker.errors.APIError:
logger.error(traceback.format_exc())
time.sleep(1)
for dep_key, dep in run_state.bundle.dependencies.items():
if not self.shared_file_system: # No dependencies if shared fs worker
self.dependency_manager.release(run_state.bundle.uuid, dep_key)
child_path = os.path.join(run_state.bundle_path, dep.child_path)
try:
remove_path(child_path)
except Exception:
logger.error(traceback.format_exc())
if not self.shared_file_system and run_state.has_contents:
# No need to upload results since results are directly written to bundle store
return run_state._replace(stage=RunStage.UPLOADING_RESULTS,
run_status='Uploading results',
container=None)
else:
return self.finalize_run(run_state)
def _transition_from_UPLOADING_RESULTS(self, run_state):
"""
If bundle not already uploading:
Use the RunManager API to upload contents at bundle_path to the server
Pass the callback to that API such that if the bundle is killed during the upload,
the callback returns false, allowing killable uploads.
If uploading and not finished:
Update run_status with upload progress
If uploading and finished:
Move to FINALIZING state
"""
def upload_results():
try:
# Upload results
logger.debug('Uploading results for run with UUID %s',
run_state.bundle.uuid)
def progress_callback(bytes_uploaded):
run_status = 'Uploading results: %s done (archived size)' % size_str(
bytes_uploaded)
self.uploading[
run_state.bundle.uuid]['run_status'] = run_status
return True
self.upload_bundle_callback(run_state.bundle.uuid,
run_state.bundle_path,
progress_callback)
self.uploading[run_state.bundle.uuid]['success'] = True
except Exception as e:
self.uploading[run_state.bundle.uuid]['run_status'] = (
"Error while uploading: %s" % e)
logger.error(traceback.format_exc())
self.uploading.add_if_new(
run_state.bundle.uuid,
threading.Thread(target=upload_results, args=[]))
if self.uploading[run_state.bundle.uuid].is_alive():
return run_state._replace(
run_status=self.uploading[run_state.bundle.uuid]['run_status'])
elif not self.uploading[run_state.bundle.uuid]['success']:
# upload failed
failure_message = run_state.failure_message
if failure_message:
run_state = run_state._replace(failure_message=(
failure_message + '. ' +
self.uploading[run_state.bundle.uuid]['run_status']))
else:
run_state = run_state._replace(failure_message=self.uploading[
run_state.bundle.uuid]['run_status'])
self.uploading.remove(run_state.bundle.uuid)
return self.finalize_run(run_state)
def finalize_run(self, run_state):
"""
Prepare the finalize message to be sent with the next checkin
"""
if run_state.is_killed:
# Append kill_message, which contains more useful info on why a run was killed, to the failure message.
failure_message = ("{}. {}".format(run_state.failure_message,
run_state.kill_message)
if run_state.failure_message else
run_state.kill_message)
run_state = run_state._replace(failure_message=failure_message)
return run_state._replace(stage=RunStage.FINALIZING,
run_status="Finalizing bundle")
def _transition_from_FINALIZING(self, run_state):
"""
If a full worker cycle has passed since we got into FINALIZING we already reported to
server so can move on to FINISHED. Can also remove bundle_path now
"""
if run_state.finalized:
if not self.shared_file_system:
remove_path(
run_state.bundle_path) # don't remove bundle if shared FS
return run_state._replace(stage=RunStage.FINISHED,
run_status='Finished')
else:
return run_state