pete_icc_profile_cache.py

# Copyright 2024 Google LLC
#
# 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.

"""Cache for ICC Profiles for DICOM Slides.

WSI DICOM is captured in a slide scanner defined color space. The ICC Profile
information is stored within the DICOM. ICC Profiles may be quite large, e.g.
the ICC Profile generated by the Leica (Aperio) scanner is ~12 MB. The
size of the ICC Profile may greatly exceed the size of the pixel image data.
Fetching the ICC Profile from the DICOM store can be expensive if the DICOM
instance is stored in archival storage tiers.

To speed ICC profile retrieval and reducte the cost, the pathology emebedding
service can cache and re-use ICC Profiles across requests. ICC profiles
describe a scanner calibration and are not indicative of patient data.

ICC Profiles bytes are cached bytes in Redis or GCS,
(keyed by the hash of the ICC Profile bytes) and in memory.
For each slide a key value pair is stored in cloud redis
which maps the slide path to the hash value of the slides ICC Profile. A small
in memory cache is used to avoid repeated Redis/GCS retrievals for the most
recently used ICC Profiles.

The storage location (Redis or GCS) for the ICC Profile bytes is determined by
the pete flags
"""

import base64
import hashlib
import io
import os
import threading
import time
from typing import Any, Mapping, MutableMapping, Optional, Union

import cachetools
from ez_wsi_dicomweb import dicom_slide
from google.api_core import exceptions
from google.cloud import storage
import google.cloud.exceptions
import redis

from serving import pete_flags
from serving.logging_lib import cloud_logging_client


_GCS_PREFIX = 'gs://'

# Cache instance index for up to 1 day (min lifetime for ILM gcs object).
_SECONDS_IN_A_DAY = 24 * 60 * 60
_LOCALMETADATA_CACHESIZE = 4  # Cache up to last 4 icc profiles.
_LOCALMETADATA_INDEX_CACHESIZE = 100  # Cache up to last 100 instance keys.
_NO_ICC_PROFILE_BYTES = b''

_local_icc_profile_bytes_index_cache = cachetools.TTLCache(
    _LOCALMETADATA_INDEX_CACHESIZE, _SECONDS_IN_A_DAY
)
_local_icc_profile_bytes_cache = cachetools.TTLCache(
    _LOCALMETADATA_CACHESIZE, _SECONDS_IN_A_DAY
)
_local_icc_profile_bytes_cache_lock = threading.Lock()

# structured logging keys
_DICOM_PATH_CACHE_KEY = 'dicom_path_cache_key'
_DICOM_PATH_KEY = 'dicom_path'
_ICC_PROFILE_SIZE_KEY = 'icc_profile_size'
_ICC_PROFILE_HASH_KEY = 'icc_profile_hash'
_ICC_PROFILE_RETRIEVAL_TIME_KEY = 'icc_profile_retrieval_time'


def _init_fork_module_state() -> None:
  global _local_icc_profile_bytes_cache
  global _local_icc_profile_bytes_index_cache
  global _local_icc_profile_bytes_cache_lock
  _local_icc_profile_bytes_cache = cachetools.TTLCache(
      _LOCALMETADATA_CACHESIZE, _SECONDS_IN_A_DAY
  )
  _local_icc_profile_bytes_index_cache = cachetools.TTLCache(
      _LOCALMETADATA_INDEX_CACHESIZE, _SECONDS_IN_A_DAY
  )
  _local_icc_profile_bytes_cache_lock = threading.Lock()


def _cache_key_gcs_path(
    gcs_icc_profile_cache_bucket: str, dicom_slide_path_cache_key: str
) -> str:
  return f'gs://{gcs_icc_profile_cache_bucket}/slide_keys/{dicom_slide_path_cache_key}'


def _get_gcs_cache_path(
    gcs_icc_profile_cache_bucket: str, icc_profile_hash: str
) -> str:
  """Returns GS style path to ICC Profile in cache."""
  return f'gs://{gcs_icc_profile_cache_bucket}/{icc_profile_hash}.icc_profile'


