Copy user-provided script to local directory

README.rst

@@ -59,12 +59,21 @@ as follows
    docker-compose build
    ./start-sge.sh
 
+If you have done this previously and need to refresh your solution you can do
+the following
+
+.. code-block:: bash
+
+   docker-compose stop
+   docker-compose build --no-cache
+   ./start-sge.sh
+
 And run tests with py.test in the master docker container
 
 .. code-block:: bash
 
   docker exec -it sge_master /bin/bash -c "cd /dask-drmaa; python setup.py develop"
-  docker exec -it sge_master py.test dask-drmaa/dask_drmaa --verbose
+  docker exec -it sge_master /bin/bash -c "cd /dask-drmaa; py.test dask_drmaa --verbose"
 
 
 Adaptive Load

dask_drmaa/core.py

@@ -1,6 +1,7 @@
 from collections import namedtuple
 import logging
 import os
+import shutil
 import socket
 import sys
 import tempfile
@@ -105,6 +106,7 @@ def __init__(self, template=None, cleanup_interval=1000, hostname=None,
                                  prefix='dask-worker-script',
                                  dir=os.path.curdir)
             self.script = fn
+            self._should_cleanup_script = True
 
             script_contents = make_job_script(executable=worker_bin_path,
                                               name='%s.%s' % (JOB_ID, TASK_ID),
@@ -120,6 +122,11 @@ def remove_script():
             os.chmod(self.script, 0o777)
 
         else:
+            self._should_cleanup_script = False
+            with ignoring(EnvironmentError):  # may be in the same path
+                script = shutil.copy(script, os.path.curdir)
+                self._should_cleanup_script = True
+            self.script = script
             assert not preexec_commands, "Cannot specify both script and preexec_commands"
 
         # TODO: check that user-provided script is executable
@@ -224,7 +231,7 @@ def close(self):
             self.stop_workers(self.workers, sync=True)
 
         self.local_cluster.close()
-        if os.path.exists(self.script):
+        if self._should_cleanup_script and os.path.exists(self.script):
             os.remove(self.script)
 
     def __enter__(self):

dask_drmaa/tests/test_core.py

@@ -9,8 +9,10 @@
 import pytest
 
 from dask_drmaa import DRMAACluster
+from dask_drmaa.core import make_job_script, worker_bin_path
 from distributed import Client
 from distributed.utils_test import loop, inc
+from distributed.utils import tmpfile
 
 
 def test_simple(loop):
@@ -184,3 +186,19 @@ def cleanup_logs():
 
     import atexit
     atexit.register(cleanup_logs)
+
+
+def test_passed_script(loop):
+    with tmpfile(extension='sh') as fn:
+        with open(fn, 'w') as f:
+            f.write(make_job_script(executable=worker_bin_path,
+                                    name='foo'))
+        os.chmod(fn, 0o777)
+        with DRMAACluster(scheduler_port=0, script=fn) as cluster:
+            tmp_script_location = cluster.script
+            assert cluster.script.split(os.path.sep)[-1] == fn.split(os.path.sep)[-1]
+            job = cluster.start_workers(1)
+            with Client(cluster, loop=loop) as client:
+                assert client.submit(lambda x: x + 1, 10).result() == 11
+        assert os.path.exists(fn)  # doesn't cleanup provided script
+        assert not os.path.exists(tmp_script_location)