diff --git a/src/rompy/backends/config.py b/src/rompy/backends/config.py index a8694c67..d9f0d74d 100644 --- a/src/rompy/backends/config.py +++ b/src/rompy/backends/config.py @@ -102,6 +102,10 @@ class LocalConfig(BaseBackendConfig): True, description="Whether to capture stdout and stderr" ) + stream_output: bool = Field( + False, description="Whether to stream output in real-time" + ) + def get_backend_class(self): """Return the LocalRunBackend class.""" from rompy.run import LocalRunBackend @@ -117,6 +121,7 @@ def get_backend_class(self): "command": "python run_model.py", }, {"timeout": 3600, "working_dir": "/path/to/model/dir"}, + {"timeout": 3600, "stream_output": True}, ] } ) @@ -287,90 +292,59 @@ def model_post_init(self, __context) -> None: class SlurmConfig(BaseBackendConfig): """Configuration for SLURM cluster execution.""" - model_type: Literal["slurm"] = Field( - "slurm", - description="The backend type." - ) + model_type: Literal["slurm"] = Field("slurm", description="The backend type.") queue: Optional[str] = Field( - None, - description="SLURM partition name (equivalent to queue)" + None, description="SLURM partition name (equivalent to queue)" ) command: str = Field( - ..., - description="Shell command to run in the workspace directory", - min_length=1 - ) - nodes: int = Field( - 1, - ge=1, - le=100, - description="Number of nodes to allocate" - ) - ntasks: int = Field( - 1, - ge=1, - description="Number of tasks (processes) to run" + ..., description="Shell command to run in the workspace directory", min_length=1 ) + nodes: int = Field(1, ge=1, le=100, description="Number of nodes to allocate") + ntasks: int = Field(1, ge=1, description="Number of tasks (processes) to run") cpus_per_task: int = Field( - 1, - ge=1, - le=128, - description="Number of CPU cores per task" - ) - time_limit: str = Field( - "1:00:00", - description="Time limit in format HH:MM:SS" + 1, ge=1, le=128, description="Number of CPU cores per task" ) + time_limit: str = Field("1:00:00", description="Time limit in format HH:MM:SS") account: Optional[str] = Field( - None, - description="Account for billing/resource tracking" - ) - qos: Optional[str] = Field( - None, - description="Quality of Service for the job" + None, description="Account for billing/resource tracking" ) + qos: Optional[str] = Field(None, description="Quality of Service for the job") reservation: Optional[str] = Field( - None, - description="Reservation name to run job under" + None, description="Reservation name to run job under" ) output_file: Optional[str] = Field( - None, - description="Output file path for job output" + None, description="Output file path for job output" ) error_file: Optional[str] = Field( - None, - description="Error file path for job errors" - ) - job_name: Optional[str] = Field( - None, - description="Name for the SLURM job" + None, description="Error file path for job errors" ) + job_name: Optional[str] = Field(None, description="Name for the SLURM job") mail_type: Optional[str] = Field( - None, - description="Type of mail to send (BEGIN, END, FAIL, ALL, etc.)" + None, description="Type of mail to send (BEGIN, END, FAIL, ALL, etc.)" ) mail_user: Optional[str] = Field( - None, - description="Email address for notifications" + None, description="Email address for notifications" ) additional_options: List[str] = Field( default_factory=list, - description="Additional SLURM options (e.g., '--gres=gpu:1')" + description="Additional SLURM options (e.g., '--gres=gpu:1')", ) - @field_validator('time_limit') + @field_validator("time_limit") @classmethod def validate_time_limit(cls, v): """Validate time limit format (HH:MM:SS).""" import re - if not re.match(r'^\d{1,4}:\d{2}:\d{2}$', v): + + if not re.match(r"^\d{1,4}:\d{2}:\d{2}$", v): raise ValueError("Time limit must be in format HH:MM:SS") return v def get_backend_class(self): """Return the SlurmRunBackend class.""" from rompy.run.slurm import SlurmRunBackend + return SlurmRunBackend model_config = ConfigDict( @@ -402,4 +376,4 @@ def get_backend_class(self): # Type alias for all backend configurations -BackendConfig = Union[LocalConfig, DockerConfig, SlurmConfig] \ No newline at end of file +BackendConfig = Union[LocalConfig, DockerConfig, SlurmConfig] diff --git a/src/rompy/run/__init__.py b/src/rompy/run/__init__.py index bab286b0..171dd316 100644 --- a/src/rompy/run/__init__.py +++ b/src/rompy/run/__init__.py @@ -7,8 +7,9 @@ import logging import os import subprocess +import threading from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Dict, Optional if TYPE_CHECKING: from rompy.backends import LocalConfig @@ -52,6 +53,7 @@ def run( exec_working_dir = config.working_dir exec_env_vars = config.env_vars exec_timeout = config.timeout + exec_stream_output = getattr(config, "stream_output", False) logger.debug( f"Using LocalConfig: timeout={exec_timeout}, env_vars={list(exec_env_vars.keys())}" @@ -96,7 +98,7 @@ def run( # Execute command or config.run() if exec_command: success = self._execute_command( - exec_command, work_dir, env, exec_timeout + exec_command, work_dir, env, exec_timeout, exec_stream_output ) else: success = self._execute_config_run(model_run, work_dir, env) @@ -118,7 +120,12 @@ def run( return False def _execute_command( - self, command: str, work_dir: Path, env: Dict[str, str], timeout: Optional[int] + self, + command: str, + work_dir: Path, + env: Dict[str, str], + timeout: Optional[int], + stream_output: bool = False, ) -> bool: """Execute a shell command. @@ -127,10 +134,20 @@ def _execute_command( work_dir: Working directory env: Environment variables timeout: Execution timeout + stream_output: Whether to stream output in real-time Returns: True if successful, False otherwise """ + if stream_output: + return self._execute_command_streaming(command, work_dir, env, timeout) + else: + return self._execute_command_buffered(command, work_dir, env, timeout) + + def _execute_command_buffered( + self, command: str, work_dir: Path, env: Dict[str, str], timeout: Optional[int] + ) -> bool: + """Execute a shell command with buffered output.""" logger.info(f"Executing command: {command}") logger.debug(f"Working directory: {work_dir}") @@ -146,7 +163,6 @@ def _execute_command( check=False, ) - # Log output if result.stdout: logger.info(f"Command stdout:\n{result.stdout}") if result.stderr: @@ -169,6 +185,80 @@ def _execute_command( logger.exception(f"Error executing command: {e}") return False + def _execute_command_streaming( + self, command: str, work_dir: Path, env: Dict[str, str], timeout: Optional[int] + ) -> bool: + """Execute a shell command with streaming output.""" + logger.info(f"Executing command: {command}") + logger.debug(f"Working directory: {work_dir}") + + try: + process = subprocess.Popen( + command, + shell=True, + cwd=work_dir, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1, + ) + + # Capture output while streaming + stdout_lines = [] + stderr_lines = [] + + def read_stream(stream, lines, log_func): + """Read from a stream and log each line.""" + for line in stream: + line = line.rstrip() + lines.append(line) + log_func(line) + + # Start threads to read stdout and stderr concurrently + stdout_thread = threading.Thread( + target=read_stream, + args=(process.stdout, stdout_lines, logger.info), + ) + stderr_thread = threading.Thread( + target=read_stream, + args=(process.stderr, stderr_lines, lambda msg: logger.warning(msg)), + ) + + stdout_thread.start() + stderr_thread.start() + + # Wait for process to complete with timeout + try: + returncode = process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + logger.error(f"Command timed out after {timeout} seconds") + raise TimeoutError( + f"Command execution timed out after {timeout} seconds" + ) + + # Wait for reader threads to finish + stdout_thread.join() + stderr_thread.join() + + # Log remaining stderr if any (after process completed) + # (Thread already handled most output, but capture any final lines) + + if returncode == 0: + logger.debug("Command completed successfully") + return True + else: + logger.error(f"Command failed with return code: {returncode}") + if stderr_lines: + logger.error("Command stderr:\n" + "\n".join(stderr_lines)) + return False + + except Exception as e: + logger.exception(f"Error executing command: {e}") + return False + def _execute_config_run( self, model_run, work_dir: Path, env: Dict[str, str] ) -> bool: