From 539bad11c4f5d4e6944a66e130de766418f85ffc Mon Sep 17 00:00:00 2001 From: Future-Outlier Date: Tue, 7 Nov 2023 09:16:46 +0800 Subject: [PATCH] Update plugins/flytekit-spark/flytekitplugins/spark/agent.py Co-authored-by: Kevin Su Signed-off-by: Future Outlier --- plugins/flytekit-spark/flytekitplugins/spark/agent.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/plugins/flytekit-spark/flytekitplugins/spark/agent.py b/plugins/flytekit-spark/flytekitplugins/spark/agent.py index adcd446fc8..9f712d58c4 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/agent.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/agent.py @@ -71,11 +71,12 @@ async def async_get(self, context: grpc.ServicerContext, resource_meta: bytes) - cur_state = PENDING message = "" - if response.get("state"): - if response["state"].get("result_state"): - cur_state = convert_to_flyte_state(response["state"]["result_state"]) - if response["state"].get("state_message"): - message = response["state"]["state_message"] + state = response.get("state") + if state: + if state.get("result_state"): + cur_state = convert_to_flyte_state(state["result_state"]) + if state.get("state_message"): + message = state["state_message"] return GetTaskResponse(resource=Resource(state=cur_state, message=message))