main.py 文件源码

def schedule(test_run, shard_count, shard):
    # read tests from body
    test_names = flask.request.get_json(force=True)['tests']

    # first see if we have a scedule already
    schedule_id = "%s-%d" % (test_run, shard_count)
    schedule = Schedule.get_by_id(schedule_id)
    if schedule is not None:
        return flask.json.jsonify(tests=schedule.shards[str(shard)])

    # if not, do simple greedy algorithm
    test_times = ndb.get_multi(
        ndb.Key(Test, test_name) for test_name in test_names)

    def avg(test):
        if test is not None:
            return test.cost()
        return 1

    test_times = [(test_name, avg(test))
                  for test_name, test in zip(test_names, test_times)]
    test_times_dict = dict(test_times)
    test_times.sort(key=operator.itemgetter(1))

    shards = {i: [] for i in xrange(shard_count)}
    while test_times:
        test_name, time = test_times.pop()

        # find shortest shard and put it in that
        s, _ = min(
            ((i, sum(test_times_dict[t] for t in shards[i]))
             for i in xrange(shard_count)),
            key=operator.itemgetter(1))

        shards[s].append(test_name)

    # atomically insert or retrieve existing schedule
    schedule = Schedule.get_or_insert(schedule_id, shards=shards)
    return flask.json.jsonify(tests=schedule.shards[str(shard)])