airflow-unicore-integration 0.1.5__tar.gz → 0.2.0__tar.gz

{airflow_unicore_integration-0.1.5/src/airflow_unicore_integration.egg-info → airflow_unicore_integration-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: airflow-unicore-integration
-Version: 0.1.5
+Version: 0.2.0
 Summary: Running Unicore Jobs from airflow DAGs.
 Author-email: Christian Böttcher <c.boettcher@fz-juelich.de>
 License-Expression: BSD-3-Clause
@@ -24,7 +24,7 @@ Unicore Airflow Integration
 
 |Generic badge|
 
-.. |Generic badge| image:: https://github.com/UNICORE-EU/airflow-unicore-integration/actions/workflows/publish-to-pypi.yml/badge.svg 
+.. |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/>`_.
@@ -34,6 +34,31 @@ 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 UnicoreExecutor
+-------------------------
+
+To use the UnicoreExecutor, this library needs to be installed in your airflow environment and then some configuration work needs to be done.
+
+How to configure these settigns is up to your deployment, as it uses the standard airflow configuration mechanism.
+In case of a helm deployemnt via the official helm chart, you will need to use environment variables, as all unicore related options are not present in the chart and will cause schema-validation to fail.
+
+All options fall under the [unicore.executor] section in airflow.cfg, or have the ``AIRFLOW__UNICORE_EXECUTOR__`` prefix as an environment variable.
+
+========================= ============================================ ===========================================================================================
+Option name               default                                      description
+========================= ============================================ ===========================================================================================
+EXECUTION_API_SERVER_URL  <The default from the airflow config>        The url to reach the airflow API server from the execution environment (e.g. compute nodes)
+AUTH_TOKEN                mandatory                                    The unicore auth token to use for job submission
+DEFAULT_URL               http://localhost:8080/DEMO-SITE/rest/core    The default unicore site to submit jobs to
+DEFAULT_ENV               mandatory                                    The default activation script for a functional airflow environment on the execution machine
+TMP_DIR                   /tmp                                         A temporary directory to store data such as GitDagBundles
+========================= ============================================ ===========================================================================================
+
+The default env is loaded via ``. default_env.sh``, and must enable an environment, where python is available in a suitable version, and the ``apache-airflow-task-sdk`` and ``apache-airflow-providers-git`` packages are available. All other dependencies depend on the dags to be run, but must already be included in the environment.
+
+A simple solution for this may be the "activate" script for a python venv. If the target systems requires additional commands to enable python (e.g. ``module load``), these may be added to the top of the activate script.
+
 ---------------------------
 Using the Unicore Operators
 ---------------------------

{airflow_unicore_integration-0.1.5 → airflow_unicore_integration-0.2.0}/README.rst

@@ -5,7 +5,7 @@ Unicore Airflow Integration
 
 |Generic badge|
 
-.. |Generic badge| image:: https://github.com/UNICORE-EU/airflow-unicore-integration/actions/workflows/publish-to-pypi.yml/badge.svg 
+.. |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/>`_.
@@ -15,6 +15,31 @@ 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 UnicoreExecutor
+-------------------------
+
+To use the UnicoreExecutor, this library needs to be installed in your airflow environment and then some configuration work needs to be done.
+
+How to configure these settigns is up to your deployment, as it uses the standard airflow configuration mechanism.
+In case of a helm deployemnt via the official helm chart, you will need to use environment variables, as all unicore related options are not present in the chart and will cause schema-validation to fail.
+
+All options fall under the [unicore.executor] section in airflow.cfg, or have the ``AIRFLOW__UNICORE_EXECUTOR__`` prefix as an environment variable.
+
+========================= ============================================ ===========================================================================================
+Option name               default                                      description
+========================= ============================================ ===========================================================================================
+EXECUTION_API_SERVER_URL  <The default from the airflow config>        The url to reach the airflow API server from the execution environment (e.g. compute nodes)
+AUTH_TOKEN                mandatory                                    The unicore auth token to use for job submission
+DEFAULT_URL               http://localhost:8080/DEMO-SITE/rest/core    The default unicore site to submit jobs to
+DEFAULT_ENV               mandatory                                    The default activation script for a functional airflow environment on the execution machine
+TMP_DIR                   /tmp                                         A temporary directory to store data such as GitDagBundles
+========================= ============================================ ===========================================================================================
+
+The default env is loaded via ``. default_env.sh``, and must enable an environment, where python is available in a suitable version, and the ``apache-airflow-task-sdk`` and ``apache-airflow-providers-git`` packages are available. All other dependencies depend on the dags to be run, but must already be included in the environment.
+
+A simple solution for this may be the "activate" script for a python venv. If the target systems requires additional commands to enable python (e.g. ``module load``), these may be added to the top of the activate script.
+
 ---------------------------
 Using the Unicore Operators
 ---------------------------

{airflow_unicore_integration-0.1.5 → airflow_unicore_integration-0.2.0}/pyproject.toml

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

{airflow_unicore_integration-0.1.5 → airflow_unicore_integration-0.2.0}/src/airflow_unicore_integration/__init__.py

@@ -9,4 +9,7 @@ def get_provider_info():
                 "hook-class-name": "airflow_unicore_integration.hooks.unicore_hooks.UnicoreHook",
             }
         ],
