caffe2/python/parallel_workers.py - platform/external/pytorch - Git at Google

 # Copyright (c) 2016-present, Facebook, Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 ##############################################################################

 # @package parallel_workers
 # Module caffe2.python.parallel_workers
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 from __future__ import unicode_literals


 '''
 This module provides a python-land multithreaded mechanism for executing work.

 Basic usage is as follows:
    coordinator = parallel_workers.init_workers(
       my_worker_fun,
       worker_name="train"
    )
    ...
    coordinator.start()

 First argument is the function to run in a loop on potentially multiple threads.
 It has the call signature
     worker_fun(worker_id)

 Argument 'worker_name' is used to distinguish different workers,
 such as workers processing train data or workers processing test data.

 Optionally, one can define an "init function" that is called once before
 threads start, and has call signature:
    my_init_fun(worker_coordinator, global_coordinator)

 Note that for data_parallel_models, init_workers will be called
 for each GPU. Note that the 'coordinator' returned by the function is same
 each time.
 '''

 import logging
 import threading
 import atexit
 import time
 import collections
 import six
 import traceback

 from abc import ABCMeta, abstractmethod

 log = logging.getLogger("parallel_workers")
 log.setLevel(logging.INFO)
 LOG_INT_SECS = 60


 def init_workers(
     worker_fun,
     num_worker_threads=2,
     worker_name="train",
     init_fun=None,
     external_loggers=None,
     shutdown_fun=None,
 ):
     global global_coordinator

     metrics = Metrics(external_loggers)

     # Create coordinator object
     coordinator = WorkerCoordinator(
         worker_name, init_fun, shutdown_fun=shutdown_fun)

     # Launch fetch worker threads
     worker_ids = [
         global_coordinator.get_new_worker_id()
         for i in range(num_worker_threads)
     ]
     workers = [
         threading.Thread(
             target=run_worker,
             name="parallel_workers worker id {}".format(worker_id),
             args=[coordinator,
                   Worker(coordinator, worker_id, worker_fun, metrics)],
         ) for worker_id in worker_ids
     ]

     coordinator._workers = workers
     global_coordinator.add(coordinator)

     return global_coordinator


 class Metrics(object):
     def __init__(self, external_loggers):
         self._metrics = collections.defaultdict(lambda: 0)
         self._external_loggers = external_loggers

     def reset_metrics(self):
         self._metrics = collections.defaultdict(lambda: 0)

     def log_metrics(self):
         if not self._external_loggers:
             return
         for logger in self._external_loggers:
             try:
                 logger.log(self._metrics)
             except Exception as e:
                 print("Failed to call ExternalLogger: {}".format(e))

     def put_metric(self, key, value, count=True):
         self._metrics[key] += value
         if count:
             count_key = '{}_count'.format(key)
             self._metrics[count_key] += 1


 class State():
     six.add_metaclass(ABCMeta)

     @abstractmethod
     def start(self):
         pass

     @abstractmethod
     def stop(self):
         pass

     @abstractmethod
     def cleanup(self):
         pass


 class WorkerCoordinator(object):
     def __init__(self, worker_name, init_fun, state=None, shutdown_fun=None):
         self._active = True
         self._started = False
         self._workers = []
         self._worker_name = worker_name
         self._init_fun = init_fun
         self._state = state
         self._shutdown_fun = shutdown_fun

     def is_active(self):
         return self._active

     def init(self, global_coordinator):
         if self._init_fun and not self._started:
             self._init_fun(self, global_coordinator)

     def _start(self):
         if self._started:
             return
         self._active = True
         self._started = True
         if self._state:
             self._state.start()

         for w in self._workers:
             w.daemon = True
             w.start()

     def _stop(self, reason=None):
         self._active = False
         if reason is not None:
             log.error("Data input failed due to an error: {}".format(reason))
         if self._shutdown_fun and self._started:
             self._shutdown_fun()
         if self._state:
             self._state.stop()

         self._started = False

     def _wait_finish(self, cleanup=None):
         print("Wait for workers to die: {}".format(self._worker_name))
         for w in self._workers:
             if w != threading.current_thread():
                 w.join(5.0)  # don't wait forever, thread may be blocked in i/o
         success = True
         for w in self._workers:
             if w.isAlive():
                 print("Worker {} failed to close while waiting".format(w))
                 success = False

         # Release memory for the scratch blobs
         if success and self._state:
             self._state.cleanup()

         print("All workers terminated: {}".format(success))
         return success


 class GlobalWorkerCoordinator(object):
     def __init__(self):
         self._coordinators = []
         self._fetcher_id_seq = 0
         self._worker_ids = []
         self.register_shutdown_handler()

     def add(self, coordinator):
         self._coordinators.append(coordinator)

     def get_new_worker_id(self):
         worker_id = self._fetcher_id_seq
         self._worker_ids.append(worker_id)
         self._fetcher_id_seq += 1
         return worker_id

     def get_worker_ids(self):
         return self._worker_ids

     def start(self):
         for c in self._coordinators:
             c.init(self)
             c._start()

     def stop(self):
         all_success = True
         for c in self._coordinators:
             c._stop()
         for c in self._coordinators:
             success = c._wait_finish()
             all_success = all_success and success
         self._coordinators = []
         return all_success

     def stop_coordinator(self, worker_name):
         '''
         Stop a specific coordinator
         '''
         for c in self._coordinators:
             if c._worker_name == worker_name:
                 c._stop()
                 c._wait_finish()
         self._coordinators = [
             c for c in self._coordinators
             if c._worker_name != worker_name
         ]

     def register_shutdown_handler(self):
         def cleanup():
             self.stop()

         atexit.register(cleanup)


 class Worker(object):
     def __init__(
         self,
         coordinator,
         worker_id,
         worker_fun=None,
         metrics=None
     ):
         self._coordinator = coordinator
         self._worker_id = worker_id
         self._worker_fun = worker_fun
         self._metrics = metrics

     def start(self):
         self._start_time = time.time()

     def run(self):
         self._worker_fun(self._worker_id)

     def handle_exception(self, e):
         traceback.print_exc()
         logging.exception("Exception in worker", e)
         self._coordinator._stop("Exception in worker {}: {}".format(
             self._worker_id, e
         ))

     def finish(self):
         self._metrics.put_metric(
             'worker_time', time.time() - self._start_time)
         self._metrics.log_metrics()


 global_coordinator = GlobalWorkerCoordinator()


 def run_worker(coordinator, worker):
     while coordinator.is_active():
         worker.start()
         try:
             worker.run()
         except Exception as e:
             worker.handle_exception(e)
         finally:
             worker.finish()
	# Copyright (c) 2016-present, Facebook, Inc.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.
	##############################################################################

	# @package parallel_workers
	# Module caffe2.python.parallel_workers
	from __future__ import absolute_import
	from __future__ import division
	from __future__ import print_function
	from __future__ import unicode_literals


	'''
	This module provides a python-land multithreaded mechanism for executing work.

	Basic usage is as follows:
	coordinator = parallel_workers.init_workers(
	my_worker_fun,
	worker_name="train"
	)
	...
	coordinator.start()

	First argument is the function to run in a loop on potentially multiple threads.
	It has the call signature
	worker_fun(worker_id)

	Argument 'worker_name' is used to distinguish different workers,
	such as workers processing train data or workers processing test data.

	Optionally, one can define an "init function" that is called once before
	threads start, and has call signature:
	my_init_fun(worker_coordinator, global_coordinator)

	Note that for data_parallel_models, init_workers will be called
	for each GPU. Note that the 'coordinator' returned by the function is same
	each time.
	'''

	import logging
	import threading
	import atexit
	import time
	import collections
	import six
	import traceback

	from abc import ABCMeta, abstractmethod

	log = logging.getLogger("parallel_workers")
	log.setLevel(logging.INFO)
	LOG_INT_SECS = 60


	def init_workers(
	worker_fun,
	num_worker_threads=2,
	worker_name="train",
	init_fun=None,
	external_loggers=None,
	shutdown_fun=None,
	):
	global global_coordinator

	metrics = Metrics(external_loggers)

	# Create coordinator object
	coordinator = WorkerCoordinator(
	worker_name, init_fun, shutdown_fun=shutdown_fun)

	# Launch fetch worker threads
	worker_ids = [
	global_coordinator.get_new_worker_id()
	for i in range(num_worker_threads)
	]
	workers = [
	threading.Thread(
	target=run_worker,
	name="parallel_workers worker id {}".format(worker_id),
	args=[coordinator,
	Worker(coordinator, worker_id, worker_fun, metrics)],
	) for worker_id in worker_ids
	]

	coordinator._workers = workers
	global_coordinator.add(coordinator)

	return global_coordinator


	class Metrics(object):
	def __init__(self, external_loggers):
	self._metrics = collections.defaultdict(lambda: 0)
	self._external_loggers = external_loggers

	def reset_metrics(self):
	self._metrics = collections.defaultdict(lambda: 0)

	def log_metrics(self):
	if not self._external_loggers:
	return
	for logger in self._external_loggers:
	try:
	logger.log(self._metrics)
	except Exception as e:
	print("Failed to call ExternalLogger: {}".format(e))

	def put_metric(self, key, value, count=True):
	self._metrics[key] += value
	if count:
	count_key = '{}_count'.format(key)
	self._metrics[count_key] += 1


	class State():
	six.add_metaclass(ABCMeta)

	@abstractmethod
	def start(self):
	pass

	@abstractmethod
	def stop(self):
	pass

	@abstractmethod
	def cleanup(self):
	pass


	class WorkerCoordinator(object):
	def __init__(self, worker_name, init_fun, state=None, shutdown_fun=None):
	self._active = True
	self._started = False
	self._workers = []
	self._worker_name = worker_name
	self._init_fun = init_fun
	self._state = state
	self._shutdown_fun = shutdown_fun

	def is_active(self):
	return self._active

	def init(self, global_coordinator):
	if self._init_fun and not self._started:
	self._init_fun(self, global_coordinator)

	def _start(self):
	if self._started:
	return
	self._active = True
	self._started = True
	if self._state:
	self._state.start()

	for w in self._workers:
	w.daemon = True
	w.start()

	def _stop(self, reason=None):
	self._active = False
	if reason is not None:
	log.error("Data input failed due to an error: {}".format(reason))
	if self._shutdown_fun and self._started:
	self._shutdown_fun()
	if self._state:
	self._state.stop()

	self._started = False

	def _wait_finish(self, cleanup=None):
	print("Wait for workers to die: {}".format(self._worker_name))
	for w in self._workers:
	if w != threading.current_thread():
	w.join(5.0) # don't wait forever, thread may be blocked in i/o
	success = True
	for w in self._workers:
	if w.isAlive():
	print("Worker {} failed to close while waiting".format(w))
	success = False

	# Release memory for the scratch blobs
	if success and self._state:
	self._state.cleanup()

	print("All workers terminated: {}".format(success))
	return success


	class GlobalWorkerCoordinator(object):
	def __init__(self):
	self._coordinators = []
	self._fetcher_id_seq = 0
	self._worker_ids = []
	self.register_shutdown_handler()

	def add(self, coordinator):
	self._coordinators.append(coordinator)

	def get_new_worker_id(self):
	worker_id = self._fetcher_id_seq
	self._worker_ids.append(worker_id)
	self._fetcher_id_seq += 1
	return worker_id

	def get_worker_ids(self):
	return self._worker_ids

	def start(self):
	for c in self._coordinators:
	c.init(self)
	c._start()

	def stop(self):
	all_success = True
	for c in self._coordinators:
	c._stop()
	for c in self._coordinators:
	success = c._wait_finish()
	all_success = all_success and success
	self._coordinators = []
	return all_success

	def stop_coordinator(self, worker_name):
	'''
	Stop a specific coordinator
	'''
	for c in self._coordinators:
	if c._worker_name == worker_name:
	c._stop()
	c._wait_finish()
	self._coordinators = [
	c for c in self._coordinators
	if c._worker_name != worker_name
	]

	def register_shutdown_handler(self):
	def cleanup():
	self.stop()

	atexit.register(cleanup)


	class Worker(object):
	def __init__(
	self,
	coordinator,
	worker_id,
	worker_fun=None,
	metrics=None
	):
	self._coordinator = coordinator
	self._worker_id = worker_id
	self._worker_fun = worker_fun
	self._metrics = metrics

	def start(self):
	self._start_time = time.time()

	def run(self):
	self._worker_fun(self._worker_id)

	def handle_exception(self, e):
	traceback.print_exc()
	logging.exception("Exception in worker", e)
	self._coordinator._stop("Exception in worker {}: {}".format(
	self._worker_id, e
	))

	def finish(self):
	self._metrics.put_metric(
	'worker_time', time.time() - self._start_time)
	self._metrics.log_metrics()


	global_coordinator = GlobalWorkerCoordinator()


	def run_worker(coordinator, worker):
	while coordinator.is_active():
	worker.start()
	try:
	worker.run()
	except Exception as e:
	worker.handle_exception(e)
	finally:
	worker.finish()