airflow-unicore-integration 0.1.8__py3-none-any.whl → 0.1.9__py3-none-any.whl

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,7 @@ 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_UNICORE_CONN_KEY = (
         "unicore_connection_id"  # alternative connection id for the Unicore connection to use
     )
@@ -33,6 +40,8 @@ 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.
     """
 
+    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 +55,7 @@ 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
         # 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,54 +66,97 @@ class NaiveJobDescriptionGenerator(JobDescriptionGenerator):
         server = conf.get(
             "unicore.executor", "execution_api_server_url", fallback=default_execution_api_server
         )
+        logger.debug(f"Server is {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")
+        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["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",
+            "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_STORAGE_PATH": f"{tmp_dir}/{workload.ti.id}/dagbundle",
         }
 
         # 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] = []
 
-        # 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]}'")
+        # 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
+
+        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_local_url = f"{tmp_dir}/{workload.ti.id}/dagmirror"
+                # add precommand to clone repo on ligon node
+                git_precommand = f". {python_env} && mkdir -p {tmp_dir}/{workload.ti.id}/dagmirror && mkdir -p {tmp_dir}/{workload.ti.id}/dagbundle && git clone {git_remote_url} {git_local_url}"
+                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 {tmp_dir}/{workload.ti.id}"
+                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 postcommand"] = ";".join(user_added_post_commands)
         job_descr_dict["User precommand"] = ";".join(user_added_pre_commands)
         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

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.8.dist-info → airflow_unicore_integration-0.1.9.dist-info}/METADATA

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

{airflow_unicore_integration-0.1.8.dist-info → airflow_unicore_integration-0.1.9.dist-info}/RECORD

@@ -7,10 +7,11 @@ airflow_unicore_integration/hooks/unicore_hooks.py,sha256=JjcjogWtN1xveagpkraQuY
 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=BSP-XuWL5-TRGACX__MavIcmPAA52gYLdwWL236cVz8,6703
-airflow_unicore_integration-0.1.8.dist-info/licenses/LICENSE,sha256=hZ5ouAedeNr8ClHrQE-RLsgMsARcmv3kSZz7tE2BTJE,1526
-airflow_unicore_integration-0.1.8.dist-info/METADATA,sha256=3vnVOs_enmL_-Q9D1S0XOhktDWICkFjHTmaxKPm16i4,11188
-airflow_unicore_integration-0.1.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-airflow_unicore_integration-0.1.8.dist-info/entry_points.txt,sha256=PzEfCLYLSawjiYR-HNBzw8-YGfJxs1nPBULevgBQjoY,147
-airflow_unicore_integration-0.1.8.dist-info/top_level.txt,sha256=j45X-uIuOk3oL78iwlpHakMWtUkg__B7zUlJLwmZx6w,28
-airflow_unicore_integration-0.1.8.dist-info/RECORD,,
+airflow_unicore_integration/util/job.py,sha256=Te3HNieipFQ3KrCzEf17wvSN5FnBgelW9Q3jMyr6Irg,9590
+airflow_unicore_integration/util/launch_script_content.py,sha256=42_aFpaCMmvFmmUxQDGcudkleX1YSK_yYWE8T41NOy0,2915
+airflow_unicore_integration-0.1.9.dist-info/licenses/LICENSE,sha256=hZ5ouAedeNr8ClHrQE-RLsgMsARcmv3kSZz7tE2BTJE,1526
+airflow_unicore_integration-0.1.9.dist-info/METADATA,sha256=b-AYw8ivzM6mpVLuMDJpl5oHsKFYaIXs3_vtb6cniWM,11188
+airflow_unicore_integration-0.1.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+airflow_unicore_integration-0.1.9.dist-info/entry_points.txt,sha256=PzEfCLYLSawjiYR-HNBzw8-YGfJxs1nPBULevgBQjoY,147
+airflow_unicore_integration-0.1.9.dist-info/top_level.txt,sha256=j45X-uIuOk3oL78iwlpHakMWtUkg__B7zUlJLwmZx6w,28
+airflow_unicore_integration-0.1.9.dist-info/RECORD,,