code: # 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 base64
  import json
  import time
  import logging
  import pickle
  import sys
  import uuid
  from io import BytesIO
  from numbers import Integral
  
  import numpy as np
  
  from ..config import options
  from ..core.operand import Fetch
  from ..errors import ResponseMalformed, ExecutionInterrupted, ExecutionFailed, \
      ExecutionStateUnknown, ExecutionNotStopped
  from ..serialize import dataserializer
  from ..serialize.dataserializer import pyarrow
  from ..tensor.core import Indexes
  from ..utils import build_tileable_graph, sort_dataframe_result, \
      numpy_dtype_from_descr_json, serialize_graph, serialize_serializable
  
  logger = logging.getLogger(__name__)
  
  
  class Session(object):
      def __init__(self, endpoint, session_id=None, req_session=None, verify_ssl=True,
                   **session_kwargs):
          self._endpoint = endpoint.rstrip('/')
          self._session_id = session_id
          self._session_kwargs = session_kwargs
          # dict structure: {tileable_key -> graph_key, tileable_ids}
          # dict value is a tuple object which records graph key and tileable id
          self._executed_tileables = dict()
  
          self._serial_type = None
          self._pickle_protocol = pickle.HIGHEST_PROTOCOL
  
          if req_session:
              self._req_session = req_session
          else:
              import requests
              from requests.adapters import HTTPAdapter
  
              self._req_session = requests.Session()
              self._req_session.mount('http://stackoverflow.com', HTTPAdapter(max_retries=5))
  
          self._req_session.verify = verify_ssl
          if not verify_ssl:
              try:
                  import urllib3
                  urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
              except ImportError:  # pragma: no cover
                  pass
  
          self._main()
  
      @property
      def session_id(self):
          return self._session_id
  
      @property
      def endpoint(self):
          return self._endpoint
  
      @endpoint.setter
      def endpoint(self, url):
          self._endpoint = url
  
      def _main(self):
          if pyarrow is None:
              self._serial_type = dataserializer.SerialType.PICKLE
          else:
              self._serial_type = dataserializer.SerialType(options.client.serial_type.lower())
  
          session_kw = self._session_kwargs.copy()
          session_kw['pyver'] = '.'.join(str(v) for v in sys.version_info[:3])
          session_kw['pickle_protocol'] = self._pickle_protocol
          if pyarrow is not None:
              session_kw['arrow_version'] = pyarrow.__version__
  
          if self._session_id is None:
              resp = self._req_session.post(self._endpoint + '/api/session', data=session_kw)
  
              if resp.status_code >= 400:
                  raise SystemError('Failed to create mars session: ' + resp.reason)
          else:
              resp = self._req_session.get(
                  self._endpoint + '/api/session/' + self._session_id, params=session_kw)
              if resp.status_code == 404:
                  raise ValueError(f'The session with id = {self._session_id} doesn\'t exist')
              if resp.status_code >= 400:
                  raise SystemError('Failed to check mars session.')
  
          content = json.loads(resp.text)
          self._session_id = content['session_id']
          self._pickle_protocol = content.get('pickle_protocol', pickle.HIGHEST_PROTOCOL)
  
          # as pyarrow will use pickle.HIGHEST_PROTOCOL to pickle, we need to use
          # SerialType.PICKLE when pickle protocol between client and server
          # does not agree with each other
          if not content.get('arrow_compatible') or self._pickle_protocol != pickle.HIGHEST_PROTOCOL:
              self._serial_type = dataserializer.SerialType.PICKLE
  
      def _get_tileable_graph_key(self, tileable_key):
          return self._executed_tileables[tileable_key][0]
  
      def _set_tileable_graph_key(self, tileable, graph_key):
          tileable_key = tileable.key
          tileable_id = tileable.id
          if tileable_key in self._executed_tileables:
              self._executed_tileables[tileable_key][1].add(tileable_id)
          else:
              self._executed_tileables[tileable_key] = graph_key, {tileable_id}
  
      @staticmethod
      def _handle_json_response(resp, allow_empty=True, raises=True):
          try:
              resp_txt = resp.text
              if allow_empty:
                  resp_txt = resp_txt or '{}'
              resp_json = json.loads(resp_txt)
          except json.JSONDecodeError:
              text_part = resp.text if len(resp.text) < 100 else resp.text[:100] + '...'
              raise ResponseMalformed(f'Failed to parse server response. Status={resp.status_code} '
                                      f'Response="{text_part}"')
  
          if raises and resp.status_code >= 400:
              exc_info = pickle.loads(base64.b64decode(resp_json['exc_info']))
              raise exc_info[1].with_traceback(exc_info[2])
          return resp_json
  
      def _check_response_finished(self, graph_url, timeout=None):
          import requests
          try:
              resp = self._req_session.get(graph_url, params={'wait_timeout': timeout})
          except requests.ConnectionError as ex:
              err_msg = str(ex)
              if 'ConnectionResetError' in err_msg or 'Connection refused' in err_msg or \
                      'Connection aborted' in err_msg:
                  return False
              raise
  
          if resp.status_code == 504:
              logging.debug('Gateway Time-out, try again')
              return False
          if resp.status_code >= 400:
              raise SystemError(f'Failed to obtain execution status. Code: {resp.status_code}, '
                                f'Reason: {resp.reason}, Content:\n{resp.text}')
  
          resp_json = self._handle_json_response(resp, raises=False)
          if resp_json['state'] == 'succeeded':
              return True
          elif resp_json['state'] in ('running', 'preparing'):
              return False
          elif resp_json['state'] in ('cancelled', 'cancelling'):
              raise ExecutionInterrupted
          elif resp_json['state'] == 'failed':
              if 'exc_info' in resp_json:
                  exc_info = pickle.loads(base64.b64decode(resp_json['exc_info']))
                  exc = exc_info[1].with_traceback(exc_info[2])
                  raise ExecutionFailed('Graph execution failed.') from exc
              else:
                  raise ExecutionFailed('Graph execution failed with unknown reason.')
          raise ExecutionStateUnknown('Unknown graph execution state ' + resp_json['state'])
  
      def run(self, *tileables, **kw):
          timeout = kw.pop('timeout', -1)
          compose = kw.pop('compose', True)
          fetch = kw.pop('fetch', True)
          name = kw.pop('name', None)
          if kw:
              raise TypeError(f'run got unexpected key arguments {kw!r}')
  
          # those executed tileables should fetch data directly, submit the others
          run_tileables = [t for t in tileables if t.key not in self._executed_tileables]
  
          if name is not None:
              if not isinstance(name, (list, tuple)):
                  name = [name]
              if len(name) != len(tileables):
                  raise TypeError('Name must match execute tileables')
              name = ','.join(name)
  
          graph = build_tileable_graph(run_tileables, set(self._executed_tileables.keys()))
          targets = [t.key for t in run_tileables]
  
          if len(graph) > 0:
              targets_join = ','.join(targets)
              session_url = self._endpoint + '/api/session/' + self._session_id
              serialized_graph = serialize_graph(graph)
  
              resp_json = self._submit_graph(serialized_graph, targets_join, names=name or '', compose=compose)
              graph_key = resp_json['graph_key']
              graph_url = f'{session_url}/graph/{graph_key}'
  
              exec_start_time = time.time()
              time_elapsed = 0
              check_interval = options.check_interval
              while timeout <= 0 or time_elapsed < timeout:
                  timeout_val = min(check_interval, timeout - time_elapsed) if timeout > 0 else check_interval
                  try:
                      if self._check_response_finished(graph_url, timeout_val):
                          break
                  except KeyboardInterrupt:
                      resp = self._req_session.delete(graph_url)
                      if resp.status_code >= 400:
                          raise ExecutionNotStopped(
                              f'Failed to stop graph execution. Code: {resp.status_code}, '
                              f'Reason: {resp.reason}, Content:\n{resp.text}')
                  finally:
                      time_elapsed = time.time() - exec_start_time
  
              if 0 < timeout < time.time() - exec_start_time:
                  raise TimeoutError
  
              for t in tileables:
                  self._set_tileable_graph_key(t, graph_key)
  
          if not fetch:
              return
          else:
              return self.fetch(*tileables)
  
      def _is_executed(self, tileable):
          # if tileble.key in executed tileables
          # or it's a fetch already
          return tileable.key in self._executed_tileables or \
                 isinstance(tileable.op, Fetch)
  
      def fetch(self, *tileables, **kw):
          from ..tensor.indexing import TensorIndex
          from ..dataframe.indexing.iloc import DataFrameIlocGetItem, SeriesIlocGetItem
  
          timeout = kw.pop('timeout', None)
          if kw:
              raise TypeError(f'fetch got unexpected key arguments {kw!r}')
  
          results = list()
          for tileable in tileables:
              if tileable.key not in self._executed_tileables and \
                      isinstance(tileable.op, (TensorIndex, DataFrameIlocGetItem, SeriesIlocGetItem)):
                  to_fetch_tileable = tileable.inputs[0]
                  indexes = tileable.op.indexes
                  if not all(isinstance(ind, (slice, Integral)) for ind in indexes):
                      raise ValueError('Only support fetch data slices')
              else:
                  to_fetch_tileable = tileable
                  indexes = []
  
              if not self._is_executed(to_fetch_tileable):
                  raise ValueError('Cannot fetch the unexecuted tileable')
  
              key = to_fetch_tileable.key
              indexes_str = base64.b64encode(
                  serialize_serializable(Indexes(indexes=indexes))).decode('ascii')
  
              session_url = f'{self._endpoint}/api/session/{self._session_id}'
              compression_str = ','.join(v.value for v in dataserializer.get_supported_compressions())
              params = dict(compressions=compression_str, slices=indexes_str,
                            serial_type=self._serial_type.value, pickle_protocol=self._pickle_protocol)
              data_url = f'{session_url}/graph/{self._get_tileable_graph_key(key)}/data/{key}'
              resp = self._req_session.get(data_url, params=params, timeout=timeout)
              if resp.status_code >= 400:
                  raise ValueError(f'Failed to fetch data from server. Code: {resp.status_code}, '
                                   f'Reason: {resp.reason}, Content:\n{resp.text}')
              result_data = dataserializer.loads(resp.content)
              results.append(sort_dataframe_result(tileable, result_data))
          return results
  
      @classmethod
      def _process_int_or_dict_argument(cls, argument, name, params):
          if argument is None:
              return
          if not isinstance(argument, dict):
              params[name] = argument
          else:
              params[name] = ','.join(f'{k}={v}' for k, v in argument.items())
  
      def fetch_tileable_op_logs(self, tileable_op_key, offsets=None, sizes=None):
          url = f'{self._endpoint}/api/session/{self._session_id}/op/{tileable_op_key}/log'
          params = dict()
          self._process_int_or_dict_argument(offsets, 'offsets', params)
          self._process_int_or_dict_argument(sizes, 'sizes', params)
          resp = self._req_session.get(url, params=params)
          if resp.status_code >= 400:
              raise ValueError(f'Failed to fetch log from server. Code: {resp.status_code}, '
                               f'Reason: {resp.reason}, Content:\n{resp.text}')
          return json.loads(resp.content)
  
      def fetch_log(self, tileables, offsets=None, sizes=None):
          from ..custom_log import fetch
  
          return fetch(tileables, self, offsets=offsets, sizes=sizes)
  
      def get_named_tileable_infos(self, name):
          from ..context import TileableInfos
  
          url = f'{self._endpoint}/api/session/{self._session_id}'
          params = dict(name=name)
          resp = self._req_session.get(url, params=params)
          if resp.status_code >= 400:  # pragma: no cover
              raise ValueError(f'Failed to get tileable key from server. Code: {resp.status_code}, '
                               f'Reason: {resp.reason}, Content:\n{resp.text}')
          tileable_key = self._handle_json_response(resp)['tileable_key']
          nsplits, extra_meta = self._get_tileable_meta(tileable_key)
          shape = tuple(sum(s) for s in nsplits)
          return TileableInfos(tileable_key, shape, extra_meta)
  
      def create_mutable_tensor(self, name, shape, dtype, fill_value=None, chunk_size=None, *_, **__):
          from ..tensor.utils import create_mutable_tensor
          session_url = f'{self._endpoint}/api/session/{self._session_id}'
          tensor_url = f'{session_url}/mutable-tensor/{name}?action=create'
          if not isinstance(dtype, np.dtype):
              dtype = np.dtype(dtype)
          # avoid built-in scalar dtypes are made into one-field record type.
          if dtype.fields:
              dtype_descr = dtype.descr
          else:
              dtype_descr = str(dtype)
          tensor_json = {
              'shape': shape,
              'dtype': dtype_descr,
              'fill_value': fill_value,
              'chunk_size': chunk_size,
          }
          resp = self._req_session.post(tensor_url, json=tensor_json)
          shape, dtype, chunk_size, chunk_keys, chunk_eps = self._handle_json_response(resp)
          return create_mutable_tensor(name, chunk_size, shape, numpy_dtype_from_descr_json(dtype),
                                       chunk_keys, chunk_eps)
  
      def get_mutable_tensor(self, name):
          from ..tensor.utils import create_mutable_tensor
          session_url = f'{self._endpoint}/api/session/{self._session_id}'
          tensor_url = f'{session_url}/mutable-tensor/{name}'
          resp = self._req_session.get(tensor_url)
          shape, dtype, chunk_size, chunk_keys, chunk_eps = self._handle_json_response(resp)
          return create_mutable_tensor(name, chunk_size, shape, numpy_dtype_from_descr_json(dtype),
                                       chunk_keys, chunk_eps)
  
      def write_mutable_tensor(self, tensor, index, value):
          """
          How to serialize index and value:
  
          1. process_index and serialize it as json
          2. the payload of POST request:
  
              * a int64 value indicate the size of index json
              * ascii-encoded bytes of index json
              * pyarrow serialized bytes of `value`
          """
          from ..tensor.core import Indexes
          from ..serialize import dataserializer
  
          index = Indexes(indexes=index)
          index_bytes = base64.b64encode(serialize_serializable(index))
          bio = BytesIO()
          bio.write(np.int64(len(index_bytes)).tobytes())
          bio.write(index_bytes)
          dataserializer.dump(value, bio)
  
          session_url = f'{self._endpoint}/api/session/{self._session_id}'
          tensor_url = f'{session_url}/mutable-tensor/{tensor.name}'
          resp = self._req_session.put(tensor_url, data=bio.getvalue(),
                                       headers={'Content-Type': 'application/octet-stream'})
          self._handle_json_response(resp)
  
      def seal(self, tensor):
          from ..tensor.utils import create_fetch_tensor
          session_url = f'{self._endpoint}/api/session/{self._session_id}'
          tensor_url = f'{session_url}/mutable-tensor/{tensor.name}?action=seal'
          resp = self._req_session.post(tensor_url)
          graph_key_hex, tileable_key, tensor_id, tensor_meta = self._handle_json_response(resp)
          self._executed_tileables[tileable_key] = uuid.UUID(graph_key_hex), {tensor_id}
  
          # # Construct Tensor on the fly.
          shape, dtype, chunk_size, chunk_keys, _ = tensor_meta
          return create_fetch_tensor(chunk_size, shape, numpy_dtype_from_descr_json(dtype),
                                     tensor_key=tileable_key, chunk_keys=chunk_keys)
  
      def _get_tileable_nsplits(self, tileable_key):
          session_url = f'{self._endpoint}/api/session/{self._session_id}'
          graph_key = self._get_tileable_graph_key(tileable_key)
          url = f'{session_url}/graph/{graph_key}/data/{tileable_key}?type=nsplits'
          resp = self._req_session.get(url)
          new_nsplits = self._handle_json_response(resp)
          return new_nsplits
  
      def _get_tileable_meta(self, tileable_key):
          session_url = f'{self._endpoint}/api/session/{self._session_id}'
          graph_key = self._get_tileable_graph_key(tileable_key)
          url = f'{session_url}/graph/{graph_key}/data/{tileable_key}?type=meta'
          resp = self._req_session.get(url)
          meta = self._handle_json_response(resp)
          return pickle.loads(base64.b64decode(meta))  # nosec
  
      def _update_tileable_shape(self, tileable):
          tileable_key = tileable.key
          new_nsplits = self._get_tileable_nsplits(tileable_key)
          tileable._update_shape(tuple(sum(nsplit) for nsplit in new_nsplits))
          tileable.nsplits = new_nsplits
  
      def decref(self, *keys):
          for tileable_key, tileable_id in keys:
              if tileable_key not in self._executed_tileables:
                  continue
              graph_key, ids = self._executed_tileables[tileable_key]
  
              if tileable_id in ids:
                  ids.remove(tileable_id)
                  # for those same key tileables, do decref only when all those tileables are garbage collected
                  if len(ids) != 0:
                      continue
                  self.delete_data(tileable_key)
  
      def delete_data(self, tileable_key, wait=False):
          if tileable_key not in self._executed_tileables:
              return
          graph_key, _ids = self._executed_tileables[tileable_key]
          data_url = f'{self._endpoint}/api/session/{self._session_id}/graph/{graph_key}' \
                     f'/data/{tileable_key}?wait={1 if wait else 0}'
          self._req_session.delete(data_url)
          self._executed_tileables.pop(tileable_key, None)
  
      def stop(self, graph_key):
          session_url = f'{self._endpoint}/api/session/{self._session_id}'
          graph_url = session_url + '/graph/' + graph_key
          resp = self._req_session.delete(graph_url)
          if resp.status_code >= 400:
              raise SystemError(f'Failed to stop graph execution. Code: {resp.status_code}, '
                                f'Reason: {resp.reason}, Content:\n{resp.text}')
  
      def _submit_graph(self, serialized_graph, targets, names=None, compose=True):
          session_url = f'{self._endpoint}/api/session/{self._session_id}'
          resp = self._req_session.post(session_url + '/graph', dict(
              graph=base64.b64encode(serialized_graph).decode('ascii'),
              target=targets,
              names=names,
              compose='1' if compose else '0'
          ))
          return self._handle_json_response(resp)
  
      def get_graph_states(self):
          resp = self._req_session.get(f'{self._endpoint}/api/session/{self._session_id}/graph')
          return self._handle_json_response(resp)
  
      def close(self):
          session_url = f'{self._endpoint}/api/session/{self._session_id}'
          for key in list(self._executed_tileables.keys()):
              self.delete_data(key, wait=True)
          resp = self._req_session.delete(session_url)
          if resp.status_code >= 400:
              raise SystemError('Failed to close mars session.')
  
      def check_service_ready(self, timeout=1):
          import requests
          try:
              resp = self._req_session.get(self._endpoint + '/api', timeout=timeout)
          except (requests.ConnectionError, requests.Timeout):
              return False
          if resp.status_code >= 400:
              return False
          return True
  
      def count_workers(self):
          resp = self._req_session.get(self._endpoint + '/api/worker?action=count', timeout=1)
          return self._handle_json_response(resp)
  
      def get_cpu_count(self):
          resp = self._req_session.get(self._endpoint + '/api/worker?action=count_cpu', timeout=1)
          return self._handle_json_response(resp)
  
      def rescale_workers(self, new_scale, min_workers=None, wait=True, timeout=None):
          data = json.dumps(dict(new_scale=new_scale, min_workers=min_workers))
          wait_req = 1 if wait else 0
          resp = self._req_session.patch(f'{self._endpoint}/api/worker?action=count&wait={wait_req}',
                                         data, timeout=timeout)
          return self._handle_json_response(resp)
  
      def get_workers_meta(self):
          resp = self._req_session.get(self._endpoint + '/api/worker', timeout=1)
          return self._handle_json_response(resp)
  
      def get_task_count(self):
          resp = self._req_session.get(f'{self._endpoint}/api/session/{self._session_id}/graph')
          return len(self._handle_json_response(resp))
  
      def __enter__(self):
          return self
  
      def __exit__(self, *_):
          self.close()
  
apis: numpy.dtype