dayhoff-tools 1.14.0__py3-none-any.whl → 1.14.2__py3-none-any.whl

dayhoff_tools/cli/batch/__init__.py

@@ -0,0 +1,85 @@
+"""Batch job management CLI for AWS Batch.
+
+This module provides a Click-based CLI for submitting and managing batch jobs
+on AWS Batch, with support for:
+- High-level pipelines (embed-t5, boltz)
+- Generic job submission
+- Job lifecycle management (status, logs, cancel, retry, finalize)
+- Local debugging and shell access
+"""
+
+import click
+
+from .commands.boltz import boltz
+from .commands.cancel import cancel
+from .commands.embed_t5 import embed_t5
+from .commands.finalize import finalize
+from .commands.list_jobs import list_jobs
+from .commands.local import local
+from .commands.logs import logs
+from .commands.retry import retry
+from .commands.status import status
+from .commands.submit import submit
+
+
+@click.group()
+def batch_cli():
+    """Manage batch jobs on AWS Batch.
+
+    \b
+    Job Management:
+      submit     Submit a custom job from config file
+      status     Show job status
+      cancel     Cancel a running job
+      logs       View job logs
+      retry      Retry failed chunks
+      finalize   Combine results and clean up
+      local      Run a chunk locally for debugging
+      list       List recent jobs
+
+    \b
+    Embedding Pipelines:
+      embed-t5   Generate T5 protein embeddings
+
+    \b
+    Structure Prediction:
+      boltz      Predict protein structures with Boltz
+
+    \b
+    Examples:
+      # Submit an embedding job
+      dh batch embed-t5 /primordial/proteins.fasta --workers 50
+
+      # Submit a structure prediction job
+      dh batch boltz /primordial/complexes/ --workers 100
+
+      # Check job status
+      dh batch status dma-embed-20260109-a3f2
+
+      # View logs for a failed chunk
+      dh batch logs dma-embed-20260109-a3f2 --index 27
+
+      # Retry failed chunks
+      dh batch retry dma-embed-20260109-a3f2
+
+      # Finalize and combine results
+      dh batch finalize dma-embed-20260109-a3f2 --output /primordial/embeddings.h5
+    """
+    pass
+
+
+# Register job management commands
+batch_cli.add_command(submit)
+batch_cli.add_command(status)
+batch_cli.add_command(cancel)
+batch_cli.add_command(logs)
+batch_cli.add_command(retry)
+batch_cli.add_command(finalize)
+batch_cli.add_command(local)
+batch_cli.add_command(list_jobs, name="list")
+
+# Register pipeline commands
+batch_cli.add_command(embed_t5, name="embed-t5")
+batch_cli.add_command(boltz)
+
+__all__ = ["batch_cli"]

dayhoff_tools/cli/batch/aws_batch.py

