airflow-unicore-integration 0.2.1__tar.gz → 0.2.3__tar.gz

{airflow_unicore_integration-0.2.1/src/airflow_unicore_integration.egg-info → airflow_unicore_integration-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: airflow-unicore-integration
-Version: 0.2.1
+Version: 0.2.3
 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.1 → airflow_unicore_integration-0.2.3}/pyproject.toml

@@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "airflow-unicore-integration"
-version = "0.2.1"
+version = "0.2.3"
 authors = [
   { name="Christian Böttcher", email="c.boettcher@fz-juelich.de" },
 ]

{airflow_unicore_integration-0.2.1 → airflow_unicore_integration-0.2.3}/src/airflow_unicore_integration/executors/unicore_executor.py

@@ -10,19 +10,19 @@ 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 ..hooks.unicore_hooks import UnicoreHook
 from ..util.job import JobDescriptionGenerator
 from ..util.job import NaiveJobDescriptionGenerator
 
@@ -40,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,33 +116,30 @@ class UnicoreExecutor(BaseExecutor):
         overwrite_unicore_credential = executor_config.get(  # type: ignore
             UnicoreExecutor.EXECUTOR_CONFIG_UNICORE_CREDENTIAL_KEY, None
         )  # task can provide a different credential to use, else default from connection is used
-        overwrite_conn_id = executor_config.get(  # type: ignore
-            UnicoreExecutor.EXECUTOR_CONFIG_UNICORE_CONN_KEY, None
-        )
         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)
-        if overwrite_conn_id is not None:
-            hook = UnicoreHook(overwrite_conn_id)
-            base_url = hook.get_base_url()
-            credential = hook.get_credential()
+        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-0.2.1 → airflow_unicore_integration-0.2.3}/src/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.1 → airflow_unicore_integration-0.2.3}/src/airflow_unicore_integration/util/job.py

@@ -26,6 +26,7 @@ class JobDescriptionGenerator:
     EXECUTOR_CONFIG_PRE_COMMANDS = "precommands"  # gets added to the unicore job description
     EXECUTOR_CONFIG_POST_COMMANDS = "postcommands"  # gets added to the unicore job descirption
     EXECUTOR_CONFIG_JOB_TYPE = "job_type"
+    EXECUTOR_CONFIG_LOGIN_NODE = "login_node"
     EXECUTOR_CONFIG_UNICORE_CONN_KEY = (
         "unicore_connection_id"  # alternative connection id for the Unicore connection to use
     )
@@ -61,6 +62,7 @@ class NaiveJobDescriptionGenerator(JobDescriptionGenerator):
         user_defined_python_env: str = workload.ti.executor_config.get(JobDescriptionGenerator.EXECUTOR_CONFIG_PYTHON_ENV_KEY, None)  # type: ignore
         user_added_post_commands: list[str] = executor_config.get(JobDescriptionGenerator.EXECUTOR_CONFIG_POST_COMMANDS, [])  # type: ignore
         user_defined_job_type: str = executor_config.get(JobDescriptionGenerator.EXECUTOR_CONFIG_JOB_TYPE, None)  # type: ignore
+        user_defined_login_node: str = executor_config.get(JobDescriptionGenerator.EXECUTOR_CONFIG_LOGIN_NODE, None)  # type: ignore
         # get local dag path from cmd and fix dag path in arguments
         dag_rel_path = str(workload.dag_rel_path)
         if dag_rel_path.startswith("DAG_FOLDER"):
@@ -76,6 +78,8 @@ class NaiveJobDescriptionGenerator(JobDescriptionGenerator):
         # set job type
         if user_defined_job_type:
             job_descr_dict["Job type"] = user_defined_job_type
+        if user_defined_login_node:
+            job_descr_dict["Login Node"] = user_defined_login_node
 
         # check which python virtualenv to use
         if user_defined_python_env:

{airflow_unicore_integration-0.2.1 → airflow_unicore_integration-0.2.3/src/airflow_unicore_integration.egg-info}/PKG-INFO

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