Update plugins/flytekit-spark/flytekitplugins/spark/agent.py

Co-authored-by: Kevin Su <[email protected]> Signed-off-by: Future Outlier <[email protected]>
flyteorg · Nov 7, 2023 · 539bad1 · 539bad1
1 parent 31b4f0d
commit 539bad1
Showing 1 changed file with 6 additions and 5 deletions.
diff --git 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))