+        "executors": [
+            "airflow_unicore_integration.executors.unicore_executor.UnicoreExecutor",
+        ],
     }

{airflow_unicore_integration-0.1.5 → airflow_unicore_integration-0.2.0}/src/airflow_unicore_integration/executors/unicore_executor.py

@@ -19,8 +19,8 @@ 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 pyunicore import client
+from pyunicore.credentials import create_credential
 
 from ..util.job import JobDescriptionGenerator
 from ..util.job import NaiveJobDescriptionGenerator
@@ -39,9 +39,14 @@ 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
+
     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()
 
@@ -67,25 +72,25 @@ class UnicoreExecutor(BaseExecutor):
         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
+        overwrite_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
+        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
-        return hook.get_conn(
-            overwrite_base_url=unicore_site, overwrite_credential=unicore_credential
+        token = conf.get("unicore.executor", "AUTH_TOKEN", fallback="")
+        base_url = conf.get(
+            "unicore.executor", "DEFAULT_URL", fallback="http://localhost:8080/DEMO-SITE/rest/core"
         )
+        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
+        if not base_url:
+            raise TypeError()
+        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)

{airflow_unicore_integration-0.1.5 → airflow_unicore_integration-0.2.0}/src/airflow_unicore_integration/hooks/unicore_hooks.py

@@ -1,8 +1,14 @@
 from __future__ import annotations
 
-from airflow.hooks.base import BaseHook
+import logging
+from typing import Any
+
+from airflow.providers.common.compat.sdk import BaseHook
 from pyunicore import client
 from pyunicore import credentials
+from wtforms import StringField
+
+logger = logging.getLogger(__name__)
 
 
 class UnicoreHook(BaseHook):
@@ -23,25 +29,59 @@ class UnicoreHook(BaseHook):
         super().__init__()
         self.uc_conn_id = uc_conn_id
 
