dayhoff-tools 1.0.0__py3-none-any.whl

dayhoff_tools/deployment/job_runner.py

@@ -0,0 +1,153 @@
+"""Job runner for container-based deployments.
+
+This module serves as the unified entry point for both setup and execution modes.
+It can be run in three ways:
+1. dh job setup - Only performs environment setup
+2. dh job execute - Only executes the command (assumes setup is done)
+3. dh job setup_and_execute - Performs setup and then executes (default)
+
+Configuration is provided through environment variables:
+- JOB_COMMAND: Command to execute (if any)
+- REPO_ROOT: Optional path to repository root directory (for container environments)
+- GOOGLE_APPLICATION_CREDENTIALS_BASE64: Enables GCP authentication when present
+- USE_DVC: Set to "true" to enable DVC (requires GCP auth)
+- USE_RXNFP: Set to "true" to enable RXNFP library
+- FAIL_WITHOUT_GPU: Set to "true" to fail if GPU is unavailable
+
+Additional environment variables are preserved and passed to the job orchestrator.
+"""
+
+import logging
+import os
+import subprocess
+import sys
+
+import typer
+from dayhoff_tools.deployment.deploy_utils import (
+    SystemMonitor,
+    authenticate_gcp,
+    move_to_repo_root,
+    setup_dvc,
+    setup_rxnfp,
+)
+from dayhoff_tools.logs import configure_logs
+
+logger = logging.getLogger(__name__)
+
+
+def run_setup() -> None:
+    """Run all enabled setup steps.
+
+    Each setup function checks its own requirements and skips if not enabled.
+    """
+    logger.info("Starting job setup")
+
+    # Only log important environment variables
+    important_vars = [
+        "JOB_COMMAND",
+        "REPO_ROOT",
+        "USE_DVC",
+        "USE_RXNFP",
+        "FAIL_WITHOUT_GPU",
+    ]
+    for key in important_vars:
+        if key in os.environ:
+            logger.info(f"{key}={os.environ[key]}")
+
+    move_to_repo_root()
+
+    # Run setup steps
+    authenticate_gcp()  # Checks for GOOGLE_APPLICATION_CREDENTIALS_BASE64
+    setup_dvc()  # Checks for USE_DVC="true"
+    setup_rxnfp()  # Checks for USE_RXNFP="true"
+    logger.info("Setup completed successfully")
+
+
+def run_command() -> None:
+    """Execute the job command if specified.
+
+    Raises:
+        ValueError: If no job command is specified
+    """
+    job_command = os.getenv("JOB_COMMAND")
+    if not job_command:
+        raise ValueError("No job command specified")
+
+    logger.info(f"Current working directory: {os.getcwd()}")
+    logger.info(f"Executing job command: {job_command}")
+
+    # Start monitoring if FAIL_WITHOUT_GPU is enabled
+    monitor = None
+    if os.getenv("FAIL_WITHOUT_GPU", "").lower() == "true":
+        logger.info("Starting system monitoring...")
+        monitor = SystemMonitor(fail_without_gpu=True)
+        monitor.start()
+
+    try:
+        # Run command directly, allowing output to flow to parent process
+        # This avoids buffering issues and simplifies logging
+        result = subprocess.run(
+            job_command,
+            shell=True,
+            check=True,
+            stdout=None,  # Use parent's stdout
+            stderr=None,  # Use parent's stderr
+        )
+
+        logger.info("Command completed successfully")
+    except subprocess.CalledProcessError as e:
+        logger.error(f"Command failed with return code: {e.returncode}")
+        raise
+    except Exception as e:
+        logger.error(f"Error executing command: {str(e)}")
+        raise
+    finally:
+        if monitor:
+            logger.info("Stopping system monitor")
+            monitor.stop()
+
+
+def run_job(
+    mode: str = typer.Argument(
+        default="setup_and_execute",
+        help="Mode to run in: setup (setup only), execute (execute only), or setup_and_execute (both)",
+    )
+) -> None:
+    """Run a job command in the specified mode.
+
+    This function executes the job command given by the JOB_COMMAND environment variable,
+    if it is present. This method is meant for use in job containers after deployment.
+
+    Args:
+        mode: The execution mode to use. One of:
+            - setup: Only performs environment setup
+            - execute: Only executes the command (assumes setup is done)
+            - setup_and_execute: Performs setup and then executes (default)
+
+    Raises:
+        ValueError: If an invalid mode is specified
+        Exception: If any step of the process fails
+    """
+    # Configure logging first thing
+    configure_logs()
+    logger = logging.getLogger(__name__)
+
+    logger.info(f"Job runner starting in mode: {mode}")
+
+    if mode not in ["setup", "execute", "setup_and_execute"]:
+        logger.error(f"Invalid mode: {mode}")
+        raise ValueError(f"Invalid mode: {mode}")
+
+    try:
+        # Run in appropriate mode
+        if mode in ["setup", "setup_and_execute"]:
+            run_setup()
+
+        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)