@@ -0,0 +1,401 @@
+"""AWS Batch client wrapper for job submission and management."""
+
+import logging
+import time
+from dataclasses import dataclass
+from typing import Any
+
+import boto3
+from botocore.exceptions import ClientError
+
+logger = logging.getLogger(__name__)
+
+
+class BatchError(Exception):
+    """Error interacting with AWS Batch."""
+
+    pass
+
+
+@dataclass
+class ArrayJobStatus:
+    """Aggregated status for an array job."""
+
+    total: int
+    pending: int
+    runnable: int
+    starting: int
+    running: int
+    succeeded: int
+    failed: int
+
+    @property
+    def completed(self) -> int:
+        return self.succeeded + self.failed
+
+    @property
+    def in_progress(self) -> int:
+        return self.pending + self.runnable + self.starting + self.running
+
+    @property
+    def is_complete(self) -> bool:
+        return self.completed == self.total
+
+    @property
+    def success_rate(self) -> float:
+        if self.completed == 0:
+            return 0.0
+        return self.succeeded / self.completed
+
+
+class BatchClient:
+    """Client for interacting with AWS Batch."""
+
+    def __init__(self, region: str = "us-east-1"):
+        """Initialize the Batch client.
+
+        Args:
+            region: AWS region
+        """
+        self.batch = boto3.client("batch", region_name=region)
+        self.logs = boto3.client("logs", region_name=region)
+        self.region = region
+
+    def submit_job(
+        self,
+        job_name: str,
+        job_definition: str,
+        job_queue: str,
+        array_size: int | None = None,
+        environment: dict[str, str] | None = None,
+        parameters: dict[str, str] | None = None,
+        timeout_seconds: int | None = None,
+        retry_attempts: int = 3,
+        depends_on: list[dict] | None = None,
+    ) -> str:
+        """Submit a job to AWS Batch.
+
+        Args:
+            job_name: Name for the job
+            job_definition: Job definition name or ARN
+            job_queue: Queue to submit to
+            array_size: Size of array job (None for single job)
+            environment: Environment variables
+            parameters: Job parameters
+            timeout_seconds: Job timeout in seconds
+            retry_attempts: Number of retry attempts
+            depends_on: Job dependencies
+
+        Returns:
+            AWS Batch job ID
+
+        Raises:
+            BatchError: If submission fails
+        """
+        try:
+            submit_args: dict[str, Any] = {
+                "jobName": job_name,
+                "jobDefinition": job_definition,
+                "jobQueue": job_queue,
+                "retryStrategy": {"attempts": retry_attempts},
+            }
+
+            if array_size and array_size > 1:
+                submit_args["arrayProperties"] = {"size": array_size}
+
+            if environment:
+                submit_args["containerOverrides"] = {
+                    "environment": [
+                        {"name": k, "value": v} for k, v in environment.items()
+                    ]
+                }
+
+            if parameters:
+                submit_args["parameters"] = parameters
+
+            if timeout_seconds:
+                submit_args["timeout"] = {"attemptDurationSeconds": timeout_seconds}
+
+            if depends_on:
+                submit_args["dependsOn"] = depends_on
+
+            response = self.batch.submit_job(**submit_args)
+            job_id = response["jobId"]
+            logger.info(f"Submitted job {job_name} with ID {job_id}")
+            return job_id
+
+        except ClientError as e:
+            raise BatchError(f"Failed to submit job: {e}")
+
+    def submit_array_job_with_indices(
+        self,
+        job_name: str,
+        job_definition: str,
+        job_queue: str,
+        indices: list[int],
+        environment: dict[str, str] | None = None,
+        timeout_seconds: int | None = None,
+        retry_attempts: int = 3,
+    ) -> str:
+        """Submit an array job for specific indices only.
+
+        Used for retrying specific failed chunks.
+
+        Args:
+            job_name: Name for the job
+            job_definition: Job definition name or ARN
+            job_queue: Queue to submit to
+            indices: Specific array indices to run
+            environment: Environment variables
+            timeout_seconds: Job timeout in seconds
+            retry_attempts: Number of retry attempts
+
+        Returns:
+            AWS Batch job ID
+        """
+        # For small number of indices, submit individual jobs
+        # For larger numbers, we could use array job with index selection
+        # AWS Batch doesn't natively support sparse array indices, so we use a workaround
+
+        if len(indices) == 1:
+            # Single job, set index via environment
+            env = environment.copy() if environment else {}
+            env["AWS_BATCH_JOB_ARRAY_INDEX"] = str(indices[0])
+            return self.submit_job(
+                job_name=job_name,
+                job_definition=job_definition,
+                job_queue=job_queue,
+                array_size=None,
+                environment=env,
+                timeout_seconds=timeout_seconds,
+                retry_attempts=retry_attempts,
+            )
+        else:
+            # For multiple indices, we pass them as a comma-separated list
+            # The worker will pick its index from this list based on array index
+            env = environment.copy() if environment else {}
+            env["BATCH_RETRY_INDICES"] = ",".join(str(i) for i in indices)
+            return self.submit_job(
+                job_name=job_name,
+                job_definition=job_definition,
+                job_queue=job_queue,
+                array_size=len(indices),
+                environment=env,
+                timeout_seconds=timeout_seconds,
+                retry_attempts=retry_attempts,
+            )
+
+    def describe_job(self, job_id: str) -> dict:
+        """Get details for a specific job.
+
+        Args:
+            job_id: AWS Batch job ID
+
+        Returns:
+            Job details dictionary
+
+        Raises:
+            BatchError: If job not found
+        """
+        try:
+            response = self.batch.describe_jobs(jobs=[job_id])
+            if not response.get("jobs"):
+                raise BatchError(f"Job not found: {job_id}")
+            return response["jobs"][0]
+        except ClientError as e:
+            raise BatchError(f"Failed to describe job: {e}")
+
+    def get_array_job_status(self, job_id: str) -> ArrayJobStatus:
+        """Get aggregated status for an array job.
+
+        Args:
+            job_id: AWS Batch job ID (parent array job)
+
+        Returns:
+            ArrayJobStatus with counts for each status
+        """
+        job = self.describe_job(job_id)
+
+        if "arrayProperties" not in job:
+            # Single job, not an array
+            status = job.get("status", "UNKNOWN")
+            return ArrayJobStatus(
+                total=1,
+                pending=1 if status == "PENDING" else 0,
+                runnable=1 if status == "RUNNABLE" else 0,
+                starting=1 if status == "STARTING" else 0,
+                running=1 if status == "RUNNING" else 0,
+                succeeded=1 if status == "SUCCEEDED" else 0,
+                failed=1 if status == "FAILED" else 0,
+            )
+
+        # Get status summary from array properties
+        status_summary = job.get("arrayProperties", {}).get("statusSummary", {})
+
+        return ArrayJobStatus(
+            total=job.get("arrayProperties", {}).get("size", 0),
+            pending=status_summary.get("PENDING", 0),
+            runnable=status_summary.get("RUNNABLE", 0),
+            starting=status_summary.get("STARTING", 0),
+            running=status_summary.get("RUNNING", 0),
+            succeeded=status_summary.get("SUCCEEDED", 0),
+            failed=status_summary.get("FAILED", 0),
+        )
+
+    def get_failed_indices(self, job_id: str) -> list[int]:
+        """Get the array indices that failed for an array job.
+
+        Args:
+            job_id: AWS Batch job ID (parent array job)
+
+        Returns:
+            List of failed array indices
+        """
+        failed_indices = []
+
+        # List child jobs with FAILED status
+        try:
+            paginator = self.batch.get_paginator("list_jobs")
+            for page in paginator.paginate(
+                arrayJobId=job_id, jobStatus="FAILED"
+            ):
+                for job_summary in page.get("jobSummaryList", []):
+                    # Extract array index from job ID (format: jobId:index)
+                    child_id = job_summary.get("jobId", "")
+                    if ":" in child_id:
+                        index = int(child_id.split(":")[-1])
+                        failed_indices.append(index)
+        except ClientError as e:
+            logger.warning(f"Failed to list child jobs: {e}")
+
+        return sorted(failed_indices)
+
+    def cancel_job(self, job_id: str, reason: str = "Cancelled by user") -> None:
+        """Cancel a job.
+
+        Args:
+            job_id: AWS Batch job ID
+            reason: Cancellation reason
+
+        Raises:
+            BatchError: If cancellation fails
+        """
+        try:
+            self.batch.cancel_job(jobId=job_id, reason=reason)
+            logger.info(f"Cancelled job {job_id}")
+        except ClientError as e:
+            raise BatchError(f"Failed to cancel job: {e}")
+
+    def terminate_job(self, job_id: str, reason: str = "Terminated by user") -> None:
+        """Terminate a running job.
+
+        Args:
+            job_id: AWS Batch job ID
+            reason: Termination reason
+
+        Raises:
+            BatchError: If termination fails
+        """
+        try:
+            self.batch.terminate_job(jobId=job_id, reason=reason)
+            logger.info(f"Terminated job {job_id}")
+        except ClientError as e:
+            raise BatchError(f"Failed to terminate job: {e}")
+
+    def get_log_stream_name(self, job_id: str) -> str | None:
+        """Get the CloudWatch log stream name for a job.
+
+        Args:
+            job_id: AWS Batch job ID
+
+        Returns:
+            Log stream name, or None if not available
+        """
+        try:
+            job = self.describe_job(job_id)
+            container = job.get("container", {})
+            return container.get("logStreamName")
+        except BatchError:
+            return None
+
+    def get_logs(
+        self,
+        job_id: str,
+        log_group: str = "/aws/batch/job",
+        tail: int = 100,
+        start_time: int | None = None,
+        follow: bool = False,
+    ) -> list[str]:
+        """Get CloudWatch logs for a job.
+
+        Args:
+            job_id: AWS Batch job ID
+            log_group: CloudWatch log group name
+            tail: Number of lines to return (from end)
+            start_time: Start time in milliseconds since epoch
+            follow: If True, continue polling for new logs
+
+        Returns:
+            List of log messages
+        """
+        log_stream = self.get_log_stream_name(job_id)
+        if not log_stream:
+            return [f"No logs available for job {job_id}"]
+
+        messages = []
+
+        try:
+            kwargs: dict[str, Any] = {
+                "logGroupName": log_group,
+                "logStreamName": log_stream,
+                "limit": tail,
+                "startFromHead": False,
+            }
+
+            if start_time:
+                kwargs["startTime"] = start_time
+
+            response = self.logs.get_log_events(**kwargs)
+
+            for event in response.get("events", []):
+                timestamp = event.get("timestamp", 0)
+                message = event.get("message", "")
+                # Format timestamp
+                dt = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(timestamp / 1000))
+                messages.append(f"[{dt}] {message}")
+
+        except ClientError as e:
+            messages.append(f"Error fetching logs: {e}")
+
+        return messages
+
+    def wait_for_job(
+        self, job_id: str, poll_interval: int = 30, timeout: int = 86400
+    ) -> str:
+        """Wait for a job to complete.
+
+        Args:
+            job_id: AWS Batch job ID
+            poll_interval: Seconds between status checks
+            timeout: Maximum seconds to wait
+
+        Returns:
+            Final job status
+
+        Raises:
+            BatchError: If timeout exceeded
+        """
+        start_time = time.time()
+
+        while True:
+            if time.time() - start_time > timeout:
+                raise BatchError(f"Timeout waiting for job {job_id}")
+
+            job = self.describe_job(job_id)
+            status = job.get("status")
+
+            if status in ("SUCCEEDED", "FAILED"):
+                return status
+
+            logger.info(f"Job {job_id} status: {status}")
+            time.sleep(poll_interval)

dayhoff_tools/cli/batch/commands/__init__.py

@@ -0,0 +1,25 @@
+"""Batch CLI commands."""
+
+from .boltz import boltz
+from .cancel import cancel
+from .embed_t5 import embed_t5
+from .finalize import finalize
+from .list_jobs import list_jobs
+from .local import local
+from .logs import logs
+from .retry import retry
+from .status import status
+from .submit import submit
+
+__all__ = [
+    "boltz",
+    "cancel",
+    "embed_t5",
+    "finalize",
+    "list_jobs",
+    "local",
+    "logs",
+    "retry",
+    "status",
+    "submit",
+]