TYPMAINT: Type fixes #1

Author: jerdra

executors/drmaa_executor.py

@@ -2,8 +2,7 @@
 import drmaa
 from drmaa_executor_plugin.drmaa_patches import PatchedSession as drmaaSession
 
-from typing import TYPE_CHECKING
-from typing import Optional, List
+from typing import TYPE_CHECKING, Optional, TypedDict, Generator, Dict, cast
 
 from functools import wraps
 
@@ -28,6 +27,16 @@
     drmaa.JobState.FAILED: State.FAILED
 }
 
+JobID = int
+_TaskInstanceKeyDict = TypedDict('_TaskInstanceKeyDict', {
+    'dag_id': str,
+    'task_id': str,
+    'run_id': str,
+    'try_number': int
+})
+
+JobTrackingType = Dict[JobID, _TaskInstanceKeyDict]
+
 
 def check_started(method):
     '''
@@ -58,49 +67,78 @@ def __init__(self, max_concurrent_jobs: Optional[int] = None):
         # Not yet implemented
         self.max_concurrent_jobs: Optional[int] = max_concurrent_jobs
 
+    def iter_scheduled_jobs(
+            self) -> Generator[tuple[JobID, TaskInstanceKey], None, None]:
+        '''
+        Iterate over scheduled jobs
+        '''
+        for job_id, instance_info in self._get_or_create_job_ids().items():
+            yield job_id, TaskInstanceKey(**instance_info)
+
     @property
     def active_jobs(self) -> int:
         return len(self._get_or_create_job_ids())
 
+    @active_jobs.setter
+    def active_jobs(self, val):
+        self.active_jobs = val
+
     # TODO: Make `scheduler_job_ids` configurable under [executor]
     @provide_session
     def _get_or_create_job_ids(self,
-                               session: Optional[Session] = None) -> List[int]:
+                               session: Optional[Session] = None
+                               ) -> JobTrackingType:
+
         current_jobs = Variable.get("scheduler_job_ids",
                                     default_var=None,
                                     deserialize_json=True)
         if not current_jobs:
+            current_jobs = {}
             self.log.info("Setting up job tracking Airflow variable...")
-            Variable.set("scheduler_job_ids", {"jobs": []},
+            Variable.set("scheduler_job_ids",
+                         current_jobs,
                          description="Scheduler Job ID tracking",
                          serialize_json=True,
                          session=session)
             self.log.info("Created `scheduler_job_ids` variable")
 
-        return current_jobs["jobs"]
+        return current_jobs
 
     @provide_session
     def _update_job_tracker(self,
-                            jobs: List[int],
+                            jobs: JobTrackingType,
                             session: Optional[Session] = None) -> None:
-        write_json = {"jobs": jobs}
         Variable.update("scheduler_job_ids",
-                        write_json,
+                        jobs,
                         serialize_json=True,
                         session=session)
 
-    def _drop_from_tracking(self, job_id: int) -> None:
+    def _drop_from_tracking(self, job_id: JobID) -> None:
         self.log.info(
-            "Removing Job {job_id} from tracking variable `scheduler_job_ids`")
-        new_state = [j for j in self._get_or_create_job_ids() if j != job_id]
-        self._update_job_tracker(new_state)
-        self.log.info("Successfully removed {job_id} from `scheduler_job_ids`")
+            f"Removing Job {job_id} from tracking variable `scheduler_job_ids`"
+        )
+
+        jobs = self._get_or_create_job_ids()
+        try:
+            jobs.pop(job_id)
+        except KeyError:
+            self.log.error(f"Failed to remove {job_id}, job was not"
+                           " being tracked by Airflow!")
+        else:
+            self._update_job_tracker(jobs)
+            self.log.info(
+                f"Successfully removed {job_id} from `scheduler_job_ids`")
 
-    def _push_to_tracking(self, job_id: int) -> None:
+    def _push_to_tracking(self, job_id: JobID, key: TaskInstanceKey) -> None:
         self.log.info(
             "Adding Job {job_id} to tracking variable `scheduler_job_ids`")
+
+        # Convert TaskInstanceKey to serializable form
+        key_dict = _taskkey_to_dict(key)
+        entry = {job_id: key_dict}
+
         current_jobs = self._get_or_create_job_ids()
-        current_jobs.append(job_id)
+        current_jobs.update(entry)
         self._update_job_tracker(current_jobs)
         self.log.info("Successfully added {job_id} to `scheduler_job_ids`")
 
@@ -114,7 +152,7 @@ def start(self) -> None:
         current_jobs = self._get_or_create_job_ids()
 
         if current_jobs:
-            print_jobs = "\n".join([j_id for j_id in current_jobs["jobs"]])
+            print_jobs = "\n".join([f"{j_id}" for j_id in current_jobs])
             self.log.info(f"Jobs from previous session:\n{print_jobs}")
         else:
             self.log.info("No jobs are currently being tracked")
@@ -127,6 +165,7 @@ def end(self) -> None:
         self.log.info("Terminating DRMAA session")
         self.session.exit()
 
+    @check_started
     def sync(self) -> None:
         """
         Called periodically by `airflow.executors.base_executor.BaseExecutor`'s
@@ -136,9 +175,8 @@ def sync(self) -> None:
         """
 
         # Go through currently running jobs and update state
-        scheduled_jobs = self._get_or_create_job_ids()
-        for job_id in scheduled_jobs:
-            drmaa_status = self.jobStatus(job_id)
+        for job_id, task_instance_key in self.iter_scheduled_jobs():
+            drmaa_status = self.session.jobStatus(job_id)
             try:
                 status = JOB_STATE_MAP[drmaa_status]
             except KeyError:
@@ -147,7 +185,8 @@ def sync(self) -> None:
                     " Cannot be mapped into an Airflow TaskInstance State"
                     " Will try again in next sync attempt...")
             else:
-                self.change_state(status, None)
+                # Need taskinstancekey
+                self.change_state(task_instance_key, status)
                 self._drop_from_tracking(job_id)
 
     @check_started
@@ -162,7 +201,7 @@ def execute_async(self,
 
         self.log.info(f"Submitting job {key} with command {command} with"
                       f" configuration options:\n{executor_config})")
-        jt = executor_config.drm2drmaa(self.session.createJobTemplate())
+        jt = executor_config.get_drmaa_config(self.session.createJobTemplate())
 
         # CommandType always begins with "airflow" binary command
         jt.remoteCommand = command[0]
@@ -174,9 +213,18 @@ def execute_async(self,
         job_id = self.session.runJob(jt)
 
         self.log.info(f"Submitted Job {job_id}")
-        self._push_to_tracking(job_id)
+        self._push_to_tracking(job_id, key)
 
         # Prevent memory leaks on C back-end, running jobs unaffected
         # https://drmaa-python.readthedocs.io/en/latest/drmaa.html
         self.session.deleteJobTemplate(jt)
         self.jobs_submitted += 1
+
+
+def _taskkey_to_dict(key: TaskInstanceKey) -> _TaskInstanceKeyDict:
+    return {
+        "dag_id": key.dag_id,
+        "task_id": key.task_id,
+        "run_id": key.run_id,
+        "try_number": key.try_number
+    }