DAG API Enhancements: Introducing Downstream Task Parsing and Explicit Flow Definition

sky/dag.py

@@ -74,13 +74,12 @@ def add(self, task: 'task.Task') -> None:
             ValueError: If the task already exists in the DAG or if its name
             is already used.
         """
-        if task.name is None:
-            task.name = common_utils.get_unique_task_name(task)
+        task.name = common_utils.get_unique_task_name(task)
         if task.name in self._task_name_lookup:
             with ux_utils.print_exception_no_traceback():
                 raise ValueError(
-                    f'Task {task.name!r} already exists in the DAG.'
-                    f' Or the task name is already used by another task.')
+                    f'Task {task.name!r} already exists in the DAG, '
+                    f'or the task name is already used by another task.')
         self.graph.add_node(task)
         self._task_name_lookup[task.name] = task
 
@@ -91,22 +90,24 @@ def remove(self, task: Union['task.Task', str]) -> None:
             task: The Task object or name of the task to remove.
 
         Raises:
-            ValueError: If the task is still being depended on by other tasks.
+            ValueError: If the task is still being used as a downstream task by
+            other tasks.
         """
         task = self._get_task(task)
 
         # TODO(andy): Stuck by optimizer's wrong way to remove dummy sources
         # and sink nodes.
-        # if dependents:
-        #     dependent_names = ', '.join([dep.name for dep in dependents])
-        #     with ux_utils.print_exception_no_traceback():
-        #         raise ValueError(f'Task {task.name} is still being depended '
-        #                          f'by tasks {dependent_names!r}. Try to '
-        #                          'remove the dependencies first.')
+        downstreams = self.get_downstream(task)
+        if downstreams:
+            downstream_names = ', '.join(
+                cast(str, downstream.name) for downstream in downstreams)
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(f'Task {task.name} is still being used as a '
+                                 f'downstream task by {downstream_names!r}. '
+                                 'Try to remove the downstream tasks first.')
         # Here's a workaround, proactively remove all downstream edges.
-        dependents = self.get_downstream(task)
-        for dependent in dependents:
-            self.remove_edge(task, dependent)
+        for downstream in downstreams:
+            self.remove_edge(task, downstream)
 
         self.edges.pop(task, None)
         self.graph.remove_node(task)
@@ -117,19 +118,20 @@ def add_edge(self, source: TaskOrName, target: TaskOrName) -> None:
         """Add an edge from source task to target task.
 
         Args:
-            source: The task that the target task depends on.
-            target: The task that depends on the source task.
+            source: The upstream task.
+            target: The downstream task to be added.
 
         Raises:
-            ValueError: If a task tries to depend on itself or if the
-            target task is not in the DAG.
+            ValueError: If a task is set as its own downstream task or if the
+            downstream task is not in the DAG.
         """
         source = self._get_task(source)
         target = self._get_task(target)
 
         if source.name == target.name:
             with ux_utils.print_exception_no_traceback():
-                raise ValueError(f'Task {source.name} cannot depend on itself.')
+                raise ValueError(f'Task {source.name} should not be its own '
+                                 'downstream task.')
         assert target.name is not None
         if target.name not in self._task_name_lookup:
             with ux_utils.print_exception_no_traceback():
@@ -145,8 +147,8 @@ def add_downstream(self, source: TaskOrName, target: TaskOrName) -> None:
         """Add downstream tasks for a source task.
 
         Args:
-            source: The task that the downstream tasks depend on.
-            target: The task(s) to add as downstream tasks.
+            source: The upstream task.
+            target: The downstream task to be added.
         """
         return self.add_edge(source, target)
 
