dayhoff-tools 1.1.10__py3-none-any.whl → 1.13.12__py3-none-any.whl

dayhoff_tools/deployment/base.py

@@ -128,8 +128,9 @@ def build_job_image(config: dict) -> str:
 
     This function handles the complete image building process:
     1. Ensures we're in the repo root
-    2. Constructs the image URI based on config
-    3. Builds the image using docker build
+    2. Cleans Docker config to avoid credential helper conflicts
+    3. Constructs the image URI based on config
+    4. Builds the image using docker build
 
     Args:
         config: Dictionary containing the configuration loaded from YAML.
@@ -146,13 +147,20 @@ def build_job_image(config: dict) -> str:
     """
     move_to_repo_root()
 
+    # Clean Docker config to avoid VS Code dev container credential helper conflicts
+    from dayhoff_tools.deployment.deploy_utils import clean_docker_config
+
+    clean_docker_config()
+
     # Get image URI
     image_uri = _build_image_uri(config)
     docker_config = config["docker"]
 
     print("\nBuilding Docker image: ", image_uri)
     print(f"Using Dockerfile: {docker_config['dockerfile']}")
-    print(f"Using shared memory: {docker_config['shared_memory']}\n")
+    print(f"Using shared memory: {docker_config['shared_memory']}")
+    platform = docker_config.get("platform", "linux/amd64")
+    print(f"Building for platform: {platform}\n")
 
     # Build the image
     build_image_command = [
@@ -163,8 +171,20 @@ def build_job_image(config: dict) -> str:
         docker_config["dockerfile"],
         "-t",
         image_uri,
-        ".",  # Use the root of the repo as image context
     ]
+
+    # Add platform specification if provided, default to linux/amd64 for cloud deployments
+    platform = docker_config.get("platform", "linux/amd64")
+    build_image_command.extend(["--platform", platform])
+
+    # Add build args if provided (for parameterized Dockerfiles)
+    build_args = docker_config.get("build_args", {})
+    for arg_name, arg_value in build_args.items():
+        build_image_command.extend(["--build-arg", f"{arg_name}={arg_value}"])
+
+    # Add build context (defaults to "." for backward compatibility)
+    build_context = docker_config.get("build_context", ".")
+    build_image_command.append(build_context)
     subprocess.run(build_image_command, check=True)
 
     # Get and print image size
@@ -281,7 +301,7 @@ def run_container(config: dict, image_uri: str, mode: str) -> None:
     4. Handles container logs for detached mode
 
     The container name is generated using:
-    - Username (from LOCAL_USER env var)
+    - Username (from LOCAL_USER or USER env var)
     - Timestamp (YYYYMMDD_HHMMSS format)
 
     Args:
@@ -299,7 +319,7 @@ def run_container(config: dict, image_uri: str, mode: str) -> None:
         )
 
     # Generate unique container name
-    username = os.getenv("LOCAL_USER", "unknown_user")
+    username = os.getenv("LOCAL_USER") or os.getenv("USER", "unknown_user")
     timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     container_name = f"{username}_job_{timestamp}"
 
@@ -394,6 +414,50 @@ def deploy(
             get_boto_session(config)
             print("AWS credentials verified.")
 
+        # Check GCP credentials early if using GCP
+        elif cloud == "gcp":
+            print("\nVerifying GCP credentials...")
+            from dayhoff_tools.cli.cloud_commands import (
+                _is_adc_authenticated,
+                _is_gcp_user_authenticated,
+            )
+
+            user_creds_valid = _is_gcp_user_authenticated()
+            adc_creds_valid = _is_adc_authenticated()
+
+            if not user_creds_valid:
+                print(
+                    "\n⚠️  Warning: Your GCP user credentials appear to be stale/expired."
+                )
+                print(
+                    "   This may cause authentication issues when deploying to GCP Batch."
+                )
+                print("   Consider running 'dh gcp login' to refresh your credentials.")
+
+            if not adc_creds_valid:
+                print(
+                    "\n⚠️  Warning: Your Application Default Credentials (ADC) appear to be stale/expired."
+                )
+                print(
+                    "   This may cause authentication issues with API client libraries."
+                )
+                print(
+                    "   Consider running 'dh gcp use-user-adc' or 'dh gcp use-devcon-adc' to refresh your ADC."
+                )
+
+            if user_creds_valid and adc_creds_valid:
+                print("GCP credentials verified.")
+            else:
+                # Ask for confirmation before proceeding with possibly invalid credentials
+                proceed = (
+                    input("\nProceed with potentially invalid credentials? (y/n): ")
+                    .lower()
+                    .strip()
+                )
+                if proceed != "y":
+                    print("Deployment aborted.")
+                    return
+
     # Track if we built a new image
     had_image_uri = bool(config["docker"]["image_uri"])
 

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,82 @@ 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']}"
+        )
+
+    # Mount Primordial Drive if explicitly enabled via feature flag
+    # Add 'mount_primordial_drive' to features list to mount shared EFS at /primordial/
+    features = config.get("features", []) or []
+    features_set = {f if isinstance(f, str) else next(iter(f)) for f in features}
+    mount_primordial = "mount_primordial_drive" in features_set
+
+    if mount_primordial:
+        primordial_fs_id = get_primordial_fs_id(session)
+        if primordial_fs_id:
+            print(f"Adding Primordial Drive configuration (fs_id: {primordial_fs_id})")
+
+            # Add volume configuration
+            efs_volume = {
+                "name": "primordial",
+                "efsVolumeConfiguration": {
+                    "fileSystemId": primordial_fs_id,
+                    "rootDirectory": "/",
+                },
+            }
+
+            if "volumes" not in container_properties:
+                container_properties["volumes"] = []
+
+            # Check if already added to avoid duplicates
+            if not any(
+                v.get("name") == "primordial" for v in container_properties["volumes"]
+            ):
+                container_properties["volumes"].append(efs_volume)
+
+            # Add mount point
+            mount_point = {
+                "sourceVolume": "primordial",
+                "containerPath": "/primordial",
+                "readOnly": False,
+            }
+
+            if "mountPoints" not in container_properties:
+                container_properties["mountPoints"] = []
+
+            # Check if already added
+            if not any(
+                mp.get("containerPath") == "/primordial"
+                for mp in container_properties["mountPoints"]
+            ):
+                container_properties["mountPoints"].append(mount_point)
+        else:
+            print(
+                "Warning: mount_primordial_drive enabled but Primordial Drive not found in this environment"
+            )
+
     # Check if job definition already exists using the session client
     try:
         existing = batch.describe_job_definitions(
@@ -330,41 +406,67 @@ 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)},
     }
 
+    # Add tags if specified in config
+    if "tags" in aws_config:
+        job_definition_args["tags"] = aws_config["tags"]
+        print(f"Adding tags to job definition: {aws_config['tags']}")
+
     # Register new revision using the session client
     response = batch.register_job_definition(**job_definition_args)
 
     return response["jobDefinitionName"]
 
 
+def get_primordial_fs_id(session: boto3.Session) -> Optional[str]:
+    """Fetch Primordial Drive EFS ID from SSM.
+
+    Args:
+        session: Boto3 session
+
+    Returns:
+        FileSystemId if found, None otherwise
+    """
+    ssm = session.client("ssm")
+
+    # Determine environment from profile name
+    # Default to dev if cannot determine
+    env = "dev"
+    if session.profile_name and "sand" in session.profile_name:
+        env = "sand"
+
+    param_name = f"/{env}/primordial/fs_id"
+
+    try:
+        response = ssm.get_parameter(Name=param_name)
+        return response["Parameter"]["Value"]
+    except ClientError as e:
+        # Silently fail if not found - Primordial might not be deployed in this env
+        # or we might not have permissions
+        # ParameterNotFound is a ClientError with error code "ParameterNotFound"
+        return None
+    except Exception as e:
+        print(f"Warning: Failed to check for Primordial Drive: {e}")
+        return None
+
+
 def submit_aws_batch_job(
     image_uri: str,
     config: dict[str, Any],
@@ -403,9 +505,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 +541,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()
             ],
         },
     }
@@ -435,6 +559,22 @@ def submit_aws_batch_job(
     print(f"Setting retry attempts to {retry_attempts}")
     job_submit_args["retryStrategy"] = {"attempts": retry_attempts}
 
+    # Automatically add User tag for cost tracking
+    username = os.getenv("LOCAL_USER", "unknown_user")
+    default_tags = {"User": username}
+
+    # Merge with any tags specified in config (config tags take precedence)
+    if "tags" in aws_config:
+        tags = {**default_tags, **aws_config["tags"]}
+    else:
+        tags = default_tags
+
+    job_submit_args["tags"] = tags
+    job_submit_args["propagateTags"] = (
+        True  # Propagate tags to ECS tasks and EC2 instances
+    )
+    print(f"Adding tags to batch job: {tags}")
+
     # Submit the job using the session client
     response = batch.submit_job(**job_submit_args)
 

dayhoff_tools/deployment/deploy_gcp.py

@@ -50,9 +50,56 @@ 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
+        "image_uri",
+        # 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"],
@@ -75,16 +122,27 @@ def create_batch_job_config(config: dict, image_uri: str) -> dict:
         "commands": ["-c", " ".join(entrypoint_command)],
     }
 
+    # Handle container options - both shared memory and any custom options
+    docker_options = []
+
     # Add shared memory option if specified
     if "shared_memory" in config.get("docker", {}):
-        container_config["options"] = f"--shm-size={config['docker']['shared_memory']}"
+        docker_options.append(f"--shm-size={config['docker']['shared_memory']}")
+
+    # Add any custom Docker options if specified
+    if "options" in config.get("docker", {}):
+        docker_options.append(config["docker"]["options"])
+
+    # Set the options field if any options were collected
+    if docker_options:
+        container_config["options"] = " ".join(docker_options)
 
     # 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}],
         },
     }
@@ -96,8 +154,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"]:
@@ -123,6 +185,17 @@ def create_batch_job_config(config: dict, image_uri: str) -> dict:
         print("Machine Type:", policy.get("machineType", "Not specified"))
         print("Accelerators:", policy.get("accelerators", "Not specified"))
         print("Environment Variables:", list(env_vars.keys()))
+        if (
+            "runnables" in task_group["taskSpec"]
+            and task_group["taskSpec"]["runnables"]
+        ):
+            print(
+                "Container Options:",
+                task_group["taskSpec"]["runnables"][0]
+                .get("container", {})
+                .get("options", "Not specified"),
+            )
+
     except KeyError as e:
         print(f"Warning: Could not find {e} in configuration")
 

dayhoff_tools/deployment/deploy_utils.py

@@ -614,6 +614,24 @@ def setup_rxnfp() -> None:
     )
 
 
+def clean_docker_config() -> None:
+    """Clean Docker configuration to avoid credential helper conflicts.
+
+    VS Code dev containers can set up credential helpers that interfere with
+    Docker builds and registry operations. This function creates a clean
+    configuration that disables problematic credential helpers.
+
+    This is automatically called before Docker builds and registry operations
+    to prevent authentication failures.
+    """
+    docker_config_dir = os.path.expanduser("~/.docker")
+    os.makedirs(docker_config_dir, exist_ok=True)
+
+    # Write a minimal config file that disables credential helpers
+    with open(os.path.join(docker_config_dir, "config.json"), "w") as f:
+        json.dump({"auths": {}, "credsStore": ""}, f)
+
+
 def docker_login(registry: str, username: str, password: str) -> None:
     """Login to a Docker registry using provided credentials.
 
