Get and parse Slurm job status JSON object (flyteorg#3192)

* refactor: Parse job info using JSON obj

Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

* refactor: Apply JSON obj parsing to SlurmFunctionAgent get

Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

* fix: Remove the downloaded stdout file

Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

* refactor: Remove redundant check and simplify file read/write logic

Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

* Remove redundant tests

Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

* Remove raise to skip stdout delay issue

Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

* fix: Fix wrong merging operations

Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

---------

Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

Author: JiangJiaWei1103

plugins/flytekit-slurm/flytekitplugins/slurm/function/connector.py

@@ -1,10 +1,13 @@
+import json
 import tempfile
 import uuid
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Tuple, Union
 
 from asyncssh import SSHClientConnection
+from asyncssh.sftp import SFTPNoSuchFile
 
+from flytekit import logger
 from flytekit.extend.backend.base_connector import AsyncConnectorBase, ConnectorRegistry, Resource, ResourceMeta
 from flytekit.extend.backend.utils import convert_to_flyte_phase
 from flytekit.models.literals import LiteralMap
@@ -80,21 +83,25 @@ async def get(self, resource_meta: SlurmJobMetadata, **kwargs) -> Resource:
         conn = await get_ssh_conn(
             ssh_config=resource_meta.ssh_config, slurm_cluster_to_ssh_conn=self.slurm_cluster_to_ssh_conn
         )
-        job_res = await conn.run(f"scontrol show job {resource_meta.job_id}", check=True)
+        job_res = await conn.run(f"scontrol --json show job {resource_meta.job_id}", check=True)
+        job_info = json.loads(job_res.stdout)["jobs"][0]
 
         # Determine the current flyte phase from Slurm job state
-        job_state = "running"
-        msg = "No stdout available"
-        for o in job_res.stdout.split(" "):
-            if "JobState" in o:
-                job_state = o.split("=")[1].strip().lower()
-            elif "StdOut" in o:
-                stdout_path = o.split("=")[1].strip()
-                msg_res = await conn.run(f"cat {stdout_path}", check=True)
-                msg = msg_res.stdout
-
+        job_state = job_info["job_state"][0].strip().lower()
         cur_phase = convert_to_flyte_phase(job_state)
 
+        # Read stdout of the Slurm job
+        msg = ""
+        async with conn.start_sftp_client() as sftp:
+            with tempfile.NamedTemporaryFile("w+") as f:
+                try:
+                    await sftp.get(job_info["standard_output"], f.name)
+
+                    msg = f.read()
+                    logger.info(f"[SLURM STDOUT] {msg}")
+                except SFTPNoSuchFile:
+                    logger.debug("Standard output file path doesn't exist on the Slurm cluster.")
+
         return Resource(phase=cur_phase, message=msg)
 
     async def delete(self, resource_meta: SlurmJobMetadata, **kwargs) -> None:

plugins/flytekit-slurm/flytekitplugins/slurm/script/connector.py

@@ -1,10 +1,11 @@
+import json
 import tempfile
 import uuid
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Tuple, Type
 
 from asyncssh import SSHClientConnection
-from asyncssh.sftp import SFTPError
+from asyncssh.sftp import SFTPError, SFTPNoSuchFile
 
 import flytekit
 from flytekit.core.type_engine import TypeEngine
@@ -111,20 +112,25 @@ async def get(self, resource_meta: SlurmJobMetadata, **kwargs) -> Resource:
         conn = await get_ssh_conn(
             ssh_config=resource_meta.ssh_config, slurm_cluster_to_ssh_conn=self.slurm_cluster_to_ssh_conn
         )
-        job_res = await conn.run(f"scontrol show job {resource_meta.job_id}", check=True)
+        job_res = await conn.run(f"scontrol --json show job {resource_meta.job_id}", check=True)
+        job_info = json.loads(job_res.stdout)["jobs"][0]
 
         # Determine the current flyte phase from Slurm job state
-        msg = ""
-        job_state = "running"
-        for o in job_res.stdout.split(" "):
-            if "JobState" in o:
-                job_state = o.split("=")[1].strip().lower()
-            elif "StdOut" in o:
-                stdout_path = o.split("=")[1].strip()
-                msg_res = await conn.run(f"cat {stdout_path}", check=True)
-                msg = msg_res.stdout
+        job_state = job_info["job_state"][0].strip().lower()
         cur_phase = convert_to_flyte_phase(job_state)
 
+        # Read stdout of the Slurm job
+        msg = ""
+        async with conn.start_sftp_client() as sftp:
+            with tempfile.NamedTemporaryFile("w+") as f:
+                try:
+                    await sftp.get(job_info["standard_output"], f.name)
+
+                    msg = f.read()
+                    logger.info(f"[SLURM STDOUT] {msg}")
+                except SFTPNoSuchFile:
+                    logger.debug("Standard output file path doesn't exist on the Slurm cluster.")
+
         return Resource(phase=cur_phase, message=msg, outputs=resource_meta.outputs)
 
     async def delete(self, resource_meta: SlurmJobMetadata, **kwargs) -> None: