Repository: twosigma/Cook (language: python)

Issue:
Adds integration with data local service
## Changes proposed in this PR
- Adds "datasets" field to jobs
- Adds support for contacting a data local service to obtain cost data for scheduling

## Why are we making these changes?
This allows users to configure a Cook cluster to use an external data local cost service to supply cook with cost information, which it can use to prioritize scheduling tasks.




Discussion / hints:
@scrosby / @dposada any other feedback?
Any other feedback @scrosby ?

Fix patch:
diff --git a/integration/bin/run-data-local-server.sh b/integration/bin/run-data-local-server.sh
deleted file mode 100755
index 1b3854c4cd..0000000000
--- a/integration/bin/run-data-local-server.sh
+++ /dev/null
@@ -1,25 +0,0 @@
-#!/usr/bin/env bash
-
-# Usage: run-data-local-server.sh
-# Builds a docker image and launches the data local service.
-
-INTEGRATION_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && cd .. && pwd )"
-SRC_DIR="${INTEGRATION_DIR}/src/data_locality"
-NAME=data-local
-
-echo "Building docker images for ${NAME}"
-docker build -t ${NAME} ${SRC_DIR}
-
-# Expose port 35847 (used for data-local service)
-docker create \
-       -i \
-       -t \
-       --rm \
-       -p 35847:35847 \
-       --name=$NAME \
-       ${NAME}:latest
-
-docker network connect bridge ${NAME}
-docker network connect cook_nw ${NAME}
-
-docker start -ai ${NAME}
diff --git a/integration/bin/run-integration.sh b/integration/bin/run-integration.sh
index f0953c50a3..c0528c78c9 100755
--- a/integration/bin/run-integration.sh
+++ b/integration/bin/run-integration.sh
@@ -74,12 +74,6 @@ then
    COOK_MULTICLUSTER_ENV="${COOK_MULTICLUSTER_ENV} -e COOK_SLAVE_URL=${COOK_SLAVE_URL} -e COOK_MASTER_SLAVE=${COOK_MASTER_SLAVE}"
 fi
 
-DATA_LOCAL_IP=$(docker inspect data-local | jq -r '.[].NetworkSettings.IPAddress')
-DATA_LOCAL_ENV=""
-if [ ! "${DATA_LOCAL_IP}" = "null" ];
-then
-    DATA_LOCAL_ENV="-e DATA_LOCAL_SERVICE=http://${DATA_LOCAL_IP}:35847"
-fi
 
 docker create \
        --rm \
@@ -90,7 +84,6 @@ docker create \
        -e "COOK_TEST_DOCKER_IMAGE=python:3.5.9-stretch" \
        -v "/tmp/cook-integration-mount:/tmp/cook-integration-mount" \
        ${COOK_MULTICLUSTER_ENV} \
-       ${DATA_LOCAL_ENV} \
        ${DOCKER_VOLUME_ARGS} \
        cook-integration:latest \
        "$@"
diff --git a/integration/src/data_locality/Dockerfile b/integration/src/data_locality/Dockerfile
deleted file mode 100644
index 177b802691..0000000000
--- a/integration/src/data_locality/Dockerfile
+++ /dev/null
@@ -1,10 +0,0 @@
-FROM python:3.6
-
-RUN pip install flask
-
-COPY . /opt/cook/integration
-ENV FLASK_APP /opt/cook/integration/service.py
-
-EXPOSE 35847
-
-ENTRYPOINT ["flask", "run", "--host=0.0.0.0", "--port=35847"]
diff --git a/integration/src/data_locality/service.py b/integration/src/data_locality/service.py
deleted file mode 100644
index 8f178b9384..0000000000
--- a/integration/src/data_locality/service.py
+++ /dev/null
@@ -1,27 +0,0 @@
-from flask import Flask, make_response, request
-import json
-
-app = Flask(__name__)
-
-costs = {}
-
-@app.route('/retrieve-costs', methods=['POST'])
-def get_costs():
-    request_data = json.loads(request.data)
-    batch = request_data['batch']
-    tasks = request_data['tasks']
-
-    response = {'batch': batch, 'costs': []}
-    for task in tasks:
-        if task['task_id'] in costs:
-            response['costs'].append({'task_id': task['task_id'], 'costs': costs[task['task_id']]})
-
-    return make_response(json.dumps(response))
-
-@app.route('/set-costs', methods=['POST'])
-def set_costs():
-    global costs
-    payload = json.loads(request.data)
-    costs = payload
-
-    return make_response('Updated cost data')
diff --git a/integration/travis/run_integration.sh b/integration/travis/run_integration.sh
index c1ec3749f6..481f745c5a 100755
--- a/integration/travis/run_integration.sh
+++ b/integration/travis/run_integration.sh
@@ -116,12 +116,7 @@ case "$COOK_POOLS" in
     exit 1
 esac
 
-pip install flask
-export DATA_LOCAL_PORT=35847
-export DATA_LOCAL_SERVICE="http://localhost:${DATA_LOCAL_PORT}"
-export DATA_LOCAL_ENDPOINT="${DATA_LOCAL_SERVICE}/retrieve-costs"
 mkdir ${SCHEDULER_DIR}/log
-FLASK_APP=${PROJECT_DIR}/src/data_locality/service.py flask run --port=${DATA_LOCAL_PORT} > ${SCHEDULER_DIR}/log/data-local.log 2>&1 &
 
 # Seed running jobs, which are used to test the task reconciler
 cd ${SCHEDULER_DIR}
diff --git a/integration/travis/scheduler_travis_config.edn b/integration/travis/scheduler_travis_config.edn
index 1f9c4723e2..adcbf8e6c3 100644
--- a/integration/travis/scheduler_travis_config.edn
+++ b/integration/travis/scheduler_travis_config.edn
@@ -26,15 +26,11 @@
                                  :container-path "/mnt/cook-integration"
                                  :mode "RW"}]}
  :cors-origins ["https?://cors.example.com"]
