Add COOK_INSTANCE_NUM environment variable

integration/tests/cook/test_basic.py

@@ -2913,3 +2913,29 @@ def test_command_length_limit(self):
                             in str(resp.content))
         finally:
             util.kill_jobs(self.cook_url, [job_uuid], assert_response=False)
+
+
+    @unittest.skipUnless(util.supports_exit_code() and not util.has_one_agent(), "Requires exit code support and multiple agents")
+    def test_cook_instance_num(self):
+        command = 'bash -c \'exit $(($COOK_INSTANCE_NUM + 1))\''
+        job_uuid, resp = util.submit_job(self.cook_url, command=command, max_retries=2)
+        try:
+            def query():
+                jobs = util.query_jobs(self.cook_url, True, uuid=[job_uuid]).json()
+                self.logger.info(f'Found jobs: {jobs}')
+                return jobs[0]
+            def predicate(job):
+                if job['status'] != 'completed':
+                    return False
+                for instance in job['instances']:
+                    if 'exit_code' not in instance:
+                        return False
+                return True
+            job = util.wait_until(query, predicate)
+            self.assertEqual(2, len(job['instances']), job)
+            exit_codes = [i['exit_code'] for i in job['instances']]
+            exit_codes.sort()
+            self.assertEqual([1, 2], exit_codes, job)
+        finally:
+            util.kill_jobs(self.cook_url, [job_uuid], assert_response=False)
+

integration/tests/cook/util.py

@@ -1453,9 +1453,7 @@ def using_mesos():
     return _get_compute_cluster_factory_fn() == 'cook.mesos.mesos-compute-cluster/factory-fn'
 
 def has_one_agent():
-    # TODO: Actually check the number of agents.
-    cook_url = retrieve_cook_url()
-    _wait_for_cook(cook_url)
-    init_cook_session(cook_url)
-    compute_clusters = settings(cook_url)['compute-clusters']
-    return compute_clusters[0]['config']['compute-cluster-name'] == 'minikube'
+    return node_count() == 1
+
+def supports_exit_code():
+    return using_kubernetes() or is_cook_executor_in_use()

scheduler/bin/run-local-kubernetes.sh

@@ -20,14 +20,6 @@ SCHEDULER_DIR="$( dirname ${DIR} )"
 
 NAME=cook-scheduler-${COOK_PORT}
 
-SCHEDULER_EXECUTOR_DIR=${SCHEDULER_DIR}/resources/public
-
-EXECUTOR_DIR="$(dirname ${SCHEDULER_DIR})/executor"
-EXECUTOR_NAME=cook-executor-local
-COOK_EXECUTOR_COMMAND="${EXECUTOR_DIR}/dist/${EXECUTOR_NAME}/${EXECUTOR_NAME}"
-SCHEDULER_EXECUTOR_DIR=${SCHEDULER_DIR}/resources/public
-
-${EXECUTOR_DIR}/bin/prepare-executor.sh local ${SCHEDULER_EXECUTOR_DIR}
 
 if [ "${COOK_ZOOKEEPER_LOCAL}" = false ] ; then
     COOK_ZOOKEEPER="${ZOOKEEPER_IP}:2181"
@@ -54,9 +46,6 @@ FLASK_APP=${INTEGRATION_DIR}/src/data_locality/service.py flask run -p 35847 &
 echo "Creating environment variables..."
 export COOK_DATOMIC_URI="${COOK_DATOMIC_URI}"
 export COOK_FRAMEWORK_ID="${COOK_FRAMEWORK_ID}"
-export COOK_EXECUTOR=""
-export COOK_EXECUTOR_COMMAND="${COOK_EXECUTOR_COMMAND}"
-export COOK_EXECUTOR_PORTION=1
 export COOK_ONE_USER_AUTH=$(whoami)
 export COOK_HOSTNAME="cook-scheduler-${COOK_PORT}"
 export COOK_LOG_FILE="log/cook-${COOK_PORT}.log"

scheduler/src/cook/mesos/task.clj

@@ -136,7 +136,9 @@
         executor (executor-key->executor executor-key)
         resources (util/job-ent->resources job-ent)
         group-uuid (util/job-ent->group-uuid job-ent)
