diff --git a/src/opengradient/client/_conversions.py b/src/opengradient/client/_conversions.py index 495f663..18493f2 100644 --- a/src/opengradient/client/_conversions.py +++ b/src/opengradient/client/_conversions.py @@ -156,7 +156,10 @@ def convert_to_model_output(event_data: AttributeDict) -> Dict[str, np.ndarray]: if isinstance(tensor, (AttributeDict, dict)): name = tensor.get("name") value = tensor.get("value") - output_dict[name] = np.array(json.loads(value)) + try: + output_dict[name] = np.array(json.loads(value)) + except (json.JSONDecodeError, TypeError, ValueError) as e: + logging.warning(f"Failed to parse JSON tensor '{name}': {e} (raw value: {value!r})") else: logging.warning(f"Unexpected tensor type: {type(tensor)}") @@ -203,7 +206,10 @@ def convert_array_to_model_output(array_data: List) -> ModelOutput: for tensor in array_data[2]: name = tensor[0] value = tensor[1] - json_data[name] = np.array(json.loads(value)) + try: + json_data[name] = np.array(json.loads(value)) + except (json.JSONDecodeError, TypeError, ValueError) as e: + logging.warning(f"Failed to parse JSON tensor '{name}': {e} (raw value: {value!r})") return ModelOutput( numbers=number_data, @@ -211,3 +217,4 @@ def convert_array_to_model_output(array_data: List) -> ModelOutput: jsons=json_data, is_simulation_result=array_data[3], ) + diff --git a/tests/utils_test.py b/tests/utils_test.py index 082cb8d..9bba5b5 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -68,8 +68,15 @@ def test_convert_array_empty(): assert result.is_simulation_result is False -def test_convert_array_invalid_json(): - """Test handling of invalid JSON data""" +def test_convert_array_invalid_json(caplog): + """Test that invalid JSON tensor data is skipped with a warning instead of crashing. + + Previously json.loads() was unguarded and would crash the entire conversion. + After the fix, malformed JSON tensors are skipped and logged as warnings so + all other valid tensors in the same response are still returned. + """ + import logging + invalid_json_data = [ [], # Empty number tensors [], # Empty string tensors @@ -77,8 +84,15 @@ def test_convert_array_invalid_json(): False, ] - with pytest.raises(json.JSONDecodeError): - utils.convert_array_to_model_output(invalid_json_data) + with caplog.at_level(logging.WARNING): + result = utils.convert_array_to_model_output(invalid_json_data) + + # Should not raise; malformed tensor is skipped + assert isinstance(result, types.ModelOutput) + assert result.jsons == {} + + # A warning should have been emitted mentioning the tensor name + assert any("invalid_json" in record.message for record in caplog.records) def test_convert_array_invalid_shape():