AD-SDL · Dozgulbas · Jun 24, 2024 · Jun 19, 2024 · Jun 19, 2024 · Jun 19, 2024
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -17,7 +17,7 @@ repos:
       - id: nbstripout
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.4.8
+    rev: v0.4.10
     hooks:
       # Run the linter.
       - id: ruff

diff --git a/.readthedocs.yaml b/.readthedocs.yaml
@@ -9,7 +9,9 @@ version: 2
 build:
   os: ubuntu-22.04
   tools:
-    python: "3.9"
+    python: "3.11"
+  apt_packages:
+    - graphviz
 
 # Build documentation in the docs/ directory with Sphinx
 sphinx:

diff --git a/pdm.lock b/pdm.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -25,6 +25,7 @@ dependencies = [
     'ulid-py>=1.1.0',
     'uvicorn[standard]>=0.21',
     "websockets>=12.0",
+    "secure-smtplib>=0.1.1",
 ]
 
 [project.optional-dependencies]
@@ -136,7 +137,8 @@ select = [
     # "RUF"
 ]
 ignore = [
-    "E501" # Line too long
+    "E501", # Line too long
+    "B006", # Do not use mutable data structures for argument defaults
 ]
 
 # Allow fix for all enabled rules (when `--fix`) is provided.

diff --git a/requirements/dev.txt b/requirements/dev.txt
@@ -46,6 +46,7 @@ pyyaml==6.0.1
 redis==4.6.0
 requests==2.31.0
 ruff==0.4.2
+secure-smtplib==0.1.1
 setuptools==69.5.1
 six==1.16.0
 sniffio==1.3.1

diff --git a/requirements/docs.txt b/requirements/docs.txt
@@ -37,6 +37,7 @@ pytz==2024.1; python_version < "3.9"
 pyyaml==6.0.1
 redis==4.6.0
 requests==2.31.0
+secure-smtplib==0.1.1
 six==1.16.0
 sniffio==1.3.1
 snowballstemmer==2.2.0

diff --git a/requirements/requirements.txt b/requirements/requirements.txt
@@ -61,6 +61,7 @@ pyzmq==26.0.2
 redis==4.6.0
 requests==2.31.0
 ruff==0.4.2
+secure-smtplib==0.1.1
 setuptools==69.5.1
 six==1.16.0
 sniffio==1.3.1

diff --git a/src/wei/core/events.py b/src/wei/core/events.py
@@ -9,6 +9,7 @@
 from wei.core.loggers import Logger
 from wei.core.state_manager import StateManager
 from wei.types import Event
+from wei.utils import threaded_task
 
 state_manager = StateManager()
 
@@ -68,6 +69,7 @@ def str_to_bytes(s):
             cls.kafka_topic = None
 
     @classmethod
