airflow-unicore-integration 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

airflow_unicore_integration/executors/unicore_executor.py

@@ -10,16 +10,17 @@ tasks should be allowed to overwrite SITE, CREDENTIALS_*, UNICORE_CONN_ID and DE
 import time
 from typing import Any
 from typing import Dict
-from typing import List
 
 import pyunicore.client as uc_client
 from airflow.configuration import conf
 from airflow.executors.base_executor import BaseExecutor
 from airflow.executors.workloads import All
 from airflow.executors.workloads import ExecuteTask
+from airflow.models.taskinstance import TaskInstance
 from airflow.models.taskinstancekey import TaskInstanceKey
 from airflow.utils.state import TaskInstanceState
 from pyunicore import client
+from pyunicore.credentials import Credential
 from pyunicore.credentials import create_credential
 
 from ..util.job import JobDescriptionGenerator
@@ -39,37 +40,74 @@ STATE_MAPPINGS: Dict[uc_client.JobStatus, TaskInstanceState] = {
 
 class UnicoreExecutor(BaseExecutor):
 
-    EXECUTOR_CONFIG_UNICORE_CONN_KEY = (
-        "unicore_connection_id"  # alternative connection id for the Unicore connection to use
-    )
     EXECUTOR_CONFIG_UNICORE_SITE_KEY = "unicore_site"  # alternative Unicore site to run at, only required if different than connection default
     EXECUTOR_CONFIG_UNICORE_CREDENTIAL_KEY = "unicore_credential"  # alternative unicore credential to use for the job, only required if different than connection default
 
+    serve_logs = True
+
     def start(self):
         self.active_jobs: Dict[TaskInstanceKey, uc_client.Job] = {}
         # TODO get job description generator class and init params from config
         self.job_descr_generator: JobDescriptionGenerator = NaiveJobDescriptionGenerator()
 
+    def _handle_used_compute_time(self, task: TaskInstanceKey, job: uc_client.Job) -> None:
+        pass
+
     def sync(self) -> None:
         # iterate through task collection and update task/ job status - delete if needed
         for task, job in list(self.active_jobs.items()):
             state = STATE_MAPPINGS[job.status]
             if state == TaskInstanceState.FAILED:
                 self.fail(task)
-                self._forward_unicore_log(task, job)
                 self.active_jobs.pop(task)
+                self._handle_used_compute_time(task, job)
             elif state == TaskInstanceState.SUCCESS:
                 self.success(task)
-                self._forward_unicore_log(task, job)
                 self.active_jobs.pop(task)
-            elif state == TaskInstanceState.RUNNING:
+                self._handle_used_compute_time(task, job)
+            elif state == TaskInstanceState.QUEUED:
                 self.running_state(task, state)
 
         return super().sync()
 
-    def _forward_unicore_log(self, task: TaskInstanceKey, job: uc_client.Job) -> List[str]:
-        # TODO retrieve unicore logs from job directory and return
-        return []
+    def get_task_log(self, ti: TaskInstance, try_number: int) -> tuple[list[str], list[str]]:
+        # if task is still active, no neeed to check other jobs for match
+        job: uc_client.Job | None = None
+        if ti.key in self.active_jobs:
+            job = self.active_jobs[ti.key]
+        # if job is finished, need to search all jobs (that are still present on the target system)
+        else:
+            job = self._get_job_from_id(ti.key)
+
+        if not job:
+            return [], []
+
+        # TODO return filecontent for stdout and stderr from jobdirectory
+        logs = job.properties["log"]  # type: ignore
+        working_dir = job.working_dir
+        stdout = working_dir.stat("/stdout")
+        stderr = working_dir.stat("/stderr")
+
+        lines_out = stdout.raw().read().decode("utf-8").split("\n")  # type: ignore
+        lines_err = stderr.raw().read().decode("utf-8").split("\n")  # type: ignore
+
+        return logs + lines_out + lines_err, ["Test", "Where is this shown??"]
+
+    def _get_job_from_id(self, key: TaskInstanceKey) -> uc_client.Job | None:
+        # iterate over jobs and check if key matches job name
+        unicore_client = self._get_unicore_client()
+        # check 50 jobs at a time to prevent server from not returning some jobs
+        offset = 0
+        number = 50
+        job_name = self.job_descr_generator.get_job_name(key)
+        while True:
+            partial_job_list = unicore_client.get_jobs(offset=offset, num=number)
+            if len(partial_job_list) == 0:
+                return None
+            offset += len(partial_job_list)
+            for job in partial_job_list:
+                if job.properties["name"] == job_name:  # type: ignore
+                    return job
 
     def _get_unicore_client(self, executor_config: dict | None = {}):
         overwrite_unicore_site = executor_config.get(  # type: ignore
@@ -79,25 +117,29 @@ class UnicoreExecutor(BaseExecutor):
             UnicoreExecutor.EXECUTOR_CONFIG_UNICORE_CREDENTIAL_KEY, None
         )  # task can provide a different credential to use, else default from connection is used
         token = conf.get("unicore.executor", "AUTH_TOKEN", fallback="")
+        if overwrite_unicore_credential is not None:
+            token = overwrite_unicore_credential
+            self.log.debug("Using user provided token.")
         base_url = conf.get(
             "unicore.executor", "DEFAULT_URL", fallback="http://localhost:8080/DEMO-SITE/rest/core"
         )
-        credential = create_credential(token=token)
+        credential: Credential = create_credential(token=token)
         if overwrite_unicore_site is not None:
             base_url = overwrite_unicore_site
-        if overwrite_unicore_credential is not None:
-            credential = overwrite_unicore_credential
+            self.log.debug("Using user provided site.")
         if not base_url:
             raise TypeError()
+        self.log.debug(f"Using site: {base_url}")
+        self.log.debug(f"Using credential: {credential}")
         conn = client.Client(credential, base_url)
         return conn
 
     def _submit_job(self, workload: ExecuteTask):
-        uc_client = self._get_unicore_client(executor_config=workload.ti.executor_config)
+        unicore_client = self._get_unicore_client(executor_config=workload.ti.executor_config)
         job_descr = self._create_job_description(workload)
         self.log.info("Generated job description")
         self.log.debug(str(job_descr))
-        job = uc_client.new_job(job_descr)
+        job = unicore_client.new_job(job_descr)
         self.log.info("Submitted unicore job")
         self.active_jobs[workload.ti.key] = job
         return job

airflow_unicore_integration/hooks/unicore_hooks.py

@@ -28,6 +28,7 @@ class UnicoreHook(BaseHook):
     def __init__(self, uc_conn_id: str = default_conn_name) -> None:
         super().__init__()
         self.uc_conn_id = uc_conn_id
+        self.get_connection(self.uc_conn_id)
 
     @classmethod
     def get_connection_form_fields(cls):
@@ -38,9 +39,7 @@ class UnicoreHook(BaseHook):
         """Return custom UI field behaviour for UNICORE connection."""
         return {
             "hidden_fields": ["schema", "port", "extra"],
-            "relabeling": {
-                "login": "Username",
-            },
+            "relabeling": {"login": "Username", "host": "Site URL"},
             "placeholder": {"auth_token": "UNICORE auth token"},
         }
 
@@ -77,6 +76,11 @@ class UnicoreHook(BaseHook):
             credential = credentials.create_credential(token=auth_token)
         return credential
 
+    def get_token(self) -> str:
+        """Returns a token if present, None if not."""
+        params = self.get_connection(self.uc_conn_id)
+        return params.extra_dejson.get("auth_token", None)
+
     def get_base_url(self) -> str:
         """Return the base url of the connection."""
         params = self.get_connection(self.uc_conn_id)

{airflow_unicore_integration-0.2.0.dist-info → airflow_unicore_integration-0.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: airflow-unicore-integration
-Version: 0.2.0
+Version: 0.2.2
 Summary: Running Unicore Jobs from airflow DAGs.
 Author-email: Christian Böttcher <c.boettcher@fz-juelich.de>
 License-Expression: BSD-3-Clause

{airflow_unicore_integration-0.2.0.dist-info → airflow_unicore_integration-0.2.2.dist-info}/RECORD

@@ -1,18 +1,18 @@
 airflow_unicore_integration/__init__.py,sha256=Qy1mlyxe2Y-PPSn0LgIW0sT6BxFuwW4_LGsPBf-Wm4s,549
 airflow_unicore_integration/executors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 airflow_unicore_integration/executors/run_task_via_supervisor.py,sha256=3ErgPf-Oy3B4Di5yNXhhPkaojIJykvCxMZ9MlKSYPI8,2756
-airflow_unicore_integration/executors/unicore_executor.py,sha256=E1nOskWSBmC-ReLRvA8E3bY-G0lpxP403tazlBNhgFQ,5919
+airflow_unicore_integration/executors/unicore_executor.py,sha256=JaR8_503WMtI_emvcrzwJkZlm2q3c_e-3aRhP70y54k,7790
 airflow_unicore_integration/hooks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-airflow_unicore_integration/hooks/unicore_hooks.py,sha256=NEJ6sgQCD9AyZO9TA3vOR8RdijqaIX_FqY-Abz74Yys,3148
+airflow_unicore_integration/hooks/unicore_hooks.py,sha256=ot9j8sG--_eTmmfoY5TpJj8BPo8nfYvjOSmqHyyQhBM,3383
 airflow_unicore_integration/operators/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 airflow_unicore_integration/operators/container.py,sha256=D-2La59XZazXvBDVn-_wFIZ0IrCYeGy18miK8S5Usmc,693
 airflow_unicore_integration/operators/unicore_operators.py,sha256=PrcB-riakRUt8DHR9Hkk0ccOe-xk6NjJmNAmDpNgpN4,16868
 airflow_unicore_integration/policies/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 airflow_unicore_integration/util/job.py,sha256=vYyaNzGw2EQWQlDTGO6q4T1gZTHoB9-RIb5KYYd6pAc,10556
 airflow_unicore_integration/util/launch_script_content.py,sha256=42_aFpaCMmvFmmUxQDGcudkleX1YSK_yYWE8T41NOy0,2915
-airflow_unicore_integration-0.2.0.dist-info/licenses/LICENSE,sha256=hZ5ouAedeNr8ClHrQE-RLsgMsARcmv3kSZz7tE2BTJE,1526
-airflow_unicore_integration-0.2.0.dist-info/METADATA,sha256=-ifD7f7RX30qQs3OMzsXKeAvo_acGHLigthKu93-rU0,13696
-airflow_unicore_integration-0.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-airflow_unicore_integration-0.2.0.dist-info/entry_points.txt,sha256=PzEfCLYLSawjiYR-HNBzw8-YGfJxs1nPBULevgBQjoY,147
-airflow_unicore_integration-0.2.0.dist-info/top_level.txt,sha256=j45X-uIuOk3oL78iwlpHakMWtUkg__B7zUlJLwmZx6w,28
-airflow_unicore_integration-0.2.0.dist-info/RECORD,,
+airflow_unicore_integration-0.2.2.dist-info/licenses/LICENSE,sha256=hZ5ouAedeNr8ClHrQE-RLsgMsARcmv3kSZz7tE2BTJE,1526
+airflow_unicore_integration-0.2.2.dist-info/METADATA,sha256=FeW3Nt-3gyVI6gWTeWd_twhzUNEIVHqEJyPP6XMzvw8,13696
+airflow_unicore_integration-0.2.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+airflow_unicore_integration-0.2.2.dist-info/entry_points.txt,sha256=PzEfCLYLSawjiYR-HNBzw8-YGfJxs1nPBULevgBQjoY,147
+airflow_unicore_integration-0.2.2.dist-info/top_level.txt,sha256=j45X-uIuOk3oL78iwlpHakMWtUkg__B7zUlJLwmZx6w,28
+airflow_unicore_integration-0.2.2.dist-info/RECORD,,