fix(agent): Hotfix for interrupts in agent v6.3.0, and updated agent

python base image

Author: riceriley59

agent/skyhook-agent/src/skyhook_agent/controller.py

@@ -442,10 +442,10 @@ def _make_interrupt_flag(interrupt_dir: str, interrupt_id: int) -> str:
         return f"{interrupt_dir}/{interrupt_id}.complete"
 
     SKYHOOK_RESOURCE_ID, _, _, _ = _get_env_config()
-
     config_data = make_config_data_from_resource_id()
 
     interrupt = interrupts.inflate(interrupt_data)
+
     # Check if the interrupt has already been run for this particular skyhook resource
     interrupt_dir = f"{get_skyhook_directory(root_mount)}/interrupts/flags/{SKYHOOK_RESOURCE_ID}"
     os.makedirs(interrupt_dir, exist_ok=True)
@@ -469,17 +469,23 @@ def _make_interrupt_flag(interrupt_dir: str, interrupt_id: int) -> str:
             f.write(str(time.time()))
 
         return_code = _run(
+            root_mount,
             cmd,
             get_log_file(f"interrupts/{interrupt_id}", copy_dir, config_data, root_mount),
             write_cmds=True,
             no_chmod=True
         )
 
         if return_code != 0:
-            print(f"INTERRUPT FAILED: {cmd} return_code: {return_code}")
-            # If this is not removed then we will skip all failing interrupts and it will look
-            # like the interrupt was successful when it was not.
-            os.remove(interrupt_flag)
+            # Special case: preserve flags only for reboot with a return code of 15 
+            # (SIGTERM signal sent to the process by OS because of reboot)
+            if not (interrupt.type == interrupts.NodeRestart._type() and return_code == 15):
+                print(f"INTERRUPT FAILED: {cmd} return_code: {return_code}")
+
+                # If this is not removed then we will skip all failing interrupts and it will look
+                # like the interrupt was successful when it was not.
+                os.remove(interrupt_flag)
+
             return True
 
     return False

agent/skyhook-agent/tests/test_controller.py

@@ -1068,8 +1068,8 @@ def test_interrupt_applies_all_commands(self, run_mock, datetime_mock):
                 "package_version": "version"
             }
             run_mock.assert_has_calls([
-                mock.call(["systemctl", "daemon-reload"], controller.get_log_file("interrupts/service_restart_0", copy_dir, config_data, root_dir), write_cmds=True, no_chmod=True),
-                mock.call(["systemctl", "restart", "containerd"], controller.get_log_file("interrupts/service_restart_1", copy_dir, config_data, root_dir), write_cmds=True, no_chmod=True)
+                mock.call(root_dir, ["systemctl", "daemon-reload"], controller.get_log_file("interrupts/service_restart_0", copy_dir, config_data, root_dir), write_cmds=True, no_chmod=True),
+                mock.call(root_dir, ["systemctl", "restart", "containerd"], controller.get_log_file("interrupts/service_restart_1", copy_dir, config_data, root_dir), write_cmds=True, no_chmod=True)
             ])
 
     @mock.patch("skyhook_agent.controller._run")
@@ -1112,6 +1112,26 @@ def test_interrupt_failures_remove_flag(self, run_mock):
             self.assertFalse(os.path.exists(f"{interrupt_dir}/{interrupt._type()}_1.complete"))
             self.assertFalse(os.path.exists(f"{interrupt_dir}/{interrupt._type()}_1.complete"))
 
+    @mock.patch("skyhook_agent.controller._run")
+    def test_interrupt_calls_run_with_correct_parameters(self, run_mock):
+        run_mock.return_value = 0
+        SKYHOOK_RESOURCE_ID = "scr-id-1_package_version"
+        
+        with (self._setup_for_main() as (container_root_dir, config_data, root_dir, copy_dir),
+            set_env(SKYHOOK_RESOURCE_ID=SKYHOOK_RESOURCE_ID)):
+            
+            interrupt = interrupts.ServiceRestart(["foo", "bar"])
+            result = controller.do_interrupt(interrupt.make_controller_input(), root_dir, copy_dir)
+            
+            self.assertEqual(result, False)
+            expected_calls = [
+                mock.call(root_dir, ["systemctl", "daemon-reload"], mock.ANY, write_cmds=True, no_chmod=True),
+                mock.call(root_dir, ["systemctl", "restart", "foo"], mock.ANY, write_cmds=True, no_chmod=True),
+                mock.call(root_dir, ["systemctl", "restart", "bar"], mock.ANY, write_cmds=True, no_chmod=True)
+            ]
+            run_mock.assert_has_calls(expected_calls)
+            self.assertEqual(run_mock.call_count, 3)
+
     @mock.patch("skyhook_agent.controller.datetime")
     @mock.patch("skyhook_agent.controller._run")
     def test_interrupt_failure_fails_controller(self, run_mock, datetime_mock):
@@ -1140,7 +1160,7 @@ def test_interrupt_failure_fails_controller(self, run_mock, datetime_mock):
                 "package_version": "version"
             }
             run_mock.assert_has_calls([
-                mock.call(["systemctl", "daemon-reload"], controller.get_log_file("interrupts/service_restart_0", "copy_dir", config_data, root_dir), write_cmds=True, no_chmod=True)
+                mock.call(root_dir, ["systemctl", "daemon-reload"], controller.get_log_file("interrupts/service_restart_0", "copy_dir", config_data, root_dir), write_cmds=True, no_chmod=True)
             ])
 
             self.assertEqual(result, True)
@@ -1173,7 +1193,7 @@ def test_interrupt_makes_config_from_skyhook_resource_id(self, run_mock, datetim
                 "package_version": "version"
             }
             run_mock.assert_has_calls([
-                mock.call(["systemctl", "daemon-reload"], controller.get_log_file("interrupts/service_restart_0", "copy_dir", config_data, root_dir), write_cmds=True, no_chmod=True)
+                mock.call(root_dir, ["systemctl", "daemon-reload"], controller.get_log_file("interrupts/service_restart_0", "copy_dir", config_data, root_dir), write_cmds=True, no_chmod=True)
             ])
 
     def test_interrupt_noop_makes_the_flag_file(self):

containers/agent.Dockerfile

@@ -36,13 +36,13 @@ RUN make build build_version=${AGENT_VERSION}
 # Install the wheel in the builder stage
 RUN python3 -m venv venv && ./venv/bin/pip install /code/skyhook-agent/dist/skyhook_agent*.whl
 
-FROM nvcr.io/nvidia/distroless/python:3.12-v3.4.13
+FROM nvcr.io/nvidia/distroless/python:3.12-v3.4.15
 
 ARG AGENT_VERSION
 ARG GIT_SHA
 
 ## https://github.com/opencontainers/image-spec/blob/main/annotations.md
-LABEL org.opencontainers.image.base.name="nvcr.io/nvidia/distroless/python:3.12-v3.4.13" \
+LABEL org.opencontainers.image.base.name="nvcr.io/nvidia/distroless/python:3.12-v3.4.15" \
       org.opencontainers.image.licenses="Apache-2.0" \
       org.opencontainers.image.title="skyhook-agent" \
       org.opencontainers.image.version="${AGENT_VERSION}" \