Class: Polyrun::Partition::Plan

Inherits:
Object
  • Object
show all
Defined in:
lib/polyrun/partition/plan.rb,
lib/polyrun/partition/plan_sharding.rb

Overview

Assigns discrete items (e.g. spec paths, or path:line example locators) to shards (spec_queue.md).

Strategies:

  • round_robin — sorted paths, assign by index mod total_shards.
  • random_round_robin — Fisher–Yates shuffle (optional seed), then same mod assignment.
  • cost_binpack (+cost+, binpack, timing) — LPT greedy binpack using per-item weights; optional Constraints for pins / serial globs before LPT on the rest. Default timing_granularity is file (one weight per spec file). Experimental :example uses path:line locators and per-example weights in the timing JSON.
  • hrw (+rendezvous+) — rendezvous hashing for minimal remapping when m changes; optional constraints.
  • weighted_hrw — rendezvous with per-shard weights (+shard_weights+); use stable_cost_binpack for path costs.
  • lazy_robin — sorted round-robin assignment with timing loaded for diagnostics and shard_seconds.
  • preserve_order_round_robin — round-robin in paths-file order (no sort); membership from paths_build only.

Constant Summary collapse

COST_STRATEGIES = 
%w[cost cost_binpack binpack timing stable_cost_binpack].freeze
HRW_STRATEGIES = 
%w[hrw rendezvous weighted_hrw].freeze
LAZY_ROBIN_STRATEGIES = 
%w[lazy_robin].freeze
MOD_STRATEGIES = 
%w[round_robin random_round_robin lazy_robin preserve_order_round_robin].freeze

Instance Attribute Summary collapse

Class Method Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(items:, total_shards:, strategy: "round_robin", seed: nil, costs: nil, constraints: nil, root: nil, timing_granularity: :file, stable_assignment: nil, stable_imbalance_threshold: 1.30, shard_weights: nil) ⇒ Plan

Returns a new instance of Plan.

Raises:



30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
 
# File 'lib/polyrun/partition/plan.rb', line 30

def initialize(items:, total_shards:, strategy: "round_robin", seed: nil, costs: nil, constraints: nil, root: nil, timing_granularity: :file, stable_assignment: nil, stable_imbalance_threshold: 1.30, shard_weights: nil)
  @timing_granularity = TimingKeys.normalize_granularity(timing_granularity)
  @root = root ? File.expand_path(root) : Dir.pwd
  @stable_assignment = normalize_stable_assignment(stable_assignment)
  @stable_imbalance_threshold = stable_imbalance_threshold.to_f
  @items = items.map do |x|
    if @timing_granularity == :example
      TimingKeys.normalize_locator(x, @root, :example)
    else
      x.to_s.strip
    end
  end.freeze
  @total_shards = Integer(total_shards)
  raise Polyrun::Error, "total_shards must be >= 1" if @total_shards < 1

  @strategy = strategy.to_s
  @seed = seed
  @constraints = normalize_constraints(constraints)
  @costs = normalize_costs(costs)
  @shard_weights = shard_weights

  validate_constraints_strategy_combo!
  if cost_strategy? && (@costs.nil? || @costs.empty?)
    raise Polyrun::Error,
      "strategy #{@strategy} requires a timing map (path => seconds or path:line => seconds), e.g. merged polyrun_timing.json"
  end
  if lazy_robin_strategy? && (@costs.nil? || @costs.empty?)
    raise Polyrun::Error,
      "strategy lazy_robin requires a timing map (path => seconds), e.g. merged polyrun_timing.json"
  end
end

Instance Attribute Details

#constraints ⇒ Object (readonly)

Returns the value of attribute constraints.



28
29
30
 
# File 'lib/polyrun/partition/plan.rb', line 28

def constraints
  @constraints
end

#items ⇒ Object (readonly)

Returns the value of attribute items.



28
29
30
 
# File 'lib/polyrun/partition/plan.rb', line 28

def items
  @items
end

#root ⇒ Object (readonly)

Returns the value of attribute root.



28
29
30
 
