airflow-unicore-integration 0.0.5__py3-none-any.whl → 0.1.0__py3-none-any.whl

airflow_unicore_integration/__init__.py

@@ -3,8 +3,10 @@ def get_provider_info():
         "package-name": "airflow-unicore-integration",
         "name": "Unicore",
         "description": "Apache Airflow Unicore provider containing Operators and hooks.",
-        "connection-types": [{
-            "connection-type": "unicore",
-            "hook-class-name": "airflow_unicore_integration.hooks.unicore_hooks.UnicoreHook",
-        }],
+        "connection-types": [
+            {
+                "connection-type": "unicore",
+                "hook-class-name": "airflow_unicore_integration.hooks.unicore_hooks.UnicoreHook",
+            }
+        ],
     }

airflow_unicore_integration/executors/run_task_via_supervisor.py

@@ -0,0 +1,85 @@
+"""
+Usage:
+
+python run_task_via_supervisor.py [--json-string <workload string> | --json-file <workload filepath>]
+
+"""
+
+import argparse
+import sys
+
+import structlog
+from airflow.configuration import conf
+from airflow.executors import workloads
+from airflow.sdk.execution_time.supervisor import supervise
+from pydantic import TypeAdapter
+from pydantic_core._pydantic_core import ValidationError
+
+log = structlog.get_logger(logger_name=__name__)
+
+
+def execute_workload_locally(workload: workloads.All):
+    if not isinstance(workload, workloads.ExecuteTask):
+        raise ValueError(f"Executor does not know how to handle {type(workload)}")
+
+    base_url = conf.get("api", "base_url", fallback="/")
+    default_execution_api_server = f"{base_url.rstrip('/')}/execution/"
+    server = conf.get("core", "execution_api_server_url", fallback=default_execution_api_server)
+    log.info("Connecting to server:", server=server)
+
+    supervise(
+        # This is the "wrong" ti type, but it duck types the same. TODO: Create a protocol for this.
+        ti=workload.ti,  # type: ignore[arg-type]
+        dag_rel_path=workload.dag_rel_path,
+        bundle_info=workload.bundle_info,
+        token=workload.token,
+        server=server,
+        log_path=workload.log_path,
+        # Include the output of the task to stdout too, so that in process logs can be read from via the
+        # kubeapi as pod logs.
+        subprocess_logs_to_stdout=True,
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Execute a workload in a Containerised executor using the task SDK."
+    )
+
+    # Create a mutually exclusive group to ensure that only one of the flags is set
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument(
+        "--json-path",
+        help="Path to the input JSON file containing the execution workload payload.",
+        type=str,
+    )
+    group.add_argument(
+        "--json-string",
+        help="The JSON string itself containing the execution workload payload.",
+        type=str,
+    )
+    args = parser.parse_args()
+
+    decoder = TypeAdapter[workloads.All](workloads.All)
+
+    if args.json_path:
+        try:
+            with open(args.json_path) as file:
+                input_data = file.read()
+                workload = decoder.validate_json(input_data)
+        except OSError as e:
+            log.error("Failed to read file", error=str(e))
+            sys.exit(1)
+
+    elif args.json_string:
+        try:
+            workload = decoder.validate_json(args.json_string)
+        except ValidationError as e:
+            log.error("Failed to parse input JSON string", error=str(e))
+            sys.exit(1)
+
+    execute_workload_locally(workload)
+
+
+if __name__ == "__main__":
+    main()

airflow_unicore_integration/executors/unicore_executor.py

@@ -0,0 +1,123 @@
+"""
+to configure for executor:
+- Connection details for unicore: conn_id AIRFLOW__UNICORE_EXECUTOR__UNICORE_CONN_ID | should be defined, can be skipped if every task provides one
+- location (path) of python virtualenv prepared on hpc system | AIRFLOW__UNICORE_EXECUTOR__DEFAULT_ENV | should be defined, can be skipped if every task provides one
+
+tasks should be allowed to overwrite SITE, CREDENTIALS_*, UNICORE_CONN_ID and DEFAULT_ENV - i.e. everything but the database connection - credentials should be given as a uc_credential object via executor_config
+
+"""
+
+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.taskinstancekey import TaskInstanceKey
+from airflow.utils.state import TaskInstanceState
+
+from airflow_unicore_integration.hooks import unicore_hooks
+
+from ..util.job import JobDescriptionGenerator
+from ..util.job import NaiveJobDescriptionGenerator
+
+STATE_MAPPINGS: Dict[uc_client.JobStatus, TaskInstanceState] = {
+    uc_client.JobStatus.UNDEFINED: TaskInstanceState.FAILED,
+    uc_client.JobStatus.READY: TaskInstanceState.QUEUED,
+    uc_client.JobStatus.STAGINGIN: TaskInstanceState.QUEUED,
+    uc_client.JobStatus.QUEUED: TaskInstanceState.QUEUED,
+    uc_client.JobStatus.RUNNING: TaskInstanceState.RUNNING,
+    uc_client.JobStatus.STAGINGOUT: TaskInstanceState.RUNNING,
+    uc_client.JobStatus.SUCCESSFUL: TaskInstanceState.SUCCESS,
+    uc_client.JobStatus.FAILED: TaskInstanceState.FAILED,
+}
+
+
+class UnicoreExecutor(BaseExecutor):
+
+    def start(self):
+        self.active_jobs: Dict[TaskInstanceKey, uc_client.Job] = {}
+        self.uc_conn = unicore_hooks.UnicoreHook().get_conn()
+        # TODO get job description generator class and init params from config
+        self.job_descr_generator: JobDescriptionGenerator = NaiveJobDescriptionGenerator()
+
+    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)
+            elif state == TaskInstanceState.SUCCESS:
+                self.success(task)
+                self._forward_unicore_log(task, job)
+                self.active_jobs.pop(task)
+            elif state == TaskInstanceState.RUNNING:
+                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_unicore_client(self, executor_config: dict | None = {}):
+        # TODO fix this only temporary solution
+        return self.uc_conn
+        # END TODO fix this
+        # include client desires from executor_config
+        unicore_conn_id = executor_config.get(  # type: ignore
+            UnicoreExecutor.EXECUTOR_CONFIG_UNICORE_CONN_KEY,
+            conf.get("unicore.executor", "UNICORE_CONN_ID"),
+        )  # task can provide a different unicore connection to use, else airflow-wide default is used
+        self.log.info(f"Using base unicore connection with id '{unicore_conn_id}'")
+        hook = unicore_hooks.UnicoreHook(uc_conn_id=unicore_conn_id)
+        unicore_site = executor_config.get(  # type: ignore
+            UnicoreExecutor.EXECUTOR_CONFIG_UNICORE_SITE_KEY, None
+        )  # task can provide a different site to run at, else default from connetion is used
+        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
+        return hook.get_conn(
+            overwrite_base_url=unicore_site, overwrite_credential=unicore_credential
+        )
+
+    def _submit_job(self, workload: ExecuteTask):
+        uc_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)
+        self.log.info("Submitted unicore job")
+        self.active_jobs[workload.ti.key] = job
+        return job
+
+    def _create_job_description(self, workload: ExecuteTask) -> Dict[str, Any]:
+        return self.job_descr_generator.create_job_description(workload)
+
+    def queue_workload(self, workload: ExecuteTask | All, session):
+        if not isinstance(workload, ExecuteTask):
+            raise TypeError(f"Don't know how to queue workload of type {type(workload).__name__}")
+
+        # submit job to unicore and add to active_jobs dict for task state management
+        job = self._submit_job(workload)
+        self.active_jobs[workload.ti.key] = job
+
+    def end(self, heartbeat_interval=10) -> None:
+        # wait for current jobs to finish, dont start any new ones
+        while True:
+            self.sync()
+            if not self.active_jobs:
+                break
+            time.sleep(heartbeat_interval)
+
+    def terminate(self):
+        # terminate all jobs
+        for task, job in list(self.active_jobs.items()):
+            job.abort()
+        self.end()

airflow_unicore_integration/hooks/unicore_hooks.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-from  pyunicore import client, credentials
-
 from airflow.hooks.base import BaseHook
+from pyunicore import client
+from pyunicore import credentials
 
 
 class UnicoreHook(BaseHook):
@@ -22,28 +22,28 @@ class UnicoreHook(BaseHook):
     def __init__(self, uc_conn_id: str = default_conn_name) -> None:
         super().__init__()
         self.uc_conn_id = uc_conn_id
-        self.conn = None
 
-    def get_conn(self, overwrite_base_url: str | None = None, overwrite_credential : credentials.Credential | None = None) -> client.Client:
+    def get_conn(
+        self,
+        overwrite_base_url: str | None = None,
+        overwrite_credential: credentials.Credential | None = None,
+    ) -> client.Client:
         """Return a Unicore Client. base_url and credentials may be overwritten."""
-        if self.conn is None or overwrite_base_url is not None or overwrite_credential is not None: # if not generated, or overwrite attribute is set crete new
-            params = self.get_connection(self.uc_conn_id)
-            base_url = params.host
-            credential = credentials.UsernamePassword(params.login, params.password)
-            if overwrite_base_url is not None:
-                base_url = overwrite_base_url
-            if overwrite_credential is not None:
-                credential = overwrite_credential
-            self.conn = client.Client(credential, base_url)
-
-        return self.conn
-
+        self.log.debug(
+            f"Gettig connection with id '{self.uc_conn_id}' from secrets backend. Will be modified with user input for UNICORE."
+        )
+        params = self.get_connection(self.uc_conn_id)
+        base_url = params.host
+        credential = credentials.UsernamePassword(params.login, params.password)
+        if overwrite_base_url is not None:
+            base_url = overwrite_base_url
+        if overwrite_credential is not None:
+            credential = overwrite_credential
+        conn = client.Client(credential, base_url)
+        return conn
 
     def test_connection(self) -> tuple[bool, str]:
         """Test the connection by sending an access_info request"""
-        try:
-            conn = self.get_conn()
-            conn.access_info()
-            return True, "Connection successfully tested"
-        except Exception as e:
-            return False, str(e)
+        conn = self.get_conn()
+        conn.access_info()
+        return True, "Connection successfully tested"

airflow_unicore_integration/operators/unicore_operators.py

@@ -1,38 +1,73 @@
-from airflow.models.baseoperator import BaseOperator
-from airflow.decorators.base import DecoratedOperator, task_decorator_factory
-from typing import Any, List, Dict
-
-from airflow.utils.context import Context
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Sequence
 
 import pyunicore.client as uc_client
 import pyunicore.credentials as uc_credentials
-from airflow_unicore_integration.hooks import unicore_hooks
+from airflow.models.baseoperator import BaseOperator
+from airflow.utils.context import Context
 
-import logging
+from airflow_unicore_integration.hooks import unicore_hooks
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_SCRIPT_NAME = 'default_script_from_job_description'
-DEFAULT_BSS_FILE = 'default_bss_file_upload'
+DEFAULT_SCRIPT_NAME = "default_script_from_job_description"
+DEFAULT_BSS_FILE = "default_bss_file_upload"
+
 
 class JobDescriptionException(BaseException):
-    """ For errors in the unicore job description that would be generated by the unicore operators."""
+    """For errors in the unicore job description that would be generated by the unicore operators."""
+
     def __init__(self, *args: object) -> None:
         super().__init__(*args)
 
+
 class UnicoreGenericOperator(BaseOperator):
 
-    def __init__(self, name: str, application_name : str | None = None, application_version: str | None = None, executable: str | None = None, arguments: List[str] | None = None, 
-                 environment: List[str] | None = None, parameters: Dict[str,str | List[str]] | None = None, stdout: str | None = None, stderr: str | None = None, stdin: str | None = None, ignore_non_zero_exit_code: bool | None = None, 
-                 user_pre_command: str | None = None, run_user_pre_command_on_login_node: bool | None = None, user_pre_command_ignore_non_zero_exit_code: bool | None = None, user_post_command: str | None = None, 
-                 run_user_post_command_on_login_node: bool | None = None, user_post_command_ignore_non_zero_exit_code: bool | None = None, resources: Dict[str, str] | None = None, project: str | None = None, 
-                 imports: List[Dict[str,str | List[str]]] | None = None, exports: List[Dict[str,str | List[str]]] | None = None, have_client_stagein: bool | None = None, job_type: str | None = None, 
-                 login_node: str | None = None, bss_file: str | None = None, tags: List[str] | None = None, notification: str | None = None, user_email: str | None = None, xcom_output_files: List[str] = ["stdout", "stderr"], 
-                 base_url: str | None = None, credential: uc_credentials.Credential | None = None, credential_username: str | None = None, credential_password: str | None = None, credential_token: str | None = None, **kwargs):
+    def __init__(
+        self,
+        name: str,
+        application_name: str | None = None,
+        application_version: str | None = None,
+        executable: str | None = None,
+        arguments: List[str] | None = None,
+        environment: List[str] | None = None,
+        parameters: Dict[str, str | List[str]] | None = None,
+        stdout: str | None = None,
+        stderr: str | None = None,
+        stdin: str | None = None,
+        ignore_non_zero_exit_code: bool | None = None,
+        user_pre_command: str | None = None,
+        run_user_pre_command_on_login_node: bool | None = None,
+        user_pre_command_ignore_non_zero_exit_code: bool | None = None,
+        user_post_command: str | None = None,
+        run_user_post_command_on_login_node: bool | None = None,
+        user_post_command_ignore_non_zero_exit_code: bool | None = None,
+        resources: Dict[str, str] | None = None,
+        project: str | None = None,
+        imports: List[Dict[str, str | Sequence[str]]] | None = None,
+        exports: List[Dict[str, str | List[str]]] | None = None,
+        have_client_stagein: bool | None = None,
+        job_type: str | None = None,
+        login_node: str | None = None,
+        bss_file: str | None = None,
+        tags: List[str] | None = None,
+        notification: str | None = None,
+        user_email: str | None = None,
+        xcom_output_files: List[str] = ["stdout", "stderr"],
+        base_url: str | None = None,
+        credential: uc_credentials.Credential | None = None,
+        credential_username: str | None = None,
+        credential_password: str | None = None,
+        credential_token: str | None = None,
+        **kwargs,
+    ):
         """
-        Initialize a Unicore Job Operator. 
-        :param name: The name parameter defines both the airflow task name and the unicore job name. 
-        :param xcom_output_files: A list of filenames to be exported from the job directory via x_coms. By default stdout and stderr. 
+        Initialize a Unicore Job Operator.
+        :param name: The name parameter defines both the airflow task name and the unicore job name.
+        :param xcom_output_files: A list of filenames to be exported from the job directory via x_coms. By default stdout and stderr.
         :param base_url: The base URL of the UNICOREX server to be used for the unicore client. Overwrites the configured airflow conenction.
         :param credential: A unicore Credential to be used for the unicore client. Overwrites the configured user-pass in the airflow connection.
         :param credential_username: Username for the unicore client credentials. Overwrites the configured user in the airflow connection.
@@ -58,7 +93,9 @@ class UnicoreGenericOperator(BaseOperator):
         self.user_pre_command_ignore_non_zero_exit_code = user_pre_command_ignore_non_zero_exit_code
         self.user_post_command = user_post_command
         self.run_user_post_command_on_login_node = run_user_post_command_on_login_node
-        self.user_post_command_ignore_non_zero_exit_code = user_post_command_ignore_non_zero_exit_code
+        self.user_post_command_ignore_non_zero_exit_code = (
+            user_post_command_ignore_non_zero_exit_code
+        )
         self.resources = resources
         self.project = project
         self.imports = imports
@@ -85,54 +122,61 @@ class UnicoreGenericOperator(BaseOperator):
         """
         Does some basic validation and checks if a proper job description can be generated. Raises a JobDescriptionException if not.
         """
-        # check for some errors in the parameters for creating the unicore job 
+        # check for some errors in the parameters for creating the unicore job
 
         # first check if application or executable been set
-        if self.application_name is None and self.executable is None: # TODO check if executable is required if bss_file is given
+        if (
+            self.application_name is None and self.executable is None
+        ):  # TODO check if executable is required if bss_file is given
             raise JobDescriptionException
-        
+
         # if bss_file is set, we need an executable
         if self.bss_file is not None:
             if self.executable is None and self.application_name is not None:
                 raise JobDescriptionException
             # TODO validate BSS file?
             logger.info("using bss file")
-            
+
         if self.credential_token is not None and self.credential is None:
             logger.info("using provided oidc token")
             self.credential = uc_credentials.OIDCToken(token=self.credential_token)
 
-        if self.credential_username is not None and self.credential_password is not None and self.credential is None:
+        if (
+            self.credential_username is not None
+            and self.credential_password is not None
+            and self.credential is None
+        ):
             logger.info("using provied user/pass credentials")
-            self.credential = uc_credentials.UsernamePassword(username=self.credential_username, password=self.credential_password)
-        
-    
-    def get_job_description(self) -> dict[str,Any]:
+            self.credential = uc_credentials.UsernamePassword(
+                username=self.credential_username, password=self.credential_password
+            )
+
+    def get_job_description(self) -> dict[str, Any]:
         """Generates the job description to be submitted to the unicore server."""
         job_description_dict: Dict = {}
 
         # now add the various simple string attribute fragments to the list, when they are not None
         if self.name is not None:
             job_description_dict["Name"] = self.name
-        
+
         if self.application_name is not None:
             job_description_dict["ApplicationName"] = self.application_name
-        
+
         if self.application_version is not None:
             job_description_dict["ApplicationVersion"] = self.application_version
-        
+
         if self.executable is not None:
             job_description_dict["Executable"] = self.executable
 
         if self.arguments is not None:
             job_description_dict["Arguments"] = self.arguments
-        
+
         if self.environment is not None:
             job_description_dict["Environment"] = self.environment
 
         if self.parameters is not None:
             job_description_dict["Parameters"] = self.parameters
-        
+
         if self.stdout is not None:
             job_description_dict["Stdout"] = self.stdout
 
@@ -144,24 +188,32 @@ class UnicoreGenericOperator(BaseOperator):
 
         if self.ignore_non_zero_exit_code is not None:
             job_description_dict["IgnoreNonZeroExitCode"] = self.ignore_non_zero_exit_code
-        
+
         if self.user_pre_command is not None:
             job_description_dict["User precommand"] = self.user_pre_command
 
         if self.run_user_pre_command_on_login_node is not None:
-            job_description_dict["RunUserPrecommandOnLoginNode"] = self.run_user_pre_command_on_login_node
-        
+            job_description_dict["RunUserPrecommandOnLoginNode"] = (
+                self.run_user_pre_command_on_login_node
+            )
+
         if self.user_pre_command_ignore_non_zero_exit_code is not None:
-            job_description_dict["UserPrecommandIgnoreNonZeroExitCode"] = self.user_pre_command_ignore_non_zero_exit_code
+            job_description_dict["UserPrecommandIgnoreNonZeroExitCode"] = (
+                self.user_pre_command_ignore_non_zero_exit_code
+            )
 
         if self.user_post_command is not None:
             job_description_dict["User postcommand"] = self.user_post_command
 
         if self.run_user_post_command_on_login_node is not None:
-            job_description_dict["RunUserPostcommandOnLoginNode"] = self.run_user_post_command_on_login_node
+            job_description_dict["RunUserPostcommandOnLoginNode"] = (
+                self.run_user_post_command_on_login_node
+            )
 
         if self.user_post_command_ignore_non_zero_exit_code is not None:
-            job_description_dict["UserPostcommandIgnoreNonZeroExitCode"] = self.user_post_command_ignore_non_zero_exit_code
+            job_description_dict["UserPostcommandIgnoreNonZeroExitCode"] = (
+                self.user_post_command_ignore_non_zero_exit_code
+            )
 
         if self.resources is not None:
             job_description_dict["Resources"] = self.resources
@@ -171,12 +223,12 @@ class UnicoreGenericOperator(BaseOperator):
 
         if self.imports is not None:
             job_description_dict["Imports"] = self.imports
-        
+
         if self.exports is not None:
             job_description_dict["Exports"] = self.exports
 
         if self.have_client_stagein is not None:
-            job_description_dict["haveClientStageIn"] =self.have_client_stagein
+            job_description_dict["haveClientStageIn"] = self.have_client_stagein
 
         if self.job_type is not None:
             job_description_dict["Job type"] = self.job_type
@@ -197,7 +249,7 @@ class UnicoreGenericOperator(BaseOperator):
             job_description_dict["Tags"] = self.tags
 
         return job_description_dict
-        
+
     def get_uc_client(self, uc_conn_id: str | None = None) -> uc_client.Client:
         """Return a unicore client for submitting the job. Will use provided base_url and credentials if possible, else it will use the default unicore connection from airflow."""
         if self.base_url is not None and self.credential is not None:
@@ -207,7 +259,7 @@ class UnicoreGenericOperator(BaseOperator):
         else:
             hook = unicore_hooks.UnicoreHook(uc_conn_id=uc_conn_id)
         return hook.get_conn(overwrite_base_url=self.base_url, overwrite_credential=self.credential)
-    
+
     def execute_async(self, context: Context) -> Any:
         """Submits the job and returns the job object without waiting for it to finish."""
         client = self.get_uc_client()
@@ -216,35 +268,41 @@ class UnicoreGenericOperator(BaseOperator):
 
     def execute(self, context: Context) -> Any:
         """
