prophecy-libs 2.0.0.dev5__tar.gz → 2.0.0.dev7__tar.gz

{prophecy-libs-2.0.0.dev5 → prophecy-libs-2.0.0.dev7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 1.0
 Name: prophecy-libs
-Version: 2.0.0.dev5
+Version: 2.0.0.dev7
 Summary: Helper library for prophecy generated code
 Home-page: https://github.com/SimpleDataLabsInc/prophecy-python-libs
 Author: UNKNOWN

{prophecy-libs-2.0.0.dev5 → prophecy-libs-2.0.0.dev7}/prophecy/executionmetrics/componentruns/component_runs_dao.py

@@ -693,8 +693,8 @@ class ComponentRunsDAO(ExecutionMetricsDAO):
     def _row_to_component_runs_with_status(self, row: Row) -> ComponentRunsWithStatus:
         """Convert Row to ComponentRunsWithStatus."""
         base = self._row_to_component_runs(row)
-        return ComponentRunsWithStatus(
-            **base.__dict__, status=row[STATUS] if STATUS in row else None
+        return ComponentRunsWithStatus.from_component_runs(
+            base, row[STATUS] if STATUS in row else None
         )
 
     def _row_to_interim_response(self, row: Row) -> InterimResponse:
@@ -817,18 +817,20 @@ class ComponentRunsDAO(ExecutionMetricsDAO):
         """Convert Row exception to StoredSerializableException."""
         return StoredSerializableException(
             exception_type=exception_row[EXCEPTION_TYPE],
-            msg=exception_row[MSG],
-            cause_msg=exception_row[CAUSE_MSG],
+            message=exception_row[MSG],
+            cause_message=exception_row[CAUSE_MSG],
             stack_trace=exception_row[STACK_TRACE],
             time=exception_row[TIME],
         )
 
     def _convert_exception_to_row(self, exception: StoredSerializableException) -> Row:
         """Convert StoredSerializableException to Row."""
-        return Row(
-            exception_type=exception.exception_type,
-            msg=exception.msg,
-            cause_msg=exception.cause_msg,
-            stack_trace=exception.stack_trace,
-            time=exception.time,
-        )
+        row_data = {
+            EXCEPTION_TYPE: exception.exception_type,
+            MSG: exception.message,
+            CAUSE_MSG: exception.cause_message,
+            STACK_TRACE: exception.stack_trace,
+            TIME: exception.time,
+        }
+
+        return Row(**row_data)

{prophecy-libs-2.0.0.dev5 → prophecy-libs-2.0.0.dev7}/prophecy/executionmetrics/evolutions/delta_metrics_storage_initializer.py

@@ -104,7 +104,7 @@ def base_location(spark: SparkSession, suffix: str) -> str:
     spark_delta_path_prefix = "spark.prophecy.delta.path.prefix"
     default_path = "/prophecy/metadata/executionmetrics"
     try:
-        base_path = spark.conf.get(spark_delta_path_prefix)
+        base_path = spark.conf.get(spark_delta_path_prefix, default_path)
     except Exception as e:
         base_path = default_path
     return f"{base_path}{suffix}"

{prophecy-libs-2.0.0.dev5 → prophecy-libs-2.0.0.dev7}/prophecy/executionmetrics/in_memory_store.py

@@ -45,70 +45,16 @@ from prophecy.executionmetrics.evolutions.metrics_storage_initializer import (
     StorageMetadata,
     create_storage_metadata,
 )
-from prophecy.executionmetrics.utils.common import is_databricks_environment
+from prophecy.executionmetrics.utils.common import is_databricks_environment, now_millis
 from prophecy.executionmetrics.utils.external import compress
-
-
-class HasProcessesConnectionsPorts(ABC):
-    """Interface for objects with processes, connections and ports."""
-
-    pass
-
-
-@dataclass
-class WorkflowNode:
-    """Workflow node representation."""
-
-    id: str
-    component: str
-    metadata: Dict[str, Any]
-    ports: Dict[str, Any]
-    properties: Dict[str, Any] = field(default_factory=dict)
-
-
-@dataclass
-class WorkflowEdge:
-    """Workflow edge representation."""
-
-    source: str
-    source_port: str
-    target: str
-    target_port: str
-
-
-@dataclass
-class WorkflowGroup(WorkflowNode, HasProcessesConnectionsPorts):
-    """Workflow group node."""
-
-    processes: Dict[str, "WorkflowNode"] = field(default_factory=dict)
-    connections: List["WorkflowEdge"] = field(default_factory=list)
-
-
-@dataclass
-class WorkflowGraph(HasProcessesConnectionsPorts):
-    """Workflow graph representation."""
-
-    processes: Dict[str, WorkflowNode]
-    connections: List[WorkflowEdge]
-    ports: Dict[str, Any] = field(default_factory=dict)
-
-
-def get_workflow_graph(code: Optional[Dict[str, str]]) -> Optional[WorkflowGraph]:
-    """Extract workflow graph from code dictionary."""
-    if not code:
-        return None
-
-    workflow_content = None
-    for path, content in code.items():
-        if path.endswith(".prophecy/workflow.latest.json"):
-            workflow_content = content
-            break
-
-    if workflow_content:
-        data = json.loads(workflow_content)
-        return _parse_workflow_graph(data)
-
-    return None
+from prophecy.executionmetrics.workflow_parser import (
+    HasProcessesConnectionsPorts,
+    WorkflowEdge,
+    WorkflowGraph,
+    WorkflowGroup,
+    WorkflowNode,
+    get_workflow_graph,
+)
 
 
 def get_slug_to_process_map(code: Optional[Dict[str, str]]) -> Optional[Dict[str, str]]:
@@ -134,61 +80,15 @@ def extract_slug_to_process_mapping(
         slug_to_process[current_slug] = node.id
 
         if isinstance(node, WorkflowGroup):
-            for process_id, child_node in node.processes.items():
+            for _, child_node in node.processes.items():
                 extract_from_node(child_node, current_slug)
 
-    for process_id, node in processes.items():
+    for _, node in processes.items():
         extract_from_node(node)
 
     return slug_to_process
 
 
-def _parse_workflow_nodes(data: Dict[str, Any]) -> Dict[str, "WorkflowNode"]:
-    processes = {}
-    for proc_id, proc_data in data.get("processes", {}).items():
-        if proc_data.get("component") == "Subgraph":
-            node = WorkflowGroup(
-                id=proc_id,
-                component=proc_data.get("component"),
-                metadata=proc_data.get("metadata", {}),
-                ports=proc_data.get("ports", {}),
-                properties=proc_data.get("properties", {}),
-                processes=_parse_workflow_nodes(proc_data),
-                connections=proc_data.get("connections", []),
-            )
-        else:
-            node = WorkflowNode(
-                id=proc_id,
-                component=proc_data.get("component"),
-                metadata=proc_data.get("metadata", {}),
-                ports=proc_data.get("ports", {}),
-                properties=proc_data.get("properties", {}),
-            )
-        processes[proc_id] = node
-
-    return processes
-
-
-def _parse_workflow_graph(data: Dict[str, Any]) -> WorkflowGraph:
-    """Parse workflow graph from dictionary."""
-    processes = _parse_workflow_nodes(data)
-
-    connections = []
-    for conn_data in data.get("connections", []):
-        connections.append(
-            WorkflowEdge(
-                source=conn_data["source"],
-                source_port=conn_data["sourcePort"],
-                target=conn_data["target"],
-                target_port=conn_data["targetPort"],
-            )
-        )
-
-    return WorkflowGraph(
-        processes=processes, connections=connections, ports=data.get("ports", {})
-    )
-
-
 class InMemoryStore:
     """
     In-memory store for pipeline execution metrics.
@@ -215,6 +115,7 @@ class InMemoryStore:
         self._status: str = PipelineStatus.STARTED
         self._fabric_uid: str = ""
         self._time_taken: Optional[int] = None
+        self._time_started: int = now_millis()
         self._rows_read: Optional[int] = None
         self._rows_written: Optional[int] = None
         self._run_type: str = ""
@@ -282,6 +183,7 @@ class InMemoryStore:
                 )
 
             # Set pipeline execution state
+            self._time_started = now_millis()
             self._pipeline_uri = pipeline_uri
             self._job_uri = job_uri
             self._fabric_uid = fabric_uid
@@ -411,59 +313,208 @@ class InMemoryStore:
         collect_processes(graph.processes)
         return all_procs
 
+    def _all_connections(self) -> List[WorkflowEdge]:
+        graph = self._graph
+        all_conns = []
+        if not graph:
+            return all_conns
+
+        def collect_connections(
+            connections: List[WorkflowEdge], processes: Dict[str, WorkflowNode]
+        ):
+            all_conns.extend(connections)
+            # Check for subgraphs and collect their connections
+            for node in processes.values():
+                if isinstance(node, WorkflowGroup):
+                    collect_connections(node.connections, node.processes)
+
+        collect_connections(graph.connections, graph.processes)
+        return all_conns
+
     def _extract_process_to_gem_map(self) -> Dict[str, GemProgress]:
         """Extract mapping from process ID to GemProgress."""
-        if not self._graph or not self._processes:
+        graph = self._graph
+        processes = self._processes
+        if graph is None or processes is None:
             return {}
 
-        process_to_gem = {}
-
-        for proc_id, node in self._processes.items():
-            # Extract gem information from node
-            gem_name = node.metadata.get("slug", "")
-            gem_type = node.component
-
-            # Extract ports
-            in_ports = [p["id"] for p in node.ports.get("inputs", [])]
-            out_ports = [p["id"] for p in node.ports.get("outputs", [])]
-
-            # Create input/output gem edges based on connections
-            input_gems = []
-            output_gems = []
-
-            for edge in self._graph.connections:
-                if edge.target == proc_id:
-                    source_node = self._processes.get(edge.source)
-                    if source_node:
-                        input_gems.append(
-                            GemEdge(
-                                gem_name=source_node.metadata.get("slug"),
-                                from_port=edge.source_port,
-                                to_port=edge.target_port,
+        slug_to_process_map = extract_slug_to_process_mapping(graph.processes)
+        process_to_slug_map = {v: k for k, v in slug_to_process_map.items()}
+
+        connections = self._all_connections()
+        target_port_to_connection_map = {c.target_port: c for c in connections}
+
+        # Group connections by source port
+        source_port_to_connection_multi_map: Dict[str, List[WorkflowEdge]] = {}
+        for conn in connections:
+            if conn.source_port not in source_port_to_connection_multi_map:
+                source_port_to_connection_multi_map[conn.source_port] = []
+            source_port_to_connection_multi_map[conn.source_port].append(conn)
+
+        def fetch_output_details(out_port: str) -> List[Tuple[str, str]]:
+            """
+            Fetch output details for a given output port.
+
+            Args:
+                out_port: The output port ID
+
+            Returns:
+                List of tuples (target process ID, target port ID)
+            """
+            connections_list = source_port_to_connection_multi_map.get(out_port, [])
+            target_connection_to_node = {
+                c: processes.get(c.target) for c in connections_list
+            }
+
+            result = []
+            for connection, node in target_connection_to_node.items():
+                if node and node.component == "Subgraph":
+                    # Recursively fetch output details for subgraph
+                    result.extend(fetch_output_details(connection.target_port))
+                elif node:
+                    result.append((connection.target, connection.target_port))
+
+            return result
+
+        def fetch_input_details(in_port: str) -> Tuple[Optional[str], Optional[str]]:
+            """
+            Fetch input details for a given input port.
+
+            Args:
+                in_port: The input port ID
+
+            Returns:
+                Tuple of (source process ID, source port ID)
+            """
+            connection = target_port_to_connection_map.get(in_port)
+            if not connection:
+                return (None, None)
+
+            source_node = processes.get(connection.source)
+            if source_node and source_node.component == "Subgraph":
+                # Recursively fetch input details for subgraph
+                return fetch_input_details(connection.source_port)
+            else:
+                return (
+                    connection.source if connection else None,
+                    connection.source_port if connection else None,
+                )
+
+        def extract_process_to_gem_map_inner() -> Dict[str, GemProgress]:
+            """
+            Internal function to extract process to gem mapping.
+
+            Returns:
+                List of tuples (process ID, GemProgress)
+            """
+            result = {}
+
+            for process_id, current_process in processes.items():
+                gem_name = process_to_slug_map.get(
+                    current_process.id
+                ) or current_process.metadata.get("slug", "")
+
+                out_ports = [port.slug for port in current_process.ports.outputs]
+                in_ports = [port.slug for port in current_process.ports.inputs]
+
+                # Get input process details
+                input_process_details = []
+                for current_process_in_port in current_process.ports.inputs:
+                    input_process, source_port_id = fetch_input_details(
+                        current_process_in_port.id
+                    )
+                    input_process_details.append(
+                        (input_process, source_port_id, current_process_in_port.slug)
+                    )
+
+                # Get output process details
+                output_process_details = []
+                for current_process_out_port in current_process.ports.outputs:
+                    output_details = fetch_output_details(current_process_out_port.id)
+                    for output_process, target_port_id in output_details:
+                        output_process_details.append(
+                            (
+                                output_process,
+                                current_process_out_port.slug,
+                                target_port_id,
                             )
                         )
-                elif edge.source == proc_id:
-                    target_node = self._processes.get(edge.target)
-                    if target_node:
-                        output_gems.append(
-                            GemEdge(
-                                gem_name=target_node.metadata.get("slug"),
-                                from_port=edge.source_port,
-                                to_port=edge.target_port,
-                            )
+
+                # Create input gem edges
+                input_gem_edges = []
+                for (
+                    process_id_opt,
+                    from_port_id_opt,
+                    to_port_slug,
+                ) in input_process_details:
+                    gem_name_opt = (
+                        process_to_slug_map.get(process_id_opt)
+                        if process_id_opt
+                        else None
+                    )
+                    input_process = (
+                        processes.get(process_id_opt) if process_id_opt else None
+                    )
+
+                    from_port_slug = None
+                    if from_port_id_opt and input_process:
+                        for port in input_process.ports.outputs:
+                            if port.id == from_port_id_opt:
+                                from_port_slug = port.slug
+                                break
+
+                    input_gem_edges.append(
+                        GemEdge(
+                            gem_name=gem_name_opt,
+                            from_port=from_port_slug,
+                            to_port=to_port_slug,
+                        )
+                    )
+
+                # Create output gem edges
+                output_gem_edges = []
+                for process_id, from_port_slug, to_port_id in output_process_details:
+                    gem_name_opt = process_to_slug_map.get(process_id)
+                    output_process = processes.get(process_id)
+
+                    to_port_slug = None
+                    if output_process:
+                        for port in output_process.ports.inputs:
+                            if port.id == to_port_id:
+                                to_port_slug = port.slug
+                                break
+
+                    output_gem_edges.append(
+                        GemEdge(
+                            gem_name=gem_name_opt,
+                            from_port=from_port_slug,
+                            to_port=to_port_slug,
                         )
+                    )
 
-            process_to_gem[proc_id] = GemProgress(
-                gem_name=gem_name,
-                process_id=proc_id,
-                gem_type=gem_type,
-                input_gems=input_gems,
-                output_gems=output_gems,
-                in_ports=in_ports,
-                out_ports=out_ports,
-            )
+                # Create GemProgress
+                gem_progress = GemProgress(
+                    gem_name=gem_name,
+                    process_id=current_process.id,
+                    gem_type=current_process.component,
+                    input_gems=input_gem_edges,
+                    output_gems=output_gem_edges,
+                    in_ports=in_ports,
+                    out_ports=out_ports,
+                    num_rows_output=None,
+                    stdout=None,
+                    stderr=None,
+                    start_time=None,
+                    end_time=None,
+                    state=None,
+                    exception=None,
+                )
 
-        return process_to_gem
+                result[current_process.id] = gem_progress
+
+            return result
+
+        return extract_process_to_gem_map_inner()
 
     def update_run_uid(self, job_run_uid: str, task_run_uid: str) -> None:
         """Update job and task run UIDs."""
@@ -516,8 +567,8 @@ class InMemoryStore:
     def update_gem_progress(
         self,
         process: str,
-        start_time: int,
-        end_time: Optional[int],
+        start_time: str,
+        end_time: Optional[str],
         stdout: Optional[List[TimestampedOutput]],
         stderr: Optional[List[TimestampedOutput]],
         state: str,
@@ -530,13 +581,15 @@ class InMemoryStore:
 
             if process in self._process_to_gem_map:
                 gem_progress = self._process_to_gem_map[process]
-                # Update gem progress fields
-                gem_progress.start_time = start_time
-                gem_progress.end_time = end_time
-                gem_progress.stdout = stdout
-                gem_progress.stderr = stderr
-                gem_progress.state = state
-                gem_progress.exception = exception
+                updated_gem_progress = gem_progress.update_progress(
+                    start_time=int(start_time),
+                    end_time=int(end_time) if end_time else None,
+                    stdout=stdout,
+                    stderr=stderr,
+                    state=state,
+                    exception=exception,
+                )
+                self._process_to_gem_map[process] = updated_gem_progress
 
     def _update_connecting_gem_rows(
         self, interim_process: str, interim: LInterimContent
@@ -609,7 +662,9 @@ class InMemoryStore:
                     gem.num_rows = num_records
 
     def offload(
-        self, interim_details: List[Tuple[InterimKey, DataFrame]] = None
+        self,
+        pipeline_status: str,
+        interim_details: List[Tuple[InterimKey, DataFrame]] = None,
     ) -> Tuple[List["ComponentRuns"], List[str]]:
         """
         Offload metrics to persistent storage.
@@ -718,10 +773,10 @@ class InMemoryStore:
                 job_uri=self._job_uri,
                 job_run_uid=self._job_run_uid,
                 task_run_uid=self._task_run_uid,
-                status=self._status,
+                status=pipeline_status,
                 submission_time=self._submission_time,
                 fabric_uid=self._fabric_uid,
-                time_taken=self._time_taken,
+                time_taken=now_millis() - self._time_started,
                 rows_read=self._rows_read,
                 rows_written=self._rows_written,
                 run_type=self._run_type,
@@ -783,7 +838,7 @@ class InMemoryStore:
                 component_name = node.metadata.get("slug", "")
                 component_type = node.component
 
-                interim_ports = [p["id"] for p in node.ports.get("outputs", [])]
+                interim_ports = [p.id for p in node.ports.outputs]
                 gem_progress = (
                     self._process_to_gem_map.get(process_id)
                     if self._process_to_gem_map
@@ -832,11 +887,11 @@ class InMemoryStore:
                 # Default interim details
                 interim_subgraph = subgraph_name
                 interim_node = node
-                interim_ports = [p["id"] for p in node.ports.get("outputs", [])]
+                interim_ports = [p.id for p in node.ports.outputs]
 
                 # Special handling for Target nodes
-                if component_type == "Target" and node.ports.get("inputs"):
-                    first_input_port = node.ports["inputs"][0]["id"]
+                if component_type == "Target" and node.ports.inputs:
+                    first_input_port = node.ports.inputs[0].id
                     connection, parent_graph, target_subgraph = (
                         self._fetch_target_component_source_port(
                             target_port_connection_map,

{prophecy-libs-2.0.0.dev5 → prophecy-libs-2.0.0.dev7}/prophecy/executionmetrics/package.py

@@ -316,6 +316,34 @@ class ComponentRunsWithStatus(ExecutionMetricsEntity):
             "branch": self.branch,
         }
 
+    @classmethod
+    def from_component_runs(
+        cls, component_run: ComponentRuns, status: Optional[str]
+    ) -> "ComponentRunsWithStatus":
+        return ComponentRunsWithStatus(
+            uid=component_run.uid,
+            component_uri=component_run.component_uri,
+            pipeline_run_uid=component_run.pipeline_run_uid,
+            pipeline_uri=component_run.pipeline_uri,
+            fabric_uid=component_run.fabric_uid,
+            component_name=component_run.component_name,
+            interim_component_name=component_run.interim_component_name,
+            component_type=component_run.component_type,
+            interim_out_port=component_run.interim_out_port,
+            interim_subgraph_name=component_run.interim_subgraph_name,
+            interim_process_id=component_run.interim_process_id,
+            created_by=component_run.created_by,
+            records=component_run.records,
+            bytes=component_run.bytes,
+            partitions=component_run.partitions,
+            created_at=component_run.created_at,
+            expired=component_run.expired,
+            status=status,  # This will be set separately
+            job_uri=component_run.job_uri,
+            run_type=component_run.run_type,
+            branch=component_run.branch,
+        )
+
 
 @dataclass
 class ComponentRunsWithStatusAndInterims(ExecutionMetricsEntity):