diff --git a/guest_tools/gpu_verifiers/local_gpu_verifier/src/verifier/config.py b/guest_tools/gpu_verifiers/local_gpu_verifier/src/verifier/config.py index 483a710a8..e6e1c9906 100755 --- a/guest_tools/gpu_verifiers/local_gpu_verifier/src/verifier/config.py +++ b/guest_tools/gpu_verifiers/local_gpu_verifier/src/verifier/config.py @@ -544,7 +544,7 @@ def check_status(self): self.check_if_attestation_report_driver_version_matches() and \ self.check_if_attestation_report_vbios_version_matches() and \ self.check_if_attestation_report_signature_verified() and \ - self.check_if_gpu_attestation_report_cert_chain_fwid_matched and \ + self.check_if_gpu_attestation_report_cert_chain_fwid_matched() and \ self.check_if_driver_rim_fetched() and \ self.check_if_gpu_driver_rim_schema_validated() and \ self.check_if_gpu_driver_rim_cert_chain_validated() and \ diff --git a/guest_tools/ppcie-verifier-sdk-cpp/ppcie/verifier/verification.py b/guest_tools/ppcie-verifier-sdk-cpp/ppcie/verifier/verification.py index e58b8a01c..29d208a24 100644 --- a/guest_tools/ppcie-verifier-sdk-cpp/ppcie/verifier/verification.py +++ b/guest_tools/ppcie-verifier-sdk-cpp/ppcie/verifier/verification.py @@ -29,6 +29,13 @@ from .src.nvml.nvml_client import NvmlClient from .src.utils.logging import setup_logging, get_logger from .src.nscq import NSCQHandler +from .src.exceptions.exception import ( + PpcieVerifierException, + GpuPreChecksException, + SwitchPreChecksException, + GpuAttestationException, + SwitchAttestationException, +) parser = argparse.ArgumentParser() logger = setup_logging() @@ -125,7 +132,10 @@ def verification(): if number_of_gpus != 8 and number_of_switches != 4: logger.error("PPCIE: Number of GPUs present are : %d and Switches are %d which do not meet the required " "configuration. Exiting..", number_of_gpus, number_of_switches) - sys.exit() + raise PpcieVerifierException( + f"PPCIE: Number of GPUs present are : {number_of_gpus} and Switches are {number_of_switches} " + "which do not meet the required configuration" + ) status = validate_gpu_pre_checks(nvml_client, logger, status) if status.gpu_pre_checks is True: @@ -159,6 +169,16 @@ def verification(): else: logger.debug("PPCIE: All stages have passed") nvml_client.__destroy__() + except PpcieVerifierException as e: + logger.error("An error occurred while using the PPCIE Verification Tool %s", e) + depth = 0 + while e: + indent = " " * depth + cause = f"Caused by:" if e.__cause__ else "" + logger.error(f"{indent}{type(e).__name__}: {str(e)}. {cause}") + e = e.__cause__ + depth += 1 + sys.exit(1) except Exception as e: logger.error("An error occurred while using the PPCIE Verification Tool %s", e) depth = 0 @@ -168,6 +188,7 @@ def verification(): logger.error(f"{indent}{type(e).__name__}: {str(e)}. {cause}") e = e.__cause__ depth += 1 + sys.exit(1) finally: status.status(logger) logger.info("PPCIE: End of PPCIE Verification Tool") @@ -217,7 +238,7 @@ def get_number_of_switches(logger, nscq_client): logger.info("PPCIE: Number of NVSwitches are: %d", len(switch_list_uuids)) if switch_list_uuids is not None and len(switch_list_uuids) == 0: logger.error("PPCIE: There are no switches available in the system. Exiting..") - sys.exit() + raise PpcieVerifierException("PPCIE: There are no switches available in the system") return switch_list_uuids @@ -233,7 +254,7 @@ def get_number_of_gpus(logger, nvml_client): logger.info("PPCIE: Number of GPUs are: %d", number_of_gpus) if number_of_gpus == 0 or number_of_gpus is None or number_of_gpus < 0: logger.error("PPCIE: There are no GPUs available in the system. Exiting..") - sys.exit() + raise PpcieVerifierException("PPCIE: There are no GPUs available in the system") return number_of_gpus def generate_nonce(): @@ -328,7 +349,7 @@ def perform_gpu_attestation(logger, status, args): if not isinstance(evidence_items, list): logger.error("PPCIE: GPU evidence file format invalid; expected list") status.gpu_attestation = False - sys.exit() + raise GpuAttestationException("PPCIE: GPU evidence file format invalid; expected list") else: logger.debug("PPCIE: Collecting GPU evidence from NVML") nvattest_collect_evidence_cmd = ["nvattest"] @@ -347,7 +368,9 @@ def perform_gpu_attestation(logger, status, args): if collect_exit_code != 0: logger.error("PPCIE: collect-evidence failed with non-zero exit code %d", collect_exit_code) status.gpu_attestation = False - sys.exit() + raise GpuAttestationException( + f"PPCIE: collect-evidence failed with non-zero exit code {collect_exit_code}" + ) try: collect_json = extract_last_json_object(collect_output) @@ -363,7 +386,7 @@ def perform_gpu_attestation(logger, status, args): except json.JSONDecodeError as jde: logger.error("PPCIE: Failed to parse collect-evidence JSON: %s", jde) status.gpu_attestation = False - sys.exit() + raise GpuAttestationException("PPCIE: Failed to parse collect-evidence JSON") from jde # Decode all items into gpu_attestation_report for item in evidence_items: @@ -373,7 +396,7 @@ def perform_gpu_attestation(logger, status, args): except Exception as decode_err: logger.error("PPCIE: Failed to decode GPU evidence: %s", decode_err) status.gpu_attestation = False - sys.exit() + raise GpuAttestationException("PPCIE: Failed to decode GPU evidence") from decode_err gpu_attestation_report.append(decoded_evidence) # ---- Step 2: attest ---- @@ -408,7 +431,7 @@ def perform_gpu_attestation(logger, status, args): except Exception as e: status.gpu_attestation = False logger.error("PPCIE: An error occurred while attesting the GPUs %s", e) - sys.exit() + raise GpuAttestationException("PPCIE: An error occurred while attesting the GPUs") from e finally: # Clean up temporary evidence file if created if temp_gpu_evidence_file and os.path.exists(temp_gpu_evidence_file): @@ -440,11 +463,13 @@ def validate_gpu_pre_checks(nvml_client, logger, status): else: logger.error("PPCIE: Terminating the process as TNVL is not enabled for all the GPUs") status.gpu_pre_checks = False - sys.exit() + raise GpuPreChecksException( + "PPCIE: Terminating the process as TNVL is not enabled for all the GPUs" + ) except Exception as e: status.gpu_pre_checks = False logger.error("PPCIE: An error occurred while checking GPU pre checks %s", e) - sys.exit() + raise GpuPreChecksException("PPCIE: An error occurred while checking GPU pre checks") from e return status @@ -482,7 +507,7 @@ def perform_switch_attestation(logger, status, args): if not isinstance(evidence_items, list): logger.error("PPCIE: Switch evidence file format invalid; expected list") status.switch_attestation = False - sys.exit() + raise SwitchAttestationException("PPCIE: Switch evidence file format invalid; expected list") else: logger.debug("PPCIE: Collecting switch evidence from NSCQ") nvattest_collect_evidence_cmd = ["nvattest"] @@ -501,7 +526,9 @@ def perform_switch_attestation(logger, status, args): if collect_exit_code != 0: logger.error("PPCIE: collect-evidence failed with non-zero exit code %d", collect_exit_code) status.switch_attestation = False - sys.exit() + raise SwitchAttestationException( + f"PPCIE: collect-evidence failed with non-zero exit code {collect_exit_code}" + ) try: collect_json = extract_last_json_object(collect_output) @@ -517,7 +544,7 @@ def perform_switch_attestation(logger, status, args): except json.JSONDecodeError as jde: logger.error("PPCIE: Failed to parse collect-evidence JSON: %s", jde) status.switch_attestation = False - sys.exit() + raise SwitchAttestationException("PPCIE: Failed to parse collect-evidence JSON") from jde # Decode all items into switch_attestation_report for item in evidence_items: @@ -527,7 +554,7 @@ def perform_switch_attestation(logger, status, args): except Exception as decode_err: logger.error("PPCIE: Failed to decode Switch evidence: %s", decode_err) status.switch_attestation = False - sys.exit() + raise SwitchAttestationException("PPCIE: Failed to decode Switch evidence") from decode_err switch_attestation_report.append(decoded_evidence) # ---- Step 2: attest ---- @@ -562,7 +589,7 @@ def perform_switch_attestation(logger, status, args): except Exception as e: status.switch_attestation = False logger.error("PPCIE: An error occurred while attesting the NVSwitches %s", e) - sys.exit() + raise SwitchAttestationException("PPCIE: An error occurred while attesting the NVSwitches") from e finally: # Clean up temporary evidence file if created if temp_switch_evidence_file and os.path.exists(temp_switch_evidence_file): @@ -600,7 +627,9 @@ def validate_switch_pre_checks(nscq_client, logger, status, switch_uuid_list): uuid, ) status.switch_pre_checks = False - sys.exit() + raise SwitchPreChecksException( + f"PPCIE: TNVL mode for the switch with id {uuid} is not enabled" + ) lock_status, rc = nscq_client.is_switch_lock_mode(uuid) logger.debug( "PPCIE: Lock Mode of the switch is: %s with return " @@ -614,7 +643,9 @@ def validate_switch_pre_checks(nscq_client, logger, status, switch_uuid_list): uuid, ) status.switch_pre_checks = False - sys.exit() + raise SwitchPreChecksException( + f"PPCIE: Lock mode for the switch with id {uuid} is not enabled" + ) status.switch_pre_checks = True except Exception as e: status.switch_pre_checks = False @@ -622,7 +653,7 @@ def validate_switch_pre_checks(nscq_client, logger, status, switch_uuid_list): "PPCIE: An exception has occurred while attempting to check switch pre checks %s", e, ) - sys.exit() + raise SwitchPreChecksException("PPCIE: An exception has occurred while attempting to check switch pre checks") from e return status diff --git a/guest_tools/ppcie-verifier-sdk-cpp/tests/test_verification.py b/guest_tools/ppcie-verifier-sdk-cpp/tests/test_verification.py index e3331cb60..10d13de9d 100644 --- a/guest_tools/ppcie-verifier-sdk-cpp/tests/test_verification.py +++ b/guest_tools/ppcie-verifier-sdk-cpp/tests/test_verification.py @@ -42,9 +42,11 @@ class TestVerification(unittest.TestCase): @patch("ppcie.verifier.src.nvml.nvml_client.nvmlDeviceGetCount") @patch("ppcie.verifier.src.nvml.nvml_client.nvmlInit") @patch("ppcie.verifier.verification.run_nvattest_command") - def test_verification(self, run_nvattest, nvml_init, gpu_device_count, nscq_handler, mock_system_configuration_compute_settings, mock_settings, gpu_topology_check, switch_topology_check, get_ready_state, set_ready_state): + @patch("ppcie.verifier.src.nvml.nvml_client.nvmlShutdown") + def test_verification(self, nvml_shutdown, run_nvattest, nvml_init, gpu_device_count, nscq_handler, mock_system_configuration_compute_settings, mock_settings, gpu_topology_check, switch_topology_check, get_ready_state, set_ready_state): status = Status() nvml_init.return_value = None + nvml_shutdown.return_value = None gpu_device_count.return_value = 8 nscq_handler.return_value = nscq_handler nscq_handler.get_all_switch_uuid.return_value = [{'switchid-1', 'switchid-2', 'switchid-3', 'switchid-4'}] @@ -91,9 +93,11 @@ def nvattest_side_effect(logger, command, description): @patch("ppcie.verifier.src.nvml.nvml_client.nvmlDeviceGetCount") @patch("ppcie.verifier.src.nvml.nvml_client.nvmlInit") @patch("ppcie.verifier.verification.run_nvattest_command") - def test_verification_with_file_evidence(self, run_nvattest, nvml_init, gpu_device_count, nscq_handler, mock_system_configuration_compute_settings, mock_settings, gpu_topology_check, switch_topology_check, get_ready_state, set_ready_state): + @patch("ppcie.verifier.src.nvml.nvml_client.nvmlShutdown") + def test_verification_with_file_evidence(self, nvml_shutdown, run_nvattest, nvml_init, gpu_device_count, nscq_handler, mock_system_configuration_compute_settings, mock_settings, gpu_topology_check, switch_topology_check, get_ready_state, set_ready_state): status = Status() nvml_init.return_value = None + nvml_shutdown.return_value = None gpu_device_count.return_value = 8 nscq_handler.return_value = nscq_handler nscq_handler.get_all_switch_uuid.return_value = [{'switchid-1', 'switchid-2', 'switchid-3', 'switchid-4'}] @@ -129,6 +133,68 @@ def nvattest_side_effect(logger, command, description): with patch.object(sys, 'argv', testargs): verification() + @patch("ppcie.verifier.verification.NSCQHandler") + @patch("ppcie.verifier.src.nvml.nvml_client.nvmlDeviceGetCount") + @patch("ppcie.verifier.src.nvml.nvml_client.nvmlInit") + def test_verification_should_exit_non_zero_if_no_gpus_are_present(self, nvml_init, gpu_device_count, nscq_handler): + nvml_init.return_value = None + gpu_device_count.return_value = 0 + nscq_handler.return_value = nscq_handler + verification_module.parser = argparse.ArgumentParser() + testargs = ["prog", "--verifier", "local", "--nonce", "1234"] + + with patch.object(sys, 'argv', testargs): + with self.assertRaises(SystemExit) as error: + verification() + + self.assertEqual(error.exception.code, 1) + + @patch("ppcie.verifier.src.nvml.nvml_client.nvmlSystemSetConfComputeGpusReadyState") + @patch("pynvml.nvmlSystemGetConfComputeGpusReadyState") + @patch("ppcie.verifier.src.topology.validate_topology.TopologyValidation.switch_topology_check") + @patch("ppcie.verifier.src.topology.validate_topology.TopologyValidation.gpu_topology_check") + @patch("ppcie.verifier.src.nvml.nvml_client.NvmlSystemConfComputeSettings") + @patch("ppcie.verifier.src.nvml.nvml_client.nvmlSystemGetConfComputeSettings") + @patch("ppcie.verifier.verification.NSCQHandler") + @patch("ppcie.verifier.src.nvml.nvml_client.nvmlDeviceGetCount") + @patch("ppcie.verifier.src.nvml.nvml_client.nvmlInit") + @patch("ppcie.verifier.verification.run_nvattest_command") + def test_verification_should_exit_non_zero_if_collect_evidence_fails(self, run_nvattest, nvml_init, gpu_device_count, nscq_handler, mock_system_configuration_compute_settings, mock_settings, gpu_topology_check, switch_topology_check, get_ready_state, set_ready_state): + status = Status() + nvml_init.return_value = None + gpu_device_count.return_value = 8 + nscq_handler.return_value = nscq_handler + nscq_handler.get_all_switch_uuid.return_value = [{'switchid-1', 'switchid-2', 'switchid-3', 'switchid-4'}] + result = pynvml.NVML_SUCCESS + settings = NvmlSystemConfComputeSettings() + settings.version = NVML_SYSTEM_CONF_COMPUTE_VERSION + settings.environment = 1 + settings.ccFeature = 1 + settings.devToolsMode = 3 + settings.multiGpuMode = 1 + mock_settings.return_value = settings + mock_system_configuration_compute_settings.return_value = result + nscq_handler.is_switch_tnvl_mode.return_value = 1, 0 + nscq_handler.is_switch_lock_mode.return_value = 1, 0 + run_nvattest.return_value = ("collect failed", 2) + status.topology_checks = True + status.gpu_pre_checks = True + status.switch_pre_checks = True + status.gpu_attestation = True + status.switch_attestation = True + gpu_topology_check.return_value = status + switch_topology_check.return_value = status + get_ready_state.return_value = 1 + set_ready_state.return_value = pynvml.NVML_SUCCESS + verification_module.parser = argparse.ArgumentParser() + testargs = ["prog", "--verifier", "local", "--log-level", "debug", "--nonce", "1234", "--relying-party-policy", "/tests/data/relying_party_policy.rego", "--rim-url", "https://rim.attestation.nvidia.com", "--ocsp-url", "https://ocsp.ndis.nvidia.com", "--nras-url", "https://nras.attestation.nvidia.com"] + + with patch.object(sys, 'argv', testargs): + with self.assertRaises(SystemExit) as error: + verification() + + self.assertEqual(error.exception.code, 1) + @patch("ppcie.verifier.verification.run_nvattest_command") def test_perform_gpu_attestation(self, run_nvattest): logger = logging.getLogger('test') @@ -343,4 +409,4 @@ def test_generate_nonce(self): hex = generate_nonce() self.assertIsNotNone(hex) self.assertTrue(len(hex) == 64) - self.assertTrue(hex.isalnum()) \ No newline at end of file + self.assertTrue(hex.isalnum()) diff --git a/guest_tools/ppcie-verifier/ppcie/verifier/verification.py b/guest_tools/ppcie-verifier/ppcie/verifier/verification.py index aba49b7b5..aa04e4552 100644 --- a/guest_tools/ppcie-verifier/ppcie/verifier/verification.py +++ b/guest_tools/ppcie-verifier/ppcie/verifier/verification.py @@ -133,12 +133,17 @@ def verification(): if number_of_gpus != 8 and number_of_switches != 4: logger.error("PPCIE: Number of GPUs present are : %d and Switches are %d which do not meet the required " "configuration. Exiting..", number_of_gpus, number_of_switches) - sys.exit() + raise PpcieVerifierException( + f"PPCIE: Number of GPUs present are : {number_of_gpus} and Switches are {number_of_switches} " + "which do not meet the required configuration" + ) set_allow_hold_cert(args["allow_hold_cert"]) if args["gpu_attestation_mode"] != args["switch_attestation_mode"]: logger.error( "PPCIE: GPU attestation mode and Switch attestation mode should be same. Exiting..") - sys.exit() + raise PpcieVerifierException( + "PPCIE: GPU attestation mode and Switch attestation mode should be same" + ) status = validate_gpu_pre_checks(nvml_client, logger, status) if status.gpu_pre_checks is True: @@ -172,6 +177,16 @@ def verification(): else: logger.debug("PPCIE: All stages have passed") nvml_client.__destroy__() + except PpcieVerifierException as e: + logger.error("An error occurred while using the PPCIE Verification Tool %s", e) + depth = 0 + while e: + indent = " " * depth + cause = f"Caused by:" if e.__cause__ else "" + logger.error(f"{indent}{type(e).__name__}: {str(e)}. {cause}") + e = e.__cause__ + depth += 1 + sys.exit(1) except Exception as e: logger.error("An error occurred while using the PPCIE Verification Tool %s", e) depth = 0 @@ -181,6 +196,7 @@ def verification(): logger.error(f"{indent}{type(e).__name__}: {str(e)}. {cause}") e = e.__cause__ depth += 1 + sys.exit(1) finally: status.status(logger) logger.info("PPCIE: End of PPCIE Verification Tool") @@ -230,7 +246,7 @@ def get_number_of_switches(logger, nscq_client): logger.info("PPCIE: Number of NVSwitches are: %d", len(switch_list_uuids)) if switch_list_uuids is not None and len(switch_list_uuids) == 0: logger.error("PPCIE: There are no switches available in the system. Exiting..") - sys.exit() + raise PpcieVerifierException("PPCIE: There are no switches available in the system") return switch_list_uuids @@ -246,7 +262,7 @@ def get_number_of_gpus(logger, nvml_client): logger.info("PPCIE: Number of GPUs are: %d", number_of_gpus) if number_of_gpus == 0 or number_of_gpus is None or number_of_gpus < 0: logger.error("PPCIE: There are no GPUs available in the system. Exiting..") - sys.exit() + raise PpcieVerifierException("PPCIE: There are no GPUs available in the system") return number_of_gpus @@ -304,7 +320,9 @@ def perform_gpu_attestation(logger, status, args): type(evidence), ) status.gpu_attestation = False - sys.exit() + raise GpuAttestationException( + f"PPCIE: Invalid/Unknown evidence type found for GPU {type(evidence)}" + ) logger.info("PPCIE: Attesting the GPUs") gpu_attestation_result = client.attest(evidence_list) @@ -327,7 +345,7 @@ def perform_gpu_attestation(logger, status, args): except Exception as e: status.gpu_attestation = False logger.error("PPCIE: An error occurred while attesting the GPUs %s", e) - sys.exit() + raise GpuAttestationException("PPCIE: An error occurred while attesting the GPUs") from e logger.info("PPCIE: GPU Attestation Completed") return status, gpu_attestation_report @@ -352,11 +370,13 @@ def validate_gpu_pre_checks(nvml_client, logger, status): "PPCIE: Terminating the process as TNVL is not enabled for all the GPUs" ) status.gpu_pre_checks = False - sys.exit() + raise GpuPreChecksException( + "PPCIE: Terminating the process as TNVL is not enabled for all the GPUs" + ) except Exception as e: status.gpu_pre_checks = False logger.error("PPCIE: An error occurred while checking GPU pre checks %s", e) - sys.exit() + raise GpuPreChecksException("PPCIE: An error occurred while checking GPU pre checks") from e return status @@ -403,7 +423,9 @@ def perform_switch_attestation(logger, status, args): type(evidence), ) status.switch_attestation = False - sys.exit() + raise SwitchAttestationException( + f"PPCIE: Invalid/Unknown evidence type found for switch {type(evidence)}" + ) logger.info("PPCIE: Attesting the switches") attestation_result = switch_attester.attest(evidence_list) @@ -427,7 +449,7 @@ def perform_switch_attestation(logger, status, args): except Exception as e: logger.error("PPCIE: An error occurred while attesting the switches %s", e) status.switch_attestation = False - sys.exit() + raise SwitchAttestationException("PPCIE: An error occurred while attesting the switches") from e return status, switch_attestation_report @@ -455,7 +477,9 @@ def validate_switch_pre_checks(nscq_client, logger, status, switch_uuid_list): uuid, ) status.switch_pre_checks = False - sys.exit() + raise SwitchPreChecksException( + f"PPCIE: TNVL mode for the switch with id {uuid} is not enabled" + ) lock_status, rc = nscq_client.is_switch_lock_mode(uuid) logger.debug( "PPCIE: Lock Mode of the switch is: %s with return " @@ -469,7 +493,9 @@ def validate_switch_pre_checks(nscq_client, logger, status, switch_uuid_list): uuid, ) status.switch_pre_checks = False - sys.exit() + raise SwitchPreChecksException( + f"PPCIE: Lock mode for the switch with id {uuid} is not enabled" + ) status.switch_pre_checks = True except Exception as e: status.switch_pre_checks = False @@ -477,7 +503,7 @@ def validate_switch_pre_checks(nscq_client, logger, status, switch_uuid_list): "PPCIE: An exception has occurred while attempting to check switch pre checks %s", e, ) - sys.exit() + raise SwitchPreChecksException("PPCIE: An exception has occurred while attempting to check switch pre checks") from e return status diff --git a/guest_tools/ppcie-verifier/tests/test_verification.py b/guest_tools/ppcie-verifier/tests/test_verification.py index b52e96cfc..bf81cbb04 100644 --- a/guest_tools/ppcie-verifier/tests/test_verification.py +++ b/guest_tools/ppcie-verifier/tests/test_verification.py @@ -12,6 +12,8 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import argparse +import base64 import logging import unittest @@ -27,11 +29,12 @@ NVML_SYSTEM_CONF_COMPUTE_VERSION from ppcie.verifier.src.utils.status import Status from ppcie.verifier.verification import verification +import ppcie.verifier.verification as verification_module import sys import os from nv_attestation_sdk import attestation -PPCIE_EVIDENCE = os.environ.get("PPCIE_EVIDENCE") +PPCIE_EVIDENCE = os.environ.get("PPCIE_EVIDENCE") or base64.b64encode(b"evidence").decode("ascii") class TestVerification(unittest.TestCase): @@ -46,9 +49,11 @@ class TestVerification(unittest.TestCase): @patch("ppcie.verifier.verification.NSCQHandler") @patch("ppcie.verifier.src.nvml.nvml_client.nvmlDeviceGetCount") @patch("ppcie.verifier.src.nvml.nvml_client.nvmlInit") - def test_verification(self, nvml_init, gpu_device_count, nscq_handler, mock_system_configuration_compute_settings, mock_settings, attestation_client, gpu_topology_check, switch_topology_check, get_ready_state, set_ready_state): + @patch("ppcie.verifier.src.nvml.nvml_client.nvmlShutdown") + def test_verification(self, nvml_shutdown, nvml_init, gpu_device_count, nscq_handler, mock_system_configuration_compute_settings, mock_settings, attestation_client, gpu_topology_check, switch_topology_check, get_ready_state, set_ready_state): status = Status() nvml_init.return_value = None + nvml_shutdown.return_value = None gpu_device_count.return_value = 8 nscq_handler.return_value = nscq_handler nscq_handler.get_all_switch_uuid.return_value = [{'switchid-1', 'switchid-2', 'switchid-3', 'switchid-4'}] @@ -83,10 +88,27 @@ def test_verification(self, nvml_init, gpu_device_count, nscq_handler, mock_syst switch_topology_check.return_value = status get_ready_state.return_value = 1 set_ready_state.return_value = pynvml.NVML_SUCCESS + verification_module.parser = argparse.ArgumentParser() testargs = ["prog", "--gpu-attestation-mode", "LOCAL", "--switch-attestation-mode", "LOCAL", "--log", "DEBUG", "--claims-version", "3.0"] with patch.object(sys, 'argv', testargs): verification() + @patch("ppcie.verifier.verification.NSCQHandler") + @patch("ppcie.verifier.src.nvml.nvml_client.nvmlDeviceGetCount") + @patch("ppcie.verifier.src.nvml.nvml_client.nvmlInit") + def test_verification_should_exit_non_zero_if_no_gpus_are_present(self, nvml_init, gpu_device_count, nscq_handler): + nvml_init.return_value = None + gpu_device_count.return_value = 0 + nscq_handler.return_value = nscq_handler + verification_module.parser = argparse.ArgumentParser() + testargs = ["prog", "--gpu-attestation-mode", "LOCAL", "--switch-attestation-mode", "LOCAL"] + + with patch.object(sys, 'argv', testargs): + with self.assertRaises(SystemExit) as error: + verification() + + self.assertEqual(error.exception.code, 1) + @patch("ppcie.verifier.src.nvml.nvml_client.nvmlSystemSetConfComputeGpusReadyState") @patch("pynvml.nvmlSystemGetConfComputeGpusReadyState") @patch("ppcie.verifier.src.nvml.nvml_client.nvmlInit")