[Optimizer] Fix NoneType for egress when pipeline is specified (#2579)

* Fix NoneType for egress cost * Fix type check
skypilot-org · Sep 19, 2023 · 617603f · 617603f
1 parent bc4acc8
commit 617603f
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 8 deletions.
diff --git a/sky/clouds/azure.py b/sky/clouds/azure.py
@@ -96,7 +96,7 @@ def accelerators_to_hourly_cost(self,
         # GPUs, while the task specifies to use 1 GPU.
         return 0
 
-    def get_egress_cost(self, num_gigabytes):
+    def get_egress_cost(self, num_gigabytes: float):
         # In general, query this from the cloud:
         #   https://azure.microsoft.com/en-us/pricing/details/bandwidth/
         # NOTE: egress from US East.

diff --git a/sky/clouds/cloud.py b/sky/clouds/cloud.py
@@ -191,7 +191,7 @@ def accelerators_to_hourly_cost(self, accelerators: Dict[str, int],
         """Returns the hourly on-demand price for accelerators."""
         raise NotImplementedError
 
-    def get_egress_cost(self, num_gigabytes):
+    def get_egress_cost(self, num_gigabytes: float):
         """Returns the egress cost.
 
         TODO: takes into account "per month" accumulation per account.

diff --git a/sky/clouds/gcp.py b/sky/clouds/gcp.py
@@ -330,7 +330,7 @@ def accelerators_to_hourly_cost(self,
                                                            zone=zone,
                                                            clouds='gcp')
 
-    def get_egress_cost(self, num_gigabytes):
+    def get_egress_cost(self, num_gigabytes: float):
         # In general, query this from the cloud:
         #   https://cloud.google.com/storage/pricing#network-pricing
         # NOTE: egress to worldwide (excl. China, Australia).

diff --git a/sky/clouds/ibm.py b/sky/clouds/ibm.py
@@ -136,10 +136,10 @@ def accelerators_to_hourly_cost(self,
         # Currently Isn't implemented in the same manner by aws and azure.
         return 0
 
-    def get_egress_cost(self, num_gigabytes):
+    def get_egress_cost(self, num_gigabytes: float):
         """Returns the egress cost. Currently true for us-south, i.e. Dallas.
         based on https://cloud.ibm.com/objectstorage/create#pricing. """
-        cost = 0
+        cost = 0.
         price_thresholds = [{
             'threshold': 150,
             'price_per_gb': 0.05

diff --git a/sky/optimizer.py b/sky/optimizer.py
@@ -209,7 +209,9 @@ def _egress_cost_or_time(minimize_cost: bool, parent: task_lib.Task,
         if not nbytes:
             # nbytes can be None, if the task has no inputs/outputs.
             return 0
-        assert src_cloud is not None and dst_cloud is not None
+        assert src_cloud is not None and dst_cloud is not None, (src_cloud,
+                                                                 dst_cloud,
+                                                                 nbytes)
 
         if minimize_cost:
             fn = Optimizer._egress_cost
@@ -636,8 +638,12 @@ def _print_egress_plan(graph, plan, minimize_cost):
         for parent, child in graph.edges():
             src_cloud, dst_cloud, nbytes = Optimizer._get_egress_info(
                 parent, plan[parent], child, plan[child])
-            if nbytes == 0:
-                continue
+            if not nbytes:
+                # nbytes can be None, if the task has no inputs/outputs.
+                return 0
+            assert src_cloud is not None and dst_cloud is not None, (src_cloud,
+                                                                     dst_cloud,
+                                                                     nbytes)
 
             if minimize_cost:
                 fn = Optimizer._egress_cost