Custom benchmarking track and script for EsRally

track.json

{% import "rally.helpers" as rally with context %}
{
  "version": 2,
  "description": "Tracker-generated track for wiki",
  "indices": [
    {
      "name": "wiki",
      "body": "wiki.json"
    }
  ],
  "corpora": [
    {
      "name": "wiki",
      "documents": [
        {
          "target-index": "wiki",
          "source-file": "wiki-documents.json.bz2",
          "document-count": 34633,
          "compressed-bytes": 169107793,
          "uncompressed-bytes": 600965180
        }
      ]
    }
  ],
  "schedule": [
    {
      "operation": "delete-index"
    },
    {
      "operation": {
        "operation-type": "create-index",
        "settings": {{index_settings | default({}) | tojson}}
      }
    },
    {
      "operation": {
        "operation-type": "cluster-health",
        "index": "wiki",
        "request-params": {
          "wait_for_status": "{{cluster_health | default('green')}}",
          "wait_for_no_relocating_shards": "true"
        },
        "retry-until-success": true
      }
    },
    {
      "operation": {
        "operation-type": "bulk",
        "bulk-size": {{bulk_size | default(5000)}},
        "ingest-percentage": {{ingest_percentage | default(100)}}
      },
      "clients": {{bulk_indexing_clients | default(8)}}
    },
    {
      "operation": {
        "operation-type": "search",
        "param-source": "my-custom-param-source"
      }
    }
  ]
}

track.py

import random
import glob
import json
import os

class QueryParamSource:
    def __init__(self, track, params, **kwargs):
        if len(track.indices) == 1:
            default_index = track.indices[0].name
        else:
            default_index = "wiki"

        self.index_name = params.get("index", default_index)
        self.infinite = True
        random.seed(4)
        cwd = os.path.dirname(__file__)
        self.queries = glob.glob(os.path.join(cwd, "queries/*"))

    def partition(self, partition_index, total_partitions):
        return self
    
    def params(self):
        query_file = random.choice(self.queries)
        with open(query_file, 'r') as f:
            query = json.load(f)
            result = {
                "body": query,
                "index": self.index_name,
            }

            return result

def register(registry):
    registry.register_param_source("my-custom-param-source", QueryParamSource)