Cherry pick resource partitioner CI fix (#4058)

Author: cehongwang

tests/py/dynamo/partitioning/test_000_resource_partitioning.py

@@ -1,3 +1,4 @@
+import psutil
 import torch
 import torch.nn as nn
 from torch.fx.passes.splitter_base import Subgraph
@@ -79,7 +80,8 @@ def forward(self, x):
             partitioner = ResourcePartitioner(
                 submodule,
                 submodule_name=name,
-                cpu_memory_budget=2 * 1024 * 1024 * 1024,
+                cpu_memory_budget=0.89 * 1024 * 1024 * 1024
+                + psutil.Process().memory_info().rss,  # 0.89GB + current memory usage
             )
             subgraphs = partitioner.put_nodes_into_subgraphs()
             new_subgraphs = []

tests/py/dynamo/partitioning/test_001_resource_partitioning.py

@@ -1,5 +1,6 @@
 from typing import Any, List
 
+import psutil
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -84,7 +85,9 @@ def forward(self, x):
         )
 
         partitioned_module = resource_partition(
-            partitioned_module, cpu_memory_budget=2 * 1024 * 1024 * 1024  # 2GB,
+            partitioned_module,
+            cpu_memory_budget=0.89 * 1024 * 1024 * 1024
+            + psutil.Process().memory_info().rss,  # 0.89GB + current memory usage,
         )
 
         self.assertEqual(
@@ -166,7 +169,9 @@ def forward(self, x):
         )
 
         partitioned_module = resource_partition(
-            partitioned_module, cpu_memory_budget=1.4 * 1024 * 1024 * 1024  # 1.4GB,
+            partitioned_module,
+            cpu_memory_budget=0.39 * 1024 * 1024 * 1024
+            + psutil.Process().memory_info().rss,  # 0.39GB + current memory usage,
         )
 
         assert (
@@ -298,7 +303,9 @@ def forward(self, x):
         )
 
         partitioned_module = resource_partition(
-            partitioned_module, cpu_memory_budget=1.4 * 1024 * 1024 * 1024  # 1.4GB,
+            partitioned_module,
+            cpu_memory_budget=0.39 * 1024 * 1024 * 1024
+            + psutil.Process().memory_info().rss,  # 0.89GB + current memory usage,
         )
 
         assert (
@@ -396,7 +403,9 @@ def forward(self, x):
         )
 
         partitioned_module = resource_partition(
-            partitioned_module, cpu_memory_budget=1.4 * 1024 * 1024 * 1024  # 1.4GB,
+            partitioned_module,
+            cpu_memory_budget=0.39 * 1024 * 1024 * 1024
+            + psutil.Process().memory_info().rss,  # 0.89GB + current memory usage,
         )
 
         assert (