airflow-unicore-integration 0.0.4__py3-none-any.whl

airflow_unicore_integration/__init__.py

@@ -0,0 +1,10 @@
+def get_provider_info():
+    return {
+        "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",
+        }],
+    }

airflow_unicore_integration/hooks/unicore_hooks.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from  pyunicore import client, credentials
+
+from airflow.hooks.base import BaseHook
+
+
+class UnicoreHook(BaseHook):
+    """
+    Interact with Unicore.
+
+    Creates Unicore Clients from airflow connections.
+
+    :param uc_conn_id: The unicore connection id - default: uc_default
+    """
+
+    conn_name_attr = "uc_conn_id"
+    default_conn_name = "uc_default"
+    conn_type = "unicore"
+    hook_name = "Unicore"
+
+    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:
+        """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
+
+
+    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)

airflow_unicore_integration/operators/unicore_operators.py

@@ -0,0 +1,324 @@
+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 pyunicore.client as uc_client
+import pyunicore.credentials as uc_credentials
+from airflow_unicore_integration.hooks import unicore_hooks
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+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."""
+    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):
+        """
+        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.
+        :param credential_password: Password the the unicore client credentials. Overwrites the configured password in the airflow connection.
+        :param credential_token: An OIDC token to be used by the unicore client. Overwrites the configured credentials in teh airflow connection.
+
+        All other parameters are parameters for the Unicore Job Description as defined [here](https://unicore-docs.readthedocs.io/en/latest/user-docs/rest-api/job-description/index.html#overview).
+        """
+        super().__init__(**kwargs)
+        self.name = name
+        self.application_name = application_name
+        self.application_version = application_version
+        self.executable = executable
+        self.arguments = arguments
+        self.environment = environment
+        self.parameters = parameters
+        self.stdout = stdout
+        self.stderr = stderr
+        self.stdin = stdin
+        self.ignore_non_zero_exit_code = ignore_non_zero_exit_code
+        self.user_pre_command = user_pre_command
+        self.run_user_pre_command_on_login_node = run_user_pre_command_on_login_node
+        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.resources = resources
+        self.project = project
+        self.imports = imports
+        self.exports = exports
+        self.have_client_stagein = have_client_stagein
+        self.job_type = job_type
+        self.login_node = login_node
+        self.bss_file = bss_file
+        self.tags = tags
+        self.notification = notification
+        self.user_email = user_email
+        self.xcom_output_files = xcom_output_files
+
+        self.base_url = base_url
+        self.credential = credential
+        self.credential_username = credential_username
+        self.credential_password = credential_password
+        self.credential_token = credential_token
+
+        self.validate_job_description()
+        logger.debug("created Unicore Job Task")
+
+    def validate_job_description(self):
+        """
+        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 
+
+        # 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
+            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:
+            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]:
+        """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
+
+        if self.stderr is not None:
+            job_description_dict["Stderr"] = self.stderr
+
+        if self.stdin is not None:
+            job_description_dict["Stdin"] = self.stdin
+
+        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
+        
+        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
+
+        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
+
+        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
+
+        if self.resources is not None:
+            job_description_dict["Resources"] = self.resources
+
+        if self.project is not None:
+            job_description_dict["Project"] = self.project
+
+        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
+
+        if self.job_type is not None:
+            job_description_dict["Job type"] = self.job_type
+
+        if self.login_node is not None:
+            job_description_dict["Login node"] = self.login_node
+
+        if self.bss_file is not None:
+            job_description_dict["BSS file"] = self.bss_file
+
+        if self.notification is not None:
+            job_description_dict["Notification"] = self.notification
+
+        if self.user_email is not None:
+            job_description_dict["User email"] = self.user_email
+
+        if self.tags is not None:
+            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:
+            return uc_client.Client(self.credential, self.base_url)
+        if uc_conn_id is None:
+            hook = unicore_hooks.UnicoreHook()
+        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()
+        job = client.new_job(job_description=self.get_job_description(), inputs=[])
+        return job
+
+    def execute(self, context: Context) -> Any:
+        """
+        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
+        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
+        logger.debug(f"Waiting for unicore job {job}")
+        job.poll() # wait for job to finish
+
+        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
+
+
+        work_dir = job.working_dir
+
+        content = work_dir.contents()['content']
+        task_instance.xcom_push(key="workdir_content", value=content)
+
+        for filename in content.keys():
+            if "/UNICORE_Job_" in filename:
+                task_instance.xcom_push(key="Unicore Job ID", value=filename[13:])
+                jobt_text = work_dir.stat(filename).raw().read().decode("utf-8")
+                task_instance.xcom_push(key="UNICORE Job", value=jobt_text)
+                continue
+            if "bss_submit_" in filename:
+                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)
+            except HTTPError as http_error:
+                logger.error(f"Error while retreiving file {file} from workdir.", http_error)
+                continue
+            except UnicodeDecodeError as unicore_error:
+                logger.error(f"Error while decoding file {file}.", unicore_error)
+
+        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
+            }
+        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
+            }
+        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)
+
+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)
+

airflow_unicore_integration-0.0.4.dist-info/METADATA

@@ -0,0 +1,17 @@
+Metadata-Version: 2.1
+Name: airflow-unicore-integration
+Version: 0.0.4
+Summary: Running Unicore Jobs from airflow DAGs.
+Author-email: Christian Böttcher <c.boettcher@fz-juelich.de>
+Project-URL: Homepage, https://gitlab.jsc.fz-juelich.de/boettcher1/airflow_unicore_integration
+Project-URL: Issues, https://gitlab.jsc.fz-juelich.de/boettcher1/airflow_unicore_integration/-/issues
+Classifier: Development Status :: 4 - Beta
+Classifier: Framework :: Apache Airflow :: Provider
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: pyunicore >=1.0.0
+Requires-Dist: apache-airflow ==2.8.4
+

airflow_unicore_integration-0.0.4.dist-info/RECORD

@@ -0,0 +1,10 @@
+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.4.dist-info/METADATA,sha256=R6KzqZ1Bng1AjIqSzsjJDK-xCAQdKm92CiymwsL-Hk0,756
+airflow_unicore_integration-0.0.4.dist-info/WHEEL,sha256=mguMlWGMX-VHnMpKOjjQidIo1ssRlCFu4a4mBpz1s2M,91
+airflow_unicore_integration-0.0.4.dist-info/entry_points.txt,sha256=bmk1b2a4vYfZ2P4fbcPVSOLEThgKHSHftgjR82XiG6A,88
+airflow_unicore_integration-0.0.4.dist-info/top_level.txt,sha256=j45X-uIuOk3oL78iwlpHakMWtUkg__B7zUlJLwmZx6w,28
+airflow_unicore_integration-0.0.4.dist-info/RECORD,,

airflow_unicore_integration-0.0.4.dist-info/WHEEL

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

airflow_unicore_integration-0.0.4.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[apache_airflow_provider]
+provider_info = airflow_unicore_integration:get_provider_info

airflow_unicore_integration-0.0.4.dist-info/top_level.txt

@@ -0,0 +1 @@
+airflow_unicore_integration