text: # Copyright 1999-2020 Alibaba Group Holding Ltd.
  #
  # 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.
  
  import json
  import logging
  import os
  import time
  from collections import defaultdict
  
  import numpy as np
  
  from .kvstore import KVStoreActor
  from .session import SessionManagerActor, SessionActor
  from .utils import SchedulerActor
  from ..config import options
  from ..utils import BlacklistSet
  
  logger = logging.getLogger(__name__)
  
  
  class ResourceHeartbeatActor(SchedulerActor):
      def __init__(self, resource_ref):
          super().__init__()
          self._resource_ref = resource_ref
  
      def post_create(self):
          super().post_create()
          self._resource_ref = self.ctx.actor_ref(self._resource_ref)
          self.ref().do_heartbeat(_tell=True)
  
      def do_heartbeat(self):
          self._resource_ref.heartbeat(_tell=True)
          self.ref().do_heartbeat(_tell=True, _delay=1)
  
  
  class ResourceActor(SchedulerActor):
      """
      Actor managing free resources on workers
      """
      def __init__(self):
          super().__init__()
          self._meta_cache = dict()
          self._worker_blacklist = BlacklistSet(options.scheduler.worker_blacklist_time)
          self._worker_allocations = dict()
  
          self._last_heartbeat_time = time.time()
          self._last_heartbeat_interval = 0
  
          self._kv_store_ref = None
          self._heartbeat_ref = None
  
      def post_create(self):
          logger.debug('Actor %s running in process %d', self.uid, os.getpid())
  
          super().post_create()
          self.set_cluster_info_ref()
  
          self._kv_store_ref = self.ctx.actor_ref(KVStoreActor.default_uid())
          if not self.ctx.has_actor(self._kv_store_ref):
              self._kv_store_ref = None
  
          try:
              # we assign the heartbeat actor into another process
              # so it can sense whether the process is busy
              heartbeat_uid = self.ctx.distributor.make_same_process(
                  ResourceHeartbeatActor.default_uid(), self.uid, delta=1)
          except AttributeError:
              heartbeat_uid = ResourceHeartbeatActor.default_uid()
          self._heartbeat_ref = self.ctx.create_actor(
              ResourceHeartbeatActor, self.ref(), uid=heartbeat_uid)
  
          self.ref().detect_dead_workers(_tell=True)
  
      def pre_destroy(self):
          self._heartbeat_ref.destroy()
          self.unset_cluster_info_ref()
          super().pre_destroy()
  
      def heartbeat(self):
          t = time.time()
          self._last_heartbeat_interval = t - self._last_heartbeat_time
          self._last_heartbeat_time = t
  
      def detect_dead_workers(self):
          """
          Remove worker when it does not update its status for a long time
          """
          # take latency of scheduler into consideration by
          # include interval of last heartbeat
          timeout = options.scheduler.status_timeout + self._last_heartbeat_interval
          dead_workers = []
  
          check_time = self._last_heartbeat_time
          for worker in list(self._meta_cache.keys()):
              worker_meta = self._meta_cache[worker]
              try:
                  if check_time - worker_meta['update_time'] > timeout:
                      dead_workers.append(worker)
              except (KeyError, TypeError, ValueError):
                  pass
  
          self.detach_dead_workers(dead_workers)
          self.ref().detect_dead_workers(_tell=True, _delay=1)
  
      def mark_workers_alive(self, workers):
          t = time.time()
          for w in workers:
              try:
                  self._meta_cache[w]['update_time'] = t
              except KeyError:  # pragma: no cover
                  pass
  
      def detach_dead_workers(self, workers):
          from ..worker.execution import ExecutionActor
  
          workers = [w for w in workers if w in self._meta_cache and
                     w not in self._worker_blacklist]
  
          if not workers:
              return
  
          logger.warning('Workers %r dead, detaching from ResourceActor.', workers)
          for w in workers:
              del self._meta_cache[w]
              self._worker_blacklist.add(w)
  
          self._broadcast_sessions(SessionActor.handle_worker_change, [], workers)
          self._broadcast_workers(ExecutionActor.handle_worker_change, [], workers)
  
      def get_worker_count(self):
          return len(self._meta_cache)
  
      def get_total_cpu_count(self):
          """Get all workers' cpu count"""
          return sum(v['hardware']['cpu_total'] for k, v in self._meta_cache.items()
                     if k not in self._worker_blacklist)
  
      def get_workers_meta(self):
          return dict((k, v) for k, v in self._meta_cache.items()
                      if k not in self._worker_blacklist)
  
      def get_worker_endpoints(self):
          return [k for k in self._meta_cache if k not in self._worker_blacklist]
  
      def get_worker_metas(self):  # pragma: no cover
          return {k: meta for k, meta in self._meta_cache.items() if k not in self._worker_blacklist}
  
      def select_worker(self, size=None):
          # NB: randomly choice the indices, rather than the worker list to avoid the `str` to `np.str_` casting.
          available_workers = [k for k in self._meta_cache if k not in self._worker_blacklist]
          idx = np.random.choice(range(len(available_workers)), size=size)
          if size is None:
              return available_workers[idx]
          else:
              return [available_workers[i] for i in idx]
  
      def set_worker_meta(self, worker, worker_meta):
          from ..worker.execution import ExecutionActor
  
          if worker in self._worker_blacklist:
              return
  
          is_new = worker not in self._meta_cache
  
          if worker_meta.get('update_time', 0) < self._meta_cache.get(worker, {}).get('update_time', 0):
              return
  
          self._meta_cache[worker] = worker_meta
          if self._kv_store_ref is not None:
              self._kv_store_ref.write(f'/workers/meta/{worker}', json.dumps(worker_meta),
                                       _tell=True, _wait=False)
              self._kv_store_ref.write('/workers/meta_timestamp', str(int(time.time())),
                                       _tell=True, _wait=False)
          if is_new:
              self._broadcast_sessions(SessionActor.handle_worker_change, [worker], [])
              self._broadcast_workers(ExecutionActor.handle_worker_change, [worker], [])
  
      def _broadcast_sessions(self, handler, *args, **kwargs):
          if hasattr(handler, '__name__'):
              handler = handler.__name__
  
          futures = []
          kwargs.update(dict(_tell=True, _wait=False))
          for ep in self.get_schedulers():
              ref = self.ctx.actor_ref(SessionManagerActor.default_uid(), address=ep)
              futures.append(ref.broadcast_sessions(handler, *args, **kwargs))
          [f.result() for f in futures]
  
      def _broadcast_workers(self, handler, *args, **kwargs):
          from ..worker.execution import ExecutionActor
  
          if hasattr(handler, '__name__'):
              handler = handler.__name__
  
          kwargs.update(dict(_tell=True, _wait=False))
          for w in self._meta_cache.keys():
              ref = self.ctx.actor_ref(ExecutionActor.default_uid(), address=w)
              getattr(ref, handler)(*args, **kwargs)
  
      def allocate_resource(self, session_id, op_key, endpoint, alloc_dict, log_fail=False):
          """
          Try allocate resource for operands
          :param session_id: session id
          :param op_key: operand key
          :param endpoint: worker endpoint
          :param alloc_dict: allocation dict, listing resources needed by the operand
          :param log_fail: log assign failure
          :return: True if allocated successfully
          """
          worker_stats = self._meta_cache[endpoint]['hardware']
          if endpoint not in self._worker_allocations:
              self._worker_allocations[endpoint] = dict()
          if session_id not in self._worker_allocations[endpoint]:
              self._worker_allocations[endpoint][session_id] = dict()
          worker_allocs = self._worker_allocations[endpoint][session_id]
  
          res_used = defaultdict(lambda: 0)
          free_alloc_keys = []
          for alloc in worker_allocs.values():
              for k, v in alloc[0].items():
                  res_used[k] += v
          for k in free_alloc_keys:
              self.deallocate_resource(session_id, k, endpoint)
  
          for k, v in alloc_dict.items():
              res_used[k] += v
  
          sufficient = True
          for k, v in res_used.items():
              if res_used[k] > worker_stats.get(k, 0):
                  sufficient = False
                  if log_fail:
                      logger.warning('Failed to assign %s resource for %s, res_used: %s, worker stat: %s',
                                     k, op_key, res_used[k], worker_stats.get(k, 0))
                  break
  
          if sufficient:
              self._worker_allocations[endpoint][session_id][op_key] = (alloc_dict, time.time())
              return True
          else:
              return False
  
      def deallocate_resource(self, session_id, op_key, endpoint):
          """
          Deallocate resource
          :param session_id: session id
          :param op_key: operand key
          :param endpoint: worker endpoint
          """
          try:
              del self._worker_allocations[endpoint][session_id][op_key]
          except KeyError:
              pass
  
      def get_allocations(self):
          return self._worker_allocations
  