Internal

PiperOrigin-RevId: 561052201 Change-Id: I60fd51bb8eb50f795a2d79ebdcb7c015f3ba2a0e GitOrigin-RevId: 820eccc7d8d0994b3420b2fa5a1a97776a81e874
google-deepmind · Sep 4, 2023 · e79bbed · e79bbed
1 parent 983b68a
commit e79bbed
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/xmanager/xm/core.py b/xmanager/xm/core.py
@@ -353,6 +353,13 @@ def add(
       An awaitable that would be fulfilled when the job is launched.
     """
     # pyformat: enable
+
+    # Prioritize the identity given directly to the work unit at work unit
+    # creation time, as opposed to the identity passed when adding jobs to it as
+    # this is more consistent between job generator work units and regular work
+    # units.
+    identity = self.identity or identity
+
     job = job_operators.shallow_copy_job_type(job)
     if args is not None:
       _apply_args(job, args)