Repository: DataBiosphere/toil

Issue:
toil-wdl-runner can run task containers on the leader
@jeizenga has a workflow where the first task doesn't ask for more than the default amount of resources. So we kick off the `WDLRootJob` locally on the leader, and it chains to a call-a-task job, which chains to the actual task job, which doesn't need to reschedule itself for more resources, and starts up the container.

The worker log goes something like:

```
[2024-02-02T11:10:19-0800] [MainThread] [D] [toil.worker] 'WDLWorkflowJob' kind-WDLRootJob/instance-87pqhjvf v4 would chain to ID kind-WDLWorkflowNodeJob/instance-e1oepvf7
[2024-02-02T11:10:19-0800] [MainThread] [I] [toil.worker] Chaining from 'WDLWorkflowJob' kind-WDLRootJob/instance-87pqhjvf v4 to 'WDLWorkflowNodeJob' call-formatAssembly kind-WDLWorkflowNodeJob/instance-e1oepvf7 v1

...

[2024-02-02T11:10:20-0800] [MainThread] [D] [toil.worker] 'WDLWorkflowNodeJob' call-formatAssembly kind-WDLRootJob/instance-87pqhjvf v6 would chain to ID kind-WDLTaskJob/instance-5ux5yawd
[2024-02-02T11:10:20-0800] [MainThread] [I] [toil.worker] Chaining from 'WDLWorkflowNodeJob' call-formatAssembly kind-WDLRootJob/instance-87pqhjvf v6 to 'WDLTaskJob' centromereAnnotation.formatAssembly kind-WDLTaskJob/instance-5ux5yawd v1
```

We probably want to make sure the WDLTaskJob itself isn't marked local and that chaining definitely refuses to go form a local job to a non-local one.

┆Issue is synchronized with this [Jira Story](https://ucsc-cgl.atlassian.net/browse/TOIL-1492)
┆Issue Number: TOIL-1492


Discussion:

Patch:
diff --git a/src/toil/worker.py b/src/toil/worker.py
index 13d9a668..2f086c9f 100644
--- a/src/toil/worker.py
+++ b/src/toil/worker.py
@@ -118,6 +118,14 @@ def nextChainable(predecessor: JobDescription, jobStore: AbstractJobStore, confi
         logger.debug("Next job is checkpoint, so finishing")
         return None
 
+    if not config.run_local_jobs_on_workers and predecessor.local and not successor.local:
+        # This job might be running on the leader, but the next job may not.
+        #
+        # TODO: Optimize by detecting whether we actually are on the leader,
+        # somehow.
+        logger.debug("Next job is not allowed to run on the leader, so finishing")
+        return None
+
     # Made it through! This job is chainable.
     return successor
 


Test patch:
diff --git a/src/toil/test/docs/scriptsTest.py b/src/toil/test/docs/scriptsTest.py
index c1316a81..c1cd599b 100644
--- a/src/toil/test/docs/scriptsTest.py
+++ b/src/toil/test/docs/scriptsTest.py
@@ -5,6 +5,8 @@ import subprocess
 import sys
 import unittest
 
+from typing import List
+
 pkg_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))  # noqa
 sys.path.insert(0, pkg_root)  # noqa
 
@@ -36,9 +38,9 @@ class ToilDocumentationTest(ToilTest):
         unittest.TestCase.tearDown(self)
 
     """Just check the exit code"""
-    def checkExitCode(self, script):
+    def checkExitCode(self, script, extra_args: List[str] = []):
         program = os.path.join(self.directory, "scripts", script)
