[iris] Return typed JobState from job_state / state_only

Author: rjpower

lib/iris/src/iris/client/client.py

@@ -224,18 +224,14 @@ def status(self) -> job_pb2.JobStatus:
         """
         return self._client._cluster_client.get_job_status(self._job_id)
 
-    def state_only(self) -> int:
+    def state_only(self) -> job_pb2.JobState:
         """Lightweight state query that avoids loading tasks/attempts/workers."""
-        states = self._client._cluster_client.get_job_states([self._job_id])
-        wire_id = self._job_id.to_wire()
-        if wire_id not in states:
-            raise KeyError(f"Job {wire_id} not found")
-        return states[wire_id]
+        return self._client.job_state(self._job_id)
 
     @property
     def state(self) -> job_pb2.JobState:
         """Get current job state via the lightweight state-only RPC."""
-        return cast(job_pb2.JobState, self.state_only())
+        return self.state_only()
 
     def tasks(self) -> list[Task]:
         """Get all tasks for this job.
@@ -706,7 +702,7 @@ def status(self, job_id: JobName) -> job_pb2.JobStatus:
         """
         return self._cluster_client.get_job_status(job_id)
 
-    def job_state(self, job_id: JobName) -> int:
+    def job_state(self, job_id: JobName) -> job_pb2.JobState:
         """Lightweight state query that avoids loading tasks/attempts/workers.
 
         Prefer this over ``status(job_id).state`` for polling loops.
@@ -715,7 +711,7 @@ def job_state(self, job_id: JobName) -> int:
         wire_id = job_id.to_wire()
         if wire_id not in states:
             raise KeyError(f"Job {wire_id} not found")
-        return states[wire_id]
+        return cast(job_pb2.JobState, states[wire_id])
 
     def terminate(self, job_id: JobName) -> None:
         """Terminate a running job.