def get_user_variant()

in lambda/api/lambda_invoke.py [0:0]

• 40 lines of code
• 18 McCabe index (conditional complexity)

def get_user_variant(endpoint_name: str, user_id: str):
    # Get the variants metrics (this will fail if endpoint doesn't exist)
    strategy, epsilon, warmup, variant_metrics = exp_metrics.get_variant_metrics(
        endpoint_name
    )

    # Get the configuration for the endpoint name
    logger.info(f"Getting variant for user: {user_id}")
    user_variant = exp_assignment.get_assignment(
        user_id=user_id, endpoint_name=endpoint_name
    )

    # Ensure that our user variant is still in current metrics
    target_variant = user_variant
    if user_variant is not None:
        user_match = [v for v in variant_metrics if v["variant_name"] == user_variant]
        if len(user_match) == 0:
            logger.info(f"User variant {user_variant} not in endpoint variants")
            target_variant = None

    # Get the new target variant if not assigned
    status_code = 200
    if target_variant is None:
        # See if all variants have invocation metrics
        with_invocations = [
            v for v in variant_metrics if v["invocation_count"] > warmup
        ]
        if len(with_invocations) < len(variant_metrics):
            strategy = WeightedSampling.STRATEGY_NAME
            algo = WeightedSampling(variant_metrics)
        elif strategy == WeightedSampling.STRATEGY_NAME:
            algo = WeightedSampling(variant_metrics)
        elif strategy == ThompsonSampling.STRATEGY_NAME:
            algo = ThompsonSampling(variant_metrics)
        elif strategy == EpsilonGreedy.STRATEGY_NAME:
            algo = EpsilonGreedy(variant_metrics, epsilon)
        elif strategy == UCB1.STRATEGY_NAME:
            algo = UCB1(variant_metrics)
        else:
            raise Exception(f"Strategy {strategy} not supported")
        target_variant = algo.select_variant()
        status_code = 201

    # Assign the target variant to the user
    if user_variant != target_variant:
        logger.info(f"Set target variant: {target_variant} for user: {user_id}")
        exp_assignment.put_assignment(
            user_id=user_id, endpoint_name=endpoint_name, variant_name=target_variant
        )

    # Return the result
    return strategy, target_variant, status_code