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

airflow_unicore_integration/util/job.py

@@ -34,10 +34,13 @@ 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"
@@ -84,7 +87,7 @@ class NaiveJobDescriptionGenerator(JobDescriptionGenerator):
             "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"] = (
             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
         )
@@ -102,7 +105,6 @@ class NaiveJobDescriptionGenerator(JobDescriptionGenerator):
             # "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
@@ -133,9 +135,14 @@ class NaiveJobDescriptionGenerator(JobDescriptionGenerator):
             ):
                 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}"
+                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
@@ -147,18 +154,26 @@ class NaiveJobDescriptionGenerator(JobDescriptionGenerator):
                 )
                 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}"
+                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}
 
-        job_descr_dict["User postcommand"] = ";".join(user_added_post_commands)
-        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"] = (
             "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.1.9.dist-info → airflow_unicore_integration-0.1.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: airflow-unicore-integration
-Version: 0.1.9
+Version: 0.1.11
 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.9.dist-info → airflow_unicore_integration-0.1.11.dist-info}/RECORD

@@ -7,11 +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=Te3HNieipFQ3KrCzEf17wvSN5FnBgelW9Q3jMyr6Irg,9590
+airflow_unicore_integration/util/job.py,sha256=IsVjDiq3Jj-Eky0mET3awC3yBloHR9Duvm1KFiAPKpU,10264
 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,,
+airflow_unicore_integration-0.1.11.dist-info/licenses/LICENSE,sha256=hZ5ouAedeNr8ClHrQE-RLsgMsARcmv3kSZz7tE2BTJE,1526
+airflow_unicore_integration-0.1.11.dist-info/METADATA,sha256=a6vLdzCg5UYaxOIsKCPfxQPDZ-sBg6rwJQWCVRoupWM,11189
+airflow_unicore_integration-0.1.11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+airflow_unicore_integration-0.1.11.dist-info/entry_points.txt,sha256=PzEfCLYLSawjiYR-HNBzw8-YGfJxs1nPBULevgBQjoY,147
+airflow_unicore_integration-0.1.11.dist-info/top_level.txt,sha256=j45X-uIuOk3oL78iwlpHakMWtUkg__B7zUlJLwmZx6w,28
+airflow_unicore_integration-0.1.11.dist-info/RECORD,,