Skip to content

Commit 59e2bca

Browse files
matthew29tangcopybara-github
authored andcommitted
fix: Handle case when no metadata is returned from model.predict
PiperOrigin-RevId: 607414090
1 parent cfd96d8 commit 59e2bca

File tree

1 file changed

+9
-3
lines changed

1 file changed

+9
-3
lines changed

google/cloud/aiplatform/models.py

+9-3
Original file line numberDiff line numberDiff line change
@@ -1564,7 +1564,7 @@ def predict(
15641564
json_response = raw_predict_response.json()
15651565
return Prediction(
15661566
predictions=json_response["predictions"],
1567-
metadata=json_response["metadata"],
1567+
metadata=json_response.get("metadata"),
15681568
deployed_model_id=raw_predict_response.headers[
15691569
_RAW_PREDICT_DEPLOYED_MODEL_ID_KEY
15701570
],
@@ -1582,7 +1582,10 @@ def predict(
15821582
parameters=parameters,
15831583
timeout=timeout,
15841584
)
1585-
metadata = json_format.MessageToDict(prediction_response._pb.metadata)
1585+
if prediction_response._pb.metadata:
1586+
metadata = json_format.MessageToDict(prediction_response._pb.metadata)
1587+
else:
1588+
metadata = None
15861589

15871590
return Prediction(
15881591
predictions=[
@@ -1643,7 +1646,10 @@ async def predict_async(
16431646
parameters=parameters,
16441647
timeout=timeout,
16451648
)
1646-
metadata = json_format.MessageToDict(prediction_response._pb.metadata)
1649+
if prediction_response._pb.metadata:
1650+
metadata = json_format.MessageToDict(prediction_response._pb.metadata)
1651+
else:
1652+
metadata = None
16471653

16481654
return Prediction(
16491655
predictions=[

0 commit comments

Comments
 (0)