def _redis_set(
    redis_client: redis.Redis, redis_cache_key: str, data: bytes
) -> None:
  """Sets value in redis cache. If cache returns OOM empty cache."""
  try:
    redis_client.set(redis_cache_key, data, nx=False, ex=_SECONDS_IN_A_DAY)
  except redis.exceptions.ConnectionError as exp:
    cloud_logging_client.warning('Unexpected connecting to Redis cache.', exp)
    return
  except redis.exceptions.ResponseError as exp:
    if 'command not allowed under OOM prevention' not in str(exp):
      return
    try:
      lazyfree_pending_objects = redis_client.info('memory').get(
          'lazyfree_pending_objects', 0
      )
      if lazyfree_pending_objects > 0:
        cloud_logging_client.info(
            f'Redis cache has N={lazyfree_pending_objects} '
            'lazyfree_pending_objects. Redis not flushed.'
        )
        return
      redis_client.flushall(asynchronous=True)
      cloud_logging_client.info('Flushed Redis cache.')
    except (
        redis.exceptions.ConnectionError,
        redis.exceptions.ResponseError,
    ) as redis_exp:
      cloud_logging_client.warning(
          'Unexpected error flushing Redis cache.', redis_exp
      )


def _upload_file_to_blob(blob: storage.Blob, open_file: io.BytesIO) -> None:
  """Uploads file to blob; also Mock target for testing."""
  blob.upload_from_file(open_file)


