SLURM Runner and cleanup in codebase

src/madengine/__init__.py

@@ -1,7 +1,7 @@
 """
 MADEngine - AI Models automation and dashboarding command-line tool.
 
-An AI Models automation and dashboarding command-line tool to run LLMs and Deep Learning 
+An AI Models automation and dashboarding command-line tool to run LLMs and Deep Learning
 models locally or remotely with CI. The MADEngine library supports AI automation with:
 - AI Models run reliably on supported platforms and drive software quality
 - Simple, minimalistic, out-of-the-box solution that enables confidence on hardware and software stack
@@ -19,4 +19,4 @@
     # Package is not installed, use a default version
     __version__ = "dev"
 
-__all__ = ["__version__"]
+__all__ = ["__version__"]

src/madengine/core/console.py

@@ -9,24 +9,22 @@
 import subprocess
 import typing
 import re
+
 # third-party modules
 import typing_extensions
 
 
 class Console:
     """Class to run console commands.
-    
+
     Attributes:
         shellVerbose (bool): The shell verbose flag.
         live_output (bool): The live output flag.
     """
-    def __init__(
-            self, 
-            shellVerbose: bool=True, 
-            live_output: bool=False
-        ) -> None:
+
+    def __init__(self, shellVerbose: bool = True, live_output: bool = False) -> None:
         """Constructor of the Console class.
-        
+
         Args:
             shellVerbose (bool): The shell verbose flag.
             live_output (bool): The live output flag.
@@ -36,19 +34,19 @@ def __init__(
 
     def _highlight_docker_operations(self, command: str) -> str:
         """Highlight docker push/pull/build/run operations for better visibility.
-        
+
         Args:
             command (str): The command to potentially highlight.
-            
+
         Returns:
             str: The highlighted command if it's a docker operation.
         """
         # Check if this is a docker operation
-        docker_push_pattern = r'^docker\s+push\s+'
-        docker_pull_pattern = r'^docker\s+pull\s+'
-        docker_build_pattern = r'^docker\s+build\s+'
-        docker_run_pattern = r'^docker\s+run\s+'
-        
+        docker_push_pattern = r"^docker\s+push\s+"
+        docker_pull_pattern = r"^docker\s+pull\s+"
+        docker_build_pattern = r"^docker\s+build\s+"
+        docker_run_pattern = r"^docker\s+run\s+"
+
         if re.match(docker_push_pattern, command, re.IGNORECASE):
             return f"\n{'='*80}\n🚀 DOCKER PUSH OPERATION: {command}\n{'='*80}"
         elif re.match(docker_pull_pattern, command, re.IGNORECASE):
@@ -57,21 +55,21 @@ def _highlight_docker_operations(self, command: str) -> str:
             return f"\n{'='*80}\n🔨 DOCKER BUILD OPERATION: {command}\n{'='*80}"
         elif re.match(docker_run_pattern, command, re.IGNORECASE):
             return f"\n{'='*80}\n🏃 DOCKER RUN OPERATION: {command}\n{'='*80}"
-        
+
         return command
 
     def _show_docker_completion(self, command: str, success: bool = True) -> None:
         """Show completion message for docker operations.
-        
+
         Args:
             command (str): The command that was executed.
             success (bool): Whether the operation was successful.
         """
-        docker_push_pattern = r'^docker\s+push\s+'
-        docker_pull_pattern = r'^docker\s+pull\s+'
-        docker_build_pattern = r'^docker\s+build\s+'
-        docker_run_pattern = r'^docker\s+run\s+'
-        
+        docker_push_pattern = r"^docker\s+push\s+"
+        docker_pull_pattern = r"^docker\s+pull\s+"
+        docker_build_pattern = r"^docker\s+build\s+"
+        docker_run_pattern = r"^docker\s+run\s+"
+
         if re.match(docker_push_pattern, command, re.IGNORECASE):
             if success:
                 print(f"✅ DOCKER PUSH COMPLETED SUCCESSFULLY")
@@ -81,7 +79,7 @@ def _show_docker_completion(self, command: str, success: bool = True) -> None:
                 print(f"{'='*80}\n")
         elif re.match(docker_pull_pattern, command, re.IGNORECASE):
             if success:
-                print(f"✅ DOCKER PULL COMPLETED SUCCESSFULLY") 
+                print(f"✅ DOCKER PULL COMPLETED SUCCESSFULLY")
                 print(f"{'='*80}\n")
             else:
                 print(f"❌ DOCKER PULL FAILED")
@@ -102,24 +100,24 @@ def _show_docker_completion(self, command: str, success: bool = True) -> None:
                 print(f"{'='*80}\n")
 
     def sh(
-            self, 
-            command: str, 
-            canFail: bool=False, 
-            timeout: int=60, 
-            secret: bool=False, 
-            prefix: str="", 
-            env: typing.Optional[typing.Dict[str, str]]=None
-        ) -> str:
+        self,
+        command: str,
+        canFail: bool = False,
+        timeout: int = 60,
+        secret: bool = False,
+        prefix: str = "",
+        env: typing.Optional[typing.Dict[str, str]] = None,
+    ) -> str:
         """Run shell command.
-        
+
         Args:
             command (str): The shell command.
             canFail (bool): The flag to allow failure.
             timeout (int): The timeout in seconds.
             secret (bool): The flag to hide the command.
             prefix (str): The prefix of the output.
             env (typing_extensions.TypedDict): The environment variables.
-        
+
         Returns:
             str: The output of the shell command.
 
@@ -149,7 +147,12 @@ def sh(
                 outs, errs = proc.communicate(timeout=timeout)
             else:
                 outs = []
-                for stdout_line in iter(lambda: proc.stdout.readline().encode('utf-8', errors='replace').decode('utf-8', errors='replace'), ""):
+                for stdout_line in iter(
+                    lambda: proc.stdout.readline()
+                    .encode("utf-8", errors="replace")
+                    .decode("utf-8", errors="replace"),
+                    "",
+                ):
                     print(prefix + stdout_line, end="")
                     outs.append(stdout_line)
                 outs = "".join(outs)
@@ -158,14 +161,14 @@ def sh(
         except subprocess.TimeoutExpired as exc:
             proc.kill()
             raise RuntimeError("Console script timeout") from exc
-        
+
         # Check for failure
         success = proc.returncode == 0
-        
+
         # Show docker operation completion status
         if not secret:
             self._show_docker_completion(command, success)
-        
+
         if proc.returncode != 0:
             if not canFail:
                 if not secret:
@@ -182,6 +185,6 @@ def sh(
                         + "' failed with exit code "
                         + str(proc.returncode)
                     )
-                
+
         # Return the output
         return outs.strip()

src/madengine/core/constants.py

@@ -8,7 +8,7 @@
     - MAD_SETUP_MODEL_DIR: Set to "true" to enable automatic MODEL_DIR setup during import
     - MODEL_DIR: Path to model directory to copy to current working directory
     - MAD_MINIO: JSON string with MinIO configuration
-    - MAD_AWS_S3: JSON string with AWS S3 configuration  
+    - MAD_AWS_S3: JSON string with AWS S3 configuration
     - NAS_NODES: JSON string with NAS nodes configuration
     - PUBLIC_GITHUB_ROCM_KEY: JSON string with GitHub token configuration
 
@@ -17,7 +17,7 @@
     1. Environment variables (as JSON strings)
     2. credential.json file
     3. Built-in defaults
-    
+
     Invalid JSON in environment variables will fall back to defaults with error logging.
 
 Copyright (c) Advanced Micro Devices, Inc. All rights reserved.
@@ -27,6 +27,7 @@
 import json
 import logging
 
+
 # Utility function for optional verbose logging of configuration
 def _log_config_info(message: str, force_print: bool = False):
     """Log configuration information either to logger or print if specified."""
@@ -35,12 +36,14 @@ def _log_config_info(message: str, force_print: bool = False):
     else:
         logging.debug(message)
 
+
 # third-party modules
 from madengine.core.console import Console
 
 # Get the model directory, if it is not set, set it to None.
 MODEL_DIR = os.environ.get("MODEL_DIR")
 
+
 def _setup_model_dir():
     """Setup model directory if MODEL_DIR environment variable is set."""
     if MODEL_DIR:
@@ -52,13 +55,15 @@ def _setup_model_dir():
         console.sh(f"cp -vLR --preserve=all {MODEL_DIR}/* {cwd_path}")
         _log_config_info(f"Model dir: {MODEL_DIR} copied to current dir: {cwd_path}")
 
+
 # Only setup model directory if explicitly requested (when not just importing for constants)
 if os.environ.get("MAD_SETUP_MODEL_DIR", "").lower() == "true":
     _setup_model_dir()
 
 # MADEngine credentials configuration
 CRED_FILE = "credential.json"
 
+
 def _load_credentials():
     """Load credentials from file with proper error handling."""
     try:
@@ -77,8 +82,10 @@ def _load_credentials():
         _log_config_info(f"Unexpected error loading {CRED_FILE}: {e}, using defaults")
         return {}
 
+
 CREDS = _load_credentials()
 
+
 def _get_nas_nodes():
     """Initialize NAS_NODES configuration."""
     if "NAS_NODES" not in os.environ:
@@ -88,29 +95,37 @@ def _get_nas_nodes():
             return CREDS["NAS_NODES"]
         else:
             _log_config_info("NAS_NODES is using default values.")
-            return [{
-                "NAME": "DEFAULT",
-                "HOST": "localhost",
-                "PORT": 22,
-                "USERNAME": "username",
-                "PASSWORD": "password",
-            }]
+            return [
+                {
+                    "NAME": "DEFAULT",
+                    "HOST": "localhost",
+                    "PORT": 22,
+                    "USERNAME": "username",
+                    "PASSWORD": "password",
+                }
+            ]
     else:
         _log_config_info("NAS_NODES is loaded from env variables.")
         try:
             return json.loads(os.environ["NAS_NODES"])
         except json.JSONDecodeError as e:
-            _log_config_info(f"Error parsing NAS_NODES environment variable: {e}, using defaults")
-            return [{
-                "NAME": "DEFAULT",
-                "HOST": "localhost",
-                "PORT": 22,
-                "USERNAME": "username",
-                "PASSWORD": "password",
-            }]
+            _log_config_info(
+                f"Error parsing NAS_NODES environment variable: {e}, using defaults"
+            )
+            return [
+                {
+                    "NAME": "DEFAULT",
+                    "HOST": "localhost",
+                    "PORT": 22,
+                    "USERNAME": "username",
+                    "PASSWORD": "password",
+                }
+            ]
+
 
 NAS_NODES = _get_nas_nodes()
 
+
 def _get_mad_aws_s3():
     """Initialize MAD_AWS_S3 configuration."""
     if "MAD_AWS_S3" not in os.environ:
@@ -129,14 +144,18 @@ def _get_mad_aws_s3():
         try:
             return json.loads(os.environ["MAD_AWS_S3"])
         except json.JSONDecodeError as e:
-            _log_config_info(f"Error parsing MAD_AWS_S3 environment variable: {e}, using defaults")
+            _log_config_info(
+                f"Error parsing MAD_AWS_S3 environment variable: {e}, using defaults"
+            )
             return {
                 "USERNAME": None,
                 "PASSWORD": None,
             }
 
+
 MAD_AWS_S3 = _get_mad_aws_s3()
 
+
 # Check the MAD_MINIO environment variable which is a dict.
 def _get_mad_minio():
     """Initialize MAD_MINIO configuration."""
@@ -150,24 +169,28 @@ def _get_mad_minio():
             return {
                 "USERNAME": None,
                 "PASSWORD": None,
-                "MINIO_ENDPOINT": "http://localhost:9000",  
+                "MINIO_ENDPOINT": "http://localhost:9000",
                 "AWS_ENDPOINT_URL_S3": "http://localhost:9000",
             }
     else:
         _log_config_info("MAD_MINIO is loaded from env variables.")
         try:
             return json.loads(os.environ["MAD_MINIO"])
         except json.JSONDecodeError as e:
-            _log_config_info(f"Error parsing MAD_MINIO environment variable: {e}, using defaults")
+            _log_config_info(
+                f"Error parsing MAD_MINIO environment variable: {e}, using defaults"
+            )
             return {
                 "USERNAME": None,
                 "PASSWORD": None,
-                "MINIO_ENDPOINT": "http://localhost:9000",  
+                "MINIO_ENDPOINT": "http://localhost:9000",
                 "AWS_ENDPOINT_URL_S3": "http://localhost:9000",
             }
 
+
 MAD_MINIO = _get_mad_minio()
 
+
 def _get_public_github_rocm_key():
     """Initialize PUBLIC_GITHUB_ROCM_KEY configuration."""
     if "PUBLIC_GITHUB_ROCM_KEY" not in os.environ:
@@ -186,10 +209,13 @@ def _get_public_github_rocm_key():
         try:
             return json.loads(os.environ["PUBLIC_GITHUB_ROCM_KEY"])
         except json.JSONDecodeError as e:
-            _log_config_info(f"Error parsing PUBLIC_GITHUB_ROCM_KEY environment variable: {e}, using defaults")
+            _log_config_info(
+                f"Error parsing PUBLIC_GITHUB_ROCM_KEY environment variable: {e}, using defaults"
+            )
             return {
                 "username": None,
                 "token": None,
             }
 
+
 PUBLIC_GITHUB_ROCM_KEY = _get_public_github_rocm_key()