- :data-local {:fitness-calculator {:cache-ttl-ms 20000
-                                   :cost-endpoint #config/env "DATA_LOCAL_ENDPOINT"
-                                   :update-interval-ms 1000}}
  :database {:datomic-uri #config/env "COOK_DATOMIC"}
  :zookeeper {:connection #config/env "COOK_ZOOKEEPER"
              :local? #config/env-bool "COOK_ZOOKEEPER_LOCAL"
              :local-port #config/env-int-default ["COOK_ZOOKEEPER_LOCAL_PORT" -1]}
- :scheduler {:fenzo-fitness-calculator "cook.scheduler.data-locality/make-data-local-fitness-calculator"
-             :offer-incubate-ms 15000
+ :scheduler {:offer-incubate-ms 15000
              :task-constraints {:timeout-hours 1
                                 :timeout-interval-minutes 1
                                 :memory-gb 48
@@ -76,7 +72,6 @@
        :levels {"datomic.db" :warn
                 "datomic.peer" :warn
                 "datomic.kv-cluster" :warn
-                "cook.scheduler.data-locality" :debug
                 "cook.mesos.fenzo-utils" :debug
                 "cook.mesos.rebalancer" :debug
                 "cook.mesos.scheduler" :debug
diff --git a/jobclient/java/src/main/java/com/twosigma/cook/jobclient/Job.java b/jobclient/java/src/main/java/com/twosigma/cook/jobclient/Job.java
index ed722624e7..e511165c9e 100644
--- a/jobclient/java/src/main/java/com/twosigma/cook/jobclient/Job.java
+++ b/jobclient/java/src/main/java/com/twosigma/cook/jobclient/Job.java
@@ -105,7 +105,6 @@ public static class Builder {
         private String _progressOutputFile;
         private String _progressRegexString;
         private String _user;
-        private JSONArray _datasets;
 
         /**
          * Prior to {@code build()}, command, memory and cpus for a job must be provided.<br>
@@ -148,7 +147,7 @@ public Job build() {
             }
             return new Job(_uuid, _name, _command, _executor, _memory, _cpus, _disk, _gpus, _retries, _maxRuntime,
                     _expectedRuntime, _status, _priority, _pool, _isMeaCulpaRetriesDisabled, _instances, _env, _uris, _container,
-                    _labels, _constraints, _groups, _application, _checkpoint, _progressOutputFile, _progressRegexString, _user, _datasets);
+                    _labels, _constraints, _groups, _application, _checkpoint, _progressOutputFile, _progressRegexString, _user);
         }
 
         /**
@@ -170,7 +169,6 @@ public Builder of(Job job) {
             setContainer(job.getContainer());
             setPool(job.getPool());
             addLabels(job.getLabels());
-            setDatasets(job.getDatasets());
             setStatus(job.getStatus());
             setPriority(job.getPriority());
             setDisk(job.getDisk());
@@ -715,17 +713,6 @@ public Builder setUser(String user) {
             return this;
         }
 
-        /**
-         * Sets the datasets for the job being built.
-         * @param datasets {@link JSONArray} the datasets
-         * @return this builder
-         */
-        public Builder setDatasets(JSONArray datasets) {
-            _datasets = datasets;
-            return this;
-        }
-
-
         /**
          * Parse a JSON object into this Builder object, e.g.
          * <p>
@@ -859,9 +846,6 @@ public Builder parseFromJSON(JSONObject json, InstanceDecorator decorator)
             if (json.has("progress_regex_string")) {
                 setProgressRegexString(json.getString("progress_regex_string"));
             }
-            if (json.has("datasets")) {
-                setDatasets(json.getJSONArray("datasets"));
-            }
             if (json.has("pool")) {
                 setPool(json.getString("pool"));
             }
@@ -910,13 +894,12 @@ public Builder parseFromJSON(JSONObject json) throws JSONException {
     final private String _progressOutputFile;
     final private String _progressRegexString;
     final private String _user;
-    final private JSONArray _datasets;
 
     private Job(UUID uuid, String name, String command, Executor executor, Double memory, Double cpus, Disk disk, Integer gpus, Integer retries,
                 Long maxRuntime, Long expectedRuntime, Status status, Integer priority, String pool, Boolean isMeaCulpaRetriesDisabled,
                 List<Instance> instances, Map<String, String> env, List<FetchableURI> uris, JSONObject container,
                 Map<String, String> labels, Set<Constraint> constraints, List<UUID> groups, Application application, Checkpoint checkpoint,
-                String progressOutputFile, String progressRegexString, String user, JSONArray datasets) {
+                String progressOutputFile, String progressRegexString, String user) {
         _uuid = uuid;
         _name = name;
         _command = command;
@@ -951,15 +934,6 @@ private Job(UUID uuid, String name, String command, Executor executor, Double me
         } else {
             _container = null;
         }
-        if (datasets != null) {
-            try {
-                _datasets = new JSONArray(datasets.toString());
-            } catch (JSONException e) {
-                throw new RuntimeException("Failed to parse datasets string", e);
-            }
-        } else {
-            _datasets = null;
-        }
         _labels = ImmutableMap.copyOf(labels);
         _constraints = ImmutableSet.copyOf(constraints);
         _groups = groups;
@@ -1165,13 +1139,6 @@ public String getProgressRegexString() {
         return _progressRegexString;
     }
 
-    /**
-     * @return the job's datasets
-     */
-    public JSONArray getDatasets() {
-        return _datasets;
-    }
-
     /**
      * @return the job instance with the running state or {@code null} if can't find one.
      */
@@ -1276,9 +1243,6 @@ public static JSONObject jsonizeJob(Job job)
         if (job._expectedRuntime != null) {
             object.put("expected_runtime", job._expectedRuntime);
         }
-        if (job._datasets != null) {
-            object.put("datasets", job._datasets);
-        }
         return object;
     }
 
diff --git a/scheduler/bin/run-docker.sh b/scheduler/bin/run-docker.sh
index 826b92d631..5e36a5c8ca 100755
--- a/scheduler/bin/run-docker.sh
+++ b/scheduler/bin/run-docker.sh
@@ -113,27 +113,6 @@ else
     COOK_ZOOKEEPER=""
 fi
 
-echo "About to: Setup data local service"
-DATA_LOCAL_IP=$(docker inspect data-local | jq -r '.[].NetworkSettings.IPAddress // empty')
-if [[ -z "${DATA_LOCAL_IP}" ]];
-then
-    echo "Starting data local server"
-    PROJECT_ROOT="$( dirname ${SCHEDULER_DIR} )"
-    ${PROJECT_ROOT}/integration/bin/run-data-local-server.sh &
-    ITERATIONS=0
-    while [[ -z "${DATA_LOCAL_IP}" && $ITERATIONS -lt 10 ]];
-    do
-        sleep 1
-        DATA_LOCAL_IP=$(docker inspect data-local | jq -r '.[].NetworkSettings.IPAddress // empty')
-        ((ITERATIONS++))
-    done
-    if [[ -z "${DATA_LOCAL_IP}" ]];
-    then
-        echo "Unable to start data local server"
-        exit 1
-    fi
-fi
-
 echo "Starting cook..."
 
 # NOTE: since the cook scheduler directory is mounted as a volume
@@ -167,7 +146,6 @@ docker create \
     -e "COOK_ONE_USER_AUTH=root" \
     -e "COOK_EXECUTOR_PORTION=${COOK_EXECUTOR_PORTION:-0}" \
     -e "COOK_KEYSTORE_PATH=/opt/ssl/cook.p12" \
-    -e "DATA_LOCAL_ENDPOINT=http://${DATA_LOCAL_IP}:35847/retrieve-costs" \
     -v ${DIR}/../log:/opt/cook/log \
     cook-scheduler:latest ${COOK_CONFIG:-}
 
diff --git a/scheduler/bin/run-local.sh b/scheduler/bin/run-local.sh
index 11eac271ac..aaf63ecef4 100755
--- a/scheduler/bin/run-local.sh
+++ b/scheduler/bin/run-local.sh
@@ -54,7 +54,6 @@ if ! [ -x "$(command -v flask)" ]; then
 fi
 
 INTEGRATION_DIR="$(dirname ${SCHEDULER_DIR})/integration"
-FLASK_APP=${INTEGRATION_DIR}/src/data_locality/service.py flask run -p 35847 &
 
 echo "Mesos Master IP is ${MASTER_IP}"
 
@@ -76,7 +75,6 @@ export MESOS_MASTER="${MASTER_IP}:5050"
 export MESOS_NATIVE_JAVA_LIBRARY="${MESOS_NATIVE_JAVA_LIBRARY}"
 export COOK_SSL_PORT="${COOK_SSL_PORT}"
 export COOK_KEYSTORE_PATH="${COOK_KEYSTORE_PATH}"
-export DATA_LOCAL_ENDPOINT="http://localhost:35847/retrieve-costs"
 
 echo "Starting cook..."
 lein run config.edn
diff --git a/scheduler/config-composite.edn b/scheduler/config-composite.edn
index c3166ead43..c6a2aa548a 100644
--- a/scheduler/config-composite.edn
+++ b/scheduler/config-composite.edn
@@ -67,9 +67,6 @@
                                                                    :bucket-size 10000
                                                                    :tokens-replenished-per-minute 5000}}}]
  :cors-origins ["https?://cors.example.com"]
- :data-local {:fitness-calculator {:cache-ttl-ms 60000
-                                   :cost-endpoint #config/env "DATA_LOCAL_ENDPOINT"
-                                   :update-interval-ms 1000}}
  :database {:datomic-uri #config/env "COOK_DATOMIC_URI"}
  :plugins {:job-submission-validator {:batch-timeout-seconds 40
                                       :factory-fns ["cook.plugins.demo-plugin/submission-factory"
@@ -85,7 +82,6 @@
        :levels {"datomic.db" :warn
                 "datomic.kv-cluster" :warn
                 "datomic.peer" :warn
-                "cook.scheduler.data-locality" :debug
                 "cook.scheduler.fenzo-utils" :debug
                 "cook.scheduler.rebalancer" :debug
                 "cook.scheduler.scheduler" :debug
@@ -131,8 +127,7 @@
               :min-dru-diff 1.0
               :safe-dru-threshold 1.0}
  :sandbox-syncer {:sync-interval-ms 1000}
- :scheduler {:fenzo-fitness-calculator "cook.scheduler.data-locality/make-data-local-fitness-calculator"
-             :offer-incubate-ms 15000
+ :scheduler {:offer-incubate-ms 15000
              :task-constraints {:command-length-limit 5000
                                 :cpus 10
                                 :memory-gb 48
diff --git a/scheduler/config-k8s.edn b/scheduler/config-k8s.edn
index b5d9ef5789..4888b5471a 100644
--- a/scheduler/config-k8s.edn
+++ b/scheduler/config-k8s.edn
@@ -103,7 +103,6 @@
        :levels {"datomic.db" :warn
                 "datomic.kv-cluster" :warn
                 "datomic.peer" :warn
-                "cook.scheduler.data-locality" :debug
                 "cook.scheduler.fenzo-utils" :debug
                 "cook.scheduler.rebalancer" :debug
                 "cook.scheduler.scheduler" :debug
diff --git a/scheduler/config.edn b/scheduler/config.edn
index e1911a6168..0bb7d7d700 100644
--- a/scheduler/config.edn
+++ b/scheduler/config.edn
@@ -26,9 +26,6 @@
                                                                    :bucket-size 10000
                                                                    :tokens-replenished-per-minute 5000}}}]
  :cors-origins ["https?://cors.example.com"]
- :data-local {:fitness-calculator {:cache-ttl-ms 60000
-                                   :cost-endpoint #config/env "DATA_LOCAL_ENDPOINT"
-                                   :update-interval-ms 1000}}
  :database {:datomic-uri #config/env "COOK_DATOMIC_URI"}
  :executor {:command #config/env "COOK_EXECUTOR_COMMAND"
             :environment {"EXECUTOR_DEFAULT_PROGRESS_OUTPUT_NAME" "stdout"}
@@ -49,7 +46,6 @@
        :levels {"datomic.db" :warn
                 "datomic.kv-cluster" :warn
                 "datomic.peer" :warn
-                "cook.scheduler.data-locality" :debug
                 "cook.scheduler.fenzo-utils" :debug
                 "cook.scheduler.rebalancer" :debug
                 "cook.scheduler.scheduler" :debug
@@ -79,8 +75,7 @@
               :min-dru-diff 1.0
               :safe-dru-threshold 1.0}
  :sandbox-syncer {:sync-interval-ms 1000}
- :scheduler {:fenzo-fitness-calculator "cook.scheduler.data-locality/make-data-local-fitness-calculator"
-             :offer-incubate-ms 15000
+ :scheduler {:offer-incubate-ms 15000
              :task-constraints {:command-length-limit 5000
                                 :cpus 10
                                 :memory-gb 48
diff --git a/scheduler/src/cook/caches.clj b/scheduler/src/cook/caches.clj
index b90c6891c8..3341d57014 100644
--- a/scheduler/src/cook/caches.clj
+++ b/scheduler/src/cook/caches.clj
@@ -34,7 +34,6 @@
 (mount/defstate ^Cache task-ent->user-cache :start (new-cache config/config))
 (mount/defstate ^Cache job-ent->user-cache :start (new-cache config/config))
 (mount/defstate ^Cache task->feature-vector-cache :start (new-cache config/config))
-(mount/defstate ^Cache job-uuid->dataset-maps-cache :start (new-cache config/config))
 (mount/defstate ^Cache user->group-ids-cache :start (new-cache config/config))
 (mount/defstate ^Cache recent-synthetic-pod-job-uuids :start
   (-> (CacheBuilder/newBuilder)
diff --git a/scheduler/src/cook/config.clj b/scheduler/src/cook/config.clj
index 33e3e16585..b3530cf219 100644
--- a/scheduler/src/cook/config.clj
+++ b/scheduler/src/cook/config.clj
@@ -471,17 +471,6 @@
      :estimated-completion-constraint (fnk [[:config {estimated-completion-constraint nil}]]
                                         (merge {:agent-start-grace-period-mins 10}
                                                estimated-completion-constraint))
-     :data-local-fitness-calculator (fnk [[:config {data-local {}}]]
-                                      (let [fitness-calculator (get data-local :fitness-calculator {})]
-                                        {:auth (get fitness-calculator :auth nil)
-                                         :base-calculator (config-string->fitness-calculator
-                                                            (get fitness-calculator :base-calculator "com.netflix.fenzo.plugins.BinPackingFitnessCalculators/cpuMemBinPacker"))
-                                         :batch-size (get fitness-calculator :batch-size 500)
-                                         :cache-ttl-ms (get fitness-calculator :cache-ttl-ms 300000)
-                                         :cost-endpoint (get fitness-calculator :cost-endpoint nil)
-                                         :data-locality-weight (get fitness-calculator :data-locality-weight 0.95)
-                                         :launch-wait-seconds (get fitness-calculator :launch-wait-seconds 60)
-                                         :update-interval-ms (get fitness-calculator :update-interval-ms nil)}))
      :plugins (fnk [[:config {plugins {}}]]
                 (let [{:keys [job-launch-filter job-routing job-submission-validator pool-selection]} plugins]
                   (merge plugins
@@ -651,10 +640,6 @@
   []
   (-> config :settings :estimated-completion-constraint))
 
-(defn data-local-fitness-config
-  []
-  (-> config :settings :data-local-fitness-calculator))
-
 (defn fitness-calculator-config
   []
   (-> config :settings :fenzo-fitness-calculator))
diff --git a/scheduler/src/cook/mesos.clj b/scheduler/src/cook/mesos.clj
index 3b0dea314b..7a6f443f1c 100644
--- a/scheduler/src/cook/mesos.clj
+++ b/scheduler/src/cook/mesos.clj
@@ -26,7 +26,6 @@
             [cook.monitor]
             [cook.queue-limit :as queue-limit]
             [cook.rebalancer]
-            [cook.scheduler.data-locality :as dl]
             [cook.scheduler.optimizer]
             [cook.scheduler.scheduler :as sched]
             [cook.tools :as tools]
@@ -95,8 +94,7 @@
    {:keys [timeout-interval-minutes]
     :or {timeout-interval-minutes 1}
     :as task-constraints}]
-  (let [{:keys [update-interval-ms]} (config/data-local-fitness-config)
-        prepare-trigger-chan (fn prepare-trigger-chan [interval]
+  (let [prepare-trigger-chan (fn prepare-trigger-chan [interval]
                                (let [ch (async/chan (async/sliding-buffer 1))]
                                  (async/pipe (chime-ch (util/time-seq (time/now) interval))
                                              ch)
@@ -109,9 +107,7 @@
          :progress-updater-trigger-chan (prepare-trigger-chan (time/millis (:publish-interval-ms progress-config)))
          :rank-trigger-chan (prepare-trigger-chan (time/seconds 5))
          :rebalancer-trigger-chan (prepare-trigger-chan (time/seconds (:interval-seconds rebalancer-config)))
-         :straggler-trigger-chan (prepare-trigger-chan (time/minutes timeout-interval-minutes))}
-      update-interval-ms
-      (assoc :update-data-local-costs-trigger-chan (prepare-trigger-chan (time/millis update-interval-ms))))))
+         :straggler-trigger-chan (prepare-trigger-chan (time/minutes timeout-interval-minutes))})))
 
 (defn make-compute-cluster-config-updater-task
   "Create a Runnable that periodically checks for updated compute cluster configurations"
@@ -276,8 +272,6 @@
                                                                                           view-incubating-offers
                                                                                           optimizer-config
                                                                                           optimizer-trigger-chan))
-                                      (when (:update-data-local-costs-trigger-chan trigger-chans)
-                                        (dl/start-update-cycles! mesos-datomic-conn (:update-data-local-costs-trigger-chan trigger-chans)))
                                       ; This counter exists so that the cook leader has a '1' and the others have a '0'.
                                       ; So we can see who the leader is in the graphs by graphing the value.
                                       (counters/inc! mesos-leader)
diff --git a/scheduler/src/cook/rest/api.clj b/scheduler/src/cook/rest/api.clj
index 1a5267eec1..f1eee03978 100644
--- a/scheduler/src/cook/rest/api.clj
+++ b/scheduler/src/cook/rest/api.clj
@@ -49,7 +49,6 @@
             [cook.quota :as quota]
             [cook.rate-limit :as rate-limit]
             [cook.scheduler.constraints :as constraints]
-            [cook.scheduler.data-locality :as dl]
             [cook.scheduler.share :as share]
             [cook.schema :refer [constraint-operators host-placement-types straggler-handling-types]]
             [cook.task :as task]
@@ -134,8 +133,6 @@
 
 (def iso-8601-format (:date-time tf/formatters))
 
-(def partition-date-format (:basic-date tf/formatters))
-
 (s/defschema CookInfo
   "Schema for the /info endpoint response"
   {:authentication-scheme s/Str
@@ -257,22 +254,6 @@
   [s]
   (re-matches #"[\.a-zA-Z0-9_-]{1,128}" s))
 
-(defn non-empty-max-128-characters
-  "String of between 1 and 128 characters"
-  [s]
-  (<= 1 (.length s) 128))
-
-(def non-empty-max-128-characters-str
-  (s/constrained s/Str non-empty-max-128-characters))
-
-(defn valid-date-str?
-  "yyyyMMdd"
-  [s]
-  (try
-    (tf/parse partition-date-format s)
-    true
-    (catch Exception e false)))
-
 (defn valid-k8s-pod-label-value?
   "Returns true if s contains only '.', '_', '-' or any word or
   number characters, has length at least 2 and at most 63, and
@@ -290,31 +271,12 @@
    (s/optional-key :workload-id) (s/constrained s/Str valid-k8s-pod-label-value?)
    (s/optional-key :workload-details) (s/constrained s/Str valid-k8s-pod-label-value?)})
 
-; Datasets represent the data dependencies of jobs, which can be used by the scheduler to schedule jobs
-; on hosts to take advantage of data locality.
-; A dataset is decribed by two fields:
-; - dataset: a dictionary from string->string describing the dataset
-;   e.g. {"kind": "trades", "market": "NYSE"}
-; - partitions (optional): a description of a subset of the dataset required
-;   e.g. [{"begin": "20180101", "end": "20180201"}, {"begin": "20180301", "end": "20180401"}]
-; If partitions are specified, the dataset field must contain the key "partition-type" to describe the type of partition scheme.
-;   e.g. "partition-type": "date"
-(def Dataset
-  "Schema for a job dataset"
-  {:dataset {non-empty-max-128-characters-str non-empty-max-128-characters-str}
-   (s/optional-key :partitions) #{{non-empty-max-128-characters-str non-empty-max-128-characters-str}}})
-
 (def Disk
   "Schema for disk limit specifications"
   {:request (s/pred #(> % 1.0) 'greater-than-one)
    (s/optional-key :limit) (s/pred #(> % 1.0) 'greater-than-one)
    (s/optional-key :type) s/Str})
 
-(def DatePartition
-  "Schema for a date partition"
-  {(s/required-key "begin") (s/constrained s/Str valid-date-str?)
-   (s/required-key "end") (s/constrained s/Str valid-date-str?)})
-
 (def Constraint
   "Schema for user defined job host constraint"
   [(s/one NonEmptyString "attribute")
@@ -370,8 +332,7 @@
    ;; a lower case character or a digit, and has length between 2 to (62 + 2).
    :user UserName
    (s/optional-key :application) Application
-   (s/optional-key :expected-runtime) PosInt
-   (s/optional-key :datasets) #{Dataset}})
+   (s/optional-key :expected-runtime) PosInt})
 
 (def Job
   "Full schema for a job"
@@ -388,8 +349,6 @@
              (s/optional-key :max-runtime) PosInt
              (s/optional-key :name) JobName
              (s/optional-key :priority) JobPriority
-             (s/optional-key :datasets) [{:dataset {s/Keyword s/Str}
-                                          (s/optional-key :partitions) [{s/Keyword s/Str}]}]
              ;; The Java job client used to send the
              ;; status field on job submission. At some
              ;; point, that changed, but we will keep
@@ -419,8 +378,6 @@
               (s/optional-key :groups) [s/Uuid]
               (s/optional-key :instances) [Instance]
               (s/optional-key :pool) s/Str
-              (s/optional-key :datasets) #{{:dataset {s/Str s/Str}
-                                            (s/optional-key :partitions) #{{s/Str s/Str}}}}
               (s/optional-key :submit-pool) s/Str})
       prepare-schema-response))
 
@@ -699,17 +656,6 @@
                       :commit-latch/uuid (d/squuid)}]
     [commit-latch-id commit-latch]))
 
-(defn make-partition-ent
-  "Makes a datomic entity for a partition"
-  [partition-type partition]
-  (let [coerce-date (fn coerce-date [date-str]
-                      (tc/to-date (tf/parse partition-date-format date-str)))]
-    (case partition-type
-      "date"
-      {:dataset.partition/begin (coerce-date (partition "begin"))
-       :dataset.partition/end (coerce-date (partition "end"))}
-      :default (throw (IllegalArgumentException. (str "Unsupported partition type " partition-type))))))
-
 (defn get-default-container-for-pool
   "Given a pool name, determine a default container that should be run on it."
   [default-containers effective-pool-name]
@@ -746,7 +692,7 @@
   [{:keys [job pool-name pool-name-from-submission]} :- {:job Job} commit-latch-id db]
   (let [{:keys [uuid command max-retries max-runtime expected-runtime priority cpus mem disk gpus
                 user name ports uris env labels container group application disable-mea-culpa-retries
-                constraints executor progress-output-file progress-regex-string datasets checkpoint]
+                constraints executor progress-output-file progress-regex-string checkpoint]
          :or {group nil
               disable-mea-culpa-retries false}} job
         db-id (d/tempid :db.part/user)
@@ -829,21 +775,6 @@
                                    {:db/id gpus-id
                                     :resource/type :resource.type/gpus
                                     :resource/amount (double gpus)}]))])
-        datasets (map (fn [{:keys [dataset partitions]}]
-                        (let [parameters (map (fn [[k v]] {:db/id (d/tempid :db.part/user)
-                                                           :dataset.parameter/key k
-                                                           :dataset.parameter/value v})
-                                              dataset)]
-                          (if (contains? dataset "partition-type")
-                            (let [partition-type (get dataset "partition-type")
-                                  partitions (map (partial make-partition-ent partition-type) partitions)]
-                              {:db/id (d/tempid :db.part/user)
-                               :dataset/parameters parameters
-                               :dataset/partition-type partition-type
-                               :dataset/partitions partitions})
-                            {:db/id (d/tempid :db.part/user)
-                             :dataset/parameters parameters})))
-                      datasets)
         job-submit-time (Date.)
         txn (cond-> {:db/id db-id
                      :job/command command
@@ -896,8 +827,7 @@
                                pool-name))
                       (assoc :job/submit-pool-name
                              pool-name-from-submission))
-                    checkpoint (assoc :job/checkpoint (build-checkpoint checkpoint))
-                    (seq datasets) (assoc :job/datasets datasets))
+                    checkpoint (assoc :job/checkpoint (build-checkpoint checkpoint)))
         txn (plugins/adjust-job adjustment/plugin txn db)]
 
     ;; TODO batch these transactions to improve performance
@@ -982,33 +912,6 @@
    :host-placement (make-default-host-placement)
    :straggler-handling default-straggler-handling})
 
-(defn validate-partitions
-  "Ensures that the given partitions are valid.
-   Currently, we only support date partitions which have a begin and end value, specified as a yyyyMMdd date."
-  [{:keys [dataset partitions] :as in}]
-  (let [partition-type (get dataset "partition-type" nil)]
-    (cond
-      (nil? partition-type)
-      (when-not (empty? partitions)
-        (throw (ex-info "Dataset with partitions must supply partition-type"
-                        {:dataset dataset})))
-      (= "date" partition-type)
-      (run! (partial s/validate DatePartition) partitions)
-      :else
-      (throw (ex-info "Dataset with unsupported partition type"
-                      {:dataset dataset}))))
-  in)
-
-(defn munge-datasets
-  "Converts dataset and partition keys (keywords) into strings"
-  [datasets]
-  (->> datasets
-       (map (fn [{:keys [dataset partitions]}]
-              (cond-> {:dataset (walk/stringify-keys dataset)}
-                partitions
-                (assoc :partitions (into #{} (walk/stringify-keys partitions))))))
-       (into #{})))
-
 (defn validate-gpu-job
   "Validates that a job requesting GPUs is supported on the pool"
   [gpu-enabled? pool-name {:keys [gpus env]}]
@@ -1051,7 +954,7 @@
   [db pool-name user task-constraints gpu-enabled? new-group-uuids
    {:keys [cpus mem disk gpus uuid command priority max-retries max-runtime expected-runtime name
            uris ports env labels container group application disable-mea-culpa-retries
-           constraints executor progress-output-file progress-regex-string datasets checkpoint]
+           constraints executor progress-output-file progress-regex-string checkpoint]
     :or {group nil
          disable-mea-culpa-retries false}
     :as job}
@@ -1096,7 +999,6 @@
                  (when progress-output-file {:progress-output-file progress-output-file})
                  (when progress-regex-string {:progress-regex-string progress-regex-string})
                  (when application {:application application})
-                 (when datasets {:datasets (munge-datasets datasets)})
                  (when checkpoint {:checkpoint (if (-> checkpoint :options :preserve-paths)
                                                  (update-in checkpoint [:options :preserve-paths] set)
                                                  checkpoint)}))
@@ -1141,8 +1043,6 @@
     (doseq [{:keys [executable? extract?] :as uri} (:uris munged)
             :when (and (not (nil? executable?)) (not (nil? extract?)))]
       (throw (ex-info "Uri cannot set executable and extract" uri)))
-    (doseq [dataset (:datasets munged)]
-      (validate-partitions dataset))
     (when (and group-uuid
                (not (valid-group-uuid? db
                                        new-group-uuids
@@ -1312,8 +1212,6 @@
                                        (map str)))))
           instances (map #(fetch-instance-map db %1) (:job/instance job))
           submit-time (util/job->submit-time job)
-          datasets (when (seq (:job/datasets job))
-                     (dl/get-dataset-maps job))
           attempts-consumed (util/job-ent->attempts-consumed db job)
           retries-remaining (- (:job/max-retries job) attempts-consumed)
           disable-mea-culpa-retries (:job/disable-mea-culpa-retries job false)
@@ -1361,7 +1259,6 @@
               pool (assoc :pool (:pool/name pool))
               container (assoc :container (container->response-map container))
               checkpoint (assoc :checkpoint (util/job-ent->checkpoint job))
-              datasets (assoc :datasets datasets)
               submit-pool-name (assoc :submit-pool submit-pool-name)))))
 
 (defn fetch-job-map
@@ -3273,46 +3170,6 @@
                     (->> (pool/all-pools db)
                          (mapv pool-entity->consumable-map))))}))
 
-(def DataLocalUpdateTimeResponse
-  {s/Uuid (s/maybe s/Str)})
-
-(def DataLocalCostResponse
-  {s/Str {:cost s/Num
-          :suitable s/Bool}})
-
-(defn data-local-update-time-handler
-  "Handler for return the last update time of data locality"
-  [conn]
-  (base-cook-handler
-    {:allowed-methods [:get]
-     ; TODO (pschorf) - cache this if it is a performance bottleneck
-     :handle-ok (fn [_]
-                  (let [uuid->datasets (->> (d/db conn)
-                                            queries/get-pending-job-ents
-                                            (filter (fn [j] (not (empty? (:job/datasets j)))))
-                                            (map (fn [j] [(:job/uuid j) (dl/get-dataset-maps j)]))
-                                            (into {}))
-                        datasets->update-time (dl/get-last-update-time)]
-                    (pc/map-vals (fn [datasets]
-                                   (when-let [update-time (get datasets->update-time datasets nil)]
-                                     (tf/unparse iso-8601-format update-time)))
-                                 uuid->datasets)))}))
-
-(defn data-local-cost-handler
-  "Handler which returns the data locality costs for a given job"
-  [conn]
-  (base-cook-handler
-    {:allowed-methods [:get]
-     :exists? (fn [ctx]
-                (let [uuid (get-in ctx [:request :params :uuid])
-                      job-ent (d/entity (d/db conn) [:job/uuid (UUID/fromString uuid)])
-                      datasets (dl/get-dataset-maps job-ent)]
-                  (if-let [costs (get (dl/get-data-local-costs) datasets nil)]
-                    {:costs costs}
-                    false)))
-     :handle-ok (fn [{:keys [costs]}]
-                  costs)}))
-
 ;;
 ;; /compute-clusters
 ;;
@@ -3961,24 +3818,6 @@
                                  (update-instance-progress-handler
                                    conn mock-auth-fn leadership-atom leader-selector progress-aggregator-chan))}}))))
 
