def _deployment_config_lru_cache()

in src/sagemaker/jumpstart/utils.py [0:0]

• 35 lines of code
• 16 McCabe index (conditional complexity)

def _deployment_config_lru_cache(_func=None, *, maxsize: int = 128, typed: bool = False):
    """LRU cache for deployment configs."""

    def has_instance_rate_metric(config: DeploymentConfigMetadata) -> bool:
        """Determines whether metadata config contains instance rate metric stat.

        Args:
            config (DeploymentConfigMetadata): Metadata config metadata.
        Returns:
            bool: Whether the metadata config contains instance rate metric stat.
        """
        if config.benchmark_metrics is None:
            return True
        for benchmark_metric_stats in config.benchmark_metrics.values():
            if not has_instance_rate_stat(benchmark_metric_stats):
                return False
        return True

    def wrapper_cache(f):
        f = lru_cache(maxsize=maxsize, typed=typed)(f)

        @wraps(f)
        def wrapped_f(*args, **kwargs):
            res = f(*args, **kwargs)

            # Clear cache on first call if
            #   - The output does not contain Instant rate metrics
            #   as this is caused by missing policy.
            if f.cache_info().hits == 0 and f.cache_info().misses == 1:
                if isinstance(res, list):
                    for item in res:
                        if isinstance(
                            item, DeploymentConfigMetadata
                        ) and not has_instance_rate_metric(item):
                            f.cache_clear()
                            break
                elif isinstance(res, dict):
                    keys = list(res.keys())
                    if len(keys) == 0 or "Instance Rate" not in keys[-1]:
                        f.cache_clear()
                    elif len(res[keys[1]]) > len(res[keys[-1]]):
                        del res[keys[-1]]
                        f.cache_clear()
            return res

        wrapped_f.cache_info = f.cache_info
        wrapped_f.cache_clear = f.cache_clear
        return wrapped_f

    if _func is None:
        return wrapper_cache
    return wrapper_cache(_func)