class LocalComputeLogSubscriptionManager(object):
def __init__(self, manager):
self._manager = manager
self._subscriptions = defaultdict(list)
self._watchers = {}
self._observer = PollingObserver(WATCHDOG_POLLING_TIMEOUT)
self._observer.start()
def _key(self, run_id, step_key):
return '{}:{}'.format(run_id, step_key)
def add_subscription(self, subscription):
check.inst_param(subscription, 'subscription', ComputeLogSubscription)
key = self._key(subscription.run_id, subscription.step_key)
self._subscriptions[key].append(subscription)
self.watch(subscription.run_id, subscription.step_key)
def remove_all_subscriptions(self, run_id, step_key):
key = self._key(run_id, step_key)
for subscription in self._subscriptions.pop(key, []):
subscription.complete()
def watch(self, run_id, step_key):
key = self._key(run_id, step_key)
if key in self._watchers:
return
update_paths = [
self._manager.get_local_path(run_id, step_key,
ComputeIOType.STDOUT),
self._manager.get_local_path(run_id, step_key,
ComputeIOType.STDERR),
]
complete_paths = [
self._manager.complete_artifact_path(run_id, step_key)
]
directory = os.path.dirname(
self._manager.get_local_path(run_id, step_key,
ComputeIOType.STDERR))
ensure_dir(directory)
self._watchers[key] = self._observer.schedule(
LocalComputeLogFilesystemEventHandler(self, run_id, step_key,
update_paths,
complete_paths),
directory,
)
def notify_subscriptions(self, run_id, step_key):
key = self._key(run_id, step_key)
for subscription in self._subscriptions[key]:
subscription.fetch()
def unwatch(self, run_id, step_key, handler):
key = self._key(run_id, step_key)
if key in self._watchers:
self._observer.remove_handler_for_watch(handler,
self._watchers[key])
del self._watchers[key]
class ComputeLogManager(object):
def __init__(self, instance):
self._instance = instance
self._subscriptions = defaultdict(list)
self._watchers = {}
self._observer = PollingObserver(POLLING_TIMEOUT)
self._observer.start()
def watch(self, run_id, step_key):
key = _run_key(run_id, step_key)
if key in self._watchers:
return
self._watchers[key] = self._observer.schedule(
ComputeLogEventHandler(self, run_id, step_key, self._instance),
os.path.dirname(
_filepath(_filebase(self._instance, run_id, step_key),
IO_TYPE_STDOUT)),
)
def unwatch(self, run_id, step_key, handler):
key = _run_key(run_id, step_key)
if key in self._watchers:
self._observer.remove_handler_for_watch(handler,
self._watchers[key])
del self._watchers[key]
def on_compute_end(self, run_id, step_key):
run_key = _run_key(run_id, step_key)
for subscription in self._subscriptions.pop(run_key, []):
subscription.on_compute_end()
def on_subscribe(self, subscription, run_id, step_key):
key = _run_key(run_id, step_key)
self._subscriptions[key].append(subscription)
self.watch(run_id, step_key)
def on_update(self, run_id, step_key):
key = _run_key(run_id, step_key)
for subscription in self._subscriptions[key]:
subscription.fetch()
def get_observable(self, run_id, step_key, cursor):
subscription = ComputeLogSubscription(self._instance, run_id, step_key,
cursor)
self.on_subscribe(subscription, run_id, step_key)
return Observable.create(subscription) # pylint: disable=E1101
class LocalComputeLogSubscriptionManager:
def __init__(self, manager):
self._manager = manager
self._subscriptions = defaultdict(list)
self._watchers = {}
self._observer = None
def _watch_key(self, run_id, key):
return "{}:{}".format(run_id, key)
def add_subscription(self, subscription):
check.inst_param(subscription, "subscription", ComputeLogSubscription)
if self._manager.is_watch_completed(subscription.run_id,
subscription.key):
subscription.fetch()
subscription.complete()
else:
watch_key = self._watch_key(subscription.run_id, subscription.key)
self._subscriptions[watch_key].append(subscription)
self.watch(subscription.run_id, subscription.key)
def remove_all_subscriptions(self, run_id, step_key):
watch_key = self._watch_key(run_id, step_key)
for subscription in self._subscriptions.pop(watch_key, []):
subscription.complete()
def watch(self, run_id, step_key):
watch_key = self._watch_key(run_id, step_key)
if watch_key in self._watchers:
return
update_paths = [
self._manager.get_local_path(run_id, step_key,
ComputeIOType.STDOUT),
self._manager.get_local_path(run_id, step_key,
ComputeIOType.STDERR),
]
complete_paths = [
self._manager.complete_artifact_path(run_id, step_key)
]
directory = os.path.dirname(
self._manager.get_local_path(run_id, step_key,
ComputeIOType.STDERR))
if not self._observer:
self._observer = PollingObserver(self._manager.polling_timeout)
self._observer.start()
ensure_dir(directory)
self._watchers[watch_key] = self._observer.schedule(
LocalComputeLogFilesystemEventHandler(self, run_id, step_key,
update_paths,
complete_paths),
str(directory),
)
def notify_subscriptions(self, run_id, step_key):
watch_key = self._watch_key(run_id, step_key)
for subscription in self._subscriptions[watch_key]:
subscription.fetch()
def unwatch(self, run_id, step_key, handler):
watch_key = self._watch_key(run_id, step_key)
if watch_key in self._watchers:
self._observer.remove_handler_for_watch(handler,
self._watchers[watch_key])
del self._watchers[watch_key]
def dispose(self):
if self._observer:
self._observer.stop()
self._observer.join(15)