-      ; Data locality debug endpoints
-      (c-api/context
-        "/data-local/:uuid" []
-        :path-params [uuid :- s/Uuid]
-        (c-api/resource
-          {:produces ["application/json"]
-           :responses {200 {:schema DataLocalCostResponse}}
-           :get {:summary "Returns summary information on the current data locality status"
-                 :handler (data-local-cost-handler conn)}}))
-
-      (c-api/context
-        "/data-local" []
-        (c-api/resource
-          {:produces ["application/json"]
-           :responses {200 {:schema DataLocalUpdateTimeResponse}}
-           :get {:summary "Returns summary information on the current data locality status"
-                 :handler (data-local-update-time-handler conn)}}))
-
       (ANY "/queue" []
         (waiting-jobs conn mesos-pending-jobs-fn is-authorized-fn leadership-atom leader-selector))
       (ANY "/running" []
diff --git a/scheduler/src/cook/scheduler/constraints.clj b/scheduler/src/cook/scheduler/constraints.clj
index ea1c13c5bd..fe0b79807a 100644
--- a/scheduler/src/cook/scheduler/constraints.clj
+++ b/scheduler/src/cook/scheduler/constraints.clj
@@ -25,7 +25,6 @@
             [cook.group :as group]
             [cook.rate-limit :as ratelimit]
             [cook.regexp-tools :as regexp-tools]
-            [cook.scheduler.data-locality :as dl]
             [cook.tools :as util]
             [swiss.arrows :refer :all])
   (:import (com.netflix.fenzo ConstraintEvaluator ConstraintEvaluator$Result TaskRequest TaskTrackerState VirtualMachineCurrentState VirtualMachineLease TaskTracker$ActiveTask)
@@ -252,35 +251,6 @@
     [this _ vm-attributes _]
     (job-constraint-evaluate this _ vm-attributes)))
 
