content: # Copyright 2017, Google Inc. All rights reserved.
  #
  # 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 time
  import numpy as np
  from helpers.gke_helper import get_cluster
  from helpers.gcs_helper import pickle_and_upload, get_uri_blob, download_uri_and_unpickle
  from helpers.kubernetes_helper import create_job, delete_jobs_pods
  from copy import deepcopy
  from itertools import product
  from sklearn.model_selection import GridSearchCV, RandomizedSearchCV
  from skopt import BayesSearchCV
  from skopt.space import Categorical, Integer, Real
  
  
  class GKEParallel(object):
      SUPPORTED_SEARCH = [
          GridSearchCV,
          RandomizedSearchCV,
          BayesSearchCV
      ]
  
      def __init__(self, search, project_id, zone, cluster_id, bucket_name, image_name, task_name=None):
          """Wraps around a SearchCV object and handles deploying `fit`
          jobs to a GKE cluster.
          """
          if type(search) not in self.SUPPORTED_SEARCH:
              raise TypeError('Search type {} not supported.  Only supporting {}.'.format(type(search), [s.__name__ for s in self.SUPPORTED_SEARCH]))
  
          self.search = search
          self.project_id = project_id
          self.cluster_id = cluster_id
          self.bucket_name = bucket_name
          self.image_name = image_name
          self.task_name = task_name
          self.gcs_uri = None
  
          self.cluster = get_cluster(project_id, zone, cluster_id)
          self.n_nodes = self.cluster['currentNodeCount']
  
          self.task_name = None
  
          # For GridSearchCV
          self.param_grids = {}
          # For RandomizedSearchCV
          self.param_distributions = None
          self.n_iter = None
          # For BayesSearchCV
          self.search_spaces = {}
  
          self.job_names = {}
          self.output_uris = {}
          self.output_without_estimator_uris = {}
          self.dones = {}
          self.results = {}
  
          self.best_estimator_ = None
          self.best_params_ = None
          self.best_score_ = None
          self.best_search_ = None
  
          self._cancelled = False
          self._done = False
  
  
      def _make_job_name(self, worker_id):
          return '{}.worker.{}'.format(self.task_name, worker_id)
  
  
      def _make_job_body(self, worker_id, X_uri, y_uri):
          body = {
              'apiVersion': 'batch/v1',
              'kind': 'Job',
              'metadata': {
                  'name': self._make_job_name(worker_id)
              },
              'spec': {
                  'template': {
                      'spec': {
                          'containers': [
                              {
                                  'image': 'gcr.io/{}/{}'.format(self.project_id, self.image_name),
                                  'command': ['python'],
                                  'args': ['worker.py', self.bucket_name, self.task_name, worker_id, X_uri, y_uri],
                                  'name': 'worker'
                              }
                          ],
                  'restartPolicy': 'OnFailure'}
                  }
              }
          }
  
          return body
  
  
      def _deploy_job(self, worker_id, X_uri, y_uri):
          job_body = self._make_job_body(worker_id, X_uri, y_uri)
  
          print('Deploying worker {}'.format(worker_id))
          create_job(job_body)
  
  
      def _partition_grid(self, param_grid_dict, partition_keys):
          _param_grid_dict = deepcopy(param_grid_dict)
  
          partition_lists = [_param_grid_dict.pop(key) for key in partition_keys]
          
          partitioned = []
          for prod in product(*partition_lists):
              lists = [[element] for element in prod]
              singleton = dict(zip(partition_keys, lists))
              singleton.update(_param_grid_dict)
  
              partitioned.append(singleton)
  
          return partitioned
  
  
      def _partition_param_grid(self, param_grid, target_n_partition=5):
          """Returns a list of param_grids whose union is the input
          param_grid.
  
          If param_grid is a dict:
  
          The implemented strategy attempts to partition the param_grid
          into at least target_n_partition smaller param_grids.
  
          NOTE: The naive strategy implemented here does not distinguish
          between different types of parameters nor their impact on the
          running time.  The user of this module is encouraged to
          implement their own paritioning strategy based on their needs.
          """
          if type(param_grid) == list:
              # If the input is already a list of param_grids then just
              # use it as is.
              return param_grid
          else:
              # The strategy is to simply expand the grid fully with
              # respect to a parameter:
              # [1, 2, 3]x[4, 5] --> [1]x[4, 5], [2]x[4, 5], [3]x[4, 5]
              # until the target number of partitions is reached.
              partition_keys = []
              n_partition = 1
              for key, lst in param_grid.items():
                  partition_keys.append(key)
                  n_partition *= len(lst)
  
                  if n_partition >= target_n_partition:
                      break
  
              partitioned = self._partition_grid(param_grid, partition_keys)
  
              return partitioned
  
  
      def _handle_grid_search(self, X_uri, y_uri):
          param_grids = self._partition_param_grid(self.search.param_grid, self.n_nodes)
  
          for i, param_grid in enumerate(param_grids):
              worker_id = str(i)
  
              self.param_grids[worker_id] = param_grid
              self.job_names[worker_id] = self._make_job_name(worker_id)
              self.output_uris[worker_id] = 'gs://{}/{}/{}/fitted_search.pkl'.format(self.bucket_name, self.task_name, worker_id)
              self.output_without_estimator_uris[worker_id] = 'gs://{}/{}/{}/fitted_search_without_estimator.pkl'.format(self.bucket_name, self.task_name, worker_id)
              self.dones[worker_id] = False
  
              pickle_and_upload(param_grid, self.bucket_name, '{}/{}/param_grid.pkl'.format(self.task_name, worker_id))
  
              self._deploy_job(worker_id, X_uri, y_uri)
  
  
      def _handle_randomized_search(self, X_uri, y_uri):
          self.param_distributions = self.search.param_distributions
          self.n_iter = self.search.n_iter
          n_iter = self.n_iter / self.n_nodes + 1
  
          for i in xrange(self.n_nodes):
              worker_id = str(i)
  
              self.job_names[worker_id] = self._make_job_name(worker_id)
              self.output_uris[worker_id] = 'gs://{}/{}/{}/fitted_search.pkl'.format(self.bucket_name, self.task_name, worker_id)
              self.output_without_estimator_uris[worker_id] = 'gs://{}/{}/{}/fitted_search_without_estimator.pkl'.format(self.bucket_name, self.task_name, worker_id)
              self.dones[worker_id] = False
  
              pickle_and_upload(self.param_distributions, self.bucket_name, '{}/{}/param_distributions.pkl'.format(self.task_name, worker_id))
              pickle_and_upload(n_iter, self.bucket_name, '{}/{}/n_iter.pkl'.format(self.task_name, worker_id))
  
              self._deploy_job(worker_id, X_uri, y_uri)
  
  
      def _partition_space(self, space):
          """Partitions the space into two subspaces.  In the case of
          Real and Integer, the subspaces are not disjoint, but
          overlapping at an endpoint.
  
          The argument `space` should be a dict whose values are
          skopt.space's Categorical, Integer, or Real.
          """
          partition_key = np.random.choice(space.keys())
          dimension = space[partition_key]
  
          if type(dimension) == Categorical:
              categories = dimension.categories
              prior = dimension.prior
              transform = dimension.transform_
  
              if len(categories) >= 2:
                  mid_index = len(categories) / 2
                  left_categories = categories[:mid_index]
                  right_categories = categories[mid_index:]
  
                  if prior is not None:
                      left_prior = prior[:mid_index]
                      left_weight = sum(left_prior)
                      left_prior = [p/left_weight for p in left_prior]
  
                      right_prior = prior[mid_index:]
                      right_weight = sum(right_prior)
                      right_prior = [p/right_weight for p in right_prior]
                  else:
                      left_prior = None
                      right_prior = None
  
                  left = Categorical(left_categories, prior=left_prior, transform=transform)
                  right = Categorical(right_categories, prior=right_prior, transform=transform)
              else:
                  return [space]
  
          elif type(dimension) == Integer:
              low = dimension.low
              high = dimension.high
              transform = dimension.transform_
  
              if low < high:
                  mid = int((high - low) / 2)
                  left = Integer(low, mid, transform=transform)
                  right = Integer(mid, high, transform=transform)
  
              else:
                  return [space]
  
          elif type(dimension) == Real:
              low = dimension.low
              high = dimension.high
              prior = dimension.prior
              transform = dimension.transform_
  
              if low < high:
                  mid = (high - low) / 2
                  left = Real(low, mid, prior=prior, transform=transform)
                  right = Real(mid, high, prior=prior, transform=transform)
  
              else:
                  return [space]
  
          left_space = deepcopy(space)
          left_space[partition_key] = left
          right_space = deepcopy(space)
          right_space[partition_key] = right
  
          return [left_space, right_space]
  
  
      def _partition_search_spaces(self, search_spaces, target_n_partition=5):
          """Returns a list of search_spaces whose union is the input
          search_spaces.
  
          If search_spaces is a dict:
  
          The implemented strategy attempts to partition the search_spaces
          into at least target_n_partition smaller search_spaces.
  
          NOTE: The naive strategy implemented here does not distinguish
          between different types of parameters nor their impact on the
          running time.  The user of this module is encouraged to
          implement their own paritioning strategy based on their needs.
          """
          if type(search_spaces[0]) == tuple:
              # If the input is already a list of search_spaces then just
              # use it as is.
              return search_spaces.values()
          else:
              result = search_spaces.values()
              while len(result) < target_n_partition:
                  space = result.pop()
                  partitioned = self._partition_space(space)
                  result.extend(partitioned)
  
              return result
  
  
      def _handle_bayes_search(self, X_uri, y_uri):
          partitioned_search_spaces = self._partition_search_spaces(self.search.search_spaces_, self.n_nodes)
  
          for i, search_spaces in enumerate(partitioned_search_spaces):
              worker_id = str(i)
  
              self.search_spaces[worker_id] = search_spaces
              self.job_names[worker_id] = self._make_job_name(worker_id)
              self.output_uris[worker_id] = 'gs://{}/{}/{}/fitted_search.pkl'.format(self.bucket_name, self.task_name, worker_id)
              self.output_without_estimator_uris[worker_id] = 'gs://{}/{}/{}/fitted_search_without_estimator.pkl'.format(self.bucket_name, self.task_name, worker_id)
              self.dones[worker_id] = False
  
              pickle_and_upload(search_spaces, self.bucket_name, '{}/{}/search_spaces.pkl'.format(self.task_name, worker_id))
  
              self._deploy_job(worker_id, X_uri, y_uri)
  
  
      def _upload_data(self, X, y):
          if type(X) == str and X.startswith('gs://'):
              X_uri = X
          else:
              X_uri = pickle_and_upload(X, self.bucket_name, '{}/X.pkl'.format(self.task_name))
  
          if type(y) == str and y.startswith('gs://'):
              y_uri = y
          else:
              y_uri = pickle_and_upload(y, self.bucket_name, '{}/y.pkl'.format(self.task_name))
  
          search_uri = pickle_and_upload(self.search, self.bucket_name, '{}/search.pkl'.format(self.task_name))
  
          return X_uri, y_uri, search_uri
  
  
      def fit(self, X, y):
          """Deploys `fit` jobs to each worker in the cluster.
          """
          timestamp = str(int(time.time()))
          self.task_name = self.task_name or '{}.{}.{}'.format(self.cluster_id, self.image_name, timestamp)
          self._done = False
          self._cancelled = False
  
          X_uri, y_uri, _ = self._upload_data(X, y)
  
          if type(self.search) == GridSearchCV:
              handler = self._handle_grid_search
          elif type(self.search) == RandomizedSearchCV:
              handler = self._handle_randomized_search
          elif type(self.search) == BayesSearchCV:
              handler = self._handle_bayes_search
  
          print('Fitting {}'.format(type(self.search)))
          handler(X_uri, y_uri)
  
          self.persist()
  
  
      def persist(self):
          """Pickle and upload self to GCS, allowing recovering of parallel
          search objects across experiments.
          """
          self.gcs_uri = pickle_and_upload(self, self.bucket_name, '{}/gke_search.pkl'.format(self.task_name))
          print('Persisted the GKEParallel instance: {}'.format(self.gcs_uri))
  
  
      # Implement part of the concurrent.future.Future interface.
      def done(self):
          if not self._done:
              for worker_id, output_uri in self.output_uris.items():
                  print('Checking if worker {} is done'.format(worker_id))
                  self.dones[worker_id] = get_uri_blob(output_uri).exists()
  
              self._done = all(self.dones.values())
  
          return self._done
  
  
      def cancel(self):
          """Deletes the kubernetes jobs.
          Persisted data and the cluster will not be deleted."""
          if not self._cancelled:
              delete_jobs_pods(self.job_names.values())
              self._cancelled = True
  
  
      def cancelled(self):
          return self._cancelled
  
  
      def result(self, download=False):
          if not self.done():
              n_done = len(d for d in self.dones.values() if d)
              print('Not done: {} out of {} workers completed.'.format(n_done, len(self.dones)))
              return None
  
          if not self.results or download:
              for worker_id, output_uri in self.output_without_estimator_uris.items():
                  print('Getting result from worker {}'.format(worker_id))
                  self.results[worker_id] = download_uri_and_unpickle(output_uri)
  
              self._aggregate_results(download)
  
              self.persist()
  
          return self.results
  
  
      def _aggregate_results(self, download):
          best_id = None
          for worker_id, result in self.results.items():
              if self.best_score_ is None or result.best_score_ > self.best_score_ or download:
  
                  self.best_score_ = result.best_score_
                  self.best_params_ = result.best_params_
  
                  best_id = worker_id
  
          if download and self.best_estimator_ is None:
              # Download only the best estimator among the workers.
              print('Downloading the best estimator (worker {}).'.format(best_id))
              output_uri = self.output_uris[best_id]
              self.best_search_ = download_uri_and_unpickle(output_uri)
              self.best_estimator_ = self.best_search_.best_estimator_
  
  
      # Implement part of SearchCV interface by delegation.
      def predict(self, *args, **kwargs):
          return self.best_estimator_.predict(*args, **kwargs)
  
  
      def predict_proba(self, *args, **kwargs):
          return self.best_estimator_.predict_proba(*args, **kwargs)
  
  
      def predict_log_proba(self, *args, **kwargs):
          return self.best_estimator_.predict_log_proba(*args, **kwargs)
  
  