@@ -625,13 +643,8 @@ def docker_login(registry: str, username: str, password: str) -> None:
     Raises:
         subprocess.CalledProcessError: If Docker login fails
     """
-    # Create .docker directory if it doesn't exist
-    docker_config_dir = os.path.expanduser("~/.docker")
-    os.makedirs(docker_config_dir, exist_ok=True)
-
-    # Write a minimal config file that disables credential helpers
-    with open(os.path.join(docker_config_dir, "config.json"), "w") as f:
-        json.dump({"auths": {}, "credsStore": ""}, f)
+    # Clean Docker config to avoid credential helper conflicts
+    clean_docker_config()
 
     # Login to Docker using the credentials
     login_process = subprocess.run(

dayhoff_tools/deployment/job_runner.py

@@ -94,9 +94,9 @@ def run_command() -> None:
             stderr=None,  # Use parent's stderr
         )
 
-        logger.info("Command completed successfully")
+        logger.info("Job command completed successfully")
     except subprocess.CalledProcessError as e:
-        logger.error(f"Command failed with return code: {e.returncode}")
+        logger.error(f"Job command failed with return code: {e.returncode}")
         raise
     except Exception as e:
         logger.error(f"Error executing command: {str(e)}")
@@ -133,6 +133,13 @@ def run_job(
     logger = logging.getLogger(__name__)
 
     logger.info(f"Job runner starting in mode: {mode}")
+    import importlib.metadata
+
+    try:
+        version = importlib.metadata.version("dayhoff-tools")
+        logger.info(f"dayhoff-tools version: {version}")
+    except importlib.metadata.PackageNotFoundError:
+        logger.warning("Could not determine dayhoff-tools version")
 
     if mode not in ["setup", "execute", "setup_and_execute"]:
         logger.error(f"Invalid mode: {mode}")
@@ -146,8 +153,6 @@ def run_job(
         if mode in ["execute", "setup_and_execute"]:
             run_command()
 
-        logger.info("Job completed successfully")
-
     except Exception as e:
         logger.error(f"Job failed with error: {str(e)}", exc_info=True)
         sys.exit(1)