dayhoff_tools/deployment/processors.py

@@ -0,0 +1,125 @@
+import logging
+import os
+import subprocess
+import shlex
+from abc import ABC, abstractmethod
+
+logger = logging.getLogger(__name__)
+
+
+class Processor(ABC):
+    """Processes data locally.  Abstract class for specific calculations.
+    Takes in a single file and produces a single file or folder of outputs."""
+
+    @abstractmethod
+    def run(self, input_file: str) -> str:
+        """Do the calculation, including reading from input_file
+        and writing to output_file"""
+        output_path = "output_file"
+
+        return output_path
+
+
+class BoltzPredictor(Processor):
+    """Processor for running Boltz docking predictions.
+
+    This class wraps the Boltz docking tool to predict protein structures
+    from sequence data.
+    """
+
+    def __init__(self, num_workers: int, boltz_options: str | None = None):
+        """Initialize the BoltzPredictor.
+
+        Args:
+            num_workers: Number of worker threads to use as a default.
+                         This can be overridden if --num_workers is present
+                         in boltz_options.
+            boltz_options: A string containing additional command-line options
+                           to pass to the Boltz predictor. Options should be
+                           space-separated (e.g., "--option1 value1 --option2").
+        """
+        self.num_workers = num_workers
+        self.boltz_options = boltz_options
+
+    def run(self, input_file: str) -> str:
+        """Run Boltz prediction on the input file.
+
+        Constructs the command using the input file, default number of workers,
+        and any additional options provided via `boltz_options`. If `--num_workers`
+        is specified in `boltz_options`, it overrides the default `num_workers`.
+
+        Args:
+            input_file: Path to the input file containing sequences
+
+        Returns:
+            Path to the output directory created by Boltz
+
+        Raises:
+            subprocess.CalledProcessError: If Boltz prediction fails
+        """
+        # Determine expected output directory name
+        input_base = os.path.splitext(os.path.basename(input_file))[0]
+        expected_output_dir = f"boltz_results_{input_base}"
+        logger.info(f"Expected output directory: {expected_output_dir}")
+
+        # Start building the command
+        cmd = ["boltz", "predict", input_file]
+
+        # Parse additional options if provided
+        additional_args = []
+        num_workers_in_opts = False
+        if self.boltz_options:
+            try:
+                parsed_opts = shlex.split(self.boltz_options)
+                additional_args.extend(parsed_opts)
+                if "--num_workers" in parsed_opts:
+                    num_workers_in_opts = True
+                    logger.info(
+                        f"Using --num_workers from BOLTZ_OPTIONS: {self.boltz_options}"
+                    )
+            except ValueError as e:
+                logger.error(f"Error parsing BOLTZ_OPTIONS '{self.boltz_options}': {e}")
+                # Decide if we should raise an error or proceed without options
+                # For now, proceed without the additional options
+                additional_args = []  # Clear potentially partially parsed args
+
+        # Add num_workers if not specified in options
+        if not num_workers_in_opts:
+            logger.info(f"Using default num_workers: {self.num_workers}")
+            cmd.extend(["--num_workers", str(self.num_workers)])
+
+        # Add the parsed additional arguments
+        cmd.extend(additional_args)
+
+        # Log the final command
+        # Use shlex.join for safer command logging, especially if paths/args have spaces
+        try:
+            safe_cmd_str = shlex.join(cmd)
+            logger.info(f"Running command: {safe_cmd_str}")
+        except AttributeError:  # shlex.join is Python 3.8+
+            logger.info(f"Running command: {' '.join(cmd)}")
+
+        # Stream output in real-time
+        process = subprocess.Popen(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            bufsize=1,
+        )
+
+        stdout = process.stdout
+        if stdout:
+            for line in iter(stdout.readline, ""):
+                logger.info(f"BOLTZ: {line.rstrip()}")
+
+        # Wait for process to complete
+        return_code = process.wait()
+        if return_code != 0:
+            logger.error(f"Boltz prediction failed with exit code {return_code}")
+            raise subprocess.CalledProcessError(return_code, cmd)
+
+        logger.info(
+            f"Boltz prediction completed successfully. Output in {expected_output_dir}"
+        )
+        return expected_output_dir