dayhoff-tools 1.14.1__py3-none-any.whl → 1.14.3__py3-none-any.whl

dayhoff_tools/cli/batch/commands/submit.py

@@ -0,0 +1,215 @@
+"""Submit command for generic batch jobs."""
+
+import os
+from pathlib import Path
+
+import click
+import yaml
+
+from ..aws_batch import BatchClient, BatchError
+from ..job_id import generate_job_id
+from ..manifest import (
+    BATCH_JOBS_BASE,
+    BatchConfig,
+    InputConfig,
+    JobManifest,
+    JobStatus,
+    create_job_directory,
+    save_manifest,
+)
+
+
+# Default job definition for generic jobs
+DEFAULT_JOB_DEFINITION = "dayhoff-batch-base"
+DEFAULT_QUEUE = "t4-1x-spot"
+
+
+@click.command()
+@click.option("-f", "--file", "config_file", type=click.Path(exists=True), help="Config file path")
+@click.option("--command", help="Command to run (alternative to config file)")
+@click.option("--queue", default=DEFAULT_QUEUE, help=f"Batch queue [default: {DEFAULT_QUEUE}]")
+@click.option("--memory", default="30G", help="Memory limit (e.g., 30G)")
+@click.option("--vcpus", default=8, type=int, help="Number of vCPUs")
+@click.option("--gpus", default=1, type=int, help="Number of GPUs")
+@click.option("--array", default=1, type=int, help="Number of array tasks")
+@click.option("--retry", default=3, type=int, help="Retry attempts")
+@click.option("--timeout", default="6h", help="Job timeout (e.g., 6h, 1d)")
+@click.option("--image", help="Pre-built image URI")
+@click.option("--env", multiple=True, help="Environment variables (KEY=VALUE)")
+@click.option("--dry-run", is_flag=True, help="Show plan without submitting")
+@click.option("--base-path", default=BATCH_JOBS_BASE, help="Base path for job data")
+def submit(
+    config_file,
+    command,
+    queue,
+    memory,
+    vcpus,
+    gpus,
+    array,
+    retry,
+    timeout,
+    image,
+    env,
+    dry_run,
+    base_path,
+):
+    """Submit a custom batch job.
+
+    Jobs can be defined via a config file (-f) or inline options.
+
+    \b
+    Examples:
+      # Submit from config file
+      dh batch submit -f config.yaml
+
+      # Submit with inline command
+      dh batch submit --command "python train.py --epochs 100" --queue a10g-1x-spot
+
+      # Array job
+      dh batch submit -f config.yaml --array 10
+
+    \b
+    Config file format (YAML):
+      command: python scripts/train.py --epochs 100
+      queue: t4-1x-spot
+      memory: 30G
+      vcpus: 8
+      gpus: 1
+      array: 10
+      retry: 3
+      timeout: 6h
+      image: custom-image:tag
+      env:
+        MY_VAR: value
+    """
+    # Parse config file if provided
+    config = {}
+    if config_file:
+        with open(config_file) as f:
+            config = yaml.safe_load(f)
+
+    # Override with command-line options
+    job_command = command or config.get("command")
+    if not job_command:
+        raise click.UsageError("Must specify --command or provide config file with 'command' field")
+
+    job_queue = queue if queue != DEFAULT_QUEUE else config.get("queue", queue)
+    job_memory = memory if memory != "30G" else config.get("memory", memory)
+    job_vcpus = vcpus if vcpus != 8 else config.get("vcpus", vcpus)
+    job_gpus = gpus if gpus != 1 else config.get("gpus", gpus)
+    job_array = array if array != 1 else config.get("array", array)
+    job_retry = retry if retry != 3 else config.get("retry", retry)
+    job_timeout = timeout if timeout != "6h" else config.get("timeout", timeout)
+    job_image = image or config.get("image")
+
+    # Parse environment variables
+    job_env = dict(config.get("env", {}))
+    for e in env:
+        if "=" in e:
+            key, value = e.split("=", 1)
+            job_env[key] = value
+
+    # Generate job ID
+    job_id = generate_job_id("batch")
+
+    # Parse timeout
+    timeout_seconds = _parse_timeout(job_timeout)
+
+    # Show plan
+    click.echo()
+    click.echo(f"Job ID:      {job_id}")
+    click.echo(f"Command:     {job_command}")
+    click.echo(f"Queue:       {job_queue}")
+    click.echo(f"Resources:   {job_vcpus} vCPUs, {job_memory} memory, {job_gpus} GPUs")
+    click.echo(f"Array Size:  {job_array}")
+    click.echo(f"Retry:       {job_retry}")
+    click.echo(f"Timeout:     {job_timeout} ({timeout_seconds}s)")
+    if job_image:
+        click.echo(f"Image:       {job_image}")
+    if job_env:
+        click.echo(f"Environment: {len(job_env)} variables")
+
+    if dry_run:
+        click.echo()
+        click.echo(click.style("Dry run - job not submitted", fg="yellow"))
+        return
+
+    click.echo()
+
+    # Create job directory and manifest
+    job_dir = create_job_directory(job_id, base_path)
+    click.echo(f"Created job directory: {job_dir}")
+
+    manifest = JobManifest(
+        job_id=job_id,
+        user=job_id.split("-")[0],  # Extract username from job ID
+        pipeline="batch",
+        status=JobStatus.PENDING,
+        command=job_command,
+        image_uri=job_image,
+        batch=BatchConfig(
+            queue=job_queue,
+            array_size=job_array if job_array > 1 else None,
+        ),
+    )
+
+    # Submit to AWS Batch
+    try:
+        client = BatchClient()
+
+        # Prepare environment
+        submit_env = {
+            "JOB_DIR": str(job_dir),
+            "JOB_ID": job_id,
+            **job_env,
+        }
+
+        batch_job_id = client.submit_job(
+            job_name=job_id,
+            job_definition=job_image or DEFAULT_JOB_DEFINITION,
+            job_queue=job_queue,
+            array_size=job_array if job_array > 1 else None,
+            environment=submit_env,
+            timeout_seconds=timeout_seconds,
+            retry_attempts=job_retry,
+        )
+
+        # Update manifest with Batch job ID
+        manifest.status = JobStatus.SUBMITTED
+        manifest.batch.job_id = batch_job_id
+        save_manifest(manifest, base_path)
+
+        click.echo(click.style("✓ Job submitted successfully!", fg="green"))
+        click.echo()
+        click.echo(f"AWS Batch Job ID: {batch_job_id}")
+        click.echo()
+        click.echo("Next steps:")
+        click.echo(f"  Check status: dh batch status {job_id}")
+        click.echo(f"  View logs:    dh batch logs {job_id}")
+        click.echo(f"  Cancel:       dh batch cancel {job_id}")
+
+    except BatchError as e:
+        manifest.status = JobStatus.FAILED
+        manifest.error_message = str(e)
+        save_manifest(manifest, base_path)
+        click.echo(click.style(f"✗ Failed to submit job: {e}", fg="red"), err=True)
+        raise SystemExit(1)
+
+
+def _parse_timeout(timeout_str: str) -> int:
+    """Parse timeout string to seconds.
+
+    Supports formats like: 6h, 1d, 30m, 3600
+    """
+    timeout_str = timeout_str.strip().lower()
+
+    if timeout_str.endswith("h"):
+        return int(timeout_str[:-1]) * 3600
+    elif timeout_str.endswith("d"):
+        return int(timeout_str[:-1]) * 86400
+    elif timeout_str.endswith("m"):
+        return int(timeout_str[:-1]) * 60
+    elif timeout_str.endswith("s"):
+        return int(timeout_str[:-1])
+    else:
+        return int(timeout_str)

dayhoff_tools/cli/batch/job_id.py

@@ -0,0 +1,151 @@
+"""Job ID generation for batch jobs.
+
+Job IDs follow the format: {username}-{pipeline}-{YYYYMMDD}-{random4}
+Examples:
+- dma-embed-20260109-a3f2
+- josh-boltz-20260109-b7c1
+- sam-batch-20260109-c9d2 (for generic submit jobs)
+"""
+
+import json
+import os
+import secrets
+import subprocess
+from datetime import datetime
+from functools import lru_cache
+
+
+class JobIdError(Exception):
+    """Error generating job ID."""
+
+    pass
+
+
+@lru_cache(maxsize=1)
+def get_aws_username() -> str:
+    """Extract username from AWS SSO session.
+
+    Attempts multiple methods in order:
+    1. AWS_SSO_USER environment variable (if set by dh aws login)
+    2. Parse from `aws sts get-caller-identity` ARN
+
+    Returns:
+        Username string (lowercase, alphanumeric only)
+
+    Raises:
+        JobIdError: If username cannot be determined
+    """
+    # Method 1: Check environment variable (fastest)
+    env_user = os.environ.get("AWS_SSO_USER")
+    if env_user:
+        return _sanitize_username(env_user)
+
+    # Method 2: Parse from STS caller identity
+    try:
+        result = subprocess.run(
+            ["aws", "sts", "get-caller-identity", "--output", "json"],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+        if result.returncode == 0:
+            identity = json.loads(result.stdout)
+            arn = identity.get("Arn", "")
+            # ARN format: arn:aws:sts::ACCOUNT:assumed-role/AWSReservedSSO_ROLE/username
+            # or: arn:aws:iam::ACCOUNT:user/username
+            if "/AWSReservedSSO_" in arn:
+                # SSO assumed role - username is last part
+                username = arn.split("/")[-1]
+                return _sanitize_username(username)
+            elif ":user/" in arn:
+                # IAM user
+                username = arn.split("/")[-1]
+                return _sanitize_username(username)
+    except (subprocess.TimeoutExpired, json.JSONDecodeError, FileNotFoundError):
+        pass
+
+    # Method 3: Fall back to system username
+    import getpass
+
+    try:
+        username = getpass.getuser()
+        return _sanitize_username(username)
+    except Exception:
+        pass
+
+    raise JobIdError(
+        "Could not determine AWS username. "
+        "Ensure you're logged in with 'dh aws login' or set AWS_SSO_USER environment variable."
+    )
+
+
+def _sanitize_username(username: str) -> str:
+    """Sanitize username to be safe for job IDs.
+
+    - Convert to lowercase
+    - Keep only alphanumeric characters
+    - Truncate to 20 characters
+    """
+    sanitized = "".join(c for c in username.lower() if c.isalnum())
+    return sanitized[:20] if sanitized else "unknown"
+
+
+def generate_job_id(pipeline: str = "batch") -> str:
+    """Generate a unique job ID.
+
+    Args:
+        pipeline: Pipeline type (e.g., 'embed', 'boltz', 'batch')
+
+    Returns:
+        Job ID in format: {username}-{pipeline}-{YYYYMMDD}-{random4}
+
+    Examples:
+        >>> generate_job_id("embed")
+        'dma-embed-20260109-a3f2'
+        >>> generate_job_id()
+        'dma-batch-20260109-b7c1'
+    """
+    username = get_aws_username()
+    date_str = datetime.now().strftime("%Y%m%d")
+    random_suffix = secrets.token_hex(2)  # 4 hex characters
+
+    # Sanitize pipeline name
+    pipeline_clean = "".join(c for c in pipeline.lower() if c.isalnum())[:10]
+
+    return f"{username}-{pipeline_clean}-{date_str}-{random_suffix}"
+
+
+def parse_job_id(job_id: str) -> dict:
+    """Parse a job ID into its components.
+
+    Args:
+        job_id: Job ID string
+
+    Returns:
+        Dictionary with keys: username, pipeline, date, suffix
+
+    Raises:
+        ValueError: If job ID format is invalid
+    """
+    parts = job_id.split("-")
+    if len(parts) < 4:
+        raise ValueError(f"Invalid job ID format: {job_id}")
+
+    # Handle usernames with dashes by taking last 3 parts as known components
+    suffix = parts[-1]
+    date_str = parts[-2]
+    pipeline = parts[-3]
+    username = "-".join(parts[:-3])
+
+    # Validate date format
+    try:
+        datetime.strptime(date_str, "%Y%m%d")
+    except ValueError:
+        raise ValueError(f"Invalid date in job ID: {date_str}")
+
+    return {
+        "username": username,
+        "pipeline": pipeline,
+        "date": date_str,
+        "suffix": suffix,
+    }

dayhoff_tools/cli/batch/manifest.py

@@ -0,0 +1,293 @@
+"""Manifest management for batch jobs.
+
+Manifests are JSON files stored in Primordial that track job metadata,
+status, and configuration. They provide the single source of truth for
+job state.
+"""
+
+import json
+import os
+import tempfile
+from datetime import datetime
+from enum import Enum
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class JobStatus(str, Enum):
+    """Possible job statuses."""
+
+    PENDING = "pending"
+    SUBMITTED = "submitted"
+    RUNNING = "running"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+    FINALIZING = "finalizing"
+    FINALIZED = "finalized"
+
+
+class InputConfig(BaseModel):
+    """Configuration for job input."""
+
+    source: str = Field(..., description="Path to input file or directory")
+    num_sequences: int | None = Field(None, description="Number of sequences (for FASTA)")
+    num_chunks: int | None = Field(None, description="Number of chunks created")
+    sequences_per_chunk: int | None = Field(None, description="Sequences per chunk")
+
+
+class BatchConfig(BaseModel):
+    """AWS Batch job configuration."""
+
+    job_id: str | None = Field(None, description="AWS Batch job ID")
+    job_definition: str | None = Field(None, description="Job definition name:revision")
+    queue: str = Field(..., description="Batch queue name")
+    array_size: int | None = Field(None, description="Array job size")
+
+
+class OutputConfig(BaseModel):
+    """Configuration for job output."""
+
+    destination: str | None = Field(None, description="Final output path")
+    finalized: bool = Field(False, description="Whether output has been finalized")
+
+
+class RetryInfo(BaseModel):
+    """Information about a retry attempt."""
+
+    retry_id: str = Field(..., description="Retry job ID")
+    indices: list[int] = Field(..., description="Array indices being retried")
+    batch_job_id: str | None = Field(None, description="AWS Batch job ID for retry")
+    created: datetime = Field(default_factory=datetime.utcnow)
+
+
+class JobManifest(BaseModel):
+    """Complete manifest for a batch job."""
+
+    job_id: str = Field(..., description="Job ID")
+    user: str = Field(..., description="Username who submitted the job")
+    pipeline: str = Field(..., description="Pipeline type (embed-t5, boltz, batch)")
+    status: JobStatus = Field(JobStatus.PENDING, description="Current job status")
+    created: datetime = Field(default_factory=datetime.utcnow)
+    updated: datetime = Field(default_factory=datetime.utcnow)
+
+    input: InputConfig | None = Field(None, description="Input configuration")
+    batch: BatchConfig | None = Field(None, description="Batch job configuration")
+    output: OutputConfig | None = Field(None, description="Output configuration")
+
+    retries: list[RetryInfo] = Field(default_factory=list, description="Retry history")
+
+    # Additional metadata
+    image_uri: str | None = Field(None, description="Container image URI")
+    command: str | None = Field(None, description="Command to run")
+    error_message: str | None = Field(None, description="Error message if failed")
+
+    class Config:
+        json_encoders = {datetime: lambda v: v.isoformat()}
+
+
+# Default base path for job data
+BATCH_JOBS_BASE = "/primordial/.batch-jobs"
+
+
+def get_job_dir(job_id: str, base_path: str = BATCH_JOBS_BASE) -> Path:
+    """Get the directory path for a job.
+
+    Args:
+        job_id: Job ID
+        base_path: Base path for batch jobs (default: /primordial/.batch-jobs)
+
+    Returns:
+        Path to job directory
+    """
+    return Path(base_path) / job_id
+
+
+def get_manifest_path(job_id: str, base_path: str = BATCH_JOBS_BASE) -> Path:
+    """Get the manifest file path for a job.
+
+    Args:
+        job_id: Job ID
+        base_path: Base path for batch jobs
+
+    Returns:
+        Path to manifest.json
+    """
+    return get_job_dir(job_id, base_path) / "manifest.json"
+
+
+def create_job_directory(job_id: str, base_path: str = BATCH_JOBS_BASE) -> Path:
+    """Create the directory structure for a new job.
+
+    Creates:
+    - {base_path}/{job_id}/
+    - {base_path}/{job_id}/input/
+    - {base_path}/{job_id}/output/
+
+    Args:
+        job_id: Job ID
+        base_path: Base path for batch jobs
+
+    Returns:
+        Path to job directory
+    """
+    job_dir = get_job_dir(job_id, base_path)
+    (job_dir / "input").mkdir(parents=True, exist_ok=True)
+    (job_dir / "output").mkdir(parents=True, exist_ok=True)
+    return job_dir
+
+
+def save_manifest(manifest: JobManifest, base_path: str = BATCH_JOBS_BASE) -> Path:
+    """Save a manifest to disk atomically.
+
+    Uses write-to-temp-then-rename for atomicity to prevent corruption
+    if interrupted.
+
+    Args:
+        manifest: JobManifest to save
+        base_path: Base path for batch jobs
+
+    Returns:
+        Path to saved manifest
+    """
+    manifest.updated = datetime.utcnow()
+    manifest_path = get_manifest_path(manifest.job_id, base_path)
+    manifest_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Write to temp file first, then rename for atomicity
+    temp_fd, temp_path = tempfile.mkstemp(
+        dir=manifest_path.parent, prefix=".manifest_", suffix=".json"
+    )
+    try:
+        with os.fdopen(temp_fd, "w") as f:
+            f.write(manifest.model_dump_json(indent=2))
+        os.rename(temp_path, manifest_path)
+    except Exception:
+        # Clean up temp file on error
+        if os.path.exists(temp_path):
+            os.unlink(temp_path)
+        raise
+
+    return manifest_path
+
+
+def load_manifest(job_id: str, base_path: str = BATCH_JOBS_BASE) -> JobManifest:
+    """Load a manifest from disk.
+
+    Args:
+        job_id: Job ID
+        base_path: Base path for batch jobs
+
+    Returns:
+        JobManifest
+
+    Raises:
+        FileNotFoundError: If manifest doesn't exist
+        ValueError: If manifest is invalid
+    """
+    manifest_path = get_manifest_path(job_id, base_path)
+    if not manifest_path.exists():
+        raise FileNotFoundError(f"Manifest not found for job: {job_id}")
+
+    with open(manifest_path) as f:
+        data = json.load(f)
+
+    return JobManifest(**data)
+
+
+def update_manifest(
+    job_id: str, updates: dict[str, Any], base_path: str = BATCH_JOBS_BASE
+) -> JobManifest:
+    """Update specific fields in a manifest.
+
+    Args:
+        job_id: Job ID
+        updates: Dictionary of fields to update
+        base_path: Base path for batch jobs
+
+    Returns:
+        Updated JobManifest
+    """
+    manifest = load_manifest(job_id, base_path)
+
+    # Apply updates
+    for key, value in updates.items():
+        if hasattr(manifest, key):
+            setattr(manifest, key, value)
+        else:
+            raise ValueError(f"Unknown manifest field: {key}")
+
+    save_manifest(manifest, base_path)
+    return manifest
+
+
+def list_jobs(
+    base_path: str = BATCH_JOBS_BASE,
+    user: str | None = None,
+    status: JobStatus | None = None,
+    pipeline: str | None = None,
+    limit: int = 50,
+) -> list[JobManifest]:
+    """List jobs from the batch jobs directory.
+
+    Args:
+        base_path: Base path for batch jobs
+        user: Filter by username
+        status: Filter by status
+        pipeline: Filter by pipeline type
+        limit: Maximum number of jobs to return
+
+    Returns:
+        List of JobManifest objects, sorted by created date (newest first)
+    """
+    base = Path(base_path)
+    if not base.exists():
+        return []
+
+    manifests = []
+    for job_dir in base.iterdir():
+        if not job_dir.is_dir():
+            continue
+
+        manifest_path = job_dir / "manifest.json"
+        if not manifest_path.exists():
+            continue
+
+        try:
+            with open(manifest_path) as f:
+                data = json.load(f)
+            manifest = JobManifest(**data)
+
+            # Apply filters
+            if user and manifest.user != user:
+                continue
+            if status and manifest.status != status:
+                continue
+            if pipeline and manifest.pipeline != pipeline:
+                continue
+
+            manifests.append(manifest)
+        except (json.JSONDecodeError, ValueError):
+            # Skip invalid manifests
+            continue
+
+    # Sort by created date, newest first
+    manifests.sort(key=lambda m: m.created, reverse=True)
+
+    return manifests[:limit]
+
+
+def delete_job_directory(job_id: str, base_path: str = BATCH_JOBS_BASE) -> None:
+    """Delete a job directory and all its contents.
+
+    Args:
+        job_id: Job ID
+        base_path: Base path for batch jobs
+    """
+    import shutil
+
+    job_dir = get_job_dir(job_id, base_path)
+    if job_dir.exists():
+        shutil.rmtree(job_dir)