-(defrecord data-locality-constraint [job launch-wait-seconds]
-  JobConstraint
-  (job-constraint-name [this] (get-class-name this))
-  (job-constraint-evaluate [this _ vm-attributes]
-    (let [{:keys [job/submit-time]} job
-          datasets (dl/get-dataset-maps job)
-          data-locality-costs (dl/get-data-local-costs)
-          launch-after-age (t/plus (tc/from-date submit-time) (t/seconds launch-wait-seconds))
-          launch-without-data? (t/after? (t/now) launch-after-age)]
-      (if (contains? data-locality-costs datasets)
-        (let [{:keys [suitable] :or {suitable true}} (get-in data-locality-costs [datasets (get vm-attributes "HOSTNAME")])]
-          (if suitable
-            [true nil]
-            [false "Host is not suitable for datasets"]))
-        (if launch-without-data?
-          [true nil]
-          [false "No data locality costs available"]))))
-  (job-constraint-evaluate [this _ vm-attributes _]
-    (job-constraint-evaluate this _ vm-attributes)))
-
-(defn build-data-locality-constraint
-  "If the job supports data local and the data local fitness calculator is in use,
-   returns a data-locality-constraint"
-  [{:keys [job/datasets] :as job}]
-  (let [{:keys [launch-wait-seconds]} (config/data-local-fitness-config)
-        fitness-calculator (config/fitness-calculator-config)]
-    (when (and (not (empty? datasets)) (= fitness-calculator dl/data-local-fitness-calculator))
-      (->data-locality-constraint job launch-wait-seconds))))
-
 (defn job->default-constraints
   "Returns the list of default constraints configured for the job's pool"
   [job]
@@ -440,7 +410,6 @@
                                   build-disk-host-constraint
                                   build-user-defined-constraint
                                   build-estimated-completion-constraint
-                                  build-data-locality-constraint
                                   build-checkpoint-locality-constraint])
 
 (defn fenzoize-job-constraint
diff --git a/scheduler/src/cook/scheduler/data_locality.clj b/scheduler/src/cook/scheduler/data_locality.clj
deleted file mode 100644
index bb59c4c1e6..0000000000
--- a/scheduler/src/cook/scheduler/data_locality.clj
+++ /dev/null
@@ -1,218 +0,0 @@
-(ns cook.scheduler.data-locality
-  (:require [cheshire.core :as cheshire]
-            [clj-http.client :as http]
-            [clj-time.coerce :as tc]
-            [clj-time.core :as t]
-            [clj-time.format :as tf]
-            [clojure.set :as set]
-            [clojure.tools.logging :as log]
-            [cook.cache :as ccache]
-            [cook.caches :as caches]
-            [cook.config :as config]
-            [cook.queries :as queries]
-            [cook.tools :as util]
-            [datomic.api :as d]
-            [metrics.histograms :as histograms]
-            [metrics.timers :as timers]
-            [plumbing.core :as pc])
-  (:import (com.netflix.fenzo VMTaskFitnessCalculator)
-           (java.util UUID)))
-
-(def partition-date-format (:basic-date tf/formatters))
-
-(defn- make-partition-map
-  [partition-type partition]
-  (let [format-date (fn format-date [date] (tf/unparse partition-date-format (tc/from-date date)))]
-    (case partition-type
-      "date" {"begin" (format-date (:dataset.partition/begin partition))
-              "end" (format-date (:dataset.partition/end partition))})))
-
-(defn- make-dataset-maps
-  "Takes a job and returns a set of datasets with the shape
-   #{{:dataset {\"key\" \"value\"}
-      :partitions [{\"key\" \"value\"}]}}"
-  [job]
-  (->> job
-       :job/datasets
-       (map (fn [{:keys [dataset/partitions dataset/partition-type dataset/parameters]}]
-              (let [partitions (when (not (empty? partitions))
-                                 (->> partitions
-                                      (map (partial make-partition-map partition-type))
-                                      (into #{})))]
-                (cond-> {:dataset (into {} (map (fn [p] [(:dataset.parameter/key p) (:dataset.parameter/value p)])
-                                                parameters))}
-                  partitions (assoc :partitions partitions)))))
-       (into #{})))
-
-(defn get-dataset-maps
-  "Returns the (possibly cached) datasets for the given job"
-  [job]
-  (ccache/lookup-cache! caches/job-uuid->dataset-maps-cache
-                        :job/uuid
-                        make-dataset-maps
-                        job))
-
-(histograms/defhistogram [cook-mesos data-locality cost-update-staleness])
-(histograms/defhistogram [cook-mesos data-locality cost-match-staleness])
-
-(let [datasets->host-name->cost-atom (atom {})
-      datasets->last-update-time-atom (atom {})]
-
-  (defn reset-data-local-costs! []
-    (reset! datasets->host-name->cost-atom {})
-    (reset! datasets->last-update-time-atom {}))
-
-  (defn update-data-local-costs
-    "Updates the current data local costs. Costs larger than the configured max-cost will be capped at max-cost
-     and negative costs will be reset to 0."
-    [datasets->host->cost stale-datasets]
-    (let [{:keys [cache-ttl-ms]} (config/data-local-fitness-config)
-          datasets->last-update-time @datasets->last-update-time-atom
-          current-time (t/now)
-          datasets-to-remove (filter (fn [datasets]
-                                       (when-let [update-time (datasets->last-update-time datasets nil)]
-                                         (< cache-ttl-ms
-                                            (t/in-millis (t/interval update-time current-time)))))
-                                     stale-datasets)
-          clean-datasets->host->cost (pc/map-vals (fn [host->cost]
-                                                    (pc/map-vals (fn [cost] (update cost :cost #(-> % (min 1.0) (max 0))))
-                                                                 host->cost))
-                                                  datasets->host->cost)]
-      (run! (fn [[_ update-time]] (histograms/update! cost-update-staleness (t/in-millis (t/interval update-time current-time))))
-            (apply dissoc datasets->last-update-time stale-datasets))
-      (swap! datasets->host-name->cost-atom (fn update-datasets->host-name->cost-atom-atom
-                                              [current]
-                                              (let [remove-old-values (apply dissoc current datasets-to-remove)]
-                                                (merge remove-old-values clean-datasets->host->cost))))
-      (swap! datasets->last-update-time-atom (fn update-last-update-time
-                                               [current]
-                                               (let [remove-old-values (apply dissoc current datasets-to-remove)
-                                                     new-values (pc/map-from-keys (constantly current-time) (keys clean-datasets->host->cost))]
-                                                 (merge remove-old-values new-values))))))
-
-  (defn get-data-local-costs
-    "Returns the current cost for jobs to run on each host"
-    []
-    @datasets->host-name->cost-atom)
-
-
-  (defn get-last-update-time
-    "Returns the last update time information"
-    []
-    @datasets->last-update-time-atom)
-
-  (defn jobs-to-update
-    "Return the list of job ids to update data locality costs for in the next iteration.
-     The number of jobs is limited by the configured batch size."
-    [db]
-    (let [pending-jobs (->> db
-                            queries/get-pending-job-ents
-                            (filter (fn [j] (not (empty? (:job/datasets j)))))
-                            (map (fn [{:keys [job/uuid job/submit-time] :as job}]
-                                   {:job/uuid uuid
-                                    :job/submit-time submit-time
-                                    :job/datasets (get-dataset-maps job)})))
-          datasets->last-update-time @datasets->last-update-time-atom
-          have-data?->job (group-by (fn [job] (contains? datasets->last-update-time (:job/datasets job)))
-                                    pending-jobs)
-          missing-data (->> (get have-data?->job false [])
-                            (sort-by :job/submit-time))
-          sorted-have-data (->> (get have-data?->job true [])
-                                (sort-by (fn [j] (datasets->last-update-time (:job/datasets j)))))
-          pending-job-datasets (->> pending-jobs
-                                    (map :job/datasets)
-                                    (into #{}))
-          no-longer-waiting (set/difference (-> datasets->last-update-time keys set)
-                                            pending-job-datasets)]
-      {:to-fetch (concat missing-data sorted-have-data)
-       :to-remove no-longer-waiting}))
-
-  (defn update-cost-staleness-metric
-    "Updates a data local cost staleness metric for the given jobs"
-    [jobs]
-    (let [job-datasets (->> jobs (map get-dataset-maps) (filter #(not (empty? %))) (into #{}))
-          update-times (->> (get-last-update-time)
-                            (filter (partial contains? job-datasets)))
-          now (t/now)]
-      (run! (fn [[_ update-time]]
-              (histograms/update! cost-match-staleness (t/in-millis (t/interval update-time now))))
-            update-times))))
-
-(defn fetch-data-local-costs
-  "Contacts the server to obtain the data local costs for the given job ids"
-  [jobs]
-  (let [{:keys [auth cost-endpoint]} (config/data-local-fitness-config)
-        job-uuid->datasets (into {} (map (fn [job] [(str (:job/uuid job)) (:job/datasets job)])
-                                         jobs))
-        batch-id (UUID/randomUUID)
-        request {:batch batch-id
-                 :tasks (map (fn [job] {:task_id (str (:job/uuid job))
-                                        :datasets (:job/datasets job)})
-                             jobs)}
-        _ (log/info "Updating data local costs for" (count jobs) "tasks with batch id" batch-id)
-        _ (log/debug "Updating data local costs :" (cheshire/generate-string request))
-        {:keys [body]} (http/post cost-endpoint (merge {:accept :json
-                                                        :as :json-string-keys
-                                                        :body (cheshire/generate-string request)
-                                                        :content-type :json}
-                                                       auth))
-        _ (log/debug "Got response:" body)]
-    (pc/for-map [{:strs [task_id costs]} (body "costs")]
-       (job-uuid->datasets task_id)
-       (pc/for-map [{:strs [node cost suitable]
-                     :or {suitable true}} costs] node {:cost cost
-                                                       :suitable suitable}))))
-
-(defn fetch-and-update-data-local-costs
-  "Determine the datasets which need to be updated, fetch the costs, and update the cache with the
-   latest costs."
-  [db]
-  (let [{:keys [batch-size]} (config/data-local-fitness-config)
-        {:keys [to-fetch to-remove]} (jobs-to-update db)
-        batch (into [] (take batch-size to-fetch))]
-    (log/debug "Updating data local costs for" (map :job/uuid batch))
-    (let [new-costs (fetch-data-local-costs batch)]
-      (log/info "Got updated costs for" (count new-costs) "jobs, attempting to remove" (count to-remove) "stale costs")
-      (log/debug "Got updated costs" new-costs)
-      (update-data-local-costs new-costs to-remove))))
-
-(timers/deftimer [cook-mesos data-locality cost-update-duration])
-
-(defn start-update-cycles!
-  "Starts a `chime-at-ch` on `trigger-chan` to update data local costs."
-  [datomic-conn trigger-chan]
-  (log/info "Starting data local service update chan")
-  (util/chime-at-ch trigger-chan
-                    (fn [] (timers/time!
-                            cost-update-duration
-                            (fetch-and-update-data-local-costs (d/db datomic-conn))))
-                    {:error-handler (fn [e]
-                                      (log/error e "Error updating data local costs"))}))
-
-(deftype DataLocalFitnessCalculator [^VMTaskFitnessCalculator base-calculator data-locality-weight]
-  VMTaskFitnessCalculator
-  (getName [this] (-> this .getClass .getSimpleName))
-  (calculateFitness [this task-request target-vm tracker-state]
-    (let [base-fitness (.calculateFitness base-calculator task-request target-vm tracker-state)
-          {:keys [job/uuid] :as job} (:job task-request)
-          datasets (get-dataset-maps job)]
-      (if-not (empty? (or datasets []))
-        (let [{:keys [cost suitable] :or {cost 1.0 suitable true}} (get-in (get-data-local-costs) [datasets (.getHostname target-vm)])
-              suitability-adjusted-cost (if suitable cost 1.0)
-              normalized-fitness (- 1.0 suitability-adjusted-cost)
-              data-local-fitness (* data-locality-weight normalized-fitness)
-              fitness (+ data-local-fitness (* (- 1 data-locality-weight) base-fitness))]
-          (log/debug "Computed data local fitness:" {:hostname (.getHostname target-vm)
-                                                     :job uuid
-                                                     :costs (get (get-data-local-costs) datasets)
-                                                     :fitness fitness})
-          fitness)
-        base-fitness))))
-
-(def data-local-fitness-calculator "cook.scheduler.data-locality/make-data-local-fitness-calculator")
-
-(defn make-data-local-fitness-calculator
-  "Loads settings from configuration to build a data local fitness calculator"
-  []
-  (let [{:keys [base-calculator data-locality-weight]} (config/data-local-fitness-config)]
-    (->DataLocalFitnessCalculator base-calculator data-locality-weight)))
diff --git a/scheduler/src/cook/scheduler/scheduler.clj b/scheduler/src/cook/scheduler/scheduler.clj
index 1285e42a5f..725c747ceb 100644
--- a/scheduler/src/cook/scheduler/scheduler.clj
+++ b/scheduler/src/cook/scheduler/scheduler.clj
@@ -41,7 +41,6 @@
             [cook.quota :as quota]
             [cook.rate-limit :as ratelimit]
             [cook.scheduler.constraints :as constraints]
-            [cook.scheduler.data-locality :as dl]
             [cook.scheduler.dru :as dru]
             [cook.scheduler.fenzo-utils :as fenzo]
             [cook.scheduler.offer :as offer]
@@ -605,7 +604,6 @@
       {:matches [] :failures []})
     (do
       (log/debug "In" pool-name "pool, tasks to scheduleOnce" considerable)
-      (dl/update-cost-staleness-metric considerable)
       (let [t (System/currentTimeMillis)
             leases (mapv #(offer/offer->lease % t) offers)
             considerable->task-id (plumbing.core/map-from-keys (fn [_] (str (d/squuid))) considerable)
diff --git a/scheduler/src/cook/schema.clj b/scheduler/src/cook/schema.clj
index 674c707a1d..9b6c1bdb5e 100644
--- a/scheduler/src/cook/schema.clj
+++ b/scheduler/src/cook/schema.clj
@@ -193,13 +193,6 @@ for a job. E.g. {:resources {:cpus 4 :mem 3} :constraints {\"unique_host_constra
     :db/valueType :db.type/ref
     :db/cardinality :db.cardinality/one
     :db.install/_attribute :db.part/db}
-   {:db/id (d/tempid :db.part/db)
-    :db/ident :job/datasets
-    :db/isComponent true
-    :db/valueType :db.type/ref
-    :db/cardinality :db.cardinality/many
-    :db.install/_attribute :db.part/db
-    :db/doc "Datasets required by a job"}
    {:db/id (d/tempid :db.part/db)
     :db/ident :job/last-waiting-start-time
     :db/doc "represents the last time the job entered the waiting state"
@@ -922,52 +915,6 @@ for a job. E.g. {:resources {:cpus 4 :mem 3} :constraints {\"unique_host_constra
     :db/doc "The DRU mode of the pool."
     :db/valueType :db.type/keyword
     :db/cardinality :db.cardinality/one
-    :db.install/_attribute :db.part/db}
-
-   ;; Dataset entity
-   {:db/id (d/tempid :db.part/db)
-    :db/ident :dataset/partition-type
-    :db/doc "The partition type of the dataset"
-    :db/valueType :db.type/string
-    :db/cardinality :db.cardinality/one
-    :db.install/_attribute :db.part/db}
-   {:db/id (d/tempid :db.part/db)
-    :db/ident :dataset/parameters
-    :db/isComponent true
-    :db/doc "Parameters defining the dataset"
-    :db/valueType :db.type/ref
-    :db/cardinality :db.cardinality/many
-    :db.install/_attribute :db.part/db}
-   {:db/id (d/tempid :db.part/db)
-    :db/ident :dataset/partitions
-    :db/isComponent true
-    :db/doc "Partitions of a dataset"
-    :db/valueType :db.type/ref
-    :db/cardinality :db.cardinality/many
-    :db.install/_attribute :db.part/db}
-   {:db/id (d/tempid :db.part/db)
-    :db/ident :dataset.parameter/key
-    :db/doc "Key for a parameter"
-    :db/valueType :db.type/string
-    :db/cardinality :db.cardinality/one
-    :db.install/_attribute :db.part/db}
-   {:db/id (d/tempid :db.part/db)
-    :db/ident :dataset.parameter/value
-    :db/doc "Value for a parameter"
-    :db/valueType :db.type/string
-    :db/cardinality :db.cardinality/one
-    :db.install/_attribute :db.part/db}
-   {:db/id (d/tempid :db.part/db)
-    :db/ident :dataset.partition/begin
-    :db/doc "Begin date for a date range partition"
-    :db/valueType :db.type/instant
-    :db/cardinality :db.cardinality/one
-    :db.install/_attribute :db.part/db}
-   {:db/id (d/tempid :db.part/db)
-    :db/ident :dataset.partition/end
-    :db/doc "End date for a date range partition"
-    :db/valueType :db.type/instant
-    :db/cardinality :db.cardinality/one
     :db.install/_attribute :db.part/db}])
 
 (def migration-add-index-to-job-state


Test patch:
diff --git a/integration/tests/cook/util.py b/integration/tests/cook/util.py
index ed8615b801..09b1f426b5 100644
--- a/integration/tests/cook/util.py
+++ b/integration/tests/cook/util.py
@@ -1656,11 +1656,6 @@ def should_expect_sandbox_directory_for_job(job):
     instance = job['instances'][0]
     return should_expect_sandbox_directory(instance)
 
-
-def data_local_service_is_set():
-    return os.getenv('DATA_LOCAL_SERVICE', None) is not None
-
-
 def demo_plugins_are_configured(cook_url):
     settings_dict = settings(cook_url)
     # Because we always create plugin configuration in config.clj, the first keys always exist.
@@ -1683,22 +1678,6 @@ def pool_mover_plugin_configured():
                  'cook.plugins.pool-mover/make-pool-mover-job-adjuster'
     return configured
 
-
-@functools.lru_cache()
-def _fenzo_fitness_calculator():
-    """Get the cook executor config from the /settings endpoint"""
-    cook_url = retrieve_cook_url()
-    _wait_for_cook(cook_url)
-    init_cook_session(cook_url)
-    fitness_calculator = get_in(settings(cook_url), 'fenzo-fitness-calculator')
-    logger.info(f"Cook's fitness calculator is {fitness_calculator}")
-    return fitness_calculator
-
-
-def using_data_local_fitness_calculator():
-    return _fenzo_fitness_calculator() == 'cook.scheduler.data-locality/make-data-local-fitness-calculator'
-
-
 def get_agent_endpoint(slaves, agent_hostname):
     agent = [agent for agent in slaves
              if agent['hostname'] == agent_hostname][0]
diff --git a/jobclient/java/src/test/java/com/twosigma/cook/jobclient/JobTest.java b/jobclient/java/src/test/java/com/twosigma/cook/jobclient/JobTest.java
index 90f129d00e..74accb5cc0 100644
--- a/jobclient/java/src/test/java/com/twosigma/cook/jobclient/JobTest.java
+++ b/jobclient/java/src/test/java/com/twosigma/cook/jobclient/JobTest.java
@@ -272,43 +272,6 @@ public void testParseFromJsonProgressRegexString() throws JSONException {
         Assert.assertEquals(progressOutputFile, actualJob.getProgressRegexString());
     }
 
-    @Test
-    public void testJsonizeDatasets() {
-        final JSONObject dataset = new JSONObject();
-        dataset.put("foo", "bar");
-        final JSONArray datasets = new JSONArray();
-        datasets.put(dataset);
-
-        final Job.Builder jobBuilder = new Job.Builder();
-        populateBuilder(jobBuilder);
-        jobBuilder.setDatasets(datasets);
-        final Job basicJob = jobBuilder.build();
-        final JSONObject jsonJob = Job.jsonizeJob(basicJob);
-
-        Assert.assertEquals(datasets.toString(), basicJob.getDatasets().toString());
-        Assert.assertEquals(datasets.toString(), jsonJob.getJSONArray("datasets").toString());
-    }
-
-    @Test
-    public void testParseDatasets() {
-        final JSONObject dataset = new JSONObject();
-        dataset.put("foo", "bar");
-        final JSONArray datasets = new JSONArray();
-        datasets.put(dataset);
-
-        final Job.Builder jobBuilder = new Job.Builder();
-        populateBuilder(jobBuilder);
-        jobBuilder.setDatasets(datasets);
-
-        final JSONObject json = Job.jsonizeJob(jobBuilder.build());
-        final String jsonString = new JSONArray().put(json).toString();
-        final List<Job> jobs = Job.parseFromJSON(jsonString);
-        Assert.assertEquals(jobs.size(), 1);
-        final Job actualJob = jobs.get(0);
-
-        Assert.assertEquals(datasets.toString(), actualJob.getDatasets().toString());
-    }
-
     @Test
     public void testParseFromJsonNoGpus() {
         final String jsonString = "[{" +
diff --git a/scheduler/src/cook/test/testutil.clj b/scheduler/src/cook/test/testutil.clj
index 9c53168667..e9fd82abc4 100644
--- a/scheduler/src/cook/test/testutil.clj
+++ b/scheduler/src/cook/test/testutil.clj
@@ -129,7 +129,6 @@
                            #'cook.plugins.pool/plugin
                            #'cook.plugins.submission/plugin-object
                            #'caches/job-ent->resources-cache
-                           #'caches/job-uuid->dataset-maps-cache
                            #'caches/job-ent->pool-cache
                            #'caches/task-ent->user-cache
                            #'caches/task->feature-vector-cache
@@ -225,21 +224,6 @@
       @(d/transact conn t))
     conn))
 
-(defn- make-dataset-entity
-  [{:keys [dataset partitions]}]
-  (let [dataset-ent {:db/id (d/tempid :db.part/user)
-                     :dataset/parameters (map (fn [[k v]]
-                                                {:db/id (d/tempid :db.part/user)
-                                                 :dataset.parameter/key k
-                                                 :dataset.parameter/value v})
-                                              dataset)}]
-    (if (contains? dataset "partition-type")
-      (let [partition-type (get dataset "partition-type")]
-        (assoc dataset-ent
-               :dataset/partition-type partition-type
-               :dataset/partitions (map (partial api/make-partition-ent partition-type) partitions)))
-      dataset-ent)))
-
 (defn create-dummy-job
   "Return the entity id for the created dummy job."
   [conn & {:keys [command committed? container custom-executor? datasets disable-mea-culpa-retries env executor gpus disk group
@@ -290,9 +274,7 @@
                         (when pool
                           {:job/pool (d/entid (d/db conn) [:pool/name pool])})
                         (when expected-runtime
-                          {:job/expected-runtime expected-runtime})
-                        (when datasets
-                          {:job/datasets (map make-dataset-entity datasets)}))
+                          {:job/expected-runtime expected-runtime}))
         job-info (if gpus
                    (update-in job-info [:job/resource] conj {:resource/type :resource.type/gpus
                                                              :resource/amount (double gpus)})
diff --git a/scheduler/test/cook/test/mesos.clj b/scheduler/test/cook/test/mesos.clj
index dc090d1b01..e593c58591 100644
--- a/scheduler/test/cook/test/mesos.clj
+++ b/scheduler/test/cook/test/mesos.clj
@@ -243,22 +243,6 @@
         (async/close! chan)
         (d/delete-database test-db-uri)))))
 
-
-(deftest test-make-trigger-chans
-  (with-redefs [config/data-local-fitness-config (constantly {:update-interval-ms nil})]
-    (let [trigger-chans (mesos/make-trigger-chans {:interval-seconds 1}
-                                                  {:publish-interval-ms 1000}
-                                                  {:optimizer-interval-seconds 1}
-                                                  {})]
-      (is (nil? (:update-data-local-costs-trigger-chan trigger-chans)))))
-
-  (with-redefs [config/data-local-fitness-config (constantly {:update-interval-ms 1000})]
-    (let [trigger-chans (mesos/make-trigger-chans {:interval-seconds 1}
-                                                  {:publish-interval-ms 1000}
-                                                  {:optimizer-interval-seconds 1}
-                                                  {})]
-      (is (:update-data-local-costs-trigger-chan trigger-chans)))))
-
 (defn dumy-new-cluster-configurations-fn [])
 (deftest test-dynamic-compute-cluster-configurations-setup
   (let [uri "datomic:mem://test-compute-cluster-config"
diff --git a/scheduler/test/cook/test/rest/api.clj b/scheduler/test/cook/test/rest/api.clj
index 9eb7daf4e7..abf553bd26 100644
--- a/scheduler/test/cook/test/rest/api.clj
+++ b/scheduler/test/cook/test/rest/api.clj
@@ -34,7 +34,6 @@
             [cook.rest.api :as api]
             [cook.rest.authorization :as auth]
             [cook.rest.impersonation :as imp]
-            [cook.scheduler.data-locality :as dl]
             [cook.scheduler.scheduler :as sched]
             [cook.task :as task]
             [cook.test.testutil :as testutil
@@ -1436,7 +1435,7 @@
           ; will have to dissoc it.
           [{:keys [mem max-retries max-runtime expected-runtime name gpus
                    command ports priority uuid user cpus application
-                   disable-mea-culpa-retries executor datasets checkpoint disk]
+                   disable-mea-culpa-retries executor checkpoint disk]
             :or {disable-mea-culpa-retries false}}]
           (cond-> {;; Fields we will fill in from the provided args:
                    :command command
@@ -1466,10 +1465,8 @@
             disk (assoc :disk disk)
             expected-runtime (assoc :expected-runtime expected-runtime)
             executor (assoc :executor executor)
-            checkpoint (assoc :checkpoint checkpoint)
-            datasets (assoc :datasets datasets)))]
-    (with-redefs [caches/job-uuid->dataset-maps-cache (testutil/new-cache)
-                  config/compute-clusters (constantly [{:factory-fn 'cook.mesos.mesos-compute-cluster/factory-fn
+            checkpoint (assoc :checkpoint checkpoint)))]
+    (with-redefs [config/compute-clusters (constantly [{:factory-fn 'cook.mesos.mesos-compute-cluster/factory-fn
                                                         :config {:framework-id "test-framework"}}])]
 
       (testing "Job creation"
@@ -1529,20 +1526,6 @@
                                   (re-pattern (str ".*:job/uuid.*" uuid ".*already exists"))
                                   (testutil/create-jobs! conn {::api/job-pool-name-maps [{:job job}]})))))
 
-        (testing "should work with datasets"
-          (let [conn (restore-fresh-database! "datomic:mem://mesos-api-test")
-                {:keys [uuid] :as job} (assoc (minimal-job) :datasets #{{:dataset {"foo" "bar"}}
-                                                                        {:dataset {"foo" "bar"
-                                                                                   "partition-type" "date"}
-                                                                         :partitions #{{"begin" "20180101"
-                                                                                        "end" "20180201"}
-                                                                                       {"begin" "20180301"
-                                                                                        "end" "20180401"}}}})]
-            (is (= {::api/results (str "submitted jobs " uuid)}
-                   (testutil/create-jobs! conn {::api/job-pool-name-maps [{:job job}]})))
-            (is (= (expected-job-map job)
-                   (dissoc (api/fetch-job-map (db conn) uuid) :submit_time)))))
-
         (testing "should work when the job specifies an application"
           (let [conn (restore-fresh-database! "datomic:mem://mesos-api-test")
                 application {:name "foo-app", :version "0.1.0"}
@@ -1610,15 +1593,6 @@
                        (dissoc :submit_time)
                        (update :executor name))))))
 
-        (testing "should work with data locality support"
-          (let [conn (restore-fresh-database! "datomic:mem://data-locality-submit")
-                {:keys [uuid] :as job} (assoc (minimal-job) :supports_data_locality true)]
-            (is (= {::api/results (str "submitted jobs " uuid)}
-                   (testutil/create-jobs! conn {::api/job-pool-name-maps [{:job job}]})))
-            (is (= (expected-job-map job)
-                   (-> (api/fetch-job-map (db conn) uuid)
-                       (dissoc :submit_time))))))
-
         (testing "should work for docker containers with"
           (let [docker-container {:docker {:force-pull-image false
                                            :image "docker-image"}
@@ -2492,31 +2466,6 @@
             (is (= 201 status) (str response))
             (is (= "large-job-pool" (-> conn d/db (d/entity [:job/uuid job-uuid]) cached-queries/job->pool-name)))))))))
 
-(deftest test-validate-partitions
-  (is (api/validate-partitions {:dataset {"foo" "bar"}}))
-  (is (thrown-with-msg? ExceptionInfo #"supply partition-type"
-                        (api/validate-partitions {:dataset {"foo" "bar"}
-                                                  :partitions [{}]})))
-  (is (api/validate-partitions {:dataset {"foo" "bar"
-                                          "partition-type" "date"}
-                                :partitions [{"begin" "20180801"
-                                              "end" "20180801"}]}))
-  (is (thrown-with-msg? ExceptionInfo #"unsupported partition type"
-                        (api/validate-partitions {:dataset {"partition-type" "foo"}})))
-  (is (thrown-with-msg? ExceptionInfo #"\"foo\" disallowed-key"
-                        (api/validate-partitions {:dataset {"partition-type" "date"}
-                                                  :partitions [{"begin" "20180101"
-                                                                "end" "20180101"
-                                                                "foo" "bar"}]})))
-  (is (thrown-with-msg? ExceptionInfo #"\"end\" missing"
-                        (api/validate-partitions {:dataset {"partition-type" "date"}
-                                                  :partitions [{"begin" "20180101"}]}))))
-
-(deftest test-date?
-  (is (api/valid-date-str? "20180101"))
-  (is (not (api/valid-date-str? "2018-01-01")))
-  (is (not (api/valid-date-str? "20180132"))))
-
 (deftest test-match-default-containers
   (let [default-containers
         [{:pool-regex "^foo$" :container {:foo 1}}
diff --git a/scheduler/test/cook/test/scheduler/constraints.clj b/scheduler/test/cook/test/scheduler/constraints.clj
index a5ce7b12ab..68612e1f2e 100644
--- a/scheduler/test/cook/test/scheduler/constraints.clj
+++ b/scheduler/test/cook/test/scheduler/constraints.clj
@@ -21,7 +21,6 @@
             [cook.caches :as caches]
             [cook.config :as config]
             [cook.scheduler.constraints :as constraints]
-            [cook.scheduler.data-locality :as dl]
             [cook.scheduler.offer :as offer]
             [cook.scheduler.scheduler :as sched]
             [cook.test.testutil :as testutil
@@ -439,75 +438,6 @@
     (is (not (first (constraints/job-constraint-evaluate constraint nil {"host-start-time" 0.0}))))
     (is (first (constraints/job-constraint-evaluate constraint nil {"host-start-time" 51.0})))))
 
-
-(deftest test-data-locality-constraint
-  (with-redefs [caches/job-uuid->dataset-maps-cache (testutil/new-cache)]
-    (testing "disabled when not using data local fitness calculator"
-      (with-redefs [config/fitness-calculator-config (constantly config/default-fitness-calculator)
-                    config/data-local-fitness-config (constantly {:launch-wait-seconds 60})]
-        (is (nil? (constraints/build-data-locality-constraint {:job/uuid (UUID/randomUUID)})))
-        (is (nil? (constraints/build-data-locality-constraint {:job/uuid (UUID/randomUUID)
-                                                               :job/datasets #{{:dataset {"a" "a"}}}})))))
-
-    (testing "disabled for non data-local jobs"
-      (with-redefs [config/fitness-calculator-config (constantly dl/data-local-fitness-calculator)
-                    config/data-local-fitness-config (constantly {:launch-wait-seconds 60})]
-        (is (nil? (constraints/build-data-locality-constraint {:job/uuid (UUID/randomUUID)})))
-        (is (not (nil? (constraints/build-data-locality-constraint {:job/uuid (UUID/randomUUID)
-                                                                    :job/datasets #{{:dataset {"a" "a"}}}}))))))
-
-    (testing "passes jobs older than launch-wait-seconds"
-      (with-redefs [config/fitness-calculator-config (constantly dl/data-local-fitness-calculator)
-                    config/data-local-fitness-config (constantly {:launch-wait-seconds 60})]
-        (dl/reset-data-local-costs!)
-        (let [submit-time (tc/to-date (t/minus (t/now) (t/seconds 61)))
-              constraint (constraints/build-data-locality-constraint {:job/uuid (UUID/randomUUID)
-                                                                      :job/datasets #{{:dataset {"a" "a"}}}
-                                                                      :job/submit-time submit-time})
-              [passes reason] (constraints/job-constraint-evaluate constraint
-                                                                   nil
-                                                                   nil)]
-          (is passes))))
-
-    (testing "requires data for newer jobs"
-      (dl/reset-data-local-costs!)
-      (with-redefs [config/fitness-calculator-config (constantly dl/data-local-fitness-calculator)
-                    config/data-local-fitness-config (constantly {:launch-wait-seconds 60})]
-        (let [with-data-datasets #{{:dataset {"a" "a"}}}
-              _ (dl/update-data-local-costs {with-data-datasets {"hostA" {:cost 0
-                                                                          :suitable true}}} [])
-              with-data-constraint (constraints/build-data-locality-constraint {:job/uuid (UUID/randomUUID)
-                                                                                :job/datasets #{{:dataset/parameters #{{:dataset.parameter/key "a" :dataset.parameter/value "a"}}}}
-                                                                                :job/submit-time (tc/to-date (t/now))})
-              without-data-constraint (constraints/build-data-locality-constraint {:job/uuid (UUID/randomUUID)
-                                                                                   :job/datasets #{{:datasets {"b" "b"}}}
-                                                                                   :job/submit-time (tc/to-date (t/now))})
-              [with-data-result _] (constraints/job-constraint-evaluate with-data-constraint nil nil)
-              [without-data-result msg] (constraints/job-constraint-evaluate without-data-constraint nil nil)]
-          (is with-data-result)
-          (is (not without-data-result))
-          (is (= "No data locality costs available" msg)))))
-
-    (testing "fails for unsuitable hosts"
-      (dl/reset-data-local-costs!)
-      (with-redefs [config/fitness-calculator-config (constantly dl/data-local-fitness-calculator)
-                    config/data-local-fitness-config (constantly {:launch-wait-seconds 60})]
-        (let [datasets #{{:dataset {"a" "a"}}}
-              _ (dl/update-data-local-costs {datasets {"hostA" {:cost 0
-                                                                :suitable true}
-                                                       "hostB" {:cost 1.0
-                                                                :suitable false}}} [])
-              constraint (constraints/build-data-locality-constraint {:job/uuid (UUID/randomUUID)
-                                                                      :job/datasets #{{:dataset/parameters #{{:dataset.parameter/key "a" :dataset.parameter/value "a"}}}}
-                                                                      :job/submit-time (tc/to-date (t/now))})]
-
-          ;; suitable
-          (is (= [true nil] (constraints/job-constraint-evaluate constraint nil {"HOSTNAME" "hostA"})))
-          ;; default allow
-          (is (= [true nil] (constraints/job-constraint-evaluate constraint nil {"HOSTNAME" "hostC"})))
-          ;; unsuitable
-          (is (= [false "Host is not suitable for datasets"] (constraints/job-constraint-evaluate constraint nil {"HOSTNAME" "hostB"}))))))))
-
 (deftest test-job->previous-hosts-to-avoid
   (testing "uniqueness"
     (let [hostnames (constraints/job->previous-hosts-to-avoid
diff --git a/scheduler/test/cook/test/scheduler/data_locality.clj b/scheduler/test/cook/test/scheduler/data_locality.clj
deleted file mode 100644
index b1d018024c..0000000000
--- a/scheduler/test/cook/test/scheduler/data_locality.clj
+++ /dev/null
@@ -1,323 +0,0 @@
-(ns cook.test.scheduler.data-locality
-  (:require [cheshire.core :as cheshire]
-            [clj-http.client :as http]
-            [clj-time.coerce :as tc]
-            [clj-time.core :as t]
-            [clojure.test :refer :all]
-            [cook.caches :as caches]
-            [cook.config :as config]
-            [cook.scheduler.data-locality :as dl]
-            [cook.test.testutil :as testutil :refer [create-dummy-job restore-fresh-database!]]
-            [cook.tools :as util]
-            [datomic.api :as d]
-            [plumbing.core :as pc])
-  (:import (com.netflix.fenzo TaskRequest VirtualMachineCurrentState VMTaskFitnessCalculator)
-           (java.util UUID)))
-
-(deftest test-update-data-local-costs
-  (with-redefs [config/data-local-fitness-config (constantly {:cache-ttl-ms 5000})
-                caches/job-uuid->dataset-maps-cache (testutil/new-cache)]
-    (testing "sanitizes input costs"
-      (let [job-1 (str (UUID/randomUUID))
-            job-2 (str (UUID/randomUUID))]
-        (dl/reset-data-local-costs!)
-        (dl/update-data-local-costs {job-1 {"hostA" {:cost 2.0
-                                                     :suitable true}
-                                            "hostB" {:cost 0.2
-                                                     :suitable true}
-                                            "hostC" {:cost -20
-                                                     :suitable false}}
-                                     job-2 {"hostB" {:cost 10.0
-                                                     :suitable false}}}
-                                    [])
-        (is (= {job-1 {"hostA" {:cost 1.0
-                                :suitable true}
-                       "hostB" {:cost 0.2
-                                :suitable true}
-                       "hostC" {:cost 0
-                                :suitable false}}
-                job-2 {"hostB" {:cost 1.0
-                                :suitable false}}}
-               (dl/get-data-local-costs)))))
-
-    (testing "correctly updates existing values"
-      (let [job-1 (UUID/randomUUID)
-            job-2 (UUID/randomUUID)
-            job-3 (UUID/randomUUID)
-            job-4 (UUID/randomUUID)]
-        (dl/reset-data-local-costs!)
-        (dl/update-data-local-costs {job-1 {"hostA" {:cost 100
-                                                     :suitable true}}
-                                     job-2 {"hostB" {:cost 200
-                                                     :suitable true}}
-                                     job-4 {"hostC" {:cost 300
-                                                     :suitable true}}}
-                                    [])
-        (dl/update-data-local-costs {job-3 {"hostB" {:cost 300
-                                                     :suitable true}}
-                                     job-4 {"hostA" {:cost 200
-                                                     :suitable true}}}
-                                    [])
-        (is (= {job-1 {"hostA" {:cost 100
-                                :suitable true}}
-                job-2 {"hostB" {:cost 200
-                                :suitable true}}
-                job-3 {"hostB" {:cost 300
-                                :suitable true}}
-                job-4 {"hostA" {:cost 200
-                                :suitable true}}})
-            (dl/get-data-local-costs))))
-
-    (testing "waits until ttl expiry to remove elements"
-      (dl/reset-data-local-costs!)
-      (let [older-datasets #{{:dataset {"a" "a"}}}
-            newer-datasets #{{:dataset {"b" "b"}}}
-            cost {"hostA" {:cost 1.0
-                           :suitable true}}]
-        (with-redefs [t/now (constantly (tc/from-long 0))]
-          (dl/update-data-local-costs {older-datasets cost}
-                                      []))
-        (with-redefs [t/now (constantly (tc/from-long 4000))]
-          (dl/update-data-local-costs {newer-datasets cost}
-                                      [older-datasets]))
-        (is (= {older-datasets cost
-                newer-datasets cost}
-               (dl/get-data-local-costs)))
-
-        (with-redefs [t/now (constantly (tc/from-long 8000))]
-          (dl/update-data-local-costs {} [newer-datasets older-datasets]))
-
-        (is (= {newer-datasets cost}
-               (dl/get-data-local-costs)))))))
-
-(deftype FixedFitnessCalculator [fitness]
-  VMTaskFitnessCalculator
-  (getName [this] "FixedFitnessCalculator")
-  (calculateFitness [this _ _ _] fitness))
-
-(defn fake-vm-for-host [hostname]
-  (reify
-    VirtualMachineCurrentState
-    (getHostname [this] hostname)))
-
-(defrecord FakeTaskRequest [job]
-    TaskRequest)
-
-(deftest test-data-local-fitness-calculator
-  (with-redefs [caches/job-uuid->dataset-maps-cache (testutil/new-cache)]
-    (let [conn (restore-fresh-database! "datomic:mem://test-data-local-fitness-calculator")
-          base-fitness 0.5
-          base-calculator (FixedFitnessCalculator. base-fitness)
-          data-locality-weight 0.9
-          base-fitness-portion (* base-fitness (- 1 data-locality-weight))
-          calculator (dl/->DataLocalFitnessCalculator base-calculator
-                                                      data-locality-weight)
-          [d1 d2] [#{{:dataset {"a" "a"}}} #{{:dataset {"b" "b"}}}]
-          j1 (create-dummy-job conn :job-state :job.state/waiting
-                               :datasets d1)
-          j2 (create-dummy-job conn :job-state :job.state/waiting)
-          job-1 (d/entity (d/db conn) j1)
-          job-2 (d/entity (d/db conn) j2)]
-      (dl/update-data-local-costs {d1 {"hostA" {:cost 0
-                                                :suitable true}
-                                       "hostB" {:cost 0.5
-                                                :suitable true}
-                                       "unsuitable" {:cost 0.25
-                                                     :suitable false}}
-                                   d2 {"hostA" {:cost 0
-                                                :suitable true}
-                                       "hostB" {:cost 0
-                                                :suitable true}}}
-                                  [])
-      (testing "uses base fitness for jobs that do not support data locality"
-        (is (= base-fitness (.calculateFitness calculator (FakeTaskRequest. job-2) (fake-vm-for-host "hostA") nil)))
-        (is (= base-fitness (.calculateFitness calculator (FakeTaskRequest. job-2) (fake-vm-for-host "hostB") nil))))
-
-      (testing "calculates fitness for jobs that support data locality"
-        (is (= (+ data-locality-weight base-fitness-portion)
-               (.calculateFitness calculator (FakeTaskRequest. job-1) (fake-vm-for-host "hostA") nil)))
-        (is (= (+ (* 0.5 data-locality-weight) base-fitness-portion)
-               (.calculateFitness calculator (FakeTaskRequest. job-1) (fake-vm-for-host "hostB") nil)))
-        (is (= base-fitness-portion
-               (.calculateFitness calculator (FakeTaskRequest. job-1) (fake-vm-for-host "hostC") nil)))
-        (is (= base-fitness-portion
-               (.calculateFitness calculator (FakeTaskRequest. job-1) (fake-vm-for-host "unsuitable") nil)))))))
-
-
-(deftest test-jobs-to-update
-  (with-redefs [config/data-local-fitness-config (constantly {:batch-size 3})
-                caches/job-uuid->dataset-maps-cache (testutil/new-cache)]
-    (dl/reset-data-local-costs!)
-    (testing "does not update data for running and completed jobs"
-      (let [conn (restore-fresh-database! "datomic:mem://test-job-ids-to-update")
-            _ (create-dummy-job conn :job-state :job.state/running
-                                :datasets #{{:dataset {"a" "a"}}})
-            _ (create-dummy-job conn :job-state :job.state/completed
-                                :datasets #{{:dataset {"b" "b"}}})
-            _ (create-dummy-job conn :job-state :job.state/running
-                                :datasets #{{:dataset {"c" "c"}}})
-            {:keys [to-fetch]} (dl/jobs-to-update (d/db conn))]
-        (is (empty? to-fetch))))
-
-    (testing "ignores jobs which don't support data locality"
-      (let [conn (restore-fresh-database! "datomic:mem://test-job-ids-to-update")
-            _ (create-dummy-job conn :job-state :job.state/waiting)
-            {:keys [to-fetch]} (dl/jobs-to-update (d/db conn))]
-        (is (empty? to-fetch))))
-
-    (testing "prefers jobs with missing data, sorted by submit time"
-      (let [conn (restore-fresh-database! "datomic:mem://test-job-ids-to-update")
-            [d1 d2 d3 d4] [#{{:dataset {"a" "a"}}} #{{:dataset {"b" "b"}}}
-                           #{{:dataset {"c" "c"}}} #{{:dataset {"d" "d"}}}]
-            j1 (create-dummy-job conn :job-state :job.state/waiting
-                                 :submit-time (java.util.Date. 0)
-                                 :datasets d1)
-            j2 (create-dummy-job conn :job-state :job.state/waiting
-                                 :submit-time (java.util.Date. 1)
-                                 :datasets d2)
-            j3 (create-dummy-job conn :job-state :job.state/waiting
-                                 :submit-time (java.util.Date. 3)
-                                 :datasets d3)
-            j4 (create-dummy-job conn :job-state :job.state/waiting
-                                 :submit-time (java.util.Date. 2)
-                                 :datasets d4)
-            db (d/db conn)
-            _ (dl/update-data-local-costs {d1 {"hostA" {:cost 100
-                                                        :suitable true}}} [])
-            {:keys [to-fetch to-remove]} (dl/jobs-to-update db)]
-        (is (= (map :job/datasets to-fetch) [d2 d4 d3 d1]))
-        (is (= #{} to-remove))))
-
-    (testing "removes jobs which are no longer waiting"
-      (dl/reset-data-local-costs!)
-      (let [conn (restore-fresh-database! "datomic:mem://test-job-ids-to-update")
-            [d1 d2 d3] [#{{:dataset {"a" "a"}}} #{{:dataset {"b" "b"}}}
-                           #{{:dataset {"c" "c"}}}]
-            j1 (create-dummy-job conn :job-state :job.state/waiting
-                                 :datasets d1)
-            j2 (create-dummy-job conn :job-state :job.state/running
-                                 :datasets d2)
-            j3 (create-dummy-job conn :job-state :job.state/completed
-                                 :datasets d3)
-            db (d/db conn)
-            _ (dl/update-data-local-costs {d1 {"hostA" {:cost 100
-                                                        :suitable true}}
-                                           d2 {"hostA" {:cost 100
-                                                        :suitable true}}
-                                           d3 {"hostA" {:cost 100
-                                                        :suitable true}}}
-                                          [])
-            {:keys [to-fetch to-remove]} (dl/jobs-to-update db)]
-        (is (= (map :job/datasets to-fetch) [d1]))
-        (is (= to-remove #{d2 d3}))))))
-
-
-(deftest test-fetch-and-update-data-local-costs
-  (let [first-cost {"hostA" {:cost 1.0
-                             :suitable true}}
-        second-cost {"hostA" {:cost 0.5
-                              :suitable true}}
-        third-cost {"hostA" {:cost 0.2
-                             :suitable true}}
-        current-cost-atom (atom first-cost)]
-    (with-redefs [config/data-local-fitness-config (constantly {:batch-size 2
-                                                                :cache-ttl-ms 30000})
-                  dl/fetch-data-local-costs (fn [jobs]
-                                              (pc/map-from-keys (fn [_] @current-cost-atom)
-                                                                (map :job/datasets jobs)))
-                  caches/job-uuid->dataset-maps-cache (testutil/new-cache)]
-      (testing "calls service and updates jobs"
-        (dl/reset-data-local-costs!)
-        (let [conn (restore-fresh-database! "datomic:mem://test-fetch-and-update-data-local-costs")
-              j1 (create-dummy-job conn :job-state :job.state/waiting
-                                   :datasets #{{:dataset {"foo" "bar"}}})
-              j2 (create-dummy-job conn :job-state :job.state/waiting
-                                   :datasets #{{:dataset {"bar" "baz"}}})
-              db (d/db conn)
-              [id1 id2] (map #(:job/uuid (d/entity db %)) [j1 j2])]
-          (dl/fetch-and-update-data-local-costs db)
-          (is (= {#{{:dataset {"foo" "bar"}}} first-cost
-                  #{{:dataset {"bar" "baz"}}} first-cost}
-                 (dl/get-data-local-costs)))))
-
-      (testing "correctly computes update order"
-        (dl/reset-data-local-costs!)
-        (let [conn (restore-fresh-database! "datomic:mem://test-fetch-and-update-data-local-costs")
-              [d1 d2 d3] [#{{:dataset {"a" "a"}}} #{{:dataset {"b" "b"}}} #{{:dataset {"c" "c"}}}]
-              j1 (create-dummy-job conn :job-state :job.state/waiting
-                                   :datasets d1
-                                   :submit-time (java.util.Date. 0))
-              j2 (create-dummy-job conn :job-state :job.state/waiting
-                                   :datasets d2
-                                   :submit-time (java.util.Date. 1))
-              j3 (create-dummy-job conn :job-state :job.state/waiting
-                                   :datasets d3
-                                   :submit-time (java.util.Date. 2))
-              db (d/db conn)
-              [id1 id2 id3] (map #(:job/uuid (d/entity db %)) [j1 j2 j3])]
-          (dl/fetch-and-update-data-local-costs db)
-          (is (= {d1 first-cost
-                  d2 first-cost}
-                 (dl/get-data-local-costs)))
-          (reset! current-cost-atom second-cost)
-          (dl/fetch-and-update-data-local-costs db)
-          (let [cost-data (dl/get-data-local-costs)
-                ; We should have updated id3 and one other uuid, let's grab that one
-                [updated-dataset] (filter #(= @current-cost-atom (cost-data %))
-                                          [d1 d2])
-                [other-dataset] (keys (dissoc cost-data d3 updated-dataset))]
-            (is (= {d3 second-cost
-                    updated-dataset second-cost
-                    other-dataset first-cost}
-                   cost-data))
-
-            (reset! current-cost-atom third-cost)
-            (dl/fetch-and-update-data-local-costs db)
-            ; Now, we should update the job we skipped last time
-            (is (= third-cost (get (dl/get-data-local-costs) other-dataset)))))))))
-
-
-(deftest test-fetch-data-local-costs
-  (testing "correctly returns costs"
-    (let [jobs [{:job/uuid (UUID/randomUUID)
-                 :job/datasets #{{:dataset {"foo" "bar"}}}}
-                {:job/uuid (UUID/randomUUID)
-                 :job/datasets #{{:dataset {"bar" "baz"}}}}]
-          costs [[{"node" "hostA"
-                   "cost" 0.0}
-                  {"node" "hostB"
-                   "cost" 1.0
-                   "suitable" false}]
-                 [{"node" "hostA"
-                   "cost" 1.0}
-                  {"node" "hostB"
-                   "cost" 0.0
-                   "suitable" true}]]
-          request-body-atom (atom {})
-          cost-endpoint "http://test.example.com/cost"]
-      (with-redefs [config/data-local-fitness-config (constantly {:cost-endpoint cost-endpoint})
-                    http/post (fn [endpoint {:keys [body]}]
-                                (is (= cost-endpoint endpoint))
-                                (reset! request-body-atom body)
-                                {:body
-                                 {"costs" [{"task_id" (-> jobs first :job/uuid str)
-                                            "costs" (first costs)}
-                                           {"task_id" (-> jobs second :job/uuid str)
-                                            "costs" (second costs)}]}})]
-        (is (= {#{{:dataset {"foo" "bar"}}} {"hostA" {:cost 0.0
-                                                      :suitable true}
-                                             "hostB" {:cost 1.0
-                                                      :suitable false}}
-                #{{:dataset {"bar" "baz"}}} {"hostA" {:cost 1.0
-                                                      :suitable true}
-                                             "hostB" {:cost 0.0
-                                                      :suitable true}}}
-               (dl/fetch-data-local-costs jobs)))
-
-        (let [{:strs [batch tasks]} (cheshire/parse-string @request-body-atom)]
-          (is batch)
-          (is (= [{"task_id" (-> jobs first :job/uuid str)
-                   "datasets" [{"dataset" {"foo" "bar"}}]}
-                  {"task_id" (-> jobs second :job/uuid str)
-                   "datasets" [{"dataset" {"bar" "baz"}}]}]
-                 tasks)))))))
diff --git a/scheduler/test/cook/test/scheduler/scheduler.clj b/scheduler/test/cook/test/scheduler/scheduler.clj
index 4d876c3a9a..a706e3e8aa 100644
--- a/scheduler/test/cook/test/scheduler/scheduler.clj
+++ b/scheduler/test/cook/test/scheduler/scheduler.clj
@@ -38,7 +38,6 @@
             [cook.progress :as progress]
             [cook.quota :as quota]
             [cook.rate-limit :as rate-limit]
-            [cook.scheduler.data-locality :as dl]
             [cook.scheduler.offer :as offer]
             [cook.scheduler.scheduler :as sched]
             [cook.scheduler.share :as share]