+    @classmethod
+    def get_connection_form_fields(cls):
+        return {"auth_token": StringField("Auth Token")}
+
+    @classmethod
+    def get_ui_field_behaviour(cls) -> dict[str, Any]:
+        """Return custom UI field behaviour for UNICORE connection."""
+        return {
+            "hidden_fields": ["schema", "port", "extra"],
+            "relabeling": {
+                "login": "Username",
+            },
+            "placeholder": {"auth_token": "UNICORE auth token"},
+        }
+
     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."""
-        self.log.debug(
+        logger.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)
+        auth_token = params.extra_dejson.get("auth_token", None)
+        if auth_token is not None:
+            credential = credentials.create_credential(token=auth_token)
         if overwrite_base_url is not None:
             base_url = overwrite_base_url
         if overwrite_credential is not None:
             credential = overwrite_credential
+        if not base_url:
+            raise TypeError()
         conn = client.Client(credential, base_url)
         return conn
 
+    def get_credential(self) -> credentials.Credential:
+        """Return the credential part of the connection as a Credential object."""
+        params = self.get_connection(self.uc_conn_id)
+        credential = credentials.UsernamePassword(params.login, params.password)
+        auth_token = params.extra_dejson.get("auth_token", None)
+        if auth_token is not None:
+            credential = credentials.create_credential(token=auth_token)
+        return credential
+
+    def get_base_url(self) -> str:
+        """Return the base url of the connection."""
+        params = self.get_connection(self.uc_conn_id)
+        return params.host
+
     def test_connection(self) -> tuple[bool, str]:
         """Test the connection by sending an access_info request"""
         conn = self.get_conn()

airflow_unicore_integration-0.2.0/src/airflow_unicore_integration/operators/container.py

@@ -0,0 +1,22 @@
+from typing import Dict
+from typing import List
+
+from airflow_unicore_integration.operators.unicore_operators import (
+    UnicoreGenericOperator,
+)
+
+
+class UnicoreContainerOperator(UnicoreGenericOperator):
+    def __init__(
+        self, name: str, docker_image_url: str, command: str, options: str | None = None, **kwargs
+    ):
+        params: Dict[str, str | List[str]] = {"COMMAND": command, "IMAGE_URL": docker_image_url}
+        if options is not None:
+            params["OPTIONS"] = options
+        super().__init__(
+            name=name,
+            application_name="CONTAINER",
+            application_version="1.0",
+            parameters=params,
+            **kwargs,
+        )

{airflow_unicore_integration-0.1.5 → airflow_unicore_integration-0.2.0}/src/airflow_unicore_integration/operators/unicore_operators.py

@@ -6,7 +6,7 @@ from typing import Sequence
 
 import pyunicore.client as uc_client
 import pyunicore.credentials as uc_credentials
-from airflow.models.baseoperator import BaseOperator
+from airflow.sdk.bases.operator import BaseOperator
 from airflow.utils.context import Context
 
 from airflow_unicore_integration.hooks import unicore_hooks
@@ -62,6 +62,7 @@ class UnicoreGenericOperator(BaseOperator):
         credential_username: str | None = None,
         credential_password: str | None = None,
         credential_token: str | None = None,
+        conn_id: str | None = None,
         **kwargs,
     ):
         """
@@ -114,6 +115,7 @@ class UnicoreGenericOperator(BaseOperator):
         self.credential_username = credential_username
         self.credential_password = credential_password
         self.credential_token = credential_token
+        self.conn_id = conn_id
 
         self.validate_job_description()
         logger.debug("created Unicore Job Task")
@@ -262,7 +264,7 @@ class UnicoreGenericOperator(BaseOperator):
 
     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()
+        client = self.get_uc_client(self.conn_id)
         job = client.new_job(job_description=self.get_job_description(), inputs=[])
         return job
 

{airflow_unicore_integration-0.1.5 → airflow_unicore_integration-0.2.0}/src/airflow_unicore_integration/util/job.py

@@ -1,10 +1,16 @@
-import os
+import json
+import logging
 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
+from airflow.providers.git.hooks.git import GitHook
+
+from .launch_script_content import LAUNCH_SCRIPT_CONTENT_STR
+
+logger = logging.getLogger(__name__)
 
 
 class JobDescriptionGenerator:
@@ -18,6 +24,8 @@ class JobDescriptionGenerator:
     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_POST_COMMANDS = "postcommands"  # gets added to the unicore job descirption
+    EXECUTOR_CONFIG_JOB_TYPE = "job_type"
     EXECUTOR_CONFIG_UNICORE_CONN_KEY = (
         "unicore_connection_id"  # alternative connection id for the Unicore connection to use
     )
