dayhoff-tools 1.14.7__py3-none-any.whl → 1.14.8__py3-none-any.whl

dayhoff_tools/cli/batch/__init__.py

@@ -12,6 +12,7 @@ import click
 
 from .commands.boltz import boltz
 from .commands.cancel import cancel
+from .commands.clean import clean
 from .commands.embed_t5 import embed_t5
 from .commands.finalize import finalize
 from .commands.list_jobs import list_jobs
@@ -36,6 +37,7 @@ def batch_cli():
       finalize   Combine results and clean up
       local      Run a chunk locally for debugging
       list       List recent jobs
+      clean      Remove old completed job directories
 
     \b
     Embedding Pipelines:
@@ -77,6 +79,7 @@ batch_cli.add_command(retry)
 batch_cli.add_command(finalize)
 batch_cli.add_command(local)
 batch_cli.add_command(list_jobs, name="list")
+batch_cli.add_command(clean)
 
 # Register pipeline commands
 batch_cli.add_command(embed_t5, name="embed-t5")

dayhoff_tools/cli/batch/aws_batch.py

@@ -242,6 +242,64 @@ class BatchClient:
             failed=status_summary.get("FAILED", 0),
         )
 
+    def get_job_statuses_batch(self, job_ids: list[str]) -> dict[str, str]:
+        """Get status for multiple jobs in a single API call.
+
+        AWS Batch allows up to 100 job IDs per describe_jobs call.
+        This method handles batching for larger lists.
+
+        Args:
+            job_ids: List of AWS Batch job IDs
+
+        Returns:
+            Dictionary mapping job_id -> status string
+            Status will be one of: SUBMITTED, PENDING, RUNNABLE, STARTING,
+            RUNNING, SUCCEEDED, FAILED, or "UNKNOWN" if not found.
+            For array jobs, derives overall status from child statuses.
+        """
+        if not job_ids:
+            return {}
+
+        results = {}
+        batch_size = 100  # AWS Batch limit
+
+        for i in range(0, len(job_ids), batch_size):
+            batch = job_ids[i : i + batch_size]
+            try:
+                response = self.batch.describe_jobs(jobs=batch)
+                for job in response.get("jobs", []):
+                    job_id = job.get("jobId")
+                    status = job.get("status", "UNKNOWN")
+
+                    # For array jobs, derive overall status from children
+                    if "arrayProperties" in job:
+                        summary = job["arrayProperties"].get("statusSummary", {})
+                        total = job["arrayProperties"].get("size", 0)
+                        succeeded = summary.get("SUCCEEDED", 0)
+                        failed = summary.get("FAILED", 0)
+
+                        if succeeded + failed == total:
+                            # All children complete
+                            status = "SUCCEEDED" if failed == 0 else "FAILED"
+                        elif summary.get("RUNNING", 0) > 0:
+                            status = "RUNNING"
+                        elif summary.get("STARTING", 0) > 0:
+                            status = "STARTING"
+                        elif summary.get("RUNNABLE", 0) > 0:
+                            status = "RUNNABLE"
+                        elif summary.get("PENDING", 0) > 0:
+                            status = "PENDING"
+
+                    results[job_id] = status
+            except ClientError as e:
+                logger.warning(f"Failed to describe batch of jobs: {e}")
+                # Mark these as unknown
+                for job_id in batch:
+                    if job_id not in results:
+                        results[job_id] = "UNKNOWN"
+
+        return results
+
     def get_failed_indices(self, job_id: str) -> list[int]:
         """Get the array indices that failed for an array job.
 

dayhoff_tools/cli/batch/commands/clean.py

@@ -0,0 +1,139 @@
+"""Clean command for removing old job directories."""
+
+import click
+
+from ..aws_batch import BatchClient, BatchError
+from ..manifest import (
+    BATCH_JOBS_BASE,
+    JobStatus,
+    delete_job_directory,
+    list_jobs as list_manifests,
+)
+from .status import format_time_ago, _aws_status_to_job_status
+
+
+@click.command("clean")
+@click.option("--user", help="Only clean jobs for this user")
+@click.option(
+    "--older-than",
+    type=int,
+    default=7,
+    help="Only clean jobs older than N days [default: 7]",
+)
+@click.option("--dry-run", is_flag=True, help="Show what would be cleaned without deleting")
+@click.option("--force", is_flag=True, help="Delete without confirmation")
+@click.option("--base-path", default=BATCH_JOBS_BASE, help="Base path for job data")
+def clean(user, older_than, dry_run, force, base_path):
+    """Remove completed job directories to free up space.
+
+    Only removes jobs that have SUCCEEDED or FAILED in AWS Batch.
+    Jobs that are still running or pending are never removed.
+
+    \b
+    Examples:
+      dh batch clean                     # Clean jobs older than 7 days
+      dh batch clean --older-than 1      # Clean jobs older than 1 day
+      dh batch clean --dry-run           # Show what would be cleaned
+      dh batch clean --user dma          # Only clean dma's jobs
+    """
+    from datetime import datetime, timedelta, timezone
+
+    cutoff = datetime.now(timezone.utc) - timedelta(days=older_than)
+
+    # Get all manifests
+    manifests = list_manifests(
+        base_path=base_path,
+        user=user,
+        status=None,
+        pipeline=None,
+        limit=500,
+    )
+
+    if not manifests:
+        click.echo("No jobs found.")
+        return
+
+    # Filter to old jobs
+    old_manifests = []
+    for m in manifests:
+        created = m.created
+        if created.tzinfo is None:
+            created = created.replace(tzinfo=timezone.utc)
+        if created < cutoff:
+            old_manifests.append(m)
+
+    if not old_manifests:
+        click.echo(f"No jobs older than {older_than} days found.")
+        return
+
+    # Get live statuses for old jobs
+    batch_job_ids = []
+    manifest_to_batch_id = {}
+    for m in old_manifests:
+        if m.batch and m.batch.job_id:
+            batch_job_ids.append(m.batch.job_id)
+            manifest_to_batch_id[m.job_id] = m.batch.job_id
+
+    live_statuses = {}
+    if batch_job_ids:
+        try:
+            client = BatchClient()
+            live_statuses = client.get_job_statuses_batch(batch_job_ids)
+        except BatchError as e:
+            click.echo(f"Error: Could not fetch status from AWS Batch: {e}")
+            click.echo("Cannot safely clean jobs without knowing their status.")
+            return
+
+    # Find jobs that are safe to clean (SUCCEEDED or FAILED)
+    safe_to_clean = []
+    for manifest in old_manifests:
+        if manifest.job_id in manifest_to_batch_id:
+            batch_id = manifest_to_batch_id[manifest.job_id]
+            aws_status = live_statuses.get(batch_id, "UNKNOWN")
+            if aws_status in ("SUCCEEDED", "FAILED"):
+                safe_to_clean.append((manifest, aws_status))
+        elif manifest.status in (JobStatus.FINALIZED, JobStatus.CANCELLED):
+            # Already finalized or cancelled - safe to clean
+            safe_to_clean.append((manifest, manifest.status.value.upper()))
+
+    if not safe_to_clean:
+        click.echo(f"No completed jobs older than {older_than} days to clean.")
+        return
+
+    # Show what will be cleaned
+    click.echo()
+    click.echo(f"{'JOB ID':<35} {'STATUS':<12} {'CREATED':<12}")
+    click.echo("-" * 65)
+
+    for manifest, status in safe_to_clean:
+        click.echo(
+            f"{manifest.job_id:<35} "
+            f"{status:<12} "
+            f"{format_time_ago(manifest.created):<12}"
+        )
+
+    click.echo()
+    click.echo(f"Found {len(safe_to_clean)} completed jobs to clean.")
+
+    if dry_run:
+        click.echo("(dry-run: no changes made)")
+        return
+
+    # Confirm before deleting
+    if not force:
+        if not click.confirm("Delete these job directories?"):
+            click.echo("Cancelled.")
+            return
+
+    # Delete job directories
+    deleted = 0
+    for manifest, _ in safe_to_clean:
+        try:
+            delete_job_directory(manifest.job_id, base_path)
+            deleted += 1
+            click.echo(f"  Deleted: {manifest.job_id}")
+        except Exception as e:
+            click.echo(f"  Failed to delete {manifest.job_id}: {e}")
+
+    click.echo()
+    click.echo(f"Cleaned {deleted} job directories.")

dayhoff_tools/cli/batch/commands/list_jobs.py

@@ -2,8 +2,9 @@
 
 import click
 
+from ..aws_batch import BatchClient, BatchError
 from ..manifest import BATCH_JOBS_BASE, JobStatus, list_jobs as list_manifests
-from .status import format_status, format_time_ago
+from .status import format_status, format_time_ago, _aws_status_to_job_status
 
 
 @click.command("list")
@@ -23,6 +24,7 @@ def list_jobs(user, status_filter, pipeline, limit, base_path):
     """List recent batch jobs.
 
     Shows a table of recent jobs with their status, pipeline type, and creation time.
+    Status is fetched live from AWS Batch.
 
     \b
     Examples:
@@ -34,12 +36,13 @@ def list_jobs(user, status_filter, pipeline, limit, base_path):
     """
     status_enum = JobStatus(status_filter) if status_filter else None
 
+    # Fetch more manifests than requested to allow filtering by live status
     manifests = list_manifests(
         base_path=base_path,
         user=user,
-        status=status_enum,
+        status=None,  # Don't filter by status yet - will filter after getting live status
         pipeline=pipeline,
-        limit=limit,
+        limit=limit * 3,  # Fetch extra to account for status filtering
     )
 
     if not manifests:
@@ -48,6 +51,51 @@ def list_jobs(user, status_filter, pipeline, limit, base_path):
             click.echo("Try removing filters to see all jobs.")
         return
 
+    # Collect AWS Batch job IDs for live status lookup
+    batch_job_ids = []
+    manifest_to_batch_id = {}
+    for m in manifests:
+        if m.batch and m.batch.job_id:
+            batch_job_ids.append(m.batch.job_id)
+            manifest_to_batch_id[m.job_id] = m.batch.job_id
+
+    # Fetch live statuses from AWS Batch
+    live_statuses = {}
+    if batch_job_ids:
+        try:
+            client = BatchClient()
+            live_statuses = client.get_job_statuses_batch(batch_job_ids)
+        except BatchError as e:
+            click.echo(f"Warning: Could not fetch live status from AWS Batch: {e}")
+
+    # Build display data with live status
+    display_data = []
+    for manifest in manifests:
+        # Use live status if available, otherwise fall back to manifest status
+        if manifest.job_id in manifest_to_batch_id:
+            batch_id = manifest_to_batch_id[manifest.job_id]
+            aws_status = live_statuses.get(batch_id)
+            if aws_status:
+                live_status = _aws_status_to_job_status(aws_status)
+            else:
+                live_status = manifest.status
+        else:
+            live_status = manifest.status
+
+        # Apply status filter if specified
+        if status_enum and live_status != status_enum:
+            continue
+
+        display_data.append((manifest, live_status))
+
+        # Stop once we have enough
+        if len(display_data) >= limit:
+            break
+
+    if not display_data:
+        click.echo("No jobs found matching filters.")
+        return
+
     # Print header
     click.echo()
     click.echo(
@@ -55,17 +103,17 @@ def list_jobs(user, status_filter, pipeline, limit, base_path):
     )
     click.echo("-" * 85)
 
-    for manifest in manifests:
+    for manifest, live_status in display_data:
         click.echo(
             f"{manifest.job_id:<35} "
-            f"{format_status(manifest.status):<21} "  # Extra space for ANSI color codes
+            f"{format_status(live_status):<21} "  # Extra space for ANSI color codes
             f"{manifest.pipeline:<12} "
             f"{manifest.user:<10} "
             f"{format_time_ago(manifest.created):<12}"
         )
 
     click.echo()
-    click.echo(f"Showing {len(manifests)} jobs.")
+    click.echo(f"Showing {len(display_data)} jobs.")
 
     # Show filter hints
     hints = []

dayhoff_tools/cli/batch/commands/status.py

@@ -80,21 +80,79 @@ def status(job_id, user, status_filter, pipeline, base_path):
         _show_job_list(user, status_filter, pipeline, base_path)
 
 
+def _aws_status_to_job_status(aws_status: str) -> JobStatus:
+    """Convert AWS Batch status to JobStatus enum."""
+    mapping = {
+        "SUBMITTED": JobStatus.SUBMITTED,
+        "PENDING": JobStatus.PENDING,
+        "RUNNABLE": JobStatus.RUNNING,  # Runnable means waiting for compute
+        "STARTING": JobStatus.RUNNING,
+        "RUNNING": JobStatus.RUNNING,
+        "SUCCEEDED": JobStatus.SUCCEEDED,
+        "FAILED": JobStatus.FAILED,
+    }
+    return mapping.get(aws_status, JobStatus.SUBMITTED)
+
+
 def _show_job_list(user, status_filter, pipeline, base_path):
     """Show a list of recent jobs."""
     status_enum = JobStatus(status_filter) if status_filter else None
     manifests = list_manifests(
         base_path=base_path,
         user=user,
-        status=status_enum,
+        status=None,  # Don't filter yet - we'll filter after getting live status
         pipeline=pipeline,
-        limit=20,
+        limit=50,  # Fetch more, filter later
     )
 
     if not manifests:
         click.echo("No jobs found.")
         return
 
+    # Collect AWS Batch job IDs for live status lookup
+    batch_job_ids = []
+    manifest_to_batch_id = {}
+    for m in manifests:
+        if m.batch and m.batch.job_id:
+            batch_job_ids.append(m.batch.job_id)
+            manifest_to_batch_id[m.job_id] = m.batch.job_id
+
+    # Fetch live statuses from AWS Batch
+    live_statuses = {}
+    if batch_job_ids:
+        try:
+            client = BatchClient()
+            live_statuses = client.get_job_statuses_batch(batch_job_ids)
+        except BatchError as e:
+            click.echo(f"Warning: Could not fetch live status from AWS Batch: {e}")
+
+    # Build display data with live status
+    display_data = []
+    for manifest in manifests:
+        # Use live status if available, otherwise fall back to manifest status
+        if manifest.job_id in manifest_to_batch_id:
+            batch_id = manifest_to_batch_id[manifest.job_id]
+            aws_status = live_statuses.get(batch_id)
+            if aws_status:
+                live_status = _aws_status_to_job_status(aws_status)
+            else:
+                live_status = manifest.status
+        else:
+            live_status = manifest.status
+
+        # Apply status filter if specified
+        if status_enum and live_status != status_enum:
+            continue
+
+        display_data.append((manifest, live_status))
+
+    if not display_data:
+        click.echo("No jobs found matching filters.")
+        return
+
+    # Limit to 20 after filtering
+    display_data = display_data[:20]
+
     # Print header
     click.echo()
     click.echo(
@@ -102,17 +160,17 @@ def _show_job_list(user, status_filter, pipeline, base_path):
     )
     click.echo("-" * 85)
 
-    for manifest in manifests:
+    for manifest, live_status in display_data:
         click.echo(
             f"{manifest.job_id:<35} "
-            f"{format_status(manifest.status):<21} "  # Extra space for color codes
+            f"{format_status(live_status):<21} "  # Extra space for color codes
             f"{manifest.pipeline:<12} "
             f"{manifest.user:<10} "
             f"{format_time_ago(manifest.created):<12}"
         )
 
     click.echo()
-    click.echo(f"Showing {len(manifests)} most recent jobs.")
+    click.echo(f"Showing {len(display_data)} most recent jobs.")
     click.echo("Use 'dh batch status <job-id>' for details.")
 
 

{dayhoff_tools-1.14.7.dist-info → dayhoff_tools-1.14.8.dist-info}/METADATA

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

{dayhoff_tools-1.14.7.dist-info → dayhoff_tools-1.14.8.dist-info}/RECORD

@@ -7,18 +7,19 @@ dayhoff_tools/batch/workers/embed_t5.py,sha256=A5WqsQa1WZ7_la5X5wt0XUP-VwOglH04W
 dayhoff_tools/chemistry/standardizer.py,sha256=uMn7VwHnx02nc404eO6fRuS4rsl4dvSPf2ElfZDXEpY,11188
 dayhoff_tools/chemistry/utils.py,sha256=jt-7JgF-GeeVC421acX-bobKbLU_X94KNOW24p_P-_M,2257
 dayhoff_tools/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-dayhoff_tools/cli/batch/__init__.py,sha256=gEuiDPWio8ihrIfF_Q6-5hKOnliPTxDKQrKN8-4Y3Ac,2320
-dayhoff_tools/cli/batch/aws_batch.py,sha256=IMnNBXmyfQ1GkBqWV9OWYNNjxFNDXE_tv8sUWq9n1E8,12727
+dayhoff_tools/cli/batch/__init__.py,sha256=0QMxkelGUdvdd7-P-WG8YX2h0qjkKKb-0JbEHGrH7pE,2437
+dayhoff_tools/cli/batch/aws_batch.py,sha256=DOp8KvmTrie15O61DP1HT13PSr3s5imxM-VZHvaLO7Q,15174
 dayhoff_tools/cli/batch/commands/__init__.py,sha256=1xRzzL_mc1hz1Pv0OWNr-g6fkL5XbEsOTGHzrqddLCA,458
 dayhoff_tools/cli/batch/commands/boltz.py,sha256=KbBxeF0HS_ysgd7MyeJgrkcGTAXChJLQxOhvJ14SxjY,11989
 dayhoff_tools/cli/batch/commands/cancel.py,sha256=ZnHAJVzMGC0_1EQGpMSdYUlzm9yi-E9NxRJKBsetYW8,3111
+dayhoff_tools/cli/batch/commands/clean.py,sha256=v8R_YtsdkfPjrFmo7a9cMdbDkWxotSmUeE26dnRlsnU,4537
 dayhoff_tools/cli/batch/commands/embed_t5.py,sha256=QXFydAw0wndevdzXF1cxikxMmvn1BuQ5p9lwutQFajU,11453
 dayhoff_tools/cli/batch/commands/finalize.py,sha256=REt-1f1Vb7OycJJ6JTUV_xJQVKBPXz0RxxhEuNfd39g,7957
-dayhoff_tools/cli/batch/commands/list_jobs.py,sha256=55KqOuhpxpuFBk95Dv2Wd79MprJh1ouy_SWDmBpT7Zg,2438
+dayhoff_tools/cli/batch/commands/list_jobs.py,sha256=COfxZddDVUAHeTayNAB3ruYNhgrE3osgFxY2qzf33cg,4284
 dayhoff_tools/cli/batch/commands/local.py,sha256=dZeKhNakaM1jS-EoByAwg1nWspRRoOmYzcwzjEKBaIA,3226
 dayhoff_tools/cli/batch/commands/logs.py,sha256=ctgJksdzFmqBdD18ePPsZe2BpuJYtHz2xAaMPnUplmQ,5293
 dayhoff_tools/cli/batch/commands/retry.py,sha256=PbstKr7iP7Uu4csisp3cvM8al9YOFw27XVlc591jLbE,4509
-dayhoff_tools/cli/batch/commands/status.py,sha256=0xDIVQzgYl5mgpHcDhIwmy44IVt0B5L33LHX91nfWuU,7684
+dayhoff_tools/cli/batch/commands/status.py,sha256=k8ltke62vKUX_4MmnFDEd6H2s8KmVFA71UUXLbleA7o,9790
 dayhoff_tools/cli/batch/commands/submit.py,sha256=AXbvSReN8fLlzR5swE81pH7yvbCC1GUMCbsDrfoHAws,6786
 dayhoff_tools/cli/batch/job_id.py,sha256=mqr8SwcPlUWIYLaR_C4kACmL2ZFK8jaddd7B-45-XaQ,4246
 dayhoff_tools/cli/batch/manifest.py,sha256=blUw2ASxFTmdLV5ikz4VD91n5CN0V37rlyX11M0H2VE,8624
@@ -70,7 +71,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=UETBtZD3r7WgvURqfGbyHlT7cxoiVq8isjzMuerKw8I,24475
-dayhoff_tools-1.14.7.dist-info/METADATA,sha256=k1Nua41haYs1TSTYGO-tgESgBcVf2hAhFAnI2WJ9Yck,3184
-dayhoff_tools-1.14.7.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-dayhoff_tools-1.14.7.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
-dayhoff_tools-1.14.7.dist-info/RECORD,,
+dayhoff_tools-1.14.8.dist-info/METADATA,sha256=86xFQ0fq394JPcdZ-o2lRsWuDoLiGDkSw5SsePljptM,3184
+dayhoff_tools-1.14.8.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+dayhoff_tools-1.14.8.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
+dayhoff_tools-1.14.8.dist-info/RECORD,,