diff --git a/tools/releasetools/ota_utils.py b/tools/releasetools/ota_utils.py
index 97a6280..2e26a05 100644
--- a/tools/releasetools/ota_utils.py
+++ b/tools/releasetools/ota_utils.py
@@ -160,6 +160,39 @@
                       is_post_build):
   """Update the fields of the DeviceState proto with build info."""
 
+  def UpdatePartitionStates(partition_states):
+    """Update the per-partition state according to its build.prop"""
+
+    build_info_set = ComputeRuntimeBuildInfos(build_info,
+                                              boot_variable_values)
+    for partition in PARTITIONS_WITH_CARE_MAP:
+      partition_prop = build_info.info_dict.get(
+          '{}.build.prop'.format(partition))
+      # Skip if the partition is missing, or it doesn't have a build.prop
+      if not partition_prop or not partition_prop.build_props:
+        continue
+
+      partition_state = partition_states.add()
+      partition_state.partition_name = partition
+      # Update the partition's runtime device names and fingerprints
+      partition_devices = set()
+      partition_fingerprints = set()
+      for runtime_build_info in build_info_set:
+        partition_devices.add(
+            runtime_build_info.GetPartitionBuildProp('ro.product.device',
+                                                     partition))
+        partition_fingerprints.add(
+            runtime_build_info.GetPartitionFingerprint(partition))
+
+      partition_state.device.extend(sorted(partition_devices))
+      partition_state.build.extend(sorted(partition_fingerprints))
+
+      # TODO(xunchang) set the boot image's version with kmi. Note the boot
+      # image doesn't have a file map.
+      partition_state.version = build_info.GetPartitionBuildProp(
+          'ro.build.date.utc', partition)
+
+  # TODO(xunchang), we can save a call to ComputeRuntimeBuildInfos.
   build_devices, build_fingerprints = \
       CalculateRuntimeDevicesAndFingerprints(build_info, boot_variable_values)
   device_state.device.extend(sorted(build_devices))
@@ -167,7 +200,7 @@
   device_state.build_incremental = build_info.GetBuildProp(
       'ro.build.version.incremental')
 
-  # TODO(xunchang) update the partition state
+  UpdatePartitionStates(device_state.partition_state)
 
   if is_post_build:
     device_state.sdk_level = build_info.GetBuildProp(
@@ -302,14 +335,12 @@
           "building the incremental." % (pre_timestamp, post_timestamp))
 
 
-def CalculateRuntimeDevicesAndFingerprints(build_info, boot_variable_values):
-  """Returns a tuple of sets for runtime devices and fingerprints"""
+def ComputeRuntimeBuildInfos(default_build_info, boot_variable_values):
+  """Returns a set of build info objects that may exist during runtime."""
 
-  device_names = {build_info.device}
-  fingerprints = {build_info.fingerprint}
-
+  build_info_set = {default_build_info}
   if not boot_variable_values:
-    return device_names, fingerprints
+    return build_info_set
 
   # Calculate all possible combinations of the values for the boot variables.
   keys = boot_variable_values.keys()
@@ -319,7 +350,7 @@
   for placeholder_values in combinations:
     # Reload the info_dict as some build properties may change their values
     # based on the value of ro.boot* properties.
-    info_dict = copy.deepcopy(build_info.info_dict)
+    info_dict = copy.deepcopy(default_build_info.info_dict)
     for partition in PARTITIONS_WITH_CARE_MAP:
       partition_prop_key = "{}.build.prop".format(partition)
       input_file = info_dict[partition_prop_key].input_file
@@ -333,10 +364,22 @@
             PartitionBuildProps.FromInputFile(input_file, partition,
                                               placeholder_values)
     info_dict["build.prop"] = info_dict["system.build.prop"]
+    build_info_set.add(BuildInfo(info_dict, default_build_info.oem_dicts))
 
-    new_build_info = BuildInfo(info_dict, build_info.oem_dicts)
-    device_names.add(new_build_info.device)
-    fingerprints.add(new_build_info.fingerprint)
+  return build_info_set
+
+
+def CalculateRuntimeDevicesAndFingerprints(default_build_info,
+                                           boot_variable_values):
+  """Returns a tuple of sets for runtime devices and fingerprints"""
+
+  device_names = set()
+  fingerprints = set()
+  build_info_set = ComputeRuntimeBuildInfos(default_build_info,
+                                            boot_variable_values)
+  for runtime_build_info in build_info_set:
+    device_names.add(runtime_build_info.device)
+    fingerprints.add(runtime_build_info.fingerprint)
   return device_names, fingerprints