diff --git a/mmc/flash.cpp b/mmc/flash.cpp
index 3bd67fa..e1c5e9c 100644
--- a/mmc/flash.cpp
+++ b/mmc/flash.cpp
@@ -3,6 +3,7 @@
 #include "flash.hpp"
 
 #include "activation.hpp"
+#include "item_updater.hpp"
 
 namespace phosphor
 {
@@ -33,7 +34,8 @@
     msg.read(newStateID, newStateObjPath, newStateUnit, newStateResult);
 
     auto mmcServiceFile = "obmc-flash-mmc@" + versionId + ".service";
-    auto mmcSetPrimary = "obmc-flash-mmc-setprimary@" + versionId + ".service";
+    auto flashId = parent.versions.find(versionId)->second->path();
+    auto mmcSetPrimary = "obmc-flash-mmc-setprimary@" + flashId + ".service";
 
     if (newStateUnit == mmcServiceFile && newStateResult == "done")
     {
diff --git a/mmc/item_updater_helper.cpp b/mmc/item_updater_helper.cpp
index d3279c0..70a94bf 100644
--- a/mmc/item_updater_helper.cpp
+++ b/mmc/item_updater_helper.cpp
@@ -34,11 +34,11 @@
     utils::execute("/sbin/fw_setenv", "rwreset", "true");
 }
 
-void Helper::removeVersion(const std::string& versionId)
+void Helper::removeVersion(const std::string& flashId)
 {
     auto method = bus.new_method_call(SYSTEMD_BUSNAME, SYSTEMD_PATH,
                                       SYSTEMD_INTERFACE, "StartUnit");
-    auto serviceFile = "obmc-flash-mmc-remove@" + versionId + ".service";
+    auto serviceFile = "obmc-flash-mmc-remove@" + flashId + ".service";
     method.append(serviceFile, "replace");
     bus.call_noreply(method);
 
@@ -48,11 +48,11 @@
     std::this_thread::sleep_for(removeWait);
 }
 
-void Helper::updateUbootVersionId(const std::string& versionId)
+void Helper::updateUbootVersionId(const std::string& flashId)
 {
     auto method = bus.new_method_call(SYSTEMD_BUSNAME, SYSTEMD_PATH,
                                       SYSTEMD_INTERFACE, "StartUnit");
-    auto serviceFile = "obmc-flash-mmc-setprimary@" + versionId + ".service";
+    auto serviceFile = "obmc-flash-mmc-setprimary@" + flashId + ".service";
     method.append(serviceFile, "replace");
     bus.call_noreply(method);
 
diff --git a/mmc/obmc-flash-mmc-setprimary@.service.in b/mmc/obmc-flash-mmc-setprimary@.service.in
index 69feafb..77c724b 100644
--- a/mmc/obmc-flash-mmc-setprimary@.service.in
+++ b/mmc/obmc-flash-mmc-setprimary@.service.in
@@ -1,5 +1,5 @@
 [Unit]
-Description=Set %I as primary version
+Description=Set %I as primary partition
 
 [Service]
 Type=oneshot
diff --git a/mmc/obmc-flash-mmc@.service.in b/mmc/obmc-flash-mmc@.service.in
index f8b4262..6cdbdb1 100644
--- a/mmc/obmc-flash-mmc@.service.in
+++ b/mmc/obmc-flash-mmc@.service.in
@@ -1,6 +1,5 @@
 [Unit]
 Description=Write image %I to BMC storage
-OnFailure=obmc-flash-mmc-remove@%i.service
 
 [Service]
 Type=oneshot