Fix StreamFlow MPI example

Prior to this commit, the MPI example was not working for these reasons:
 - When scheduling a job on multiple locations the STREAMFLOW_HOSTS variable
   was not correctly populated;
 - Helm connector had some issues with file transfers and tar streaming;
 - DockerCompose connector had issues with parsing available locations.
All these issues are solved by this fix, which also fixes #82.

Author: GlassOfWhiskey

examples/mpi/streamflow.yml

@@ -28,5 +28,6 @@ deployments:
     type: helm
     config:
       chart: environment/helm/openmpi
-      kubeconfig: /home/glassofwhiskey/.kube/config-streamflow
+      kubeconfig: ~/.kube/config-streamflow
       releaseName: openmpi-rel
+    workdir: /tmp

streamflow/core/scheduling.py

@@ -129,13 +129,13 @@ def __init__(
         self,
         job: str,
         target: Target,
-        locations: MutableSequence[Location],
+        locations: MutableSequence[AvailableLocation],
         status: Status,
         hardware: Hardware,
     ):
         self.job: str = job
         self.target: Target = target
-        self.locations: MutableSequence[Location] = locations
+        self.locations: MutableSequence[AvailableLocation] = locations
         self.status: Status = status
         self.hardware: Hardware = hardware
 
@@ -232,7 +232,7 @@ def get_connector(self, job_name: str) -> Connector | None:
 
     def get_locations(
         self, job_name: str, statuses: MutableSequence[Status] | None = None
-    ) -> MutableSequence[Location]:
+    ) -> MutableSequence[AvailableLocation]:
         allocation = self.get_allocation(job_name)
         return (
             allocation.locations

streamflow/cwl/command.py

@@ -778,14 +778,9 @@ async def execute(self, job: Job) -> CWLCommandOutput:
             ]
         # If step is assigned to multiple locations, add the STREAMFLOW_HOSTS environment variable
         if len(locations) > 1:
-            service = self.step.workflow.context.scheduler.get_service(job.name)
-            available_locations = await connector.get_available_locations(
-                service=service
+            parsed_env["STREAMFLOW_HOSTS"] = ",".join(
+                [loc.hostname for loc in locations]
             )
-            hosts = {
-                k: v.hostname for k, v in available_locations.items() if k in locations
-            }
-            parsed_env["STREAMFLOW_HOSTS"] = ",".join(hosts.values())
         # Process streams
         stdin = utils.eval_expression(
             expression=self.stdin,

streamflow/cwl/main.py

@@ -20,7 +20,7 @@
 def _parse_arg(path: str, context: StreamFlowContext):
     if "://" in path:
         return path
-    elif os.path.isabs(path):
+    elif not os.path.isabs(path):
         return os.path.join(os.path.dirname(context.config["path"]), path)
     else:
         return path

streamflow/cwl/utils.py

@@ -189,7 +189,7 @@ async def _register_path(
                 ):
                     data_location = data_locations[0]
                 else:
-                    raise WorkflowExecutionException(f"Error registering path {path}")
+                    return None
             link_location = context.data_manager.register_path(
                 location=location,
                 path=path,

streamflow/deployment/connector/container.py

@@ -760,8 +760,16 @@ async def deploy(self, external: bool) -> None:
             )
             if logger.isEnabledFor(logging.DEBUG):
                 logger.debug(f"EXECUTING command {deploy_command}")
-            proc = await asyncio.create_subprocess_exec(*shlex.split(deploy_command))
-            await proc.wait()
+            proc = await asyncio.create_subprocess_exec(
+                *shlex.split(deploy_command),
+                stderr=asyncio.subprocess.STDOUT,
+                stdout=asyncio.subprocess.PIPE,
+            )
+            stdout, _ = await proc.communicate()
+            if proc.returncode != 0:
+                raise WorkflowExecutionException(
+                    f"FAILED Deployment of {self.deployment_name} environment:\n\t{stdout.decode().strip()}"
+                )
 
     @cachedmethod(lambda self: self.locationsCache)
     async def get_available_locations(
@@ -771,24 +779,30 @@ async def get_available_locations(
         output_directory: str | None = None,
         tmp_directory: str | None = None,
     ) -> MutableMapping[str, AvailableLocation]:
-        ps_command = self.base_command() + "".join(["ps ", service or ""])
+        ps_command = self.base_command() + "".join(
+            ["ps ", "--format ", "json ", service or ""]
+        )
         if logger.isEnabledFor(logging.DEBUG):
             logger.debug(f"EXECUTING command {ps_command}")
         proc = await asyncio.create_subprocess_exec(
             *shlex.split(ps_command),
             stdout=asyncio.subprocess.PIPE,
-            stderr=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
         )
         stdout, _ = await proc.communicate()
-        lines = (line for line in stdout.decode().strip().split("\n"))
-        locations = {}
-        for line in lines:
-            if line.startswith("---------"):
-                break
-        for line in lines:
-            location_name = line.split()[0].strip()
-            locations[location_name] = await self._get_location(location_name)
-        return locations
+        locations = json.loads(stdout.decode().strip())
+        return {
+            loc["Name"]: v
+            for loc, v in zip(
+                locations,
+                await asyncio.gather(
+                    *(
+                        asyncio.create_task(self._get_location(location["Name"]))
+                        for location in locations
+                    )
+                ),
+            )
+        }
 
     @classmethod
     def get_schema(cls) -> str:
@@ -804,8 +818,16 @@ async def undeploy(self, external: bool) -> None:
             )
             if logger.isEnabledFor(logging.DEBUG):
                 logger.debug(f"EXECUTING command {undeploy_command}")
-            proc = await asyncio.create_subprocess_exec(*shlex.split(undeploy_command))
-            await proc.wait()
+            proc = await asyncio.create_subprocess_exec(
+                *shlex.split(undeploy_command),
+                stderr=asyncio.subprocess.STDOUT,
+                stdout=asyncio.subprocess.PIPE,
+            )
+            stdout, _ = await proc.communicate()
+            if proc.returncode != 0:
+                raise WorkflowExecutionException(
+                    f"FAILED Undeployment of {self.deployment_name} environment:\n\t{stdout.decode().strip()}"
+                )
 
 
 class SingularityBaseConnector(ContainerConnector, ABC):
@@ -1073,7 +1095,7 @@ async def undeploy(self, external: bool) -> None:
                 proc = await asyncio.create_subprocess_exec(
                     *shlex.split(undeploy_command),
                     stdout=asyncio.subprocess.PIPE,
-                    stderr=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.STDOUT,
                 )
-                await proc.wait()
+                stdout, _ = await proc.communicate()
             self.instanceNames = []

streamflow/deployment/connector/kubernetes.py

@@ -64,14 +64,36 @@ async def _get_helm_version():
 
 
 class KubernetesResponseWrapper(BaseStreamWrapper):
+    def __init__(self, stream):
+        super().__init__(stream)
+        self.msg: bytes = bytes()
+
     async def read(self, size: int | None = None):
-        while not self.stream.closed:
+        if len(self.msg) > 0:
+            if len(self.msg) > size:
+                data = self.msg[0:size]
+                self.msg = self.msg[size:]
+                return data
+            else:
+                data = self.msg
+                size -= len(self.msg)
+                self.msg = bytes()
+        else:
+            data = bytes()
+        while size > 0 and not self.stream.closed:
             async for msg in self.stream:
                 channel = msg.data[0]
-                data = msg.data[1:]
-                if data and channel == ws_client.STDOUT_CHANNEL:
-                    return data
-        return None
+                self.msg = msg.data[1:]
+                if self.msg and channel == ws_client.STDOUT_CHANNEL:
+                    if len(self.msg) > size:
+                        data += self.msg[0:size]
+                        self.msg = self.msg[size:]
+                        return data
+                    else:
+                        data += self.msg
+                        size -= len(self.msg)
+                        self.msg = bytes()
+        return data if len(data) > 0 else None
 
     async def write(self, data: Any):
         channel_prefix = bytes(chr(ws_client.STDIN_CHANNEL), "ascii")
@@ -116,7 +138,7 @@ def __init__(
         )
         self.inCluster = inCluster
         self.kubeconfig = (
-            kubeconfig
+            str(Path(kubeconfig).expanduser())
             if kubeconfig is not None
             else os.path.join(str(Path.home()), ".kube", "config")
         )
@@ -207,33 +229,18 @@ async def _copy_local_to_remote_single(
     async def _copy_remote_to_local(
         self, src: str, dst: str, location: Location, read_only: bool = False
     ):
-        pod, container = location.name.split(":")
-        command = ["tar", "chf", "-", "-C", "/", posixpath.relpath(src, "/")]
-        # noinspection PyUnresolvedReferences
-        response = await self.client_ws.connect_get_namespaced_pod_exec(
-            name=pod,
-            namespace=self.namespace or "default",
-            container=container,
-            command=command,
-            stderr=True,
-            stdin=False,
-            stdout=True,
-            tty=False,
-            _preload_content=False,
-        )
-        try:
-            async with aiotarstream.open(
-                stream=KubernetesResponseWrapper(response),
-                mode="r",
-                copybufsize=self.transferBufferSize,
-            ) as tar:
-                await extract_tar_stream(tar, src, dst, self.transferBufferSize)
-        except tarfile.TarError as e:
-            raise WorkflowExecutionException(
-                f"Error copying {src} from location {location} to {dst}: {e}"
-            ) from e
-        finally:
-            await response.close()
+        async with self._get_stream_reader(location, src) as reader:
+            try:
+                async with aiotarstream.open(
+                    stream=reader,
+                    mode="r",
+                    copybufsize=self.transferBufferSize,
+                ) as tar:
+                    await extract_tar_stream(tar, src, dst, self.transferBufferSize)
+            except tarfile.TarError as e:
+                raise WorkflowExecutionException(
+                    f"Error copying {src} from location {location} to {dst}: {e}"
+                ) from e
 
     async def _copy_remote_to_remote(
         self,
@@ -277,9 +284,9 @@ async def _copy_remote_to_remote(
                                         namespace=self.namespace or "default",
                                         container=location.name.split(":")[1],
                                         command=write_command,
-                                        stderr=True,
-                                        stdin=False,
-                                        stdout=True,
+                                        stderr=False,
+                                        stdin=True,
+                                        stdout=False,
                                         tty=False,
                                         _preload_content=False,
                                     ),
@@ -383,9 +390,9 @@ def _get_stream_reader(self, location: Location, src: str) -> StreamWrapperConte
                     namespace=self.namespace or "default",
                     container=container,
                     command=["tar", "chf", "-", "-C", dirname, basename],
-                    stderr=False,
-                    stdin=True,
-                    stdout=False,
+                    stderr=True,
+                    stdin=False,
+                    stdout=True,
                     tty=False,
                     _preload_content=False,
                 ),
@@ -565,20 +572,20 @@ def __init__(
         self.stringValues: MutableSequence[str] | None = stringValues
         self.skipCrds: bool = skipCrds
         self.registryConfig = (
-            registryConfig
+            str(Path(registryConfig).expanduser())
             if registryConfig is not None
             else os.path.join(str(Path.home()), ".config/helm/registry.json")
         )
         self.releaseName: str = (
             releaseName if releaseName is not None else f"release-{uuid.uuid1()}"
         )
         self.repositoryCache = (
-            repositoryCache
+            str(Path(repositoryCache).expanduser())
             if repositoryCache is not None
             else os.path.join(str(Path.home()), ".cache/helm/repository")
         )
         self.repositoryConfig = (
-            repositoryConfig
+            str(Path(repositoryConfig).expanduser())
             if repositoryConfig is not None
             else os.path.join(str(Path.home()), ".config/helm/repositories.yaml")
         )
@@ -589,7 +596,7 @@ def __init__(
         self.chartVersion: str | None = chartVersion
         self.wait: bool = wait
 
-    def base_command(self) -> str:
+    def _get_base_command(self) -> str:
         return "".join(
             [
                 "helm ",
@@ -616,7 +623,7 @@ async def deploy(self, external: bool) -> None:
                     f"Helm {version} is not compatible with Helm3Connector"
                 )
             # Deploy Helm charts
-            deploy_command = self.base_command() + "".join(
+            deploy_command = self._get_base_command() + "".join(
                 [
                     "install ",
                     self.get_option("atomic", self.atomic),
@@ -649,10 +656,14 @@ async def deploy(self, external: bool) -> None:
                 logger.debug(f"EXECUTING {deploy_command}")
             proc = await asyncio.create_subprocess_exec(
                 *shlex.split(deploy_command),
-                stderr=asyncio.subprocess.DEVNULL,
-                stdout=asyncio.subprocess.DEVNULL,
+                stderr=asyncio.subprocess.STDOUT,
+                stdout=asyncio.subprocess.PIPE,
             )
-            await proc.wait()
+            stdout, _ = await proc.communicate()
+            if proc.returncode != 0:
+                raise WorkflowExecutionException(
+                    f"FAILED Deployment of {self.deployment_name} environment:\n\t{stdout.decode().strip()}"
+                )
 
     @cachedmethod(lambda self: self.locationsCache)
     async def get_available_locations(
@@ -698,7 +709,7 @@ def get_schema(cls) -> str:
     async def undeploy(self, external: bool) -> None:
         if not external:
             # Undeploy
-            undeploy_command = self.base_command() + "".join(
+            undeploy_command = self._get_base_command() + "".join(
                 [
                     "uninstall ",
                     self.get_option("keep-history", self.keepHistory),
@@ -709,7 +720,15 @@ async def undeploy(self, external: bool) -> None:
             )
             if logger.isEnabledFor(logging.DEBUG):
                 logger.debug(f"EXECUTING {undeploy_command}")
-            proc = await asyncio.create_subprocess_exec(*shlex.split(undeploy_command))
-            await proc.wait()
+            proc = await asyncio.create_subprocess_exec(
+                *shlex.split(undeploy_command),
+                stderr=asyncio.subprocess.STDOUT,
+                stdout=asyncio.subprocess.PIPE,
+            )
+            stdout, _ = await proc.communicate()
+            if proc.returncode != 0:
+                raise WorkflowExecutionException(
+                    f"FAILED Undeployment of {self.deployment_name} environment:\n\t{stdout.decode().strip()}"
+                )
         # Close connections
         await super().undeploy(external)

streamflow/scheduling/scheduler.py

@@ -78,7 +78,7 @@ def _allocate_job(
             else:
                 logger.debug(
                     f"Job {job.name} allocated on locations "
-                    ", ".join([str(loc) for loc in selected_locations])
+                    f"{', '.join([str(loc) for loc in selected_locations])}"
                 )
         self.job_allocations[job.name] = JobAllocation(
             job=job.name,