+        instance-num (str (count (:job/instance job-ent)))
         environment (cond-> (assoc (util/job-ent->env job-ent)
+                              "COOK_INSTANCE_NUM" instance-num
                               "COOK_INSTANCE_UUID" task-id
                               "COOK_JOB_UUID" (-> job-ent :job/uuid str))
                             group-uuid (assoc "COOK_JOB_GROUP_UUID" (str group-uuid))
@@ -154,9 +156,7 @@
         data (.getBytes
                (if cook-executor?
                  (json/write-str {"command" (:job/command job-ent)})
-                 (pr-str
-                   ;;TODO this data is a race-condition
-                   {:instance (str (count (:job/instance job-ent)))}))
+                 (pr-str {:instance instance-num}))
                "UTF-8")]
     (when (and (= :executor/cook (:job/executor job-ent))
                (not= executor-key :cook-executor))

scheduler/test/cook/test/mesos/task.clj

@@ -353,7 +353,8 @@
               db (d/db conn)
               job-ent (d/entity db job)
               framework-id {:value "framework-id"}
-              environment {"COOK_INSTANCE_UUID" task-id
+              environment {"COOK_INSTANCE_NUM" "0"
+                           "COOK_INSTANCE_UUID" task-id
                            "COOK_JOB_CPUS" "1.0"
                            "COOK_JOB_MEM_MB" "10.0"
                            "COOK_JOB_UUID" (-> job-ent :job/uuid str)}]
@@ -398,7 +399,8 @@
               job-ent (d/entity db job)
               group-ent (d/entity db group-ent-id)
               framework-id {:value "framework-id"}
-              environment {"COOK_INSTANCE_UUID" task-id
+              environment {"COOK_INSTANCE_NUM" "0"
+                           "COOK_INSTANCE_UUID" task-id
                            "COOK_JOB_CPUS" "1.0"
                            "COOK_JOB_GROUP_UUID" (-> group-ent :group/uuid str)
                            "COOK_JOB_MEM_MB" "10.0"
@@ -424,7 +426,8 @@
               db (d/db conn)
               job-ent (d/entity db job)
               framework-id {:value "framework-id"}
-              environment {"COOK_INSTANCE_UUID" task-id
+              environment {"COOK_INSTANCE_NUM" "0"
+                           "COOK_INSTANCE_UUID" task-id
                            "COOK_JOB_CPUS" "1.0"
                            "COOK_JOB_MEM_MB" "10.0"
                            "COOK_JOB_UUID" (-> job-ent :job/uuid str)}
@@ -449,7 +452,8 @@
               db (d/db conn)
               job-ent (d/entity db job)
               framework-id {:value "framework-id"}
-              environment {"COOK_INSTANCE_UUID" task-id
+              environment {"COOK_INSTANCE_NUM" "0"
+                           "COOK_INSTANCE_UUID" task-id
                            "COOK_JOB_CPUS" "1.0"
                            "COOK_JOB_MEM_MB" "10.0"
                            "COOK_JOB_UUID" (-> job-ent :job/uuid str)
@@ -484,7 +488,8 @@
               group-ent (d/entity db group-ent-id)
               job-ent (d/entity db job)
               framework-id {:value "framework-id"}
-              environment {"COOK_INSTANCE_UUID" task-id
+              environment {"COOK_INSTANCE_NUM" "0"
+                           "COOK_INSTANCE_UUID" task-id
                            "COOK_JOB_CPUS" "1.0"
                            "COOK_JOB_GROUP_UUID" (-> group-ent :group/uuid str)
                            "COOK_JOB_GPUS" "1000.0"
@@ -519,7 +524,8 @@
               db (d/db conn)
               job-ent (d/entity db job)
               framework-id {:value "framework-id"}
-              environment {"COOK_INSTANCE_UUID" task-id
+              environment {"COOK_INSTANCE_NUM" "0"
+                           "COOK_INSTANCE_UUID" task-id
                            "COOK_JOB_CPUS" "1.0"
                            "COOK_JOB_MEM_MB" "10.0"
                            "COOK_JOB_UUID" (-> job-ent :job/uuid str)}
@@ -555,7 +561,8 @@
               db (d/db conn)
               job-ent (d/entity db job)
               framework-id {:value "framework-id"}
-              environment {"COOK_INSTANCE_UUID" task-id
+              environment {"COOK_INSTANCE_NUM" "0"
+                           "COOK_INSTANCE_UUID" task-id
                            "COOK_JOB_CPUS" "1.0"
                            "COOK_JOB_MEM_MB" "10.0"
                            "COOK_JOB_UUID" (-> job-ent :job/uuid str)}
@@ -594,7 +601,8 @@
               db (d/db conn)
               job-ent (d/entity db job)
               framework-id {:value "framework-id"}
-              environment {"COOK_INSTANCE_UUID" task-id
+              environment {"COOK_INSTANCE_NUM" "0"
+                           "COOK_INSTANCE_UUID" task-id
                            "COOK_JOB_CPUS" "1.0"
                            "COOK_JOB_MEM_MB" "200.0"
                            "COOK_JOB_UUID" (-> job-ent :job/uuid str)}
@@ -634,7 +642,8 @@
               db (d/db conn)
               job-ent (d/entity db job)
               framework-id {:value "framework-id"}
-              environment {"COOK_INSTANCE_UUID" task-id
+              environment {"COOK_INSTANCE_NUM" "0"
+                           "COOK_INSTANCE_UUID" task-id
                            "COOK_JOB_CPUS" "1.0"
                            "COOK_JOB_MEM_MB" "200.0"
                            "COOK_JOB_UUID" (-> job-ent :job/uuid str)