+    @threaded_task
     def log_event(cls, event: Event) -> None:
         """logs an event in the proper place for the given experiment
 

diff --git a/src/wei/core/experiment.py b/src/wei/core/experiment.py
@@ -11,6 +11,7 @@
     search_for_experiment_directory,
 )
 from wei.types.experiment_types import Campaign, Experiment, ExperimentDesign
+from wei.utils import threaded_task
 
 state_manager = StateManager()
 
@@ -62,6 +63,7 @@ def get_experiment(experiment_id: str) -> Experiment:
     return experiment
 
 
+@threaded_task
 def parse_experiments_from_disk():
     """Scans the experiments directory and pulls in any experiments that are not in the state_manager."""
     experiments_dir = get_experiments_directory()
@@ -87,3 +89,5 @@ def parse_experiments_from_disk():
                 ),
             )
             state_manager.set_experiment(experiment)
+        except Exception:
+            continue
diff --git a/src/wei/core/notifications.py b/src/wei/core/notifications.py
@@ -0,0 +1,89 @@
+"""
+Code for sending notifications to workcell users
+"""
+
+import smtplib
+from email.mime.multipart import MIMEMultipart
+from email.mime.text import MIMEText
+
+from wei.config import Config
+from wei.core.state_manager import StateManager
+from wei.types import Step
+from wei.types.workflow_types import WorkflowRun
+from wei.utils import threaded_task
+
+state_manager = StateManager()
+
+
+@threaded_task
+def send_email(subject: str, email_address: str, body: str):
+    """Sends an email with the provided subject and body to the specified email address, using the configured SMTP relay server"""
+    smtp_server = Config.smtp_server
+    smtp_port = Config.smtp_port
+    sender = "no-reply-rpl@anl.gov"
+
+    try:
+        # Create the MIMEText objects for the email content
+        msg = MIMEMultipart("alternative")
+        msg["Subject"] = subject
+        msg["From"] = sender
+        msg["To"] = email_address
+
+        # Attach both plain text and HTML versions
+        import re
+
+        part1 = MIMEText(
+            re.sub("<[^<]+?>", "", body), "plain"
+        )  # * Strip HTML tags for plaintext
+        part2 = MIMEText(body, "html")
+        msg.attach(part1)
+        msg.attach(part2)
+
+        # Send the email via the SMTP server
+        with smtplib.SMTP(smtp_server, smtp_port) as server:
+            server.sendmail(sender, email_address, msg.as_string())
+        print(f"Email sent to {email_address}")
+        return True
+    except Exception as e:
+        print(f"Error sending email to {email_address}: {e}")
+        return False
+
+
+def send_failed_step_notification(
+    workflow_run: WorkflowRun,
+    step: Step,
+) -> None:
+    """Send notifications using the configuration defined in the Workcell definition/cli args."""
+
+    experiment = state_manager.get_experiment(workflow_run.experiment_id)
+    for email_address in experiment.email_addresses:
+        # * Email Content
+        subject = f"STEP FAILED: {step.name}"
+        body_html = f"""\
+        <html>
+        <body>
+            <h1>Step '{step.name}' Failed</h1>
+            <ul>
+                <li>Step: {step.name} ({step.id})</li>
+                <li>Module: {step.module}</li>
+                <li>Workflow: {workflow_run.name} ({workflow_run.run_id})</li>
+                <li>Experiment: {experiment.experiment_name} ({experiment.experiment_id})</li>
+            </ul>
+            See below for more info.
+            <h2>Step Response</h2>
+            <pre>{step.result.model_dump_json(indent=2)}</pre>
+            <h2>Step Info</h2>
+            <pre>{step.model_dump_json(indent=2)}</pre>
+            <h2>Workflow Info</h2>
+            <pre>{workflow_run.model_dump_json(indent=2)}</pre>
+            <h2>Experiment Info</h2>
+            <pre>{experiment.model_dump_json(indent=2)}</pre>
+        </body>
+        </html>
+        """
+
+        send_email(
+            subject=subject,
+            email_address=email_address,
+            body=body_html,
+        )
diff --git a/src/wei/core/step.py b/src/wei/core/step.py
@@ -9,6 +9,7 @@
 from wei.core.location import free_source_and_target, update_source_and_target
 from wei.core.loggers import Logger
 from wei.core.module import clear_module_reservation, get_module_about
+from wei.core.notifications import send_failed_step_notification
 from wei.core.state_manager import StateManager
 from wei.core.storage import get_workflow_run_directory
 from wei.types import (
@@ -56,6 +57,7 @@ def validate_step(step: Step) -> Tuple[bool, str]:
                             f"Step '{step.name}': Module {step.module}'s action, '{step.action}', is missing file '{action_file.name}'",
                         )
                 return True, f"Step '{step.name}': Validated successfully"
+
         return (
             False,
             f"Step '{step.name}': Module {step.module} has no action '{step.action}'",
@@ -139,6 +141,8 @@ def run_step(
     step.end_time = datetime.now()
     step.duration = step.end_time - step.start_time
     step.result = step_response
+    if step.result.action_response == StepStatus.FAILED:
+        send_failed_step_notification(wf_run, step)
     send_event(WorkflowStepEvent.from_wf_run(wf_run=wf_run, step=step))
     wf_run.hist[step.name] = step_response
     if step_response.action_response == StepStatus.FAILED:

diff --git a/src/wei/engine.py b/src/wei/engine.py
@@ -2,7 +2,6 @@
 Engine Class and associated helpers and data
 """
 
-import threading
 import time
 import traceback
 
@@ -30,8 +29,7 @@ def __init__(self) -> None:
         """Initialize the scheduler."""
         self.state_manager = StateManager()
         initialize_storage()
-        disk_scan_thread = threading.Thread(target=parse_experiments_from_disk)
-        disk_scan_thread.start()
+        parse_experiments_from_disk()
         self.state_manager.clear_state(
             reset_locations=Config.reset_locations,
             clear_workflow_runs=Config.clear_workflow_runs,

diff --git a/src/wei/experiment_client.py b/src/wei/experiment_client.py
@@ -36,6 +36,7 @@ def __init__(
         campaign_id: Optional[str] = None,
         description: Optional[str] = None,
         working_dir: Optional[PathLike] = None,
+        email_addresses: List[str] = [],
     ) -> None:
         """Initializes an Experiment, and creates its log files
 
@@ -56,12 +57,15 @@ def __init__(
         working_dir: Optional[Union[str, Path]]
             The directory to resolve relative paths from. Defaults to the current working directory.
 
+        email_addresses: Optional[List[str]]
+            List of email addresses to send notifications at the end of the experiment
         """
 
         self.server_host = server_host
         self.server_port = server_port
         self.url = f"http://{self.server_host}:{self.server_port}"
         self.experiment_id = experiment_id
+        self.email_addresses = email_addresses
 
         if experiment_name is None:
             assert (
@@ -72,6 +76,7 @@ def __init__(
             experiment_name=experiment_name,
             campaign_id=campaign_id,
             description=description,
+            email_addresses=email_addresses,
         )
 
         if working_dir is None:

diff --git a/src/wei/routers/event_routes.py b/src/wei/routers/event_routes.py
@@ -2,7 +2,6 @@
 Router for the "events" endpoints
 """
 
-import threading
 from typing import Any
 
 from fastapi import APIRouter
@@ -19,8 +18,7 @@
 @router.post("/")
 def log_event(event: Event) -> Any:
     """Logs a value to the log file for a given experiment"""
-    thread = threading.Thread(target=EventHandler.log_event(event))
-    thread.start()
+    EventHandler.log_event(event)
 
     return event
 

diff --git a/src/wei/types/experiment_types.py b/src/wei/types/experiment_types.py
@@ -27,6 +27,8 @@ class ExperimentDesign(BaseModel):
     """ID of the campaign this experiment should be associated with (note: this campaign must already exist)"""
     description: Optional[str] = None
     """Description of the experiment"""
+    email_addresses: List[str] = []
+    """List of email addresses to send notifications"""
 
 
 class Experiment(ExperimentDesign):

diff --git a/src/wei/types/workcell_types.py b/src/wei/types/workcell_types.py
@@ -88,6 +88,10 @@ class WorkcellConfig(BaseModel, extra="allow"):
     cold_start_delay: int = Field(
         default=2, description="Delay before starting the engine"
     )
+    smtp_server: str = Field(
+        default="mailgateway.anl.gov", description="Hostname for the SMTP server"
+    )
+    smtp_port: int = Field(default=25, description="Port number for the SMTP server")
 
     # Validators
     @field_validator("data_directory")

diff --git a/src/wei/utils.py b/src/wei/utils.py
@@ -119,3 +119,18 @@ def json_to_csv(json_data, csv_file_path):
         writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
         writer.writeheader()
         writer.writerows(flattened_data)
+
+
+def threaded_task(func):
+    """Mark a function as a threaded task, to be run without awaiting. Returns the thread object, so you _can_ await if needed."""
+
+    import functools
+    import threading
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs) -> threading.Thread:
+        thread = threading.Thread(target=func, args=args, kwargs=kwargs)
+        thread.start()
+        return thread
+
+    return wrapper
diff --git a/tests/test_base.py b/tests/test_base.py
@@ -27,6 +27,7 @@ def __init__(self, *args, **kwargs):
             server_port=self.server_port,
             experiment_name="Test_Experiment",
             working_dir=Path(__file__).resolve().parent,
+            email_addresses=["ryan.lewis@anl.gov"],
         )
         self.url = f"http://{self.server_host}:{self.server_port}"
         self.redis_host = self.workcell.config.redis_host