Merge branch 'main' into ir_extraction_backup

components/clp-package-utils/clp_package_utils/general.py

@@ -106,6 +106,10 @@ def get_clp_home():
 
 
 def generate_container_name(job_type: JobType) -> str:
+    """
+    :param job_type:
+    :return: A unique container name for the given job type.
+    """
     return f"clp-{job_type}-{str(uuid.uuid4())[-4:]}"
 
 
@@ -195,12 +199,15 @@ def is_path_already_mounted(
     return host_path_relative_to_mounted_root == container_path_relative_to_mounted_root
 
 
-def generate_container_config(clp_config: CLPConfig, clp_home: pathlib.Path):
+def generate_container_config(
+    clp_config: CLPConfig, clp_home: pathlib.Path
+) -> Tuple[CLPConfig, CLPDockerMounts]:
     """
     Copies the given config and sets up mounts mapping the relevant host paths into the container
 
     :param clp_config:
     :param clp_home:
+    :return: The container config and the mounts.
     """
     container_clp_config = clp_config.copy(deep=True)
 
@@ -260,8 +267,15 @@ def generate_container_config(clp_config: CLPConfig, clp_home: pathlib.Path):
 
 
 def dump_container_config(
-    clp_config: CLPConfig, container_clp_config, container_name: str
+    container_clp_config: CLPConfig, clp_config: CLPConfig, container_name: str
 ) -> Tuple[pathlib.Path, pathlib.Path]:
+    """
+    Writes the given config to the logs directory so that it's accessible in the container.
+    :param container_clp_config: The config to write.
+    :param clp_config: The corresponding config on the host (used to determine the logs directory).
+    :param container_name:
+    :return: The path to the config file in the container and on the host.
+    """
     container_config_filename = f".{container_name}-config.yml"
     config_file_path_on_host = clp_config.logs_directory / container_config_filename
     config_file_path_on_container = container_clp_config.logs_directory / container_config_filename
@@ -272,8 +286,15 @@ def dump_container_config(
 
 
 def generate_container_start_cmd(
-    container_name: str, container_mounts: List[DockerMount], execution_container: str
-):
+    container_name: str, container_mounts: List[CLPDockerMounts], container_image: str
+) -> List[str]:
+    """
+    Generates the command to start a container with the given mounts and name.
+    :param container_name:
+    :param container_mounts:
+    :param container_image:
+    :return: The command.
+    """
     clp_site_packages_dir = CONTAINER_CLP_HOME / "lib" / "python3" / "site-packages"
     # fmt: off
     container_start_cmd = [
@@ -291,7 +312,7 @@ def generate_container_start_cmd(
         if mount:
             container_start_cmd.append("--mount")
             container_start_cmd.append(str(mount))
-    container_start_cmd.append(execution_container)
+    container_start_cmd.append(container_image)
 
     return container_start_cmd
 
@@ -304,7 +325,7 @@ def validate_config_key_existence(config, key):
     return value
 
 
-def validate_and_load_config_file(
+def load_config_file(
     config_file_path: pathlib.Path, default_config_file_path: pathlib.Path, clp_home: pathlib.Path
 ):
     if config_file_path.exists():

components/clp-package-utils/clp_package_utils/scripts/compress.py

@@ -14,7 +14,7 @@
     generate_container_start_cmd,
     get_clp_home,
     JobType,
-    validate_and_load_config_file,
+    load_config_file,
     validate_and_load_db_credentials_file,
 )
 
@@ -57,11 +57,10 @@ def main(argv):
     # Validate and load config file
     try:
         config_file_path = pathlib.Path(parsed_args.config)
-        clp_config = validate_and_load_config_file(
-            config_file_path, default_config_file_path, clp_home
-        )
+        clp_config = load_config_file(config_file_path, default_config_file_path, clp_home)
         clp_config.validate_logs_dir()
 
+        # Validate and load necessary credentials
         validate_and_load_db_credentials_file(clp_config, clp_home, False)
     except:
         logger.exception("Failed to load config.")
@@ -70,8 +69,8 @@ def main(argv):
     container_name = generate_container_name(JobType.COMPRESSION)
 
     container_clp_config, mounts = generate_container_config(clp_config, clp_home)
-    config_file_path_on_container, config_file_path_on_host = dump_container_config(
-        clp_config, container_clp_config, container_name
+    generated_config_path_on_container, generated_config_path_on_host = dump_container_config(
+        container_clp_config, clp_config, container_name
     )
 
     necessary_mounts = [mounts.clp_home, mounts.input_logs_dir, mounts.data_dir, mounts.logs_dir]
@@ -83,7 +82,7 @@ def main(argv):
     compress_cmd = [
         "python3",
         "-m", "clp_package_utils.scripts.native.compress",
-        "--config", str(config_file_path_on_container),
+        "--config", str(generated_config_path_on_container),
         "--remove-path-prefix", str(CONTAINER_INPUT_LOGS_ROOT_DIR),
     ]
     # fmt: on
@@ -122,7 +121,7 @@ def main(argv):
     subprocess.run(cmd, check=True)
 
     # Remove generated files
-    config_file_path_on_host.unlink()
+    generated_config_path_on_host.unlink()
 
     return 0
 

components/clp-package-utils/clp_package_utils/scripts/decompress.py

@@ -4,6 +4,8 @@
 import subprocess
 import sys
 
+from typing import Optional
+
 from clp_py_utils.clp_config import CLPConfig
 
 from clp_package_utils.general import (
@@ -18,7 +20,7 @@
     get_clp_home,
     IR_EXTRACTION_COMMAND,
     JobType,
-    validate_and_load_config_file,
+    load_config_file,
     validate_and_load_db_credentials_file,
     validate_path_could_be_dir,
 )
@@ -34,7 +36,31 @@
 logger.addHandler(logging_console_handler)
 
 
-def handle_decompression(parsed_args, clp_config: CLPConfig, clp_home: pathlib.Path):
+def validate_and_load_config(
+    clp_home: pathlib.Path,
+    config_file_path: pathlib.Path,
+    default_config_file_path: pathlib.Path,
+) -> Optional[CLPConfig]:
+    """
+    Validates and loads the config file.
+    :param clp_home:
+    :param config_file_path:
+    :param default_config_file_path:
+    :return: clp_config on success, None otherwise.
+    """
+    try:
+        clp_config = load_config_file(config_file_path, default_config_file_path, clp_home)
+        clp_config.validate_logs_dir()
+
+        # Validate and load necessary credentials
+        validate_and_load_db_credentials_file(clp_config, clp_home, False)
+        return clp_config
+    except:
+        logger.exception("Failed to load config.")
+        return None
+
+
+def handle_decompression_command(parsed_args, clp_home: pathlib.Path, default_config_file_path: pathlib.Path):
     paths_to_decompress_file_path = None
     if parsed_args.files_from:
         paths_to_decompress_file_path = pathlib.Path(parsed_args.files_from)
@@ -48,11 +74,17 @@ def handle_decompression(parsed_args, clp_config: CLPConfig, clp_home: pathlib.P
         return -1
     extraction_dir.mkdir(exist_ok=True)
 
+    # Validate and load config file
+    clp_config = validate_and_load_config(clp_home, pathlib.Path(parsed_args.config), clp_home)
+    if not clp_config:
+        return -1
+
     container_name = generate_container_name(JobType.DECOMPRESSION)
     container_clp_config, mounts = generate_container_config(clp_config, clp_home)
-    config_file_path_on_container, config_file_path_on_host = dump_container_config(
-        clp_config, container_clp_config, container_name
+    generated_config_path_on_container, generated_config_path_on_host = dump_container_config(
+        container_clp_config, clp_config, container_name
     )
+
     # Set up mounts
     container_extraction_dir = pathlib.Path("/") / "mnt" / "extraction-dir"
     necessary_mounts = [
@@ -82,7 +114,7 @@ def handle_decompression(parsed_args, clp_config: CLPConfig, clp_home: pathlib.P
     decompress_cmd = [
         "python3",
         "-m", "clp_package_utils.scripts.native.decompress",
-        "--config", str(config_file_path_on_container),
+        "--config", str(generated_config_path_on_container),
         DECOMPRESSION_COMMAND,
         "-d", str(container_extraction_dir),
     ]
@@ -97,14 +129,19 @@ def handle_decompression(parsed_args, clp_config: CLPConfig, clp_home: pathlib.P
     subprocess.run(cmd, check=True)
 
     # Remove generated files
-    config_file_path_on_host.unlink()
+    generated_config_path_on_host.unlink()
 
 
-def handle_extraction(parsed_args, clp_config: CLPConfig, clp_home: pathlib.Path):
+def handle_extraction(parsed_args, clp_home: pathlib.Path, default_config_file_path: pathlib.Path):
+    # Validate and load config file
+    clp_config = validate_and_load_config(clp_home, pathlib.Path(parsed_args.config), clp_home)
+    if not clp_config:
+        return -1
+
     container_name = generate_container_name(JobType.IR_EXTRACTION)
     container_clp_config, mounts = generate_container_config(clp_config, clp_home)
-    config_file_path_on_container, config_file_path_on_host = dump_container_config(
-        clp_config, container_clp_config, container_name
+    generated_config_path_on_container, generated_config_path_on_host = dump_container_config(
+        container_clp_config, clp_config, container_name
     )
     necessary_mounts = [mounts.clp_home, mounts.logs_dir]
     container_start_cmd = generate_container_start_cmd(
@@ -116,7 +153,7 @@ def handle_extraction(parsed_args, clp_config: CLPConfig, clp_home: pathlib.Path
         "-m",
         "clp_package_utils.scripts.native.decompress",
         "--config",
-        str(config_file_path_on_container),
+        str(generated_config_path_on_container),
         IR_EXTRACTION_COMMAND,
         str(parsed_args.msg_ix),
     ]
@@ -134,7 +171,7 @@ def handle_extraction(parsed_args, clp_config: CLPConfig, clp_home: pathlib.Path
     subprocess.run(cmd, check=True)
 
     # Remove generated files
-    config_file_path_on_host.unlink()
+    generated_config_path_on_host.unlink()
 
 
 def main(argv):
@@ -172,23 +209,11 @@ def main(argv):
 
     parsed_args = args_parser.parse_args(argv[1:])
 
-    # Validate and load config file
-    try:
-        config_file_path = pathlib.Path(parsed_args.config)
-        clp_config = validate_and_load_config_file(
-            config_file_path, default_config_file_path, clp_home
-        )
-        clp_config.validate_logs_dir()
-
-        validate_and_load_db_credentials_file(clp_config, clp_home, False)
-    except:
-        logger.exception("Failed to load config.")
-        return -1
     command = parsed_args.command
     if DECOMPRESSION_COMMAND == command:
-        return handle_decompression(parsed_args, clp_config, clp_home)
+        return handle_decompression_command(parsed_args, clp_home, default_config_file_path)
     elif IR_EXTRACTION_COMMAND == command:
-        return handle_extraction(parsed_args, clp_config, clp_home)
+        return handle_extraction(parsed_args, clp_home, default_config_file_path)
     else:
         logger.exception(f"Unexpected command: {command}")
         return -1

components/clp-package-utils/clp_package_utils/scripts/native/compress.py

@@ -20,7 +20,7 @@
 from clp_package_utils.general import (
     CLP_DEFAULT_CONFIG_FILE_RELATIVE_PATH,
     get_clp_home,
-    validate_and_load_config_file,
+    load_config_file,
 )
 
 # Setup logging
@@ -170,9 +170,7 @@ def main(argv):
     # Validate and load config file
     try:
         config_file_path = pathlib.Path(parsed_args.config)
-        clp_config = validate_and_load_config_file(
-            config_file_path, default_config_file_path, clp_home
-        )
+        clp_config = load_config_file(config_file_path, default_config_file_path, clp_home)
         clp_config.validate_input_logs_dir()
         clp_config.validate_logs_dir()
     except: