dayhoff-tools 1.0.0__py3-none-any.whl

dayhoff_tools/deployment/deploy_aws.py

@@ -0,0 +1,458 @@
+"""AWS deployment functionality for running jobs on AWS Batch."""
+
+import base64
+import datetime
+import os
+import re
+import subprocess
+from pathlib import Path
+from typing import Any, Optional
+
+import boto3
+import yaml
+from botocore.exceptions import (
+    ClientError,
+    NoCredentialsError,
+    ProfileNotFound,
+    SSOTokenLoadError,
+)
+from dayhoff_tools.deployment.deploy_utils import docker_login, get_container_env_vars
+
+
+def get_boto_session(config: dict[str, Any]) -> boto3.Session:
+    """Creates a Boto3 session using the profile specified in the config.
+
+    Args:
+        config: Dictionary containing the configuration loaded from YAML
+
+    Returns:
+        A Boto3 session object
+
+    Raises:
+        RuntimeError: If the profile is not specified in the config or if credentials cannot be loaded
+    """
+    aws_config = config.get("aws", {})
+    profile_name = aws_config.get("aws_profile")
+    region = aws_config.get("region", "us-east-1")
+
+    if not profile_name:
+        print(
+            "Warning: aws.aws_profile not specified in config. Using default AWS credential chain in ~/.aws/config."
+        )
+        return boto3.Session(region_name=region)
+
+    try:
+        print(f"Using AWS profile: {profile_name}")
+        session = boto3.Session(profile_name=profile_name, region_name=region)
+        sts = session.client("sts")
+        sts.get_caller_identity()
+        return session
+    except ProfileNotFound:
+        raise RuntimeError(
+            f"AWS profile '{profile_name}' not found in `~/.aws/config`."
+        )
+    except (NoCredentialsError, ClientError, SSOTokenLoadError) as e:
+        raise RuntimeError(
+            f"Could not load credentials for AWS profile '{profile_name}'. "
+            f"Ensure you are logged in via SSO ('aws sso login --profile {profile_name}') "
+            f"or have valid credentials configured. Original error: {e}"
+        ) from e
+
+
+def _extract_run_name_from_config(
+    config: dict[str, Any], _test_file_content: Optional[str] = None
+) -> Optional[str]:
+    """Extract run name from the config file referenced in job_command.
+
+    Args:
+        config: Dictionary containing the configuration loaded from YAML
+        _test_file_content: Optional parameter for testing to override file content
+
+    Returns:
+        Run name if found, None otherwise
+    """
+    # Check if features and job_command exist in config
+    if "features" not in config:
+        return None
+
+    # Find job_command in features
+    job_command = None
+    for feature in config["features"]:
+        if isinstance(feature, dict) and "job_command" in feature:
+            job_command = feature["job_command"]
+            break
+        elif isinstance(feature, str) and feature.startswith("job_command:"):
+            job_command = feature.split(":", 1)[1].strip()
+            break
+
+    if not job_command:
+        return None
+
+    # Extract config file path using regex
+    config_match = re.search(r'--config=([^\s"]+)', job_command)
+    if not config_match:
+        return None
+
+    config_path = config_match.group(1)
+
+    # For testing, we can bypass file operations
+    if _test_file_content is not None:
+        run_config = yaml.safe_load(_test_file_content)
+    else:
+        # Resolve path relative to repo root
+        # Assuming we're in the repo root when this function is called
+        full_config_path = Path(config_path)
+
+        # Check if file exists
+        if not full_config_path.exists():
+            print(f"Warning: Config file {full_config_path} not found")
+            return None
+
+        try:
+            # Load the config file
+            with open(full_config_path, "r") as f:
+                run_config = yaml.safe_load(f)
+        except Exception as e:
+            print(f"Warning: Failed to extract run_name from {full_config_path}: {e}")
+            return None
+
+    # Extract run_name from wandb section
+    if (
+        run_config
+        and "init" in run_config
+        and "wandb" in run_config["init"]
+        and "run_name" in run_config["init"]["wandb"]
+    ):
+        return run_config["init"]["wandb"]["run_name"]
+
+    return None
+
+
+def _extract_job_name_from_uri(image_uri: str, config: dict[str, Any]) -> str:
+    """Extract job name from image URI and config.
+
+    Args:
+        image_uri: Full URI of the container image
+        config: Dictionary containing the configuration loaded from YAML
+
+    Returns:
+        Job name in format: username__jobname__uniquehex
+
+        The job name components are:
+        - username: The LOCAL_USER environment variable or "unknown_user"
+        - jobname: Either from aws.job_name config, wandb run_name, or omitted
+        - uniquehex: A hexadecimal string (typically 7 characters)
+          representing seconds since January 1, 2020
+
+        If aws.job_name is specified in the config:
+        - If aws.job_name is "use_wandb_run_name", the run_name from the wandb config will be used
+        - Otherwise, the specified job_name will be used
+
+        If no job_name is specified or if "use_wandb_run_name" is specified but no run_name is found,
+        the format will be username__uniquehex (without the middle component).
+    """
+    # Get username from environment
+    username = os.getenv("LOCAL_USER", "unknown_user")
+
+    # Generate a unique hex string based on seconds since 2020-01-01
+    epoch_2020 = datetime.datetime(2020, 1, 1).timestamp()
+    current_time = datetime.datetime.now().timestamp()
+    seconds_since_2020 = int(current_time - epoch_2020)
+    unique_hex = format(
+        seconds_since_2020, "x"
+    )  # Simple hex representation, typically 7 chars
+
+    # Check if job_name is specified in AWS config
+    if "aws" in config and "job_name" in config["aws"]:
+        job_name = config["aws"]["job_name"]
+
+        # Special handling for "use_wandb_run_name"
+        if job_name == "use_wandb_run_name":
+            # Get run name from config
+            run_name = _extract_run_name_from_config(config)
+            if run_name:
+                # Return username__runname__uniquehex
+                return f"{username}__{run_name}__{unique_hex}"
+        else:
+            # Use the specified job_name but format it as username__jobname__uniquehex
+            return f"{username}__{job_name}__{unique_hex}"
+
+    # Default behavior if job_name is not specified or "use_wandb_run_name" with no run_name
+    # Get run name from config
+    run_name = _extract_run_name_from_config(config)
+
+    # Build job name based on available components
+    job_name_parts = []
+
+    # Add username (use default if None)
+    job_name_parts.append(username or "unknown_user")
+
+    # Add run name if available
+    if run_name:
+        job_name_parts.append(run_name)
+
+    # Add the unique hex string
+    job_name_parts.append(unique_hex)
+
+    # Join parts with double underscores
+    return "__".join(job_name_parts)
+
+
+def push_image_to_ecr(image_uri: str, config: dict[str, Any]) -> None:
+    """Push a Docker image to Amazon ECR.
+
+    Args:
+        image_uri: Full URI of the image to push
+        config: Dictionary containing the configuration loaded from YAML
+
+    Raises:
+        subprocess.CalledProcessError: If ECR login or push fails
+        ClientError: If ECR authentication fails
+    """
+    print("\nPushing image to ECR")
+
+    session = get_boto_session(config)
+    aws_config = config["aws"]
+    registry = aws_config["registry_uri"]
+
+    # Get ECR login token using the specific session client
+    ecr_client = session.client("ecr")
+    try:
+        token = ecr_client.get_authorization_token()
+        username, password = (
+            base64.b64decode(token["authorizationData"][0]["authorizationToken"])
+            .decode()
+            .split(":")
+        )
+
+        # Login to Docker registry
+        docker_login(registry, username, password)
+
+        # Push the image
+        print(f"\nPushing image: {image_uri}")
+        subprocess.run(["docker", "push", image_uri], check=True)
+
+        print(f"\nSuccessfully pushed image to ECR: {image_uri}")
+
+    except ClientError as e:
+        print(f"AWS ECR error using profile '{session.profile_name}': {str(e)}")
+        raise
+    except subprocess.CalledProcessError as e:
+        print(f"Docker command failed: {e.stderr.decode() if e.stderr else str(e)}")
+        raise
+
+
+def create_or_update_job_definition(
+    image_uri: str,
+    config: dict[str, Any],
+) -> str:
+    """Create or update a job definition for the container.
+
+    Args:
+        image_uri: Full URI of the container image
+        config: Dictionary containing the configuration loaded from YAML
+
+    Returns:
+        Name of the created/updated job definition
+
+    Raises:
+        ValueError: If job_role_arn is not specified in AWS configuration
+    """
+    session = get_boto_session(config)
+    aws_config = config["aws"]
+
+    # Verify that job_role_arn is present
+    if "job_role_arn" not in aws_config:
+        raise ValueError(
+            "job_role_arn must be specified in AWS configuration. "
+            "This role is required for your container to access AWS resources after the "
+            "initial 15-minute temporary credential window."
+        )
+
+    batch = session.client("batch")
+    base_name = config["docker"]["base_name"]
+    job_def_name = f"job-def-{base_name}"
+
+    # Get compute specs from config
+    compute_specs = aws_config["batch_job"]["compute_specs"]
+    gpu_requirements = (
+        [{"type": "GPU", "value": str(compute_specs["gpus"])}]
+        if compute_specs.get("gpus", 0) > 0
+        else []
+    )
+
+    entrypoint_command = config["docker"].get("container_entrypoint")
+    if entrypoint_command is None:
+        raise ValueError("docker.container_entrypoint is required in configuration")
+
+    # Create linux parameters with devices
+    linux_params: dict[str, Any] = {
+        "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", {}):
+        shared_mem = config["docker"]["shared_memory"]
+        # Convert to MiB (e.g., "16g" -> 16384 MiB)
+        if isinstance(shared_mem, str):
+            if shared_mem.endswith("g"):
+                # Convert GB to MiB (1G = 1024 MiB)
+                shared_memory_mib = int(float(shared_mem[:-1]) * 1024)
+            elif shared_mem.endswith("m"):
+                # Convert MB to MiB (approximate conversion)
+                shared_memory_mib = int(float(shared_mem[:-1]))
+            else:
+                # Assume the value is already in MiB
+                shared_memory_mib = int(float(shared_mem))
+        else:
+            # Assume the value is already in MiB if not a string
+            shared_memory_mib = int(shared_mem)
+
+        # Add shared memory size to linux parameters
+        linux_params["sharedMemorySize"] = shared_memory_mib
+        print(f"Setting shared memory size to {shared_memory_mib} MiB")
+
+    # Check if job definition already exists using the session client
+    try:
+        existing = batch.describe_job_definitions(
+            jobDefinitionName=job_def_name, status="ACTIVE"
+        )["jobDefinitions"]
+
+        if existing:
+            print(f"\nUpdating existing job definition: {job_def_name}")
+        else:
+            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
+            print(f"\nCreating new job definition: {job_def_name}")
+        else:
+            # Re-raise unexpected client errors
+            raise
+
+    # Prepare job definition properties
+    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 {}),
+        },
+        "platformCapabilities": ["EC2"],
+        "timeout": {"attemptDurationSeconds": aws_config.get("timeout_seconds", 86400)},
+    }
+
+    # Register new revision using the session client
+    response = batch.register_job_definition(**job_definition_args)
+
+    return response["jobDefinitionName"]
+
+
+def submit_aws_batch_job(
+    image_uri: str,
+    config: dict[str, Any],
+) -> tuple[str, str]:
+    """Submit a job to AWS Batch.
+
+    Args:
+        image_uri: Full URI of the container image
+        config: Dictionary containing the configuration loaded from YAML
+
+    Returns:
+        Tuple containing (job_id, job_name) of the submitted job
+
+    Raises:
+        ValueError: If job_role_arn is not present in AWS configuration
+    """
+    session = get_boto_session(config)
+    aws_config = config["aws"]
+    region = session.region_name or aws_config["region"]
+    batch = session.client("batch")
+
+    # Generate job name (already includes unique hex string)
+    job_name = _extract_job_name_from_uri(image_uri, config)
+    print(f"\nGenerated job name: {job_name}")
+
+    # Log the job submission details
+    print("\nSubmitting job with configuration:")
+    print(f"Job Name: {job_name}")
+    print(f"Queue: {aws_config['job_queue']}")
+    print("Container Configuration:")
+    print(f"- Image: {image_uri}")
+    print(f"- vCPUs: {aws_config['batch_job']['compute_specs']['vcpus']}")
+    print(f"- Memory: {aws_config['batch_job']['compute_specs']['memory']} MiB")
+    print(f"- GPUs: {aws_config['batch_job']['compute_specs'].get('gpus', 0)}")
+    print(f"- Timeout: {aws_config.get('timeout_seconds', 86400)} seconds")
+    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)
+
+    print("Environment Variables:", list(env_vars.keys()))
+
+    # Create/Update Job Definition using the config (now implicitly uses the correct session)
+    job_definition = create_or_update_job_definition(image_uri, config)
+    print(f"\nUsing job definition: {job_definition}")
+
+    # Prepare job submission arguments
+    job_submit_args = {
+        "jobName": job_name,
+        "jobQueue": aws_config["job_queue"],
+        "jobDefinition": job_definition,
+        "containerOverrides": {
+            "environment": [
+                {"name": key, "value": str(value)} for key, value in env_vars.items()
+            ],
+        },
+    }
+
+    # Add array job configuration if specified
+    if "array_size" in aws_config:
+        array_size = aws_config["array_size"]
+        if array_size > 1:
+            print(f"\nConfiguring as array job with {array_size} instances")
+            job_submit_args["arrayProperties"] = {"size": array_size}
+
+    # Configure retry strategy for array jobs
+    retry_attempts = aws_config.get("retry_attempts", 2)
+    print(f"Setting retry attempts to {retry_attempts}")
+    job_submit_args["retryStrategy"] = {"attempts": retry_attempts}
+
+    # Submit the job using the session client
+    response = batch.submit_job(**job_submit_args)
+
+    job_id = response["jobId"]
+    print(f"\nJob submitted with ID: {job_id}")
+
+    # Print instructions for monitoring
+    print("\nTo monitor your job:")
+    print(
+        f"  1. AWS Console: https://{region}.console.aws.amazon.com/batch/home?region={region}#jobs/detail/{job_id}"
+    )
+    print(f"  2. CloudWatch Logs: Check logs for job {job_name} (ID: {job_id})")
+
+    # For array jobs, provide additional monitoring info
+    if "array_size" in aws_config and aws_config["array_size"] > 1:
+        print(f"  3. This is an array job with {aws_config['array_size']} child jobs")
+        print(
+            f"     Child jobs: https://{region}.console.aws.amazon.com/batch/home?region={region}#jobs/array-jobs/{job_id}"
+        )
+
+    return job_id, job_name

