Sending result_push_socket

vatlab · Aug 10, 2019 · e9d00b8 · e9d00b8
1 parent c9aaec8
commit e9d00b8
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/src/sos/task_executor.py b/src/sos/task_executor.py
@@ -117,7 +117,8 @@ def execute_single_task(self,
                             params,
                             runtime,
                             sig_content,
-                            quiet=False):
+                            quiet=False,
+                            **kwargs):
         '''
         Execute a single task, with
 

diff --git a/src/sos/workers.py b/src/sos/workers.py
@@ -411,14 +411,15 @@ def run_task(self, work):
         from .task_executor import BaseTaskExecutor
         executor = BaseTaskExecutor()
 
-        if env.result_socket_port is not None and env.result_socket_port != config[
-            "sockets"]["result_push_socket"]:
+        result_socket = work['config']['sockets']['result_push_socket']
+
+        if env.result_socket_port is not None and env.result_socket_port != result_socket:
             close_socket(env.result_socket)
             env.result_socket = None
 
         if env.result_socket is None:
             env.result_socket = create_socket(env.zmq_context, zmq.PUSH)
-            env.result_socket_port = config["sockets"]["result_push_socket"]
+            env.result_socket_port = result_socket
             # the result_socket_port contains IP of the worker that request the substep
             env.result_socket.connect(env.result_socket_port)