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)