Branch-redistribution feature about jacquard HOT 1 OPEN

import random
import itertools
import jacquard.config

from jacquard.buckets import Bucket, NUM_BUCKETS
from jacquard.experiments.constraints import Constraints

EXPERIMENT_NAME = '<redacted>'
REMOVING_BRANCH = '<redacted>'
EXPERIMENT_CONSTRAINTS = Constraints()

REDISTRIBUTE_TO = (
    {
        'branch': '<redacted>',
        'settings': {'<redacted>': 2},
    },
    {
        'branch': 'control',
        'settings': {'<redacted>': 1},
    },
)

config = jacquard.config.load_config('/etc/jacquard/config.cfg')

redistributions = itertools.cycle(REDISTRIBUTE_TO)

removal_key = [EXPERIMENT_NAME, REMOVING_BRANCH]

with config.storage.transaction() as store:
    bucket_order = list(range(NUM_BUCKETS))
    # Permute buckets to avoid bias
    random.shuffle(bucket_order)

    for bucket_number in bucket_order:
        store_key = 'buckets/%s' % bucket_number
        bucket = Bucket.from_json(
            store.get(store_key, '{}'),
        )

        if not bucket.covers(removal_key):
            continue

        print("Modifying bucket %s" % bucket_number)

        bucket.remove(removal_key)

        new_settings = next(redistributions)
        bucket.add(
            [EXPERIMENT_NAME, new_settings['branch']],
            new_settings['settings'],
            EXPERIMENT_CONSTRAINTS,
        )

        store[store_key] = bucket.to_json()