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
 
 
diff --git a/tools/releasetools/test_ota_from_target_files.py b/tools/releasetools/test_ota_from_target_files.py
index 045191c..6f5e78f 100644
--- a/tools/releasetools/test_ota_from_target_files.py
+++ b/tools/releasetools/test_ota_from_target_files.py
@@ -146,7 +146,7 @@
       ),
       'vendor.build.prop': common.PartitionBuildProps.FromDictionary(
           'vendor', {
-               'ro.vendor.build.fingerprint': 'vendor-build-fingerprint'}
+              'ro.vendor.build.fingerprint': 'vendor-build-fingerprint'}
       ),
       'property1': 'value1',
       'property2': 4096,
@@ -1402,7 +1402,7 @@
                      int(metadata_dict.get('ota-required-cache', 0)))
     self.assertEqual(metadata_proto.retrofit_dynamic_partitions,
                      metadata_dict.get(
-                        'ota-retrofit-dynamic-partitions') == 'yes')
+                         'ota-retrofit-dynamic-partitions') == 'yes')
 
   def test_GetPackageMetadata_incremental_package(self):
     vendor_build_prop = copy.deepcopy(self.VENDOR_BUILD_PROP)
@@ -1463,15 +1463,15 @@
     self.assertEqual(
         'vendor-device-pro|vendor-device-std|vendor-product-device',
         metadata_dict['pre-device'])
-    suffix = ':source-version-release/source-build-id/' \
-             'source-version-incremental:build-type/build-tags'
+    source_suffix = ':source-version-release/source-build-id/' \
+                    'source-version-incremental:build-type/build-tags'
     pre_fingerprints = [
         'vendor-product-brand/vendor-product-name/vendor-device-pro'
-        '{}'.format(suffix),
+        '{}'.format(source_suffix),
         'vendor-product-brand/vendor-product-name/vendor-device-std'
-        '{}'.format(suffix),
+        '{}'.format(source_suffix),
         'vendor-product-brand/vendor-product-name/vendor-product-device'
-        '{}'.format(suffix),
+        '{}'.format(source_suffix),
     ]
     self.assertEqual('|'.join(pre_fingerprints), metadata_dict['pre-build'])
 
@@ -1486,3 +1486,28 @@
     self.assertEqual('|'.join(post_fingerprints), metadata_dict['post-build'])
 
     self.CheckMetadataEqual(metadata_dict, metadata_proto)
+
+    pre_partition_states = metadata_proto.precondition.partition_state
+    self.assertEqual(2, len(pre_partition_states))
+    self.assertEqual('system', pre_partition_states[0].partition_name)
+    self.assertEqual(['generic'], pre_partition_states[0].device)
+    self.assertEqual(['generic/generic/generic{}'.format(source_suffix)],
+                     pre_partition_states[0].build)
+
+    self.assertEqual('vendor', pre_partition_states[1].partition_name)
+    self.assertEqual(['vendor-device-pro', 'vendor-device-std',
+                      'vendor-product-device'], pre_partition_states[1].device)
+    vendor_fingerprints = post_fingerprints
+    self.assertEqual(vendor_fingerprints, pre_partition_states[1].build)
+
+    post_partition_states = metadata_proto.postcondition.partition_state
+    self.assertEqual(2, len(post_partition_states))
+    self.assertEqual('system', post_partition_states[0].partition_name)
+    self.assertEqual(['generic'], post_partition_states[0].device)
+    self.assertEqual([self.constructFingerprint('generic/generic/generic')],
+                     post_partition_states[0].build)
+
+    self.assertEqual('vendor', post_partition_states[1].partition_name)
+    self.assertEqual(['vendor-device-pro', 'vendor-device-std',
+                      'vendor-product-device'], post_partition_states[1].device)
+    self.assertEqual(vendor_fingerprints, post_partition_states[1].build)