pulse-engine 0.2.7__tar.gz → 0.2.9__tar.gz

{pulse_engine-0.2.7 → pulse_engine-0.2.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pulse-engine
-Version: 0.2.7
+Version: 0.2.9
 Summary: Pulse Engine — Hybrid framework for building Pulse products
 Author: Pulse Team
 Requires-Python: >=3.11,<3.13

{pulse_engine-0.2.7 → pulse_engine-0.2.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pulse-engine"
-version = "0.2.7"
+version = "0.2.9"
 description = "Pulse Engine — Hybrid framework for building Pulse products"
 authors = ["Pulse Team"]
 readme = "README.md"

{pulse_engine-0.2.7 → pulse_engine-0.2.9}/src/pulse_engine/runners/prefect_pipeline_flow.py

@@ -143,6 +143,40 @@ def _is_s3_ref(value: Any) -> bool:
     return isinstance(value, dict) and value.get("_s3_ref") is True
 
 
+def _fetch_s3_ref(ref: dict[str, Any]) -> Any:
+    """Download and deserialise a single S3 ref. Used for when-expression evaluation."""
+    import boto3
+
+    region = boto3.session.Session().region_name or "us-east-1"
+    s3 = boto3.client("s3", region_name=region)
+    obj = s3.get_object(Bucket=ref["bucket"], Key=ref["key"])
+    return json.loads(obj["Body"].read().decode())
+
+
+def _resolve_step_results_for_when(
+    when_expr: str, step_results: dict[str, Any]
+) -> dict[str, Any]:
+    """Copy of step_results with S3 refs expanded for steps in the when expression."""
+    import re
+
+    refs = re.findall(r"\$steps\.([\w-]+)\.", when_expr)
+    if not refs:
+        return step_results
+    resolved = dict(step_results)
+    for step in set(refs):
+        val = resolved.get(step)
+        if isinstance(val, dict) and _is_s3_ref(val):
+            try:
+                resolved[step] = _fetch_s3_ref(val)
+            except Exception:
+                logger.warning(
+                    "Could not fetch S3 ref for step '%s' during when evaluation",
+                    step,
+                    exc_info=True,
+                )
+    return resolved
+
+
 def _merge_s3_refs(
     refs: list[dict[str, Any]],
     s3_bucket: str,
@@ -546,7 +580,10 @@ async def _dispatch_vm(
 
     # Embed all values as Python literals to avoid shell quoting issues entirely.
     remote_script = (
-        f"import subprocess, sys, json\n"
+        f"import os, subprocess, sys, json\n"
+        # Inject resolved engine credentials into the VM process environment so that
+        # subprocess calls (aws ecr get-login-password, docker pull) inherit them.
+        f"os.environ.update({repr(_aws_env_inject)})\n"
         f"env = {repr(env)}\n"
         f"image = {repr(image)}\n"
         f"ecr_registry = {repr(ecr_registry)}\n"
@@ -1147,7 +1184,8 @@ def _execute_layer(
 
         # Gate 2: step-level when (only when expression doesn't reference $item.*)
         if when_expr and "$item." not in when_expr:
-            if not evaluate_when(when_expr, step_results, step_statuses):
+            _when_results = _resolve_step_results_for_when(when_expr, step_results)
+            if not evaluate_when(when_expr, _when_results, step_statuses):
                 logger.info("Step '%s' skipped — when condition false", step_name)
                 step_statuses[step_name] = "skipped"
                 step_results[step_name] = {}
@@ -1155,8 +1193,60 @@ def _execute_layer(
 
         if step_config.get("for_each"):
             expr = parse_expression(step_config["for_each"])
+            # Fetch S3 refs before fan-out so field access (e.g. .batches) works
+            resolved_for_each = step_results
+            if (
+                expr.source == "steps"
+                and expr.step
+                and _is_s3_ref(step_results.get(expr.step))
+            ):
+                try:
+                    resolved_for_each = {
+                        **step_results,
+                        expr.step: _fetch_s3_ref(step_results[expr.step]),
+                    }
+                except Exception:
+                    logger.warning(
+                        "Could not fetch S3 ref for for_each step '%s'",
+                        expr.step,
+                        exc_info=True,
+                    )
+            elif expr.source == "collect" and collect_from:
+                # $collect.field — resolve S3 refs for all collect_from steps
+                overrides: dict[str, Any] = {}
+                for _cname in collect_from:
+                    _val = step_results.get(_cname)
+                    if _is_s3_ref(_val) and _val is not None:
+                        try:
+                            overrides[_cname] = _fetch_s3_ref(_val)
+                        except Exception:
+                            logger.warning(
+                                "Could not fetch S3 ref for collect step '%s'",
+                                _cname,
+                                exc_info=True,
+                            )
+                    elif isinstance(_val, list):
+                        # fan-out: list of S3 refs — fetch and flatten
+                        fetched: list[Any] = []
+                        for _ref in _val:
+                            if _is_s3_ref(_ref):
+                                try:
+                                    fetched.append(_fetch_s3_ref(_ref))
+                                except Exception:
+                                    logger.warning(
+                                        "Could not fetch S3 ref in collect list"
+                                        " for step '%s'",
+                                        _cname,
+                                        exc_info=True,
+                                    )
+                            else:
+                                fetched.append(_ref)
+                        if fetched:
+                            overrides[_cname] = fetched
+                if overrides:
+                    resolved_for_each = {**step_results, **overrides}
             items = resolve_for_each_items(
-                expr, module_args, step_results, collect_from or None
+                expr, module_args, resolved_for_each, collect_from or None
             )
 
             # Item-level when filter (fan-out only)