diff --git a/src/sagemaker/hyperpod/common/telemetry/telemetry_logging.py b/src/sagemaker/hyperpod/common/telemetry/telemetry_logging.py index 79eb2d29..e4891fb2 100644 --- a/src/sagemaker/hyperpod/common/telemetry/telemetry_logging.py +++ b/src/sagemaker/hyperpod/common/telemetry/telemetry_logging.py @@ -160,7 +160,7 @@ def wrapper(*args, **kwargs): duration = round(perf_counter() - start, 2) extra += f"&x-latency={duration}" _send_telemetry_request( - Status.SUCCESS, + STATUS_TO_CODE[str(Status.SUCCESS)], [FEATURE_TO_CODE[str(feature)]], None, None, @@ -172,7 +172,7 @@ def wrapper(*args, **kwargs): duration = round(perf_counter() - start, 2) extra += f"&x-latency={duration}" _send_telemetry_request( - Status.FAILURE, + STATUS_TO_CODE[str(Status.FAILURE)], [FEATURE_TO_CODE[str(feature)]], None, str(e), diff --git a/test/unit_tests/common/telemetry/test_telemetry_logging.py b/test/unit_tests/common/telemetry/test_telemetry_logging.py index 12939bdc..a54e36c5 100644 --- a/test/unit_tests/common/telemetry/test_telemetry_logging.py +++ b/test/unit_tests/common/telemetry/test_telemetry_logging.py @@ -17,6 +17,8 @@ import requests import logging +from src.sagemaker.hyperpod.common.telemetry.telemetry_logging import STATUS_TO_CODE + # Test data MOCK_CONTEXTS = { "eks_arn": "arn:aws:eks:us-west-2:123456789012:cluster/my-cluster", @@ -163,7 +165,7 @@ def sample_function(): args = mock_telemetry.call_args[0] # Check status - assert args[0] == Status.SUCCESS + assert args[0] == STATUS_TO_CODE[str(Status.SUCCESS)] # Check feature code assert args[1] == [FEATURE_TO_CODE[str(Feature.HYPERPOD)]] @@ -198,11 +200,11 @@ def sample_function(succeed: bool): # Check success call success_call = mock_telemetry.call_args_list[0] - assert success_call[0][0] == Status.SUCCESS + assert success_call[0][0] == STATUS_TO_CODE[str(Status.SUCCESS)] # Check failure call failure_call = mock_telemetry.call_args_list[1] - assert failure_call[0][0] == Status.FAILURE + assert failure_call[0][0] == STATUS_TO_CODE[str(Status.FAILURE)] # Test _requests_helper