-        Submits the job description to the unicore server. 
+        Submits the job description to the unicore server.
         After the job is finished (failed or successful), some basic output such as logs, status messages, job id or file content will be retreived and returned via x_coms or airflow logs.
         The exit code of this functions is the exit code of the unicore job, to ensure proper handling of failed or aborted jobs.
         """
-        import logging        
-        from pyunicore.client import JobStatus, Job
+        import logging
+
+        from pyunicore.client import Job
+        from pyunicore.client import JobStatus
+
         logger = logging.getLogger(__name__)
-        
-        job: Job = self.execute_async(context) # TODO depending on params this may spawn multiple jobs -> in those cases, e.g. output needs to be handled differently
+
+        job: Job = self.execute_async(
+            context
+        )  # TODO depending on params this may spawn multiple jobs -> in those cases, e.g. output needs to be handled differently
         logger.debug(f"Waiting for unicore job {job}")
-        job.poll() # wait for job to finish
+        job.poll()  # wait for job to finish
 
-        task_instance = context['task_instance']
+        task_instance = context["task_instance"]
 
-        
         task_instance.xcom_push(key="status_message", value=job.properties["statusMessage"])
         task_instance.xcom_push(key="log", value=job.properties["log"])
         for line in job.properties["log"]:
             logger.info(f"UNICORE LOGS: {line}")
-        
+
         if job.status is not JobStatus.SUCCESSFUL:
             from airflow.exceptions import AirflowFailException
-            logger.error(f"Unicore job not successful. Job state is {job.status}. Aborting this task.")
-            raise AirflowFailException
 
+            logger.error(
+                f"Unicore job not successful. Job state is {job.status}. Aborting this task."
+            )
+            raise AirflowFailException
 
         work_dir = job.working_dir
 
-        content = work_dir.contents()['content']
+        content = work_dir.contents()["content"]
         task_instance.xcom_push(key="workdir_content", value=content)
 
         for filename in content.keys():
@@ -257,14 +315,15 @@ class UnicoreGenericOperator(BaseOperator):
                 bss_submit_text = work_dir.stat(filename).raw().read().decode("utf-8")
                 task_instance.xcom_push(key="BSS_SUBMIT", value=bss_submit_text)
                 continue
-        
+
         from requests.exceptions import HTTPError
+
         for file in self.xcom_output_files:
             try:
                 logger.debug(f"Retreiving file {file} from workdir.")
                 remote_file = work_dir.stat(file)
                 content = remote_file.raw().read().decode("utf-8")
-                task_instance.xcom_push(key=file,value=content)
+                task_instance.xcom_push(key=file, value=content)
             except HTTPError as http_error:
                 logger.error(f"Error while retreiving file {file} from workdir.", http_error)
                 continue
@@ -274,51 +333,63 @@ class UnicoreGenericOperator(BaseOperator):
         exit_code = job.properties["exitCode"]
         return exit_code
 
+
 class UnicoreScriptOperator(UnicoreGenericOperator):
     """
     This operator uplaods and submits a script to the unicore server as a job. Behaviour and parameters are otherwise the same as the UnicoreBaseOpertor.
     """
+
     def __init__(self, name: str, script_content: str, **kwargs):
         """
         :param script_content: the content of the script to be executed as a string
         """
         super().__init__(name=name, executable=DEFAULT_SCRIPT_NAME, **kwargs)
-        lines = script_content.split('\n')
-        script_stagein = {
-            "To":   DEFAULT_SCRIPT_NAME,
-            "Data": lines
-            }
+        lines = script_content.split("\n")
+        script_stagein = {"To": DEFAULT_SCRIPT_NAME, "Data": lines}
         if self.imports is not None:
             self.imports.append(script_stagein)
         else:
             self.imports = [script_stagein]
 
+
 class UnicoreBSSOperator(UnicoreGenericOperator):
     """
     This operator uplaods and submits a BSS-script to the unicore server as a job. Behaviour and parameters are otherwise the same as the UnicoreBaseOpertor.
     """
+
     def __init__(self, name: str, bss_file_content: str, **kwargs):
         """
         :param bss_file_content: the content of the BSS-script to be executed as a string
         """
-        super().__init__(name=name, bss_file=DEFAULT_BSS_FILE, job_type="raw", xcom_output_files=[], **kwargs)
-        lines = bss_file_content.split('\n')
-        bss_stagein = {
-            "To":   DEFAULT_BSS_FILE,
-            "Data": lines
-            }
+        super().__init__(name=name, **kwargs)
+        lines = bss_file_content.split("\n")
+        bss_stagein = {"To": DEFAULT_BSS_FILE, "Data": lines}
         if self.imports is not None:
             self.imports.append(bss_stagein)
         else:
             self.imports = [bss_stagein]
 
+
 class UnicoreExecutableOperator(UnicoreGenericOperator):
     """A simplified operator for just executing an executable. Still supports all parameters fo the UnicoreBaseOperator."""
-    def __init__(self, name: str, executable: str, xcom_output_files : List[str] = ["stdout","stderr"], **kwargs) -> None:
-        super().__init__(name=name, executable=executable, xcom_output_files=xcom_output_files, **kwargs)
+
+    def __init__(
+        self,
+        name: str,
+        executable: str,
+        xcom_output_files: List[str] = ["stdout", "stderr"],
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            name=name,
+            executable=executable,
+            xcom_output_files=xcom_output_files,
+            **kwargs,
+        )
+
 
 class UnicoreDateOperator(UnicoreExecutableOperator):
     """An operator for executing the 'date' executable. Mostly for testing purposes. Still supports all parameters fo the UnicoreBaseOperator."""
-    def __init__(self, name: str, **kwargs) -> None:
-        super().__init__(name=name, executable='date',**kwargs)
 
+    def __init__(self, name: str, **kwargs) -> None:
+        super().__init__(name=name, executable="date", **kwargs)

airflow_unicore_integration/util/job.py

@@ -0,0 +1,101 @@
+from typing import Any
+from typing import Dict
+
+from airflow.configuration import conf
+from airflow.executors.workloads import ExecuteTask
+from airflow.models.taskinstancekey import TaskInstanceKey
+
+
+class JobDescriptionGenerator:
+    """
+    A generator class for generating unicore jhob descriptions that may supprot different kinds of systems and/ or environments.
+    """
+
+    EXECUTOR_CONFIG_PYTHON_ENV_KEY = "python_env"  # full path to a python virtualenv that includes airflow and all required libraries for the task (without the .../bin/activate part)
+    EXECUTOR_CONFIG_RESOURCES = "Resources"  # gets added to the unicore job description
+    EXECUTOR_CONFIG_ENVIRONMENT = "Environment"  # gets added to the unicore job description
+    EXECUTOR_CONFIG_PARAMETERS = "Parameters"  # gets added to the unicore job description
+    EXECUTOR_CONFIG_PROJECT = "Project"  # gets added to the unicore job description
+    EXECUTOR_CONFIG_PRE_COMMANDS = "precommands"  # gets added to the unicore job description
+    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
+
+    def create_job_description(self, workload: ExecuteTask) -> Dict[str, Any]:
+        raise NotImplementedError()
+
+
+class NaiveJobDescriptionGenerator(JobDescriptionGenerator):
+    """
+    This class generates a naive unicore job, that expects there to be a working python env containign airflow and any other required dependencies on the executing system.
+    """
+
+    def create_job_description(self, workload: ExecuteTask) -> Dict[str, Any]:
+        key: TaskInstanceKey = workload.ti.key
+        executor_config = workload.ti.executor_config
+        if not executor_config:
+            executor_config = {}
+        job_descr_dict: Dict[str, Any] = {}
+        # get user config from executor_config
+        user_added_env: Dict[str, str] = executor_config.get(JobDescriptionGenerator.EXECUTOR_CONFIG_ENVIRONMENT, None)  # type: ignore
+        user_added_params: Dict[str, str] = executor_config.get(JobDescriptionGenerator.EXECUTOR_CONFIG_PARAMETERS, None)  # type: ignore
+        user_added_project: str = executor_config.get(JobDescriptionGenerator.EXECUTOR_CONFIG_PROJECT, None)  # type: ignore
+        user_added_resources: Dict[str, str] = executor_config.get(JobDescriptionGenerator.EXECUTOR_CONFIG_RESOURCES, None)  # type: ignore
+        user_added_pre_commands: list[str] = executor_config.get(JobDescriptionGenerator.EXECUTOR_CONFIG_PRE_COMMANDS, [])  # type: ignore
+        user_defined_python_env: str = workload.ti.executor_config.get(JobDescriptionGenerator.EXECUTOR_CONFIG_PYTHON_ENV_KEY, 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"):
+            dag_rel_path = dag_rel_path[10:]
+        local_dag_path = conf.get("core", "DAGS_FOLDER") + "/" + dag_rel_path
+        base_url = conf.get("api", "base_url", fallback="/")
+        default_execution_api_server = f"{base_url.rstrip('/')}/execution/"
+        server = conf.get("core", "execution_api_server_url", fallback=default_execution_api_server)
+
+        # check which python virtualenv to use
+        if user_defined_python_env:
+            python_env = user_defined_python_env
+        else:
+            python_env = conf.get("unicore.executor", "DEFAULT_ENV")
+        # prepare dag file to be uploaded via unicore
+        dag_file = open(local_dag_path)
+        dag_content = dag_file.readlines()
+        dag_import = {"To": dag_rel_path, "Data": dag_content}
+        worker_script_import = {
+            "To": "run_task_via_supervisor.py",
+            "From": "https://gist.githubusercontent.com/cboettcher/3f1101a1d1b67e7944d17c02ecd69930/raw/6da9ec16ba598ddda9cf288900498fab5e226788/run_task_via_supervisor.py",
+        }
+        # start filling the actual job description
+        job_descr_dict["Name"] = f"{key.dag_id} - {key.task_id} - {key.run_id} - {key.try_number}"
+        job_descr_dict["Executable"] = (
+            "python"  # TODO may require module load to be setup for some systems
+        )
+        job_descr_dict["Arguments"] = [
+            "run_task_via_supervisor.py",
+            f"--json-string '{workload.model_dump_json()}'",
+        ]
+        job_descr_dict["Environment"] = {
+            "AIRFLOW__CORE__EXECUTION_API_SERVER_URL": server,
+            "AIRFLOW__CORE__DAGS_FOLDER": "./",
+            "AIRFLOW__LOGGING__LOGGING_LEVEL": "DEBUG",
+            "AIRFLOW__CORE__EXECUTOR": "LocalExecutor,airflow_unicore_integration.executors.unicore_executor.UnicoreExecutor",
+        }
+        user_added_pre_commands.append(f"source {python_env}/bin/activate")
+        job_descr_dict["User precommand"] = ";".join(user_added_pre_commands)
+        job_descr_dict["RunUserPrecommandOnLoginNode"] = (
+            "false"  # precommand is activating the python env, this can also be done on compute node right before running the job
+        )
+        job_descr_dict["Imports"] = [dag_import, worker_script_import]
+        # add user defined options to description
+        if user_added_env:
+            job_descr_dict["Environment"].update(user_added_env)
+        if user_added_params:
+            job_descr_dict["Parameters"] = user_added_params
+        if user_added_project:
+            job_descr_dict["Project"] = user_added_project
+        if user_added_resources:
+            job_descr_dict["Resources"] = user_added_resources
+
+        return job_descr_dict

{airflow_unicore_integration-0.0.5.dist-info → airflow_unicore_integration-0.1.0.dist-info}/METADATA

@@ -1,26 +1,39 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: airflow-unicore-integration
-Version: 0.0.5
+Version: 0.1.0
 Summary: Running Unicore Jobs from airflow DAGs.
 Author-email: Christian Böttcher <c.boettcher@fz-juelich.de>
+License-Expression: BSD-3-Clause
 Project-URL: Homepage, https://github.com/UNICORE-EU/airflow-unicore-integration
 Project-URL: Issues, https://github.com/UNICORE-EU/airflow-unicore-integration/issues
 Classifier: Development Status :: 4 - Beta
 Classifier: Framework :: Apache Airflow :: Provider
 Classifier: Programming Language :: Python :: 3
-Classifier: License :: OSI Approved :: BSD License
 Classifier: Operating System :: OS Independent
-Requires-Python: >=3.10
+Requires-Python: >=3.9
 Description-Content-Type: text/x-rst
 License-File: LICENSE
-Requires-Dist: pyunicore >=1.0.0
-Requires-Dist: apache-airflow ==2.8.4
+Requires-Dist: pyunicore>=1.0.0
+Requires-Dist: apache-airflow>=3.0.0
+Dynamic: license-file
 
 ===========================
 Unicore Airflow Integration
 ===========================
 
 
+|Generic badge|
+
+.. |Generic badge| image:: https://github.com/UNICORE-EU/airflow-unicore-integration/actions/workflows/publish-to-pypi.yml/badge.svg 
+   :target: https://github.com/UNICORE-EU/airflow-unicore-integration/actions/workflows/publish-to-pypi.yml
+
+This project integrates `UNICORE <https://github.com/UNICORE-EU>`_ and `Apache Airflow <https://airflow.apache.org/>`_.
+UNICORE is a software suite that, among other functions, provides seamless access to high-performance compute and data resources.
+Airflow is a platform to programmatically author, schedule and monitor workflows.
+
+In the current state, this projects provides a set of airflow `operators <https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/operators.html>`_, which can be used as part of airflow workflows to submit jobs to Unicore.
+The UnicoreExecutor only offers experimental support for airflow 3 so far. Further support is currently being worked on.
+
 ---------------------------
 Using the Unicore Operators
 ---------------------------
@@ -47,7 +60,7 @@ exports                 List(exports)           None
 For imports and exports go `here <https://unicore-docs.readthedocs.io/en/latest/user-docs/rest-api/job-description/index.html#importing-files-into-the-job-workspace>`_ for details.
 
 
-The ``UnicoreGenericOperator`` supports the following additional parameters: 
+The ``UnicoreGenericOperator`` supports the following additional parameters:
 
 ======================= ======================= =========================================== ====================
 parameter name          type                    default                                     description
@@ -98,7 +111,7 @@ Behaviour on Errors and Success
 The Unicore Operators do not do a lot of error and exception handling, and mostly just forward any problems to be handled by airflow.
 All of the Unicore logic is handled by the `pyunicore library <https://github.com/HumanBrainProject/pyunicore>`_.
 
-While some validation of the resulting Unicore job description is done automatically, it may still be possible to build an invalid job description with the operators. 
+While some validation of the resulting Unicore job description is done automatically, it may still be possible to build an invalid job description with the operators.
 This may lead to a submission failure with Unicore. In this case, an exception is thrown to be handled by airflow.
 
 
@@ -128,10 +141,10 @@ There are some example DAGs in this repository under ``project-dir/dags``.
 - ``unicore-test-2.py`` has some basic examples for the generic operator.
 - ``unicore-test-3.py`` also includes script-operator examples.
 - ``unicore-test-4.py`` has some examples with more arguments.
-- ``unicore-test-bss.py`` shows how bss submission can be done (very simple example). 
-- ``unicore-test-credentials.py`` demonstrates that not only the credentials from the airflow connections backend can be used, but they can also be provided in the constructor of the o`perator.
-- ``unicore-test-import-export.py`` gives shprt examples for the imports and exports usage.
- 
+- ``unicore-test-bss.py`` shows how bss submission can be done (very simple example).
+- ``unicore-test-credentials.py`` demonstrates that not only the credentials from the airflow connections backend can be used, but they can also be provided in the constructor of the operator.
+- ``unicore-test-import-export.py`` gives short examples for the imports and exports usage.
+
 
 -----------------
 Setup testing env

airflow_unicore_integration-0.1.0.dist-info/RECORD

@@ -0,0 +1,16 @@
+airflow_unicore_integration/__init__.py,sha256=sibcOlQKBCUKfTbZlAl0L7r1d2nxf0nGAQys612vzeA,429
+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=f2PKj66tZlWb39xF5YacEBoZ2-eZ_JR9pnKQwNHAybc,5619
+airflow_unicore_integration/hooks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+airflow_unicore_integration/hooks/unicore_hooks.py,sha256=JjcjogWtN1xveagpkraQuYOdXjkp2lSnEdQc0waqhU4,1662
+airflow_unicore_integration/operators/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+airflow_unicore_integration/operators/unicore_operators.py,sha256=xYX1t_QzjpqHZDoEu6jJNNXVmBEnjIlApFWvgYoYoB0,16790
+airflow_unicore_integration/policies/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+airflow_unicore_integration/util/job.py,sha256=AqbtV0rOCzvpW_7eqwMuMvUzIn0BDbuZVts-ju49zJE,5868
+airflow_unicore_integration-0.1.0.dist-info/licenses/LICENSE,sha256=hZ5ouAedeNr8ClHrQE-RLsgMsARcmv3kSZz7tE2BTJE,1526
+airflow_unicore_integration-0.1.0.dist-info/METADATA,sha256=K5ks_88YYhpL8L8s23tNH7l-v0BMQF_MdPydk11NntA,11188
+airflow_unicore_integration-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+airflow_unicore_integration-0.1.0.dist-info/entry_points.txt,sha256=PzEfCLYLSawjiYR-HNBzw8-YGfJxs1nPBULevgBQjoY,147
+airflow_unicore_integration-0.1.0.dist-info/top_level.txt,sha256=j45X-uIuOk3oL78iwlpHakMWtUkg__B7zUlJLwmZx6w,28
+airflow_unicore_integration-0.1.0.dist-info/RECORD,,

{airflow_unicore_integration-0.0.5.dist-info → airflow_unicore_integration-0.1.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (70.1.1)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

{airflow_unicore_integration-0.0.5.dist-info → airflow_unicore_integration-0.1.0.dist-info}/entry_points.txt

@@ -1,2 +1,5 @@
+[airflow.policy]
+_ = airflow_unicore_integration.policies
+
 [apache_airflow_provider]
 provider_info = airflow_unicore_integration:get_provider_info

{airflow_unicore_integration-0.0.5.dist-info → airflow_unicore_integration-0.1.0.dist-info/licenses}/LICENSE

@@ -26,4 +26,4 @@ DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

airflow_unicore_integration-0.0.5.dist-info/RECORD

@@ -1,11 +0,0 @@
-airflow_unicore_integration/__init__.py,sha256=bh8G7mutfVy8umYxhoMZLBLufNxZCp44RxekRw2y4cw,395
-airflow_unicore_integration/hooks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-airflow_unicore_integration/hooks/unicore_hooks.py,sha256=3sXaVYNhd6w-DdWS0VmtM0e037lUcRvUOuNRWz_bCYk,1768
-airflow_unicore_integration/operators/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-airflow_unicore_integration/operators/unicore_operators.py,sha256=1o9hhEJTcrSYuQ4AexWLXodylZwNerP9zpqgakJ7ZIk,16445
-airflow_unicore_integration-0.0.5.dist-info/LICENSE,sha256=FgrHcXlmWUhvWaEmE1-OJIR_mJIwlSBBlvPy2PLh83I,1525
-airflow_unicore_integration-0.0.5.dist-info/METADATA,sha256=kvCgHJkbdYSZxAzIwNmRLwksFKt6x4Io6UNrdj7E0Bo,10242
-airflow_unicore_integration-0.0.5.dist-info/WHEEL,sha256=mguMlWGMX-VHnMpKOjjQidIo1ssRlCFu4a4mBpz1s2M,91
-airflow_unicore_integration-0.0.5.dist-info/entry_points.txt,sha256=bmk1b2a4vYfZ2P4fbcPVSOLEThgKHSHftgjR82XiG6A,88
-airflow_unicore_integration-0.0.5.dist-info/top_level.txt,sha256=j45X-uIuOk3oL78iwlpHakMWtUkg__B7zUlJLwmZx6w,28
-airflow_unicore_integration-0.0.5.dist-info/RECORD,,