def _update_profile_cache(
    dicom_slide_path_cache_key: str,
    icc_profile: bytes,
    struct_log: MutableMapping[str, Any],
):
  """Updates ICC Profile cache in memory, redis (instance index), and in GCS."""
  cloud_logging_client.debug('Saving ICC profile to cache.')
  gcs_icc_profile_cache_bucket = (
      pete_flags.ICC_PROFILE_CACHE_GCS_BUCKET_FLAG.value
  )
  redis_host = pete_flags.ICC_PROFILE_CACHE_REDIS_IP_FLAG.value
  redis_port = pete_flags.ICC_PROFILE_CACHE_REDIS_PORT_FLAG.value
  if icc_profile:
    # if icc profile bytes compute hash and icc_profile storage locations.
    icc_profile_hash = hashlib.sha3_512(icc_profile).hexdigest()
    icc_profile_gcs_cache_path = _get_gcs_cache_path(
        gcs_icc_profile_cache_bucket, icc_profile_hash
    )
    icc_profile_hash = icc_profile_hash.encode('utf-8')
  else:
    # unused place holder values
    icc_profile_gcs_cache_path = ''
    icc_profile_hash = _NO_ICC_PROFILE_BYTES
  struct_log[_ICC_PROFILE_HASH_KEY] = icc_profile_hash
  struct_log[_ICC_PROFILE_SIZE_KEY] = len(icc_profile)
  with _local_icc_profile_bytes_cache_lock:
    # update local in memory cache
    if not icc_profile:
      # no icc profile just set place holder in local in memory cache to show
      # hash value has no associated icc profile
      cloud_logging_client.debug(
          'DICOM does not have ICC profile storing state in local cache.',
          struct_log,
      )
      _local_icc_profile_bytes_index_cache[dicom_slide_path_cache_key] = (
          _NO_ICC_PROFILE_BYTES
      )
    else:
      cloud_logging_client.debug(
          'DICOM has ICC profile storing in local cache.',
          struct_log,
      )
      # set icc profile hash and cache icc profile by hash value in the
      # in memory cache.
      _local_icc_profile_bytes_index_cache[dicom_slide_path_cache_key] = (
          icc_profile_hash
      )
      _local_icc_profile_bytes_cache[icc_profile_hash] = icc_profile
  client = None
  if redis_host:
    # if redis host is defined cache the key and optionally icc profile bytes
    # in redis.
    with redis.Redis(host=redis_host, port=redis_port) as redis_client:
      if not icc_profile:
        # if no icc profile set cache to place holder indicating no icc profile
        # bytes stored and return.
        cloud_logging_client.debug(
            'DICOM does not have ICC profile storing state in redis.',
            struct_log,
        )
        _redis_set(
            redis_client, dicom_slide_path_cache_key, _NO_ICC_PROFILE_BYTES
        )
        return
      # set cache key (dicon path) to point to ICC profile hash.
      cloud_logging_client.debug(
          'Storing slide path hashed reference to ICC profile bytes in Redis.',
          struct_log,
      )
      _redis_set(redis_client, dicom_slide_path_cache_key, icc_profile_hash)
      if pete_flags.STORE_ICC_PROFILE_BYTES_IN_REDIS_FLAG.value:
        cloud_logging_client.debug(
            'Storing ICC profile bytes in Redis.',
            struct_log,
        )
        # if icc profile bytes are stored in redis set key associated with
        # the icc profile hash to the icc profile bytes and return.
        _redis_set(redis_client, icc_profile_gcs_cache_path, icc_profile)
        return
  elif gcs_icc_profile_cache_bucket:
    # store key in gcs
    cloud_logging_client.debug(
        'Storing slide path hashed reference to ICC profile bytes in GCS.',
        struct_log,
    )
    client = storage.Client()
    blob = storage.Blob.from_string(
        _cache_key_gcs_path(
            gcs_icc_profile_cache_bucket, dicom_slide_path_cache_key
        ),
        client=storage.Client(),
    )
    if not icc_profile:
      # store place holder value indicating no stored bytes for dicom path
      # and return
      with io.BytesIO(_NO_ICC_PROFILE_BYTES) as upload_bytes:
        _upload_file_to_blob(blob, upload_bytes)
      return
    # upload icc profile hash to blob providing reference from the DICOM
    # path to the stored icc profile.
    with io.BytesIO(icc_profile_hash) as upload_bytes:
      _upload_file_to_blob(blob, upload_bytes)
  if not gcs_icc_profile_cache_bucket:
    return
  if client is None:
    client = storage.Client()
  cloud_logging_client.debug(
      'Storing ICC profile bytes in GCS.',
      struct_log,
  )
  # store icc profile bytes in GCS.
  try:
    blob = storage.Blob.from_string(icc_profile_gcs_cache_path, client=client)
    if blob.exists():
      try:
        blob.reload()
        # Its possible cache is updated due to redis Hash key being lost but
        # the data remains in GCS. If blob exists and is identical to the
        # version in the container, don't upload it again. Check against size
        #  and GCS generated hash to validate actual bytes stored in GCS.
        if blob.size == len(icc_profile):
          hash_val = base64.b64encode(
              bytes.fromhex(hashlib.md5(icc_profile).hexdigest())
          ).decode('utf-8')
          if blob.md5_hash == hash_val:
            return
      except exceptions.NotFound:
        pass
    with io.BytesIO(icc_profile) as icc_profile_file:
      _upload_file_to_blob(blob, icc_profile_file)
  except google.cloud.exceptions.GoogleCloudError as exp:
    cloud_logging_client.warning('Unexpected uploading blob to GCS.', exp)
    # running in thread eat exception.


def _run_cache_update_thread(th: threading.Thread) -> None:
  """Starts cache update thread; also mock target for testing."""
  th.start()


def _download_blob_as_bytes(blob: storage.Blob) -> bytes:
  """Downloads blob as bytes; also Mock target for testing."""
  return blob.download_as_bytes()


def _download_icc_profile(
    gcs_cache_path: str,
    redis_host: str,
    redis_port: int,
    struct_log: Mapping[str, Any],
) -> Optional[bytes]:
  """Downloads ICC Profile from GCS or Redis."""
  if redis_host and pete_flags.STORE_ICC_PROFILE_BYTES_IN_REDIS_FLAG.value:
    try:
      cloud_logging_client.debug(
          'Searching Redis for cached ICC profile.', struct_log
      )
      with redis.Redis(host=redis_host, port=redis_port) as redis_client:
        return redis_client.get(gcs_cache_path)
    except redis.exceptions.ConnectionError as exp:
      cloud_logging_client.warning('Unexpected connecting to redis cache.', exp)
      return None
  if not gcs_cache_path:
    return None
  cloud_logging_client.debug(
      'Searching GCS for cached ICC profile.', struct_log
  )
  client = storage.Client()
  blob = storage.Blob.from_string(gcs_cache_path, client=client)
  try:
    return _download_blob_as_bytes(blob)
  except exceptions.NotFound:
    return None
  except google.cloud.exceptions.GoogleCloudError as exp:
    cloud_logging_client.warning('Unexpected downloading blob from GCS.', exp)
    return None