dayhoff_tools/deployment/deploy_gcp.py

@@ -0,0 +1,176 @@
+"""GCP-specific deployment functionality."""
+
+import json
+import os
+import subprocess
+import tempfile
+
+from dayhoff_tools.deployment.deploy_utils import get_container_env_vars
+
+
+def check_job_exists(job_name: str, region: str) -> bool:
+    """Check if a job with the given name already exists in GCP Batch.
+
+    Args:
+        job_name: Name of the job to check
+        region: GCP region to check in
+
+    Returns:
+        bool: True if the job exists, False otherwise
+
+    Note:
+        This uses gcloud batch jobs describe, which will return a non-zero
+        exit code if the job doesn't exist.
+    """
+    try:
+        subprocess.run(
+            [
+                "gcloud",
+                "batch",
+                "jobs",
+                "describe",
+                job_name,
+                "--location",
+                region,
+            ],
+            check=True,
+            capture_output=True,  # Suppress output
+        )
+        return True
+    except subprocess.CalledProcessError:
+        return False
+
+
+def create_batch_job_config(config: dict, image_uri: str) -> dict:
+    """Create a GCP Batch job configuration from YAML config.
+
+    Args:
+        config: Dictionary containing the configuration loaded from YAML
+        image_uri: URI of the Docker image to use
+
+    Returns:
+        Dictionary containing GCP Batch job configuration
+    """
+    gcp_config = config["gcp"]
+
+    # Start with the allocation and logs policies
+    batch_config = {
+        "allocationPolicy": gcp_config["allocation_policy"],
+        "logsPolicy": gcp_config["logs_policy"],
+    }
+
+    entrypoint_command = config["docker"].get("container_entrypoint")
+    if entrypoint_command is None:
+        raise ValueError("docker.container_entrypoint is required in configuration")
+
+    if not isinstance(entrypoint_command, list) or not all(
+        isinstance(x, str) for x in entrypoint_command
+    ):
+        raise ValueError("docker.container_entrypoint must be a list of strings")
+
+    # Build the container configuration with bash entrypoint
+    container_config = {
+        "imageUri": image_uri,
+        "entrypoint": "/bin/bash",
+        "commands": ["-c", " ".join(entrypoint_command)],
+    }
+
+    # Add shared memory option if specified
+    if "shared_memory" in config.get("docker", {}):
+        container_config["options"] = f"--shm-size={config['docker']['shared_memory']}"
+
+    # Build the task group configuration
+    task_group = {
+        "taskCount": gcp_config["batch_job"]["taskCount"],
+        "parallelism": gcp_config["batch_job"]["parallelism"],
+        "taskSpec": {
+            "computeResource": gcp_config["batch_job"]["computeResource"],
+            "runnables": [{"container": container_config}],
+        },
+    }
+
+    # Get all environment variables, including special ones like WANDB_API_KEY and GCP credentials
+    env_vars = get_container_env_vars(config)
+
+    # Add environment variables if any exist
+    if env_vars:
+        task_group["taskSpec"]["runnables"][0]["environment"] = {"variables": env_vars}
+
+    # Add machine type and optional accelerators from instance config
+    instance_config = gcp_config["batch_job"]["instance"]
+    if "machineType" in instance_config:
+        # Add machine type to the allocation policy
+        if "policy" not in batch_config["allocationPolicy"]["instances"]:
+            batch_config["allocationPolicy"]["instances"]["policy"] = {}
+        batch_config["allocationPolicy"]["instances"]["policy"]["machineType"] = (
+            instance_config["machineType"]
+        )
+
+    # Add accelerators if present (optional)
+    if "accelerators" in instance_config:
+        batch_config["allocationPolicy"]["instances"]["policy"]["accelerators"] = (
+            instance_config["accelerators"]
+        )
+
+    # Add the task group to the configuration
+    batch_config["taskGroups"] = [task_group]
+
+    # Debug logging to verify configuration
+    print("\nGCP Batch Configuration:")
+    print("------------------------")
+    try:
+        policy = batch_config["allocationPolicy"]["instances"]["policy"]
+        print("Machine Type:", policy.get("machineType", "Not specified"))
+        print("Accelerators:", policy.get("accelerators", "Not specified"))
+        print("Environment Variables:", list(env_vars.keys()))
+    except KeyError as e:
+        print(f"Warning: Could not find {e} in configuration")
+
+    return batch_config
+
+
+def submit_gcp_batch_job(config: dict, image_uri: str) -> None:
+    """Submit a job to GCP Batch.
+
+    Args:
+        config: Dictionary containing the configuration loaded from YAML
+        image_uri: URI of the Docker image to use
+
+    Raises:
+        ValueError: If a job with the same name already exists
+    """
+    job_name = config["gcp"]["job_name"]
+    region = config["gcp"]["region"]
+
+    # Check if job already exists
+    if check_job_exists(job_name, region):
+        raise ValueError(
+            f"Job '{job_name}' already exists in region {region}. "
+            "Please choose a different job name or delete the existing job first."
+        )
+
+    # Create GCP Batch job configuration
+    batch_config = create_batch_job_config(config, image_uri)
+
+    # Write the configuration to a temporary file
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file:
+        json.dump(batch_config, temp_file, indent=2)
+        temp_config_path = temp_file.name
+
+    try:
+        # Submit the job using gcloud
+        command = [
+            "gcloud",
+            "batch",
+            "jobs",
+            "submit",
+            job_name,
+            "--location",
+            region,
+            "--config",
+            temp_config_path,
+        ]
+        subprocess.run(command, check=True)
+    finally:
+        # Clean up the temporary file
+        os.unlink(temp_config_path)