diff --git a/roles/build/templates/nifi/Dockerfile.j2 b/roles/build/templates/nifi/Dockerfile.j2
index 63c51747ea349daa1509f556ee4c323043635c3c..f85cf4788654c054d76d7f85697d46387ad60c9d 100644
--- a/roles/build/templates/nifi/Dockerfile.j2
+++ b/roles/build/templates/nifi/Dockerfile.j2
@@ -97,7 +97,7 @@ WORKDIR ${NIFI_HOME}
 # thus normal shell processing does not happen:
 # https://docs.docker.com/engine/reference/builder/#exec-form-entrypoint-example
 USER root
-RUN yum install -y supervisor rsync
+RUN yum install -y supervisor rsync python2-psutil
 RUN yum clean all
 COPY nifisupervisord.conf /etc/supervisord.conf
 ENTRYPOINT ["/usr/bin/supervisord", "-c", "/etc/supervisord.conf"]
diff --git a/roles/nifi/tasks/restart.yml b/roles/nifi/tasks/restart.yml
index d5434f2a368c8bfa115e96c39708d4c6764438b3..23b79832b515c82dc2583e8a4ede8025e75af0ff 100644
--- a/roles/nifi/tasks/restart.yml
+++ b/roles/nifi/tasks/restart.yml
@@ -1,6 +1,22 @@
 ---
 
-- name: Restart NiFi
+- name: Getting pids for nifi
+  pids:
+      name: nifi
+  register: pids_of_nifi
+
+- name: Stop NiFi
+  remote_user: root
+  command: "supervisorctl stop nifi"
+
+- name: Check nifi processes
+  remote_user: root
+  wait_for:
+    path: /proc/{{ item }}/status
+    state: absent
+  with_items: "{{ pids_of_nifi.pids }}"
+
+- name: Start NiFi
   remote_user: root
-  command: "supervisorctl restart nifi"
+  command: "supervisorctl start nifi"