"One job many slices"

PiperOrigin-RevId: 622858672
Change-Id: Ia0537bb1e2761f78d73584c75b707e1d2c02ba8b
GitOrigin-RevId: 311a93fdb2c80c538f8335f5d993b24868f988b9

xmanager/xm/resources.py

@@ -456,7 +456,6 @@ def __init__(
       replicas = self.topology.dimensions[1]
 
     self.replicas = replicas or 1
-    self._validate_replicas()
 
   @property
   def service_tier(self):
@@ -466,13 +465,6 @@ def service_tier(self):
   def service_tier(self, new_service_tier):
     self._service_tier = new_service_tier
 
-  def _validate_replicas(self) -> None:
-    """Raises ValueError if replication is not supported."""
-    if self.replicas > 1 and self.accelerator in TpuType:
-      raise ValueError(
-          f'Replicated jobs are not supported for {self.accelerator}.'
-      )
-
   def __repr__(self) -> str:
     """Returns string representation of the requirements."""
     args = []

xmanager/xm/resources_test.py

@@ -157,8 +157,10 @@ def test_replicas(self):
     requirements = resources.JobRequirements(replicas=2)
     self.assertEqual(requirements.replicas, 2)
 
-    with self.assertRaises(ValueError):
-      resources.JobRequirements(replicas=2, tpu_v3='4x4')
+    requirements = resources.JobRequirements(replicas=2, tpu_v3='4x4')
+    self.assertEqual(requirements.replicas, 2)
+    self.assertEqual(requirements.accelerator, resources.ResourceType.TPU_V3)
+    self.assertEqual(requirements.topology.name, '4x4')
 
     resources.JobRequirements(replicas=2, v100='4x2')
     resources.JobRequirements(v100='4x2')