# File 'lib/polyrun/partition/plan.rb', line 28

def root
  @root
end

#seed ⇒ Object (readonly)

Returns the value of attribute seed.



28
29
30
 
# File 'lib/polyrun/partition/plan.rb', line 28

def seed
  @seed
end

#stable_imbalance_threshold ⇒ Object (readonly)

Returns the value of attribute stable_imbalance_threshold.



125
126
127
 
# File 'lib/polyrun/partition/plan.rb', line 125

def stable_imbalance_threshold
  @stable_imbalance_threshold
end

#strategy ⇒ Object (readonly)

Returns the value of attribute strategy.



28
29
30
 
# File 'lib/polyrun/partition/plan.rb', line 28

def strategy
  @strategy
end

#timing_granularity ⇒ Object (readonly)

Returns the value of attribute timing_granularity.



28
29
30
 
# File 'lib/polyrun/partition/plan.rb', line 28

def timing_granularity
  @timing_granularity
end

#total_shards ⇒ Object (readonly)

Returns the value of attribute total_shards.



28
29
30
 
# File 'lib/polyrun/partition/plan.rb', line 28

def total_shards
  @total_shards
end

Class Method Details

.cost_strategy?(name) ⇒ Boolean

Returns:

  • (Boolean) 


149
150
151
 
# File 'lib/polyrun/partition/plan.rb', line 149

def self.cost_strategy?(name)
  COST_STRATEGIES.include?(name.to_s)
end

.hrw_strategy?(name) ⇒ Boolean

Returns:

  • (Boolean) 


153
154
155
 
# File 'lib/polyrun/partition/plan.rb', line 153

def self.hrw_strategy?(name)
  HRW_STRATEGIES.include?(name.to_s)
end

.lazy_robin_strategy?(name) ⇒ Boolean

Returns:

  • (Boolean) 


157
158
159
 
# File 'lib/polyrun/partition/plan.rb', line 157

def self.lazy_robin_strategy?(name)
  LAZY_ROBIN_STRATEGIES.include?(name.to_s)
end

.load_timing_costs(path, granularity: :file, root: nil) ⇒ Object 



145
146
147
 
# File 'lib/polyrun/partition/plan.rb', line 145

def self.load_timing_costs(path, granularity: :file, root: nil)
  TimingKeys.load_costs_json_file(path, granularity, root: root)
end

.timing_load_strategy?(name) ⇒ Boolean

Returns:

  • (Boolean) 


161
162
163
 
# File 'lib/polyrun/partition/plan.rb', line 161

def self.timing_load_strategy?(name)
  cost_strategy?(name) || hrw_strategy?(name) || lazy_robin_strategy?(name)
end

Instance Method Details

#default_weight ⇒ Object 



112
113
114
115
116
117
118
119
 
# File 'lib/polyrun/partition/plan.rb', line 112

def default_weight
  vals = @costs&.values || []
  if vals.empty?
    1.0
  else
    vals.sum / vals.size
  end
end

#file_weight(path) ⇒ Object 



104
105
106
 
# File 'lib/polyrun/partition/plan.rb', line 104

def file_weight(path)
  (lazy_robin_strategy? || cost_strategy?) ? weight_for(path) : weight_for_optional(path)
end

#hrw_salt ⇒ Object 



43
44
45
46
 
# File 'lib/polyrun/partition/plan_sharding.rb', line 43

def hrw_salt
  s = seed
  (s.nil? || s.to_s.empty?) ? "polyrun-hrw" : s.to_s
end

#hrw_shards ⇒ Object 



6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
 
# File 'lib/polyrun/partition/plan_sharding.rb', line 6

def hrw_shards
  @hrw_shards ||= begin
    buckets = Array.new(total_shards) { [] }
    salt = hrw_salt
    weighted = strategy == "weighted_hrw"
    items.each do |path|
      j =
        if @constraints && (fj = @constraints.forced_shard_for(path))
          Integer(fj)
        elsif weighted
          Hrw.weighted_shard_for(
            path: path,
            total_shards: total_shards,
            seed: salt,
            shard_weights: @shard_weights
          )
        else
          Hrw.shard_for(path: path, total_shards: total_shards, seed: salt)
        end
      raise Polyrun::Error, "constraint shard out of range" if j < 0 || j >= total_shards

      buckets[j] << path
    end
    buckets
  end