def _get_dicom_slide_icc_profile(
    slide: Union[dicom_slide.DicomSlide, dicom_slide.DicomMicroscopeImage],
    level: Optional[Union[dicom_slide.Level, dicom_slide.ResizedLevel]] = None,
) -> bytes:
  if isinstance(slide, dicom_slide.DicomSlide):
    return slide.get_icc_profile_bytes()
  if isinstance(slide, dicom_slide.DicomMicroscopeImage):
    return slide.get_level_icc_profile_bytes(level)
  raise ValueError('Unexpected object')


def _normalize_bucket_name(bucket_name: str) -> str:
  bucket_name = bucket_name.rstrip('/')
  if bucket_name.lower().startswith(_GCS_PREFIX):
    return bucket_name[len(_GCS_PREFIX) :]
  return bucket_name


def get_dicom_icc_profile(
    slide: Union[dicom_slide.DicomSlide, dicom_slide.DicomMicroscopeImage],
    slide_level: Optional[
        Union[dicom_slide.Level, dicom_slide.ResizedLevel]
    ] = None,
) -> bytes:
  """Returns the ICC Profile bytes a DICOM Slide."""
  start_time = time.time()
  gcs_icc_profile_cache_bucket = (
      pete_flags.ICC_PROFILE_CACHE_GCS_BUCKET_FLAG.value
  )
  redis_host = pete_flags.ICC_PROFILE_CACHE_REDIS_IP_FLAG.value
  redis_port = pete_flags.ICC_PROFILE_CACHE_REDIS_PORT_FLAG.value

  gcs_icc_profile_cache_bucket = _normalize_bucket_name(
      gcs_icc_profile_cache_bucket
  )
  # Retrieve HASH of icc profile bytes from redis cache.
  # ICC Profiles stored in by HASH of bytes to enable commonly used profiles
  # to be shared across instances.
  # use hash of slide path as key
  dicom_slide_string_path = str(slide.path)
  dicom_slide_path = dicom_slide_string_path.encode('utf-8')
  dicom_slide_path_cache_key = hashlib.sha3_512(dicom_slide_path).hexdigest()
  struct_log = {
      _DICOM_PATH_CACHE_KEY: dicom_slide_path_cache_key,
      _DICOM_PATH_KEY: dicom_slide_string_path,
  }

  cloud_logging_client.debug(
      'Generating hash key for path to identify cached ICC profile.',
      struct_log,
  )
  with _local_icc_profile_bytes_cache_lock:
    if pete_flags.IS_DEBUGGING_FLAG.value:
      icc_profile_hash = None
    else:
      cloud_logging_client.debug(
          'Searching local cache for slide ICC profile hash.', struct_log
      )
      icc_profile_hash = _local_icc_profile_bytes_index_cache.get(
          dicom_slide_path_cache_key
      )
  if not gcs_icc_profile_cache_bucket and (
      not redis_host
      or not pete_flags.STORE_ICC_PROFILE_BYTES_IN_REDIS_FLAG.value
  ):
    cloud_logging_client.warning(
        'Performance warning: No Redis server or GCS bucket defined to back'
        ' DICOM ICC profile cache.'
    )
  if icc_profile_hash is None:
    if redis_host:
      cloud_logging_client.debug(
          'Searching Redis for cached DICOM path ICC profile hash.', struct_log
      )
      try:
        with redis.Redis(host=redis_host, port=redis_port) as redis_client:
          icc_profile_hash = redis_client.get(dicom_slide_path_cache_key)
          if isinstance(icc_profile_hash, str):
            icc_profile_hash = icc_profile_hash.encode('utf-8')
      except redis.exceptions.ConnectionError as exp:
        cloud_logging_client.warning(
            'Unexpected connecting to redis cache.', exp
        )
        icc_profile_hash = None
    elif gcs_icc_profile_cache_bucket:
      cloud_logging_client.debug(
          'Searching GCS for cached DICOM path ICC profile hash.', struct_log
      )
      try:
        blob = storage.Blob.from_string(
            _cache_key_gcs_path(
                gcs_icc_profile_cache_bucket, dicom_slide_path_cache_key
            ),
            client=storage.Client(),
        )
        icc_profile_hash = _download_blob_as_bytes(blob)
      except (
          exceptions.NotFound,
          google.cloud.exceptions.GoogleCloudError,
      ) as _:
        icc_profile_hash = None
  if icc_profile_hash is not None:  # if hash entry is found
    struct_log[_ICC_PROFILE_HASH_KEY] = icc_profile_hash
    cloud_logging_client.debug(
        'Found cached DICOM path icc profile hash.', struct_log
    )
    if icc_profile_hash == _NO_ICC_PROFILE_BYTES:
      elapsed_time = time.time() - start_time
      struct_log[_ICC_PROFILE_RETRIEVAL_TIME_KEY] = elapsed_time
      cloud_logging_client.info(
          f'Slide has no ICC profile bytes (local cache); {elapsed_time} sec.',
          struct_log,
      )
      return _NO_ICC_PROFILE_BYTES  # no icc profile for dicom series.
    # test local small in memory cache for hashed bytes.
    with _local_icc_profile_bytes_cache_lock:
      _local_icc_profile_bytes_index_cache[dicom_slide_path_cache_key] = (
          icc_profile_hash
      )
      icc_profile = _local_icc_profile_bytes_cache.get(icc_profile_hash)
      if not pete_flags.IS_DEBUGGING_FLAG.value and icc_profile is not None:
        struct_log[_ICC_PROFILE_SIZE_KEY] = len(icc_profile)
        elapsed_time = time.time() - start_time
        struct_log[_ICC_PROFILE_RETRIEVAL_TIME_KEY] = elapsed_time
        cloud_logging_client.info(
            f'Returning ICC profile bytes (local cache); {elapsed_time} sec.',
            struct_log,
        )
        return icc_profile
    gcs_cache_path = _get_gcs_cache_path(
        gcs_icc_profile_cache_bucket, icc_profile_hash.decode('utf-8')
    )
    # try to retireve from GCS.
    icc_profile_bytes = _download_icc_profile(
        gcs_cache_path, redis_host, redis_port, struct_log
    )
    if icc_profile_bytes is not None:
      with _local_icc_profile_bytes_cache_lock:
        _local_icc_profile_bytes_cache[icc_profile_hash] = icc_profile_bytes
      struct_log[_ICC_PROFILE_SIZE_KEY] = (
          len(icc_profile) if icc_profile is not None else 'NONE'
      )
      elapsed_time = time.time() - start_time
      struct_log[_ICC_PROFILE_RETRIEVAL_TIME_KEY] = elapsed_time
      cloud_logging_client.info(
          f'Returning cached slide ICC profile; {elapsed_time} sec.', struct_log
      )
      return icc_profile_bytes
  # Could not find ICC Profile in cache. Fetch from DICOM store and update
  # cache.
  icc_profile = _get_dicom_slide_icc_profile(slide, slide_level)
  struct_log[_ICC_PROFILE_SIZE_KEY] = len(icc_profile)
  cloud_logging_client.debug(
      'Retrieved ICC profile from DICOM store.', struct_log
  )
  th = threading.Thread(
      target=_update_profile_cache,
      args=(dicom_slide_path_cache_key, icc_profile, struct_log),
      daemon=True,
  )
  _run_cache_update_thread(th)
  elapsed_time = time.time() - start_time
  struct_log[_ICC_PROFILE_RETRIEVAL_TIME_KEY] = elapsed_time
  cloud_logging_client.info(
      f'Returning ICC profile loaded from the DICOM store; {elapsed_time} sec.',
      struct_log,
  )
  return icc_profile


os.register_at_fork(after_in_child=_init_fork_module_state)