diff --git a/src/main/clojure/clojask/classes/DataFrame.clj b/src/main/clojure/clojask/classes/DataFrame.clj index 70e7969..b552979 100644 --- a/src/main/clojure/clojask/classes/DataFrame.clj +++ b/src/main/clojure/clojask/classes/DataFrame.clj @@ -5,6 +5,7 @@ [clojask.classes.MGroup] [clojure.java.io :as io] [clojask.utils :as u] + [clojask.engine :as engine] [clojask.onyx-comps :refer [start-onyx start-onyx-aggre-only start-onyx-groupby]] [clojask.sort :as sort] [clojask.aggregate.aggre-onyx-comps :refer [start-onyx-aggre]] @@ -126,7 +127,9 @@ index (if (= select [nil]) (take (count key-index) (iterate inc 0)) (keep key-index select))] (assert (or (= (count select) (count index)) (= select [nil])) "Must select existing columns; get-col-names lists them.") (if (= ifheader true) (.printCol this output-dir index out)) - (let [res (start-onyx num-worker batch-size this output-dir exception order index melt out)] + (let [res (if (= (engine/engine) :local) + (engine/start-compute num-worker batch-size this output-dir exception index melt out) + (start-onyx num-worker batch-size this output-dir exception order index melt out))] (if (= res "success") "success" "failed")))) diff --git a/src/main/clojure/clojask/engine.clj b/src/main/clojure/clojask/engine.clj new file mode 100644 index 0000000..c5cc5e1 --- /dev/null +++ b/src/main/clojure/clojask/engine.clj @@ -0,0 +1,130 @@ +(ns clojask.engine + "A thread-pool engine for the jobs clojask runs: one input, num-work + identical workers and one output, all in this JVM. It replaces Onyx one + driver at a time. The system property clojask.engine, else the + environment variable CLOJASK_ENGINE, chooses the engine: \"onyx\" (the + default) or \"local\". Drivers not ported yet always run on Onyx." + (:require [clojure.java.io :as io] + [clojask.utils :as u]) + (:import [com.clojask.exception ExecutionException] + [java.util.concurrent ArrayBlockingQueue Callable ExecutorService + Executors Future ThreadFactory TimeUnit])) + +(defn engine + "The engine for the next compute, :onyx or :local." + [] + (let [v (or (System/getProperty "clojask.engine") (System/getenv "CLOJASK_ENGINE") "onyx")] + (case v + "onyx" :onyx + "local" :local + (throw (IllegalArgumentException. (str "Unknown clojask engine \"" v "\"; use \"onyx\" or \"local\".")))))) + +(defn- thread-factory + "Daemon threads named prefix-1, prefix-2, ..., so a job that is never + shut down cannot keep the JVM alive." + [prefix] + (let [n (atom 0)] + (reify ThreadFactory + (newThread [_ r] + (doto (Thread. ^Runnable r (str prefix "-" (swap! n inc))) + (.setDaemon true)))))) + +(defn- feed + "Submit (work batch) to pool for every batch, in order, and put the + futures on queue. batches is not held anywhere else, so a lazy input is + not retained as it is consumed." + [^ArrayBlockingQueue queue ^ExecutorService pool work batches] + (doseq [batch batches] + (.put queue (.submit pool ^Callable (fn [] (work batch)))))) + +(defn run + "Run one job and return nil once every result has been handed to sink. + + open returns {:data batches :close f}, :close optional. A reader thread + walks batches and submits (work batch) to a pool of num-work threads; + the calling thread hands the results to sink in input order. At most + 2 * num-work + 2 batches are in flight. work must return realized data, + otherwise the work moves to the calling thread. + + The first exception from open, the batches, work or sink is rethrown + unwrapped, after the threads are stopped and the input is closed." + [num-work open work sink] + (let [pool (Executors/newFixedThreadPool num-work (thread-factory "clojask-worker")) + reader-pool (Executors/newSingleThreadExecutor (thread-factory "clojask-reader")) + queue (ArrayBlockingQueue. (* 2 num-work)) + end (Object.)] + (try + (.submit reader-pool + ^Callable (fn [] + (try + (let [{:keys [data close]} (open)] + (try + (feed queue pool work data) + (finally + (when close (close))))) + (.put queue end) + ;; the calling thread has given up on the job + (catch InterruptedException _) + (catch Throwable t + (.put queue t))))) + (loop [] + (let [item (.take queue)] + (cond + (identical? item end) nil + (instance? Throwable item) (throw item) + :else (do (sink (try + (.get ^Future item) + (catch java.util.concurrent.ExecutionException e + (throw (.getCause e))))) + (recur))))) + (finally + (.shutdownNow pool) + (.shutdownNow reader-pool) + ;; let the reader run its finally, so the input is closed on return + (.awaitTermination reader-pool 10 TimeUnit/SECONDS))))) + +(defn- stage-error + [stage ^Throwable e] + (ExecutionException. (format "[%s] (original error: %s)" stage (.getMessage e)) e)) + +(defn- compute-worker + "The work on one batch of a plain compute: the selected columns of every + row that passes the filters, evaluated and formatted. Rows that fail a + filter are dropped. Same rows as onyx-comps/worker-func-gen-format." + [df exception index] + (let [col-info (:col-info df) + operations (.getDesc col-info) + types (.getType col-info) + formats (.getFormatter col-info) + filters (.getFilters (:row-info df)) + index (vec index) + eval-res (if exception u/eval-res u/eval-res-ne)] + (fn [batch] + (into [] + (keep (fn [row] + (when (u/filter-check filters types row) + (mapv (fn [i] ((or (get formats i) str) (eval-res row types formats operations i))) index)))) + batch)))) + +(defn start-compute + "The local engine's start-onyx: compute df into dist, appending after + the header that DataFrame.compute may have written. Output is always in + input order." + [num-work batch-size df dist exception index melt out] + (let [work (try + (compute-worker df exception index) + (catch Throwable e + (throw (stage-error "preparing stage (compute)" e))))] + (try + (with-open [wtr (io/writer dist :append true)] + (run num-work + (fn [] + (let [{:keys [data close]} ((.getFunc df) :map)] + {:data (partition-all batch-size data) :close close})) + work + (fn [rows] + (doseq [row rows] + (out wtr (melt row)))))) + (catch Throwable e + (throw (stage-error "run stage (compute)" e)))) + "success")) diff --git a/test/clojask/shutdown_test.clj b/test/clojask/shutdown_test.clj index cf9697f..bc6e1c6 100644 --- a/test/clojask/shutdown_test.clj +++ b/test/clojask/shutdown_test.clj @@ -10,6 +10,21 @@ (def input "test/clojask/Employees-example.csv") +(defn- on-onyx + "These tests exercise the Onyx environment, whichever engine the rest of + the suite runs on." + [t] + (let [old (System/getProperty "clojask.engine")] + (System/setProperty "clojask.engine" "onyx") + (try + (t) + (finally + (if old + (System/setProperty "clojask.engine" old) + (System/clearProperty "clojask.engine")))))) + +(use-fixtures :each on-onyx) + (defn- env-is-down [] (is (nil? oc/env)) (is (nil? oc/peer-group))