dayhoff-tools 1.1.33__py3-none-any.whl → 1.1.35__py3-none-any.whl

dayhoff_tools/deployment/deploy_aws.py

@@ -286,15 +286,15 @@ def create_or_update_job_definition(
         raise ValueError("docker.container_entrypoint is required in configuration")
 
     # Create linux parameters with devices
-    linux_params: dict[str, Any] = {
-        "devices": [
+    linux_params: dict[str, Any] = {}
+    if compute_specs.get("gpus", 0) > 0:
+        linux_params["devices"] = [
             {
                 "hostPath": "/dev/nvidia0",
                 "containerPath": "/dev/nvidia0",
                 "permissions": ["READ", "WRITE"],
             },
-        ],
-    }
+        ]
 
     # Add shared memory configuration if specified in docker config
     if "shared_memory" in config.get("docker", {}):
@@ -318,6 +318,32 @@ def create_or_update_job_definition(
         linux_params["sharedMemorySize"] = shared_memory_mib
         print(f"Setting shared memory size to {shared_memory_mib} MiB")
 
+    # Prepare containerProperties
+    container_properties = {
+        "image": image_uri,
+        "vcpus": compute_specs["vcpus"],
+        "memory": compute_specs["memory"],
+        "resourceRequirements": gpu_requirements,
+        "executionRoleArn": aws_config["execution_role_arn"],
+        "jobRoleArn": aws_config["job_role_arn"],
+        "privileged": compute_specs.get("gpus", 0) > 0,
+        "command": entrypoint_command,
+    }
+
+    if linux_params:
+        container_properties["linuxParameters"] = linux_params
+
+    # Add volumes and mount points if defined in AWS batch_job config
+    batch_job_config = aws_config.get("batch_job", {})
+    if "volumes" in batch_job_config:
+        container_properties["volumes"] = batch_job_config["volumes"]
+        print(f"Adding volumes to job definition: {batch_job_config['volumes']}")
+    if "mountPoints" in batch_job_config:
+        container_properties["mountPoints"] = batch_job_config["mountPoints"]
+        print(
+            f"Adding mount points to job definition: {batch_job_config['mountPoints']}"
+        )
+
     # Check if job definition already exists using the session client
     try:
         existing = batch.describe_job_definitions(
@@ -330,31 +356,20 @@ def create_or_update_job_definition(
             print(f"\nCreating new job definition: {job_def_name}")
 
     except batch.exceptions.ClientError as e:
-        # Handle case where the error is specifically 'JobDefinitionNotFoundException'
-        # Boto3 typically includes error codes in the response
-        if (
-            e.response.get("Error", {}).get("Code") == "ClientError"
-        ):  # Simple check, might need refinement
+        if e.response.get("Error", {}).get(
+            "Code"
+        ) == "ClientError" and "JobDefinitionNotFoundException" in str(
+            e
+        ):  # More specific check for not found
             print(f"\nCreating new job definition: {job_def_name}")
         else:
-            # Re-raise unexpected client errors
             raise
 
-    # Prepare job definition properties
+    # Prepare job definition arguments
     job_definition_args = {
         "jobDefinitionName": job_def_name,
         "type": "container",
-        "containerProperties": {
-            "image": image_uri,
-            "vcpus": compute_specs["vcpus"],
-            "memory": compute_specs["memory"],
-            "resourceRequirements": gpu_requirements,
-            "executionRoleArn": aws_config["execution_role_arn"],
-            "jobRoleArn": aws_config["job_role_arn"],
-            "privileged": compute_specs.get("gpus", 0) > 0,
-            "command": entrypoint_command,
-            **({"linuxParameters": linux_params} if linux_params else {}),
-        },
+        "containerProperties": container_properties,
         "platformCapabilities": ["EC2"],
         "timeout": {"attemptDurationSeconds": aws_config.get("timeout_seconds", 86400)},
     }
@@ -403,9 +418,30 @@ def submit_aws_batch_job(
     print(f"- Job Role: {aws_config['job_role_arn']}")
 
     # Get all environment variables, including special ones like WANDB_API_KEY and GCP credentials
-    env_vars = get_container_env_vars(config)
+    env_vars_map = get_container_env_vars(config)  # This returns a dict
+
+    # If EFS is configured for InterProScan, override INTERPROSCAN_INSTALL_DIR
+    # Check based on the conventional volume name used in interp_bulk.yaml
+    efs_interproscan_mount_path = None
+    aws_batch_job_config = aws_config.get("batch_job", {})
+    if "mountPoints" in aws_batch_job_config:
+        for mp in aws_batch_job_config["mountPoints"]:
+            if (
+                mp.get("sourceVolume") == "interproscan-efs-volume"
+            ):  # Convention from YAML
+                efs_interproscan_mount_path = mp.get("containerPath")
+                break
+
+    if efs_interproscan_mount_path:
+        env_vars_map["INTERPROSCAN_INSTALL_DIR"] = efs_interproscan_mount_path
+        print(
+            f"INTERPROSCAN_INSTALL_DIR overridden to EFS mount path: {efs_interproscan_mount_path}"
+        )
 
-    print("Environment Variables:", list(env_vars.keys()))
+    print(
+        "Environment Variables (after potential EFS override):",
+        list(env_vars_map.keys()),
+    )
 
     # Create/Update Job Definition using the config (now implicitly uses the correct session)
     job_definition = create_or_update_job_definition(image_uri, config)
@@ -418,7 +454,8 @@ def submit_aws_batch_job(
         "jobDefinition": job_definition,
         "containerOverrides": {
             "environment": [
-                {"name": key, "value": str(value)} for key, value in env_vars.items()
+                {"name": key, "value": str(value)}
+                for key, value in env_vars_map.items()
             ],
         },
     }

dayhoff_tools/deployment/deploy_gcp.py

@@ -50,9 +50,55 @@ def create_batch_job_config(config: dict, image_uri: str) -> dict:
 
     Returns:
         Dictionary containing GCP Batch job configuration
+
+    Raises:
+        ValueError: If the configuration contains unexpected keys.
     """
     gcp_config = config["gcp"]
 
+    # Validate top-level gcp_config keys used for Batch job JSON construction
+    EXPECTED_GCP_CONFIG_KEYS = {
+        "allocation_policy",  # Goes into batch_config.allocationPolicy
+        "logs_policy",  # Goes into batch_config.logsPolicy
+        "batch_job",  # Contains detailed task and resource specs
+        # Keys like job_name, region, registry_uri, repository are used by other functions
+        # or for other purposes, not directly for constructing the core batch_config JSON here.
+    }
+    actual_gcp_keys = set(gcp_config.keys())
+    # Filter out keys not relevant to this function's direct Batch config construction
+    # These keys are used by the calling context or other parts of the deployment.
+    keys_to_ignore_for_this_check = {"job_name", "region", "registry_uri", "repository"}
+    relevant_gcp_keys = {
+        key for key in actual_gcp_keys if key not in keys_to_ignore_for_this_check
+    }
+
+    unhandled_gcp_keys = relevant_gcp_keys - EXPECTED_GCP_CONFIG_KEYS
+    if unhandled_gcp_keys:
+        raise ValueError(
+            f"Unexpected keys in 'gcp' configuration section: {unhandled_gcp_keys}. "
+            f"Expected keys for Batch job JSON construction are: {EXPECTED_GCP_CONFIG_KEYS}"
+        )
+
+    # Validate keys within gcp_config["batch_job"]
+    if "batch_job" not in gcp_config:
+        raise ValueError("Missing 'batch_job' section in 'gcp' configuration.")
+
+    gcp_batch_job_config = gcp_config["batch_job"]
+    EXPECTED_GCP_BATCH_JOB_KEYS = {
+        "taskCount",
+        "parallelism",
+        "computeResource",
+        "instance",  # Contains machineType, accelerators
+        "volumes",
+    }
+    actual_batch_job_keys = set(gcp_batch_job_config.keys())
+    unhandled_batch_job_keys = actual_batch_job_keys - EXPECTED_GCP_BATCH_JOB_KEYS
+    if unhandled_batch_job_keys:
+        raise ValueError(
+            f"Unexpected keys in 'gcp.batch_job' configuration section: {unhandled_batch_job_keys}. "
+            f"Expected keys are: {EXPECTED_GCP_BATCH_JOB_KEYS}"
+        )
+
     # Start with the allocation and logs policies
     batch_config = {
         "allocationPolicy": gcp_config["allocation_policy"],
@@ -92,10 +138,10 @@ def create_batch_job_config(config: dict, image_uri: str) -> dict:
 
     # Build the task group configuration
     task_group = {
-        "taskCount": gcp_config["batch_job"]["taskCount"],
-        "parallelism": gcp_config["batch_job"]["parallelism"],
+        "taskCount": gcp_batch_job_config["taskCount"],
+        "parallelism": gcp_batch_job_config["parallelism"],
         "taskSpec": {
-            "computeResource": gcp_config["batch_job"]["computeResource"],
+            "computeResource": gcp_batch_job_config["computeResource"],
             "runnables": [{"container": container_config}],
         },
     }
@@ -107,8 +153,12 @@ def create_batch_job_config(config: dict, image_uri: str) -> dict:
     if env_vars:
         task_group["taskSpec"]["runnables"][0]["environment"] = {"variables": env_vars}
 
+    # Add volumes to the taskSpec if specified in the config
+    if "volumes" in gcp_batch_job_config and gcp_batch_job_config["volumes"]:
+        task_group["taskSpec"]["volumes"] = gcp_batch_job_config["volumes"]
+
     # Add machine type and optional accelerators from instance config
-    instance_config = gcp_config["batch_job"]["instance"]
+    instance_config = gcp_batch_job_config["instance"]
     if "machineType" in instance_config:
         # Add machine type to the allocation policy
         if "policy" not in batch_config["allocationPolicy"]["instances"]:

{dayhoff_tools-1.1.33.dist-info → dayhoff_tools-1.1.35.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: dayhoff-tools
-Version: 1.1.33
+Version: 1.1.35
 Summary: Common tools for all the repos at Dayhoff Labs
 Author: Daniel Martin-Alarcon
 Author-email: dma@dayhofflabs.com

{dayhoff_tools-1.1.33.dist-info → dayhoff_tools-1.1.35.dist-info}/RECORD

@@ -7,8 +7,8 @@ dayhoff_tools/cli/main.py,sha256=47EGb28ALaYFc7oAUGlY1D66AIDmc4RZiXxN-gPVrpQ,451
 dayhoff_tools/cli/swarm_commands.py,sha256=5EyKj8yietvT5lfoz8Zx0iQvVaNgc3SJX1z2zQR6o6M,5614
 dayhoff_tools/cli/utility_commands.py,sha256=ER4VrJt4hu904MwrcltUXjwBWT4uFrP-aPXjdXyT3F8,24685
 dayhoff_tools/deployment/base.py,sha256=8tXwsPYvRo-zV-aNhHw1c7Rji-KWg8S5xoCCznFnVVI,17412
-dayhoff_tools/deployment/deploy_aws.py,sha256=O0gQxHioSU_sNU8T8MD4wSOPvWc--V8eRRZzlRu035I,16446
-dayhoff_tools/deployment/deploy_gcp.py,sha256=s3nt-VsF0NrEA-J9Ire8G9zsFgtp8A6dHWT7IYQL5uQ,6529
+dayhoff_tools/deployment/deploy_aws.py,sha256=jQyQ0fbm2793jEHFO84lr5tNqiOpdBg6U0S5zCVJr1M,17884
+dayhoff_tools/deployment/deploy_gcp.py,sha256=jiEE_tBVeSavAI8o_6qPDPpaoXKexcaNIa4uXcv3y0M,8839
 dayhoff_tools/deployment/deploy_utils.py,sha256=StFwbqnr2_FWiKVg3xnJF4kagTHzndqqDkpaIOaAn_4,26027
 dayhoff_tools/deployment/job_runner.py,sha256=hljvFpH2Bw96uYyUup5Ths72PZRL_X27KxlYzBMgguo,5086
 dayhoff_tools/deployment/processors.py,sha256=q2io07xO6f6twEH1iLz9GFoGaKh76qC4kcv519Q2G7g,34583
@@ -26,7 +26,7 @@ dayhoff_tools/intake/uniprot.py,sha256=BZYJQF63OtPcBBnQ7_P9gulxzJtqyorgyuDiPeOJq
 dayhoff_tools/logs.py,sha256=DKdeP0k0kliRcilwvX0mUB2eipO5BdWUeHwh-VnsICs,838
 dayhoff_tools/sqlite.py,sha256=jV55ikF8VpTfeQqqlHSbY8OgfyfHj8zgHNpZjBLos_E,18672
 dayhoff_tools/warehouse.py,sha256=TqV8nex1AluNaL4JuXH5zuu9P7qmE89lSo6f_oViy6U,14965
-dayhoff_tools-1.1.33.dist-info/METADATA,sha256=8v8Kj6e0GdUHhtDx3l1huf7t1YwcnA0vlbbaKWwpBAw,2843
-dayhoff_tools-1.1.33.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-dayhoff_tools-1.1.33.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
-dayhoff_tools-1.1.33.dist-info/RECORD,,
+dayhoff_tools-1.1.35.dist-info/METADATA,sha256=lula1chZ1oJh4kiICRO_J1Lh4nhG49THMge3MBuM428,2843
+dayhoff_tools-1.1.35.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+dayhoff_tools-1.1.35.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
+dayhoff_tools-1.1.35.dist-info/RECORD,,