<commit_before>(ns onyx.scheduling.percentage-task-scheduler
  (:require [onyx.scheduling.common-task-scheduler :as cts]
            [onyx.log.commands.common :as common]))

(defn sort-tasks-by-pct [replica job tasks]
  (let [indexed
        (map-indexed
         (fn [k t]
           {:position k :task t :pct (get-in replica [:task-percentages job t])})
         (reverse tasks))]
    (reverse (sort-by (juxt :pct :position) indexed))))

(defn min-task-allocations [replica job tasks n-peers]
  (mapv
   (fn [task]
     (let [n (int (Math/floor (* (* 0.01 (:pct task)) n-peers)))]
       (assoc task :allocation n)))
   tasks))

(defn percentage-balanced-taskload [replica job candidate-tasks n-peers]
  (let [sorted-tasks (sort-tasks-by-pct replica job candidate-tasks)
        init-allocations (min-task-allocations replica job sorted-tasks n-peers)
        init-usage (apply + (map :allocation init-allocations))
        left-over-peers (- n-peers init-usage)
        with-leftovers (update-in init-allocations [0 :allocation] + left-over-peers)]
    (into {} (map (fn [t] {(:task t) t}) with-leftovers))))

(defmethod cts/drop-peers :onyx.task-scheduler/percentage
  [replica job n]
  (let [tasks (keys (get-in replica [:allocations job]))
        balanced (percentage-balanced-taskload replica job tasks n)]
    (mapcat
     (fn [[task {:keys [allocation]}]]
       (drop-last allocation (get-in replica [:allocations job task])))
     balanced)))

(defn reuse-spare-peers [replica job tasks spare-peers]
  (loop [[head & tail :as task-seq] (get-in replica [:tasks job])
         results tasks
         capacity spare-peers]
    (let [tail (vec tail)]
      (cond (or (<= capacity 0) (not (seq task-seq)))
            results
            (< (get results head) (or (get-in replica [:task-saturation job head] Double/POSITIVE_INFINITY)))
            (recur (conj tail head) (update-in results [head] inc) (dec capacity))
            :else
            (recur tail results capacity)))))

(defmethod cts/task-distribute-peer-count :onyx.task-scheduler/percentage
  [replica job n]
  (let [tasks (get-in replica [:tasks job])
        init
        (reduce
         (fn [all [task k]]
           (assoc all task (min (get-in replica [:task-saturation job task] Double/POSITIVE_INFINITY)
                                (int (* n (get-in replica [:task-percentages job task]) 0.01)))))
         {}
         (map vector tasks (range)))
        spare-peers (- n (apply + (vals init)))]
    (reuse-spare-peers replica job init spare-peers)))<commit_msg>Fix pct task scheduler misallocation. Use largest remainders method to allocate remaining peers. Ensure peers on oversaturated tasks are reallocated to unsaturated tasks.<commit_after>(ns onyx.scheduling.percentage-task-scheduler
  (:require [onyx.scheduling.common-task-scheduler :as cts]
            [onyx.log.commands.common :as common]))

(defn tasks-by-pct [replica job tasks]
  (map
    (fn [t]
      {:task t :pct (get-in replica [:task-percentages job t])})
    tasks))

(defn rescale-task-percentages 
  "Rescale task percentages after saturated tasks were removed"
  [tasks]
  (let [total (/ (apply + (map :pct tasks)) 100)]
    (map (fn [task]
           (update-in task [:pct] / total))
         tasks))) 

(defn largest-remainder-allocations 
  "Allocates remaining peers to the tasks with the largest remainder.
  e.g. 3 tasks pct allocated 3.5, 1.75, 1.75 -> 3, 2, 2"
  [tasks n-peers]
  (let [tasks* (rescale-task-percentages tasks)
        unrounded (map (fn [task] 
                         (* 0.01 (:pct task) n-peers))
                       tasks*)
        full (map int unrounded) 
        taken (apply + full)
        remaining (- n-peers taken)
        full-allocated (zipmap tasks* full)
        remainders (->> (map (fn [task v] 
                               (vector task (- v (int v)))) 
                             tasks*
                             unrounded)
                        (sort-by second)
                        (reverse)
                        (take remaining)
                        (map (juxt first (constantly 1)))
                        (into {}))
        final-allocations (merge-with + full-allocated remainders)]
    (mapv (fn [[task allocation]] 
            (assoc task :allocation (int allocation)))
          final-allocations)))

(defn percentage-balanced-taskload 
  "Percentage balance taskload by allocating via largest remainders.
  If a task becomes oversaturated, take it out of the pool, fully allocated,
  and restart the process with the remaining peers and tasks"
  [replica job candidate-tasks n-peers]
  {:post [(>= n-peers 0)
          (= n-peers (reduce + (map :allocation (vals %))))]}
  (let [sorted-tasks (tasks-by-pct replica job candidate-tasks)
        allocations (largest-remainder-allocations sorted-tasks n-peers)
        oversaturated (filter (fn [{:keys [task allocation]}]
                                (> allocation (get-in replica [:task-saturation job task])))
                              allocations)
        cutoff-oversaturated (->> oversaturated
                                  (map (fn [{:keys [task] :as t}]
                                         [task (assoc t :allocation (get-in replica [:task-saturation job task]))]))
                                  (into {}))]
    (if (empty? cutoff-oversaturated)
      (into {} (map (fn [t] {(:task t) t}) allocations))
      (let [n-peers-fully-saturated (apply + (map :allocation (vals cutoff-oversaturated)))
            n-remaining-peers (- n-peers n-peers-fully-saturated)
            unallocated-tasks (remove cutoff-oversaturated candidate-tasks)] 
        (merge (percentage-balanced-taskload replica job unallocated-tasks n-remaining-peers)
               cutoff-oversaturated)))))

(defmethod cts/drop-peers :onyx.task-scheduler/percentage
  [replica job n]
  (let [tasks (keys (get-in replica [:allocations job]))
        balanced (percentage-balanced-taskload replica job tasks n)]
    (mapcat
      (fn [[task {:keys [allocation]}]]
        (drop-last allocation (get-in replica [:allocations job task])))
      balanced)))

(defmethod cts/task-distribute-peer-count :onyx.task-scheduler/percentage
  [replica job n]
  (->> (percentage-balanced-taskload replica job (get-in replica [:tasks job]) n)
       vals
       (map (juxt :task :allocation))
       (into {})))
{'fraction_non_alphanumeric': 0.12809658200062715, 'fraction_numerical': 0.004390090937597993, 'mean_word_length': 2.3275952008346374, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1556109', 'n_tokens_mistral': 2043, 'n_tokens_neox': 1938, 'n_words': 570}