mset wrapper that batches keys per shard and execute as few
        msets as necessary to set the keys in all the shards involved.

        This method should be invoked on a TwemRedis instance as if it
        were being invoked directly on a StrictRedis instance.

def mset(self, args):
        """
        mset wrapper that batches keys per shard and execute as few
        msets as necessary to set the keys in all the shards involved.

        This method should be invoked on a TwemRedis instance as if it
        were being invoked directly on a StrictRedis instance.
        """
        key_map = collections.defaultdict(dict)
        result_count = 0
        for key in args.keys():
            value = args[key]
            shard_num = self.get_shard_num_by_key(key)
            key_map[shard_num][key] = value

        # TODO: parallelize
        for shard_num in key_map.keys():
            shard = self.get_shard_by_num(shard_num)
            result_count += shard.mset(key_map[shard_num])

        return result_count