yarn: use system wide install of worker script

joblib · Jun 5, 2017 · 7cb7e85 · 7cb7e85
1 parent 900abbb
commit 7cb7e85
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 16 deletions.
diff --git a/joblibhadoop/yarn/remoteworker.py → bin/joblib-yarn-worker b/joblibhadoop/yarn/remoteworker.py → bin/joblib-yarn-worker
@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+
 """Remote worker script."""
 
 import sys

diff --git a/joblibhadoop/yarn/pool.py b/joblibhadoop/yarn/pool.py
@@ -18,13 +18,12 @@ def __init__(self, processes=None, port=0, authkey=None):
         self.stopping = False
         self.knit = Knit(autodetect=True)
 
-        cmd = ('python remoteworker.py --host {} --port {} --key {}'
+        cmd = ('joblib-yarn-worker --host {} --port {} --key {}'
                .format(socket.gethostname(),
                        self.server.address[1],
                        self.authkey))
-        self.app_id = self.knit.start(
-            cmd, num_containers=self._processes,
-            files=['joblibhadoop/yarn/remoteworker.py', ])
+        self.app_id = self.knit.start(cmd,
+                                      num_containers=self._processes)
         self.thread = Thread(target=self._monitor_appid)
         self.thread.deamon = True
         self.thread.start()

diff --git a/joblibhadoop/yarn/remotepool.py b/joblibhadoop/yarn/remotepool.py
@@ -112,15 +112,11 @@ def _start_remote_worker(self, pid):
         remote_worker = RemoteWorker(pid)
         debug('starting remote worker %d', pid)
 
-        args = ['python', self.workerscript]
-        args.append('--host')
-        args.append(socket.gethostname())
-        args.append('--port')
-        args.append(str(self.server.address[1]))
-        args.append('--workerid')
-        args.append(str(pid))
-        args.append('--key')
-        args.append(self.authkey.decode())
+        args = [self.workerscript,
+                '--host', socket.gethostname(),
+                '--port', str(self.server.address[1]),
+                '--workerid', str(pid),
+                '--key', self.authkey.decode()]
 
         remote_worker.proc = subprocess.Popen(args)
         self._pool.append(remote_worker)

diff --git a/setup.py b/setup.py
@@ -1,5 +1,6 @@
 """joblib-hadoop package installation module."""
 
+from os.path import join as pjoin
 from setuptools import setup, find_packages
 
 if __name__ == '__main__':
@@ -14,9 +15,10 @@
           license='BSD',
           platforms='any',
           packages=find_packages(),
+          scripts=[pjoin('bin', 'joblib-yarn-worker')],
           install_requires=[
-            'joblib>=0.10',
-            'knit>=0.2'
+              'joblib>=0.10',
+              'knit>=0.2'
           ],
           zip_safe=False,
-          )
+         )