From d576f11dd7d32299bd105e665694e54fc1c7055e Mon Sep 17 00:00:00 2001
From: Jim Crist-Harif <jcrist@users.noreply.github.com>
Date: Mon, 16 Mar 2020 16:13:43 -0500
Subject: [PATCH] Use the full Python path when using archived environment
 (#115)

Previously, if an archived environment was provided, we'd generate the
following script for running a dask worker:

```
source environment/bin/activate
dask-yarn services worker
```

This works fine on most systems, but for some reason fails on some
systems. The failure is odd - the `dask-yarn` CLI is properly found, but
the `dask_yarn` library fails to be on Python path. This may be due to
the shell not being rehashed properly, so the wrong python is used, I'm
not sure. It's likely a bug in conda-pack.

Either way, specifying the full path to the proper Python fixes things,
so we use that here. We now generate the following script:

```
source environment/bin/activate
environment/bin/python -m dask_yarn.cli services worker
```

This should work in all cases, even in the presence of a conda-pack bug
(which we should still fix).
---
 dask_yarn/core.py            | 24 +++++++++++-------------
 dask_yarn/tests/test_core.py | 23 +++++++++++------------
 2 files changed, 22 insertions(+), 25 deletions(-)

diff --git a/dask_yarn/core.py b/dask_yarn/core.py
index 39b18b5..c241bbb 100644
--- a/dask_yarn/core.py
+++ b/dask_yarn/core.py
@@ -122,21 +122,19 @@ def _files_and_build_script(environment):
     if scheme in {"conda", "venv", "python"}:
         path = environment[len(scheme) + 3 :]
         files = {}
+        if scheme == "conda":
+            setup = "conda activate %s" % path
+            cli = "dask-yarn"
+        elif scheme == "venv":
+            setup = "source %s/bin/activate" % path
+            cli = "dask-yarn"
+        else:
+            setup = ""
+            cli = "%s -m dask_yarn.cli" % path
     else:
-        # Treat archived environments the same as venvs
-        scheme = "venv"
-        path = "environment"
         files = {"environment": environment}
-
-    if scheme == "conda":
-        setup = "conda activate %s" % path
-        cli = "dask-yarn"
-    elif scheme == "venv":
-        setup = "source %s/bin/activate" % path
-        cli = "dask-yarn"
-    else:
-        setup = ""
-        cli = "%s -m dask_yarn.cli" % path
+        setup = "source environment/bin/activate"
+        cli = "environment/bin/python -m dask_yarn.cli"
 
     def build_script(cmd):
         command = "%s %s" % (cli, cmd)
diff --git a/dask_yarn/tests/test_core.py b/dask_yarn/tests/test_core.py
index 812077c..82164a6 100644
--- a/dask_yarn/tests/test_core.py
+++ b/dask_yarn/tests/test_core.py
@@ -318,17 +318,16 @@ def test_environment_conda(env, path):
     scheduler = spec.services["dask.scheduler"]
     assert not scheduler.files
     assert (
-        scheduler.script
-        == ("conda activate %s\n" "dask-yarn services scheduler") % path
+        scheduler.script == ("conda activate %s\ndask-yarn services scheduler") % path
     )
     worker = spec.services["dask.worker"]
     assert not worker.files
-    assert worker.script == ("conda activate %s\n" "dask-yarn services worker") % path
+    assert worker.script == ("conda activate %s\ndask-yarn services worker") % path
     client = spec.services["dask.client"]
     assert set(client.files) == {"script.py"}
     assert (
         client.script
-        == ("conda activate %s\n" "dask-yarn services client script.py foo bar") % path
+        == ("conda activate %s\ndask-yarn services client script.py foo bar") % path
     )
 
 
@@ -342,18 +341,16 @@ def test_environment_venv():
     assert not scheduler.files
     assert (
         scheduler.script
-        == ("source %s/bin/activate\n" "dask-yarn services scheduler") % path
+        == ("source %s/bin/activate\ndask-yarn services scheduler") % path
     )
     worker = spec.services["dask.worker"]
     assert not worker.files
-    assert (
-        worker.script == ("source %s/bin/activate\n" "dask-yarn services worker") % path
-    )
+    assert worker.script == ("source %s/bin/activate\ndask-yarn services worker") % path
     client = spec.services["dask.client"]
     assert set(client.files) == {"script.py"}
     assert (
         client.script
-        == ("source %s/bin/activate\n" "dask-yarn services client script.py foo bar")
+        == ("source %s/bin/activate\ndask-yarn services client script.py foo bar")
         % path
     )
 
@@ -383,18 +380,20 @@ def test_environment_archive():
     scheduler = spec.services["dask.scheduler"]
     assert set(scheduler.files) == {"environment"}
     assert scheduler.script == (
-        "source environment/bin/activate\n" "dask-yarn services scheduler"
+        "source environment/bin/activate\n"
+        "environment/bin/python -m dask_yarn.cli services scheduler"
     )
     worker = spec.services["dask.worker"]
     assert set(worker.files) == {"environment"}
     assert worker.script == (
-        "source environment/bin/activate\n" "dask-yarn services worker"
+        "source environment/bin/activate\n"
+        "environment/bin/python -m dask_yarn.cli services worker"
     )
     client = spec.services["dask.client"]
     assert set(client.files) == {"environment", "script.py"}
     assert client.script == (
         "source environment/bin/activate\n"
-        "dask-yarn services client script.py foo bar"
+        "environment/bin/python -m dask_yarn.cli services client script.py foo bar"
     )