@@ -27,12 +35,17 @@ class JobDescriptionGenerator:
     def create_job_description(self, workload: ExecuteTask) -> Dict[str, Any]:
         raise NotImplementedError()
 
+    def get_job_name(self, key: TaskInstanceKey) -> str:
+        return f"{key.dag_id} - {key.task_id} - {key.run_id} - {key.try_number}"
+
 
 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.
+    This class generates a naive unicore job, that expects there to be a working python env containing airflow and any other required dependencies on the executing system.
     """
 
+    GIT_DAG_BUNDLE_CLASSPATH = "airflow.providers.git.bundles.git.GitDagBundle"
+
     def create_job_description(self, workload: ExecuteTask) -> Dict[str, Any]:
         key: TaskInstanceKey = workload.ti.key
         executor_config = workload.ti.executor_config
@@ -46,6 +59,8 @@ class NaiveJobDescriptionGenerator(JobDescriptionGenerator):
         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
+        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
         # 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"):
@@ -56,56 +71,115 @@ class NaiveJobDescriptionGenerator(JobDescriptionGenerator):
         server = conf.get(
             "unicore.executor", "execution_api_server_url", fallback=default_execution_api_server
         )
+        logger.debug(f"Server is {server}")
+
+        # set job type
+        if user_defined_job_type:
+            job_descr_dict["Job type"] = user_defined_job_type
 
         # 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")
+        tmp_dir = conf.get("unicore.executor", "TMP_DIR", "/tmp")
         # prepare dag file to be uploaded via unicore
         # dag_file = open("/tmp/test")
         # 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/1d90bf38199d8c0adf47a79c8840c3e3ddf57462/run_task_via_supervisor.py",
+            # "From": "https://gist.githubusercontent.com/cboettcher/3f1101a1d1b67e7944d17c02ecd69930/raw/1d90bf38199d8c0adf47a79c8840c3e3ddf57462/run_task_via_supervisor.py",
+            "Data": LAUNCH_SCRIPT_CONTENT_STR,
         }
         # 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["Name"] = self.get_job_name(key)
         job_descr_dict["Executable"] = (
-            "python"  # TODO may require module load to be setup for some systems
+            f". airflow_config.env && . {python_env} && python run_task_via_supervisor.py --json-string '{workload.model_dump_json()}'"  # 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["Arguments"] = [
+        #    "-c",
+        #    "source airflow_config.env",
+        #    "source {python_env}/bin/activate",
+        #    "python",
+        #    "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__CORE__DAGS_FOLDER": "./",
             "AIRFLOW__LOGGING__LOGGING_LEVEL": "DEBUG",
             "AIRFLOW__CORE__EXECUTOR": "LocalExecutor,airflow_unicore_integration.executors.unicore_executor.UnicoreExecutor",
         }
 
         # build filecontent string for importing in the job | this is needed to avoid confusing nested quotes and trying to escape them properly when using unicore env vars directly
-        env_file_content: list[str] = [
-            f"export AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST='{os.environ.get("AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST", "")}'"
-        ]
+        env_file_content: list[str] = []
+
+        # transmit needed dag bundle information (and possibly files) to job directory
+        bundle_str = conf.get("dag.processor", "dag_bundle_config_list")
+        logger.debug(f"Dag Bundle config is: {bundle_str}")
+        bundle_dict = json.loads(bundle_str)
+        conn_id_to_transmit = None
+        bundle_type = None
 
-        # insert connection details that are provided via env vars to get bundles
-        for env_key in os.environ.keys():
-            if env_key.startswith("AIRFLOW_CONN_"):
-                env_file_content.append(f"export {env_key}='{os.environ[env_key]}'")
+        for bundle in bundle_dict:
+            if bundle["name"] == workload.bundle_info.name:
+                if bundle["classpath"] == NaiveJobDescriptionGenerator.GIT_DAG_BUNDLE_CLASSPATH:
+                    bundle_type = NaiveJobDescriptionGenerator.GIT_DAG_BUNDLE_CLASSPATH
+                    env_file_content.append(
+                        f"export AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST='[{json.dumps(bundle)}]'"
+                    )
+                    conn_id_to_transmit = bundle["kwargs"]["git_conn_id"]
+                    break
+                # TODO handle other bundle types
+
+        if bundle_type:
+            if (
+                bundle_type == NaiveJobDescriptionGenerator.GIT_DAG_BUNDLE_CLASSPATH
+                and conn_id_to_transmit
+            ):
+                git_hook = GitHook(conn_id_to_transmit)
+                git_remote_url = git_hook.repo_url
+                git_dir_prefix = f"{tmp_dir}/{workload.ti.dag_id}/{workload.ti.task_id}/{workload.ti.run_id}/{workload.ti.try_number}"
+                git_local_url = f"{git_dir_prefix}/dagmirror"
+                dag_bundle_path = f"{git_dir_prefix}/dagbundle"
+                # add precommand to clone repo on login node
+                git_precommand = f". {python_env} && mkdir -p {git_local_url} && mkdir -p {dag_bundle_path} && git clone {git_remote_url} {git_local_url}"
+                job_descr_dict["Environment"][
+                    "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_STORAGE_PATH"
+                ] = f"{dag_bundle_path}"
+                logger.info(f"git precommand is {git_precommand}")
+                user_added_pre_commands.append(git_precommand)
+                # add connection to local clone to env of job
+                airflow_conn_string = json.dumps(
+                    {"conn_type": "git", "host": f"file://{git_local_url}"}
+                )
+                env_file_content.append(
+                    f"export AIRFLOW_CONN_{str(conn_id_to_transmit).upper()}='{airflow_conn_string}'"
+                )
+                logger.info(f"connection is '{airflow_conn_string}'")
+                # add cleanup of local git repo to job description
+                git_cleanup_command = f"rm -r {git_dir_prefix}"
+                logger.info(f"git cleanup is {git_cleanup_command}")
+                user_added_post_commands.append(git_cleanup_command)
 
         airflow_env_import = {"To": "airflow_config.env", "Data": env_file_content}
 
-        user_added_pre_commands.append(
-            f"source airflow_config.env && source {python_env}/bin/activate"
-        )
-        job_descr_dict["User precommand"] = ";".join(user_added_pre_commands)
+        job_descr_dict["Imports"] = [worker_script_import, airflow_env_import]
+
+        if len(user_added_pre_commands) > 0:
+            precommand_import = {"To": "precommand.sh", "Data": user_added_pre_commands}
+            job_descr_dict["Imports"].append(precommand_import)
+            job_descr_dict["User precommand"] = "bash precommand.sh"
+        if len(user_added_post_commands) > 0:
+            postcommand_import = {"To": "postcommand.sh", "Data": user_added_post_commands}
+            job_descr_dict["Imports"].append(postcommand_import)
+            job_descr_dict["User postcommand"] = "bash postcommand.sh"
+
         job_descr_dict["RunUserPrecommandOnLoginNode"] = (
-            "false"  # precommand includes activating the python env, this should be done on compute node right before running the job
+            "true"  # precommand needs public internet access to clone dag repos
         )
-        job_descr_dict["Imports"] = [worker_script_import, airflow_env_import]
         # add user defined options to description
         if user_added_env:
             job_descr_dict["Environment"].update(user_added_env)

airflow_unicore_integration-0.2.0/src/airflow_unicore_integration/util/launch_script_content.py

@@ -0,0 +1,90 @@
+LAUNCH_SCRIPT_CONTENT_STR = '''# get some debug info if before anything can fail
+import sys
+print(sys.executable)
+
+from airflow.executors import workloads
+from airflow.sdk.execution_time.supervisor import supervise
+from airflow.configuration import conf
+from pydantic import TypeAdapter
+import argparse
+import structlog
+import sys
+log = structlog.get_logger(logger_name=__name__)
+
+"""
+Usage:
+
+python run_task_via_supervisor.py [--json-string <workload string> | --json-file <workload filepath>]
+
+"""
+
+
+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(f"Connecting to server:{server}" )
+
+    log.debug(f"Workload is: {workload}")
+    log.debug(f"Dag Bundle is: {workload.bundle_info}")
+
+    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
+        # unicore as job 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 Exception as e: # noqa: B902
+            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 Exception 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-0.1.5 → airflow_unicore_integration-0.2.0/src/airflow_unicore_integration.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: airflow-unicore-integration
-Version: 0.1.5
+Version: 0.2.0
 Summary: Running Unicore Jobs from airflow DAGs.
 Author-email: Christian Böttcher <c.boettcher@fz-juelich.de>
 License-Expression: BSD-3-Clause
@@ -24,7 +24,7 @@ Unicore Airflow Integration
 
 |Generic badge|
 
-.. |Generic badge| image:: https://github.com/UNICORE-EU/airflow-unicore-integration/actions/workflows/publish-to-pypi.yml/badge.svg 
+.. |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/>`_.
@@ -34,6 +34,31 @@ 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 UnicoreExecutor
+-------------------------
+
+To use the UnicoreExecutor, this library needs to be installed in your airflow environment and then some configuration work needs to be done.
+
+How to configure these settigns is up to your deployment, as it uses the standard airflow configuration mechanism.
+In case of a helm deployemnt via the official helm chart, you will need to use environment variables, as all unicore related options are not present in the chart and will cause schema-validation to fail.
+
+All options fall under the [unicore.executor] section in airflow.cfg, or have the ``AIRFLOW__UNICORE_EXECUTOR__`` prefix as an environment variable.
+
+========================= ============================================ ===========================================================================================
+Option name               default                                      description
+========================= ============================================ ===========================================================================================
+EXECUTION_API_SERVER_URL  <The default from the airflow config>        The url to reach the airflow API server from the execution environment (e.g. compute nodes)
+AUTH_TOKEN                mandatory                                    The unicore auth token to use for job submission
+DEFAULT_URL               http://localhost:8080/DEMO-SITE/rest/core    The default unicore site to submit jobs to
+DEFAULT_ENV               mandatory                                    The default activation script for a functional airflow environment on the execution machine
+TMP_DIR                   /tmp                                         A temporary directory to store data such as GitDagBundles
+========================= ============================================ ===========================================================================================
+
+The default env is loaded via ``. default_env.sh``, and must enable an environment, where python is available in a suitable version, and the ``apache-airflow-task-sdk`` and ``apache-airflow-providers-git`` packages are available. All other dependencies depend on the dags to be run, but must already be included in the environment.
+
+A simple solution for this may be the "activate" script for a python venv. If the target systems requires additional commands to enable python (e.g. ``module load``), these may be added to the top of the activate script.
+
 ---------------------------
 Using the Unicore Operators
 ---------------------------

{airflow_unicore_integration-0.1.5 → airflow_unicore_integration-0.2.0}/src/airflow_unicore_integration.egg-info/SOURCES.txt

@@ -14,6 +14,8 @@ src/airflow_unicore_integration/executors/unicore_executor.py
 src/airflow_unicore_integration/hooks/__init__.py
 src/airflow_unicore_integration/hooks/unicore_hooks.py
 src/airflow_unicore_integration/operators/__init__.py
+src/airflow_unicore_integration/operators/container.py
 src/airflow_unicore_integration/operators/unicore_operators.py
 src/airflow_unicore_integration/policies/__init__.py
-src/airflow_unicore_integration/util/job.py
+src/airflow_unicore_integration/util/job.py
+src/airflow_unicore_integration/util/launch_script_content.py