From d7adf40c8ae8cac3f85732ee5b7ab688a3ebbcec Mon Sep 17 00:00:00 2001
From: Luis Medel <luis@luismedel.com>
Date: Thu, 12 Dec 2024 01:36:19 +0100
Subject: [PATCH] Fix workflow matrix

---
 .bluish/bluish.yaml          |  6 +--
 .bluish/ci.yaml              |  4 --
 src/bluish/nodes/__init__.py |  6 +--
 src/bluish/nodes/workflow.py | 85 +++++++++++++++++-------------------
 4 files changed, 46 insertions(+), 55 deletions(-)

diff --git a/.bluish/bluish.yaml b/.bluish/bluish.yaml
index 663fc62..59aedde 100644
--- a/.bluish/bluish.yaml
+++ b/.bluish/bluish.yaml
@@ -1,15 +1,15 @@
 
 var:
-  project_version: "0.9.1"
+  project_version: "0.9.2"
   python_version: "3.12"
 
+runs_on: docker://python:${{ python_version }}-alpine
+
 jobs:
   publish:
 
     name: Publish the latest release of Bluish 
 
-    runs_on: docker://python:${{ python_version }}-alpine
-
     steps:
       - name: Install required packages
         uses: linux/install-packages
diff --git a/.bluish/ci.yaml b/.bluish/ci.yaml
index 8319869..b6eba8c 100644
--- a/.bluish/ci.yaml
+++ b/.bluish/ci.yaml
@@ -6,10 +6,6 @@ inputs:
   - name: test
     default: 1
 
-runs_on:
-  host: docker://python:${{ python_version }}-slim
-  automount: true
-
 jobs:
   _prepare_lint:
     name: Prepare linters
diff --git a/src/bluish/nodes/__init__.py b/src/bluish/nodes/__init__.py
index 93fb9dc..0829b26 100644
--- a/src/bluish/nodes/__init__.py
+++ b/src/bluish/nodes/__init__.py
@@ -328,9 +328,9 @@ def prepare_value(value: Any) -> Any:
             raise ValueError(f"Step {step_id} not found")
         return _try_get_value(step, varname, raw)
     elif root == "matrix":
-        job = cast(bluish.nodes.job.Job, _job(ctx))
-        if varname in job.matrix:
-            return prepare_value(job.matrix[varname])
+        matrix = getattr(ctx, "matrix", None) or getattr(ctx.parent, "matrix", None)
+        if matrix and (varname in matrix):
+            return prepare_value(matrix[varname])
     elif root == "step":
         return _try_get_value(_step(ctx), varname, raw)
     elif root == "inputs":
diff --git a/src/bluish/nodes/workflow.py b/src/bluish/nodes/workflow.py
index bf3fedb..18949d4 100644
--- a/src/bluish/nodes/workflow.py
+++ b/src/bluish/nodes/workflow.py
@@ -86,13 +86,6 @@ def is_true(v: Any) -> bool:
     def dispatch(self) -> bluish.process.ProcessResult:
         self.reset()
 
-        cleanup_host = False
-        if not self.runs_on_host:
-            self.runs_on_host = self.runs_on_host or bluish.process.prepare_host(
-                self.expand_expr(self.attrs.runs_on)
-            )
-            cleanup_host = True
-
         self.status = bluish.core.ExecutionStatus.RUNNING
 
         bluish.nodes.log_dict(
@@ -119,26 +112,10 @@ def dispatch(self) -> bluish.process.ProcessResult:
             if self.status == bluish.core.ExecutionStatus.RUNNING:
                 self.status = bluish.core.ExecutionStatus.FINISHED
 
-            if cleanup_host:
-                bluish.process.cleanup_host(self.runs_on_host)
-                self.runs_on_host = None
-
     def dispatch_job(
         self, job: bluish.nodes.job.Job, no_deps: bool
     ) -> bluish.process.ProcessResult | None:
-        cleanup_host = False
-        if not self.runs_on_host:
-            self.runs_on_host = self.runs_on_host or bluish.process.prepare_host(
-                self.expand_expr(self.attrs.runs_on)
-            )
-            cleanup_host = True
-        result = self.__dispatch_job(job, no_deps, set())
-
-        if cleanup_host:
-            bluish.process.cleanup_host(self.runs_on_host)
-            self.runs_on_host = None
-
-        return result
+        return self.__dispatch_job(job, no_deps, set())
 
     def __dispatch_job(
         self, job: bluish.nodes.job.Job, no_deps: bool, visited_jobs: set[str]
@@ -167,27 +144,45 @@ def __dispatch_job(
                     return result
 
         for wf_matrix in bluish.nodes._generate_matrices(self):
-            for job_matrix in bluish.nodes._generate_matrices(job):
-                job.reset()
-
-                if job.attrs.runs_on:
-                    job.runs_on_host = bluish.process.prepare_host(
-                        self.expand_expr(job.attrs.runs_on)
-                    )
-                else:
-                    job.runs_on_host = self.runs_on_host
-                
-                job.matrix = {**wf_matrix, **job_matrix}
+
+            self.matrix = wf_matrix
+
+            cleanup_host = False
+            if not self.runs_on_host:
+                self.runs_on_host = self.runs_on_host or bluish.process.prepare_host(
+                    self.expand_expr(self.attrs.runs_on)
+                )
+                cleanup_host = True
+
+            try:
+                for job_matrix in bluish.nodes._generate_matrices(job):
+                    job.reset()
+
+                    job.matrix = {**wf_matrix, **job_matrix}
+
+                    if job.attrs.runs_on:
+                        job.runs_on_host = bluish.process.prepare_host(
+                            self.expand_expr(job.attrs.runs_on)
+                        )
+                    else:
+                        job.runs_on_host = self.runs_on_host
+
                 
-                try:
-                    result = job.dispatch()
-                    if result and result.failed:
-                        return result
-                except:
-                    raise
-                finally:
-                    if job.runs_on_host and job.runs_on_host is not self.runs_on_host:
-                        bluish.process.cleanup_host(job.runs_on_host)
-                        job.runs_on_host = None
+                    try:
+                        result = job.dispatch()
+                        if result and result.failed:
+                            return result
+                    except:
+                        raise
+                    finally:
+                        if job.runs_on_host and job.runs_on_host is not self.runs_on_host:
+                            bluish.process.cleanup_host(job.runs_on_host)
+                            job.runs_on_host = None
+            except:
+                raise
+            finally:
+                if cleanup_host:
+                    bluish.process.cleanup_host(self.runs_on_host)
+                    self.runs_on_host = None
 
         return bluish.process.ProcessResult()