end

#manifest(shard_index) ⇒ Object 



131
132
133
134
135
136
137
138
139
140
141
142
143
 
# File 'lib/polyrun/partition/plan.rb', line 131

def manifest(shard_index)
  m = {
    "shard_index" => Integer(shard_index),
    "shard_total" => total_shards,
    "strategy" => strategy,
    "seed" => seed,
    "paths" => shard(shard_index)
  }
  m["timing_granularity"] = timing_granularity.to_s if timing_granularity == :example
  secs = shard_weight_totals
  m["shard_seconds"] = secs if emit_shard_seconds?(secs)
  m
end

#mod_shards ⇒ Object

One pass over ordered_items (round_robin / random_round_robin); avoids O(workers × n) rescans in shard.



34
35
36
37
38
39
40
41
 
# File 'lib/polyrun/partition/plan_sharding.rb', line 34

def mod_shards
  @mod_shards ||= begin
    list = ordered_items
    buckets = Array.new(total_shards) { [] }
    list.each_with_index { |path, i| buckets[i % total_shards] << path }
    buckets
  end
end

#ordered_items ⇒ Object 



62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
 
# File 'lib/polyrun/partition/plan.rb', line 62

def ordered_items
  @ordered_items ||= case strategy
  when "round_robin", "lazy_robin"
    items.sort
  when "preserve_order_round_robin"
    items.dup
  when "random_round_robin"
    StableShuffle.call(items.sort, random_seed)
  when "cost", "cost_binpack", "binpack", "timing"
    items.sort
  when "hrw", "rendezvous", "weighted_hrw"
    items.sort
  else
    raise Polyrun::Error, "unknown partition strategy: #{strategy}"
  end
end

#random_seed ⇒ Object 



48
49
50
51
52
53
 
# File 'lib/polyrun/partition/plan_sharding.rb', line 48

def random_seed
  s = seed
  return Integer(s) if s && s != ""

  0
end

#shard(shard_index) ⇒ Object

Raises:



79
80
81
82
83
84
85
86
87
88
89
90
 
# File 'lib/polyrun/partition/plan.rb', line 79

def shard(shard_index)
  idx = Integer(shard_index)
  raise Polyrun::Error, "shard_index out of range" if idx < 0 || idx >= total_shards

  if cost_strategy?
    cost_shards[idx]
  elsif hrw_strategy?
    hrw_shards[idx]
  else
    mod_shards[idx]
  end
end

#shard_file_weights(shard_index) ⇒ Object 



108
109
110
 
# File 'lib/polyrun/partition/plan.rb', line 108

def shard_file_weights(shard_index)
  shard(shard_index).map { |p| [p, file_weight(p)] }.sort_by { |(_, w)| [-w, p] }
end

#shard_weight_totals ⇒ Object 



92
93
94
95
96
97
98
99
100
101
102
 
# File 'lib/polyrun/partition/plan.rb', line 92

def shard_weight_totals
  if cost_strategy?
    cost_shards.map { |paths| paths.sum { |p| weight_for(p) } }
  elsif hrw_strategy?
    hrw_shards.map { |paths| paths.sum { |p| weight_for_optional(p) } }
  elsif lazy_robin_strategy? && @costs&.any?
    mod_shards.map { |paths| paths.sum { |p| weight_for(p) } }
  else
    []
  end
end

#stable_assignment_map ⇒ Object 



127
128
129
 
# File 'lib/polyrun/partition/plan.rb', line 127

def stable_assignment_map
  @stable_assignment
end

#stable_strategy? ⇒ Boolean

Returns:

  • (Boolean) 


121
122
123
 
# File 'lib/polyrun/partition/plan.rb', line 121

def stable_strategy?
  strategy == "stable_cost_binpack"
end