@@ -157,8 +159,8 @@ def set_downstream(self, source: TaskOrName,
         This replaces any existing downstream tasks for the given source.
 
         Args:
-            source: The task that the downstream tasks depend on.
-            targets: The task(s) that depend on the source task.
+            source: The upstream task.
+            targets: The downstream task(s) to be added.
         """
         source = self._get_task(source)
         if not isinstance(targets, list):
@@ -171,8 +173,8 @@ def remove_edge(self, source: TaskOrName, target: TaskOrName) -> None:
         """Remove an edge between two tasks.
 
         Args:
-            source: The source task to remove the edge from.
-            target: The target task to remove the edge to.
+            source: The upstream task.
+            target: The downstream task to remove the edge to.
         """
         source = self._get_task(source)
         target = self._get_task(target)
@@ -200,7 +202,7 @@ def get_downstream(self, task: TaskOrName) -> Set['task.Task']:
             task: The task to get downstream tasks for.
 
         Returns:
-            A set of tasks that depend on the given task.
+            A set of downstream tasks.
         """
         task = self._get_task(task)
         return self.edges.get(task, set())

sky/serve/serve_utils.py

@@ -37,7 +37,6 @@
 if typing.TYPE_CHECKING:
     import fastapi
 
-    from sky import task as task_lib
     from sky.serve import replica_managers
 
 SKY_SERVE_CONTROLLER_NAME: str = (

sky/task.py

@@ -1161,7 +1161,7 @@ def __rshift__(self, b):
         if current_dag is None:
             with ux_utils.print_exception_no_traceback():
                 raise ValueError('No current DAG context found. '
-                                 'Use `with sky.dag() as dag: ...` '
+                                 'Use `with sky.Dag() as dag: ...` '
                                  'to define a DAG.')
         current_dag.add_edge(self, b)
 

sky/utils/common_utils.py

@@ -232,11 +232,10 @@ def get_global_job_id(job_timestamp: str,
     return global_job_id
 
 
-def get_unique_task_name(_: 'task_lib.Task') -> str:
-    timestamp = int(time.time())
+def get_unique_task_name(task: 'task_lib.Task') -> str:
+    name = task.name or f'task_{int(time.time())}'
     unique_suffix = uuid.uuid4().hex[:6]
-    name = f'task_{timestamp}_{unique_suffix}'
-    return name
+    return f'{name}_{unique_suffix}'
 
 
 class Backoff:

sky/utils/dag_utils.py

@@ -71,9 +71,10 @@ def load_dag_from_yaml(
     """Loads a DAG from a YAML file.
 
     Supports various formats:
-    1. Single task without separators
-    2. Multiple tasks with separators (linear dependency by default)
-    3. DAG with explicit 'dependencies' field
+    1. Tasks without explicit dependencies:
+       - Single task
+       - Multiple tasks separated, with implicit linear dependency
+    2. DAG with explicit 'dependencies' field
 
     Has special handling for an initial section in YAML that contains only the
     'name' field, which is the DAG name.
@@ -102,29 +103,24 @@ def load_dag_from_yaml(
         dag.name = header.get('name')
         downstream = header.get('downstream', {})
     else:
-        multi_tasks = len(configs) > 1
-        if multi_tasks:
-            with ux_utils.print_exception_no_traceback():
-                raise ValueError('Multiple task definitions without a valid'
-                                 'header.')
-        # Single task without header
         dag.name = configs[0].get('name')
         downstream = {}
 
     # Handle YAML with only 'name'
     if not configs:
         configs = [{'name': dag.name}]
 
-    tasks = []
-
     # Create tasks
-    for config in configs:
-        if not isinstance(config, dict):
-            with ux_utils.print_exception_no_traceback():
-                raise ValueError(f'Invalid task configuration: {config}')
+    tasks = [
+        task_lib.Task.from_yaml_config(config, env_overrides)
+        for config in configs
+    ]
+
+    if not tasks:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError('No tasks defined in the YAML file')
 
-        task = task_lib.Task.from_yaml_config(config, env_overrides)
-        tasks.append(task)
+    for task in tasks:
         dag.add(task)
 
     # Handle dependencies
@@ -134,15 +130,11 @@ def load_dag_from_yaml(
             downs = [downs] if not isinstance(downs, list) else downs
             for down in downs:
                 dag.add_edge(upstream, down)
-    elif len(tasks) > 1:
-        # Implicit linear dependency
+    else:
+        # Implicit dependency
         for i in range(len(tasks) - 1):
             dag.add_edge(tasks[i], tasks[i + 1])
 
-    if not tasks:
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError('No tasks defined in the YAML file')
-
     return dag