-        process = subprocess.Popen([python, program, "file:my-jobstore", "--clean=always"],
+        process = subprocess.Popen([python, program, "file:my-jobstore", "--clean=always"] + extra_args,
                                    stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         stdout, stderr = process.communicate()
         if isinstance(stdout, bytes):
@@ -70,7 +72,8 @@ class ToilDocumentationTest(ToilTest):
         self.checkExitCode("tutorial_cwlexample.py")
 
     def testStats(self):
-        self.checkExitCode("tutorial_stats.py")
+        # This script asks for 4 cores but we might need to run the tests in only 3.
+        self.checkExitCode("tutorial_stats.py", ["--scale=0.5"])
 
     def testDynamic(self):
         self.checkExitCode("tutorial_dynamic.py")
diff --git a/src/toil/test/src/workerTest.py b/src/toil/test/src/workerTest.py
index dcecace1..6df4f447 100644
--- a/src/toil/test/src/workerTest.py
+++ b/src/toil/test/src/workerTest.py
@@ -18,6 +18,8 @@ from toil.jobStores.fileJobStore import FileJobStore
 from toil.test import ToilTest
 from toil.worker import nextChainable
 
+from typing import Optional
+
 
 class WorkerTests(ToilTest):
     """Test miscellaneous units of the worker."""
@@ -32,7 +34,7 @@ class WorkerTests(ToilTest):
 
     def testNextChainable(self):
         """Make sure chainable/non-chainable jobs are identified correctly."""
-        def createTestJobDesc(memory, cores, disk, preemptible, checkpoint):
+        def createTestJobDesc(memory, cores, disk, preemptible: bool = True, checkpoint: bool = False, local: Optional[bool] = None):
             """
             Create a JobDescription with no command (representing a Job that
             has already run) and return the JobDescription.
@@ -41,7 +43,16 @@ class WorkerTests(ToilTest):
             self.jobNumber += 1
 
             descClass = CheckpointJobDescription if checkpoint else JobDescription
-            jobDesc = descClass(requirements={'memory': memory, 'cores': cores, 'disk': disk, 'preemptible': preemptible}, jobName=name)
+            jobDesc = descClass(
+                requirements={
+                    'memory': memory,
+                    'cores': cores,
+                    'disk': disk,
+                    'preemptible': preemptible
+                }, 
+                jobName=name, 
+                local=local
+            )
 
             # Assign an ID
             self.jobStore.assign_job_id(jobDesc)
@@ -53,42 +64,64 @@ class WorkerTests(ToilTest):
             # Try with the branch point at both child and follow-on stages
 
             # Identical non-checkpoint jobs should be chainable.
-            jobDesc1 = createTestJobDesc(1, 2, 3, True, False)
-            jobDesc2 = createTestJobDesc(1, 2, 3, True, False)
+            jobDesc1 = createTestJobDesc(1, 2, 3)
+            jobDesc2 = createTestJobDesc(1, 2, 3)
             getattr(jobDesc1, successorType)(jobDesc2.jobStoreID)
             chainable = nextChainable(jobDesc1, self.jobStore, self.config)
             self.assertNotEqual(chainable, None)
-            self.assertEqual(jobDesc2.jobStoreID, chainable.jobStoreID)
+            self.assertEqual(chainable.jobStoreID, jobDesc2.jobStoreID)
 
             # Identical checkpoint jobs should not be chainable.
-            jobDesc1 = createTestJobDesc(1, 2, 3, True, False)
-            jobDesc2 = createTestJobDesc(1, 2, 3, True, True)
+            jobDesc1 = createTestJobDesc(1, 2, 3, checkpoint=True)
+            jobDesc2 = createTestJobDesc(1, 2, 3, checkpoint=True)
+            getattr(jobDesc1, successorType)(jobDesc2.jobStoreID)
+            self.assertEqual(nextChainable(jobDesc1, self.jobStore, self.config), None)
+
+            # Changing checkpoint from false to true should make it not chainable.
+            jobDesc1 = createTestJobDesc(1, 2, 3, checkpoint=False)
+            jobDesc2 = createTestJobDesc(1, 2, 3, checkpoint=True)
             getattr(jobDesc1, successorType)(jobDesc2.jobStoreID)
-            self.assertEqual(None, nextChainable(jobDesc1, self.jobStore, self.config))
+            self.assertEqual(nextChainable(jobDesc1, self.jobStore, self.config), None)
 
             # If there is no child we should get nothing to chain.
-            jobDesc1 = createTestJobDesc(1, 2, 3, True, False)
-            self.assertEqual(None, nextChainable(jobDesc1, self.jobStore, self.config))
+            jobDesc1 = createTestJobDesc(1, 2, 3)
+            self.assertEqual(nextChainable(jobDesc1, self.jobStore, self.config), None)
 
             # If there are 2 or more children we should get nothing to chain.
-            jobDesc1 = createTestJobDesc(1, 2, 3, True, False)
-            jobDesc2 = createTestJobDesc(1, 2, 3, True, False)
-            jobDesc3 = createTestJobDesc(1, 2, 3, True, False)
+            jobDesc1 = createTestJobDesc(1, 2, 3)
+            jobDesc2 = createTestJobDesc(1, 2, 3)
+            jobDesc3 = createTestJobDesc(1, 2, 3)
             getattr(jobDesc1, successorType)(jobDesc2.jobStoreID)
             getattr(jobDesc1, successorType)(jobDesc3.jobStoreID)
-            self.assertEqual(None, nextChainable(jobDesc1, self.jobStore, self.config))
+            self.assertEqual(nextChainable(jobDesc1, self.jobStore, self.config), None)
 
             # If there is an increase in resource requirements we should get nothing to chain.
-            reqs = {'memory': 1, 'cores': 2, 'disk': 3, 'preemptible': True, 'checkpoint': False}
+            base_reqs = {'memory': 1, 'cores': 2, 'disk': 3, 'preemptible': True, 'checkpoint': False}
             for increased_attribute in ('memory', 'cores', 'disk'):
+                reqs = dict(base_reqs)
                 jobDesc1 = createTestJobDesc(**reqs)
                 reqs[increased_attribute] += 1
                 jobDesc2 = createTestJobDesc(**reqs)
                 getattr(jobDesc1, successorType)(jobDesc2.jobStoreID)
-                self.assertEqual(None, nextChainable(jobDesc1, self.jobStore, self.config))
+                self.assertEqual(nextChainable(jobDesc1, self.jobStore, self.config), None)
 
             # A change in preemptability from True to False should be disallowed.
-            jobDesc1 = createTestJobDesc(1, 2, 3, True, False)
-            jobDesc2 = createTestJobDesc(1, 2, 3, False, True)
+            jobDesc1 = createTestJobDesc(1, 2, 3, preemptible=True)
+            jobDesc2 = createTestJobDesc(1, 2, 3, preemptible=False)
+            getattr(jobDesc1, successorType)(jobDesc2.jobStoreID)
+            self.assertEqual(nextChainable(jobDesc1, self.jobStore, self.config), None)
+
+            # A change in local-ness from True to False should be disallowed.
+            jobDesc1 = createTestJobDesc(1, 2, 3, local=True)
+            jobDesc2 = createTestJobDesc(1, 2, 3, local=False)
             getattr(jobDesc1, successorType)(jobDesc2.jobStoreID)
-            self.assertEqual(None, nextChainable(jobDesc1, self.jobStore, self.config))
+            self.assertEqual(nextChainable(jobDesc1, self.jobStore, self.config), None)
+
+            # A change in local-ness from False to True should be allowed,
+            # since running locally is an optional optimization.
+            jobDesc1 = createTestJobDesc(1, 2, 3, local=False)
+            jobDesc2 = createTestJobDesc(1, 2, 3, local=True)
+            getattr(jobDesc1, successorType)(jobDesc2.jobStoreID)
+            chainable = nextChainable(jobDesc1, self.jobStore, self.config)
+            self.assertNotEqual(chainable, None)
+            self.assertEqual(chainable.jobStoreID, jobDesc2.jobStoreID)