Issue #845 also sleep after UDF dep ZIP setup

Open-EO · Oct 30, 2024 · 89a1fdc · 89a1fdc
1 parent 20c731e
commit 89a1fdc
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/openeogeotrellis/deploy/batch_job.py b/openeogeotrellis/deploy/batch_job.py
@@ -1,3 +1,5 @@
+import time
+
 import json
 import logging
 import os
@@ -629,6 +631,13 @@ def _extract_and_install_udf_dependencies(process_graph: dict):
         logger.warning("Merging dependencies from multiple UDF runtimes/versions")
     udf_deps = set(d for ds in udf_dep_map.values() for d in ds)
     if udf_deps:
+
+        def sleep_after_udf_dep_setup():
+            delay = get_backend_config().udf_dependencies_sleep_after_install
+            if delay:
+                logger.info(f"Sleeping after UDF dependency setup ({delay}s)")
+                time.sleep(delay)
+
         udf_deps_install_mode = get_backend_config().udf_dependencies_install_mode
         if udf_deps_install_mode == UDF_DEPENDENCIES_INSTALL_MODE.DISABLED:
             raise ValueError("No UDF dependency handling")
@@ -639,6 +648,7 @@ def _extract_and_install_udf_dependencies(process_graph: dict):
             install_python_udf_dependencies(
                 dependencies=udf_deps, target=udf_python_dependencies_folder_path, timeout=20
             )
+            sleep_after_udf_dep_setup()
         elif udf_deps_install_mode == UDF_DEPENDENCIES_INSTALL_MODE.ZIP:
             udf_python_dependencies_archive_path = _get_env_var_or_fail("UDF_PYTHON_DEPENDENCIES_ARCHIVE_PATH")
             logger.info(f"UDF dep handling with {udf_deps_install_mode=} {udf_python_dependencies_archive_path=}")
@@ -649,10 +659,12 @@ def _extract_and_install_udf_dependencies(process_graph: dict):
                 format="zip",
                 timeout=20,
             )
+            sleep_after_udf_dep_setup()
         else:
             raise ValueError(f"Unsupported UDF dependencies install mode: {udf_deps_install_mode}")
 
 
+
 def start_main():
     setup_logging(
         get_logging_config(

diff --git a/openeogeotrellis/udf.py b/openeogeotrellis/udf.py
@@ -152,10 +152,6 @@ def install_python_udf_dependencies(
         if exit_code != 0:
             raise RuntimeError(f"pip install of UDF dependencies failed with {exit_code=}")
 
-        sleep_after_install = get_backend_config().udf_dependencies_sleep_after_install
-        if sleep_after_install:
-            _log.info(f"Sleeping after pip install ({sleep_after_install}s)")
-            time.sleep(sleep_after_install)
 
 
 def build_python_udf_dependencies_archive(