Merge branch 'migration-sequencing-phase-1' into migration-sequencing-phase-2

# Conflicts:
#	src/databricks/labs/ucx/assessment/sequencing.py
#	tests/unit/assessment/test_sequencing.py

Author: ericvergnaud

src/databricks/labs/ucx/sequencing/sequencing.py renamed to src/databricks/labs/ucx/assessment/sequencing.py

@@ -62,7 +62,6 @@ def __init__(self, ws: WorkspaceClient, path_lookup: PathLookup, admin_locator:
         self._admin_locator = admin_locator
         self._last_node_id = 0
         self._nodes: dict[tuple[str, str], MigrationNode] = {}
-        self._incoming: dict[tuple[str, str], set[tuple[str, str]]] = defaultdict(set)
         self._outgoing: dict[tuple[str, str], set[tuple[str, str]]] = defaultdict(set)
 
     def register_workflow_task(self, task: jobs.Task, job: jobs.Job, graph: DependencyGraph) -> MigrationNode:
@@ -80,15 +79,12 @@ def register_workflow_task(self, task: jobs.Task, job: jobs.Job, graph: Dependen
             object_owner=job_node.object_owner,  # no task owner so use job one
         )
         self._nodes[task_node.key] = task_node
-        self._incoming[job_node.key].add(task_node.key)
         self._outgoing[task_node.key].add(job_node.key)
         if task.existing_cluster_id:
             cluster_node = self.register_cluster(task.existing_cluster_id)
             if cluster_node:
-                self._incoming[cluster_node.key].add(task_node.key)
                 self._outgoing[task_node.key].add(cluster_node.key)
                 # also make the cluster dependent on the job
-                self._incoming[cluster_node.key].add(job_node.key)
                 self._outgoing[job_node.key].add(cluster_node.key)
         graph.visit(self._visit_dependency, None)
         return task_node
@@ -107,7 +103,6 @@ def register_dependency(self, parent_node: MigrationNode, object_type: str, obje
         if not dependency_node:
             dependency_node = self._create_dependency_node(object_type, object_id)
         if parent_node:
-            self._incoming[parent_node.key].add(dependency_node.key)
             self._outgoing[dependency_node.key].add(parent_node.key)
         return dependency_node
 
@@ -153,7 +148,6 @@ def register_workflow_job(self, job: jobs.Job) -> MigrationNode:
             for job_cluster in job.settings.job_clusters:
                 cluster_node = self.register_job_cluster(job_cluster)
                 if cluster_node:
-                    self._incoming[cluster_node.key].add(job_node.key)
                     self._outgoing[job_node.key].add(cluster_node.key)
         return job_node
 
@@ -182,7 +176,7 @@ def register_cluster(self, cluster_id: str) -> MigrationNode:
 
     def generate_steps(self) -> Iterable[MigrationStep]:
         """The below algo is adapted from Kahn's topological sort.
-        The main differences are as follows:
+        The differences are as follows:
         1) we want the same step number for all nodes with same dependency depth
             so instead of pushing 'leaf' nodes to a queue, we fetch them again once all current 'leaf' nodes are processed
             (these are transient 'leaf' nodes i.e. they only become 'leaf' during processing)
@@ -191,14 +185,17 @@ def generate_steps(self) -> Iterable[MigrationStep]:
             to avoid an infinite loop. We also avoid side effects (such as negative counts).
             This algo works correctly for simple cases, but is not tested on large trees.
         """
-        incoming_counts = self._populate_incoming_counts()
+        incoming_keys = self._collect_incoming_keys()
+        incoming_counts = self._compute_incoming_counts(incoming_keys)
         step_number = 1
         sorted_steps: list[MigrationStep] = []
         while len(incoming_counts) > 0:
             leaf_keys = self._get_leaf_keys(incoming_counts)
             for leaf_key in leaf_keys:
                 del incoming_counts[leaf_key]
-                sorted_steps.append(self._nodes[leaf_key].as_step(step_number, list(self._required_step_ids(leaf_key))))
+                sorted_steps.append(
+                    self._nodes[leaf_key].as_step(step_number, list(self._required_step_ids(incoming_keys[leaf_key])))
+                )
                 self._on_leaf_key_processed(leaf_key, incoming_counts)
             step_number += 1
         return sorted_steps
@@ -212,14 +209,23 @@ def _on_leaf_key_processed(self, leaf_key: tuple[str, str], incoming_counts: dic
             if incoming_counts[dependency_key] > 0:
                 incoming_counts[dependency_key] -= 1
 
-    def _required_step_ids(self, node_key: tuple[str, str]) -> Iterable[int]:
-        for leaf_key in self._incoming[node_key]:
-            yield self._nodes[leaf_key].node_id
+    def _collect_incoming_keys(self) -> dict[tuple[str, str], set[tuple[str, str]]]:
+        result: dict[tuple[str, str], set[tuple[str, str]]] = defaultdict(set)
+        for source, outgoing in self._outgoing.items():
+            for target in outgoing:
+                result[target].add(source)
+        return result
+
+    def _required_step_ids(self, required_step_keys: set[tuple[str, str]]) -> Iterable[int]:
+        for source_key in required_step_keys:
+            yield self._nodes[source_key].node_id
 
-    def _populate_incoming_counts(self) -> dict[tuple[str, str], int]:
+    def _compute_incoming_counts(
+        self, incoming: dict[tuple[str, str], set[tuple[str, str]]]
+    ) -> dict[tuple[str, str], int]:
         result = defaultdict(int)
         for node_key in self._nodes:
-            result[node_key] = len(self._incoming[node_key])
+            result[node_key] = len(incoming[node_key])
         return result
 
     @classmethod

src/databricks/labs/ucx/sequencing/__init__.py

@@ -1,3 +1,4 @@
+import dataclasses
 from unittest.mock import create_autospec
 
 from pathlib import Path
@@ -7,9 +8,9 @@
 from databricks.sdk.service.compute import ClusterDetails
 from databricks.sdk.service.jobs import NotebookTask
 
+from databricks.labs.ucx.assessment.sequencing import MigrationSequencer, MigrationStep
 from databricks.labs.ucx.framework.owners import AdministratorLocator, AdministratorFinder
 from databricks.labs.ucx.mixins.cached_workspace_path import WorkspaceCache
-from databricks.labs.ucx.sequencing.sequencing import MigrationSequencer
 from databricks.labs.ucx.source_code.base import CurrentSessionState
 from databricks.labs.ucx.source_code.graph import DependencyGraph, Dependency
 from databricks.labs.ucx.source_code.jobs import WorkflowTask
@@ -35,13 +36,17 @@ def test_sequencer_builds_cluster_and_children_from_task(ws, simple_dependency_r
     sequencer.register_workflow_task(task, job, graph)
     steps = list(sequencer.generate_steps())
     step = steps[-1]
-    assert step.step_id
-    assert step.object_type == "CLUSTER"
-    assert step.object_id == "cluster-123"
-    assert step.object_name == "my-cluster"
-    assert step.object_owner == "John Doe"
-    assert step.step_number == 3
-    assert len(step.required_step_ids) == 2
+    # we don't know the ids of the steps, se let's zero them
+    step = dataclasses.replace(step, step_id=0, required_step_ids=[0] * len(step.required_step_ids))
+    assert step == MigrationStep(
+        step_id=0,
+        step_number=3,
+        object_type="CLUSTER",
+        object_id="cluster-123",
+        object_name="my-cluster",
+        object_owner="John Doe",
+        required_step_ids=[0, 0],
+    )
 
 
 def test_sequencer_builds_steps_from_dependency_graph(ws, simple_dependency_resolver, mock_path_lookup):