Add email callback on train complete

src/lightning/pytorch/callbacks/__init__.py

@@ -16,6 +16,7 @@
 from lightning.pytorch.callbacks.checkpoint import Checkpoint
 from lightning.pytorch.callbacks.device_stats_monitor import DeviceStatsMonitor
 from lightning.pytorch.callbacks.early_stopping import EarlyStopping
+from lightning.pytorch.callbacks.email_callback import EmailCallback
 from lightning.pytorch.callbacks.finetuning import BackboneFinetuning, BaseFinetuning
 from lightning.pytorch.callbacks.gradient_accumulation_scheduler import GradientAccumulationScheduler
 from lightning.pytorch.callbacks.lambda_function import LambdaCallback
@@ -42,6 +43,7 @@
  "Checkpoint",
  "DeviceStatsMonitor",
  "EarlyStopping",
+ "EmailCallback",
  "GradientAccumulationScheduler",
  "LambdaCallback",
  "LearningRateFinder",

src/lightning/pytorch/callbacks/email_callback.py

@@ -0,0 +1,166 @@
+# Copyright The Lightning AI team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+r"""
+EmailCallback
+===============
+
+Sends an email to a list of emails when training is complete.
+"""
+
+import logging
+import smtplib
+import textwrap
+from email.mime.multipart import MIMEMultipart
+from email.mime.text import MIMEText
+from enum import Enum
+from typing import List, Optional
+
+from typing_extensions import override
+
+import lightning.pytorch as pl
+from lightning.pytorch.callbacks.callback import Callback
+
+log = logging.getLogger(__name__)
+
+
+class SMTPProvider(Enum):
+ """Enum representing different SMTP providers with their server address and port.
+
+ Attributes:
+ GMAIL (tuple): Gmail SMTP server address and port.
+
+ """
+
+ GMAIL = ("smtp.gmail.com", 587)
+ # YAHOO = ("smtp.mail.yahoo.com", 587)
+ # OUTLOOK = ("smtp.office365.com", 587)
+ # ZOHO = ("smtp.zoho.com", 587)
+ # Add more providers as needed
+
+
+class EmailCallback(Callback):
+ r"""Send an email notification when training is complete.
+
+ Args:
+ sender_email: Email address of the sender.
+ password: Password for the sender's email.
+ receiver_emails: List of email addresses to send the notification to. Defaults to sender_email if None.
+ smtp_provider: SMTP provider to use for sending the email. Defaults to SMTPProvider.GMAIL.
+ metric_precision: Number of decimal places to use for metric values in the email. Defaults to 5.
+
+ Example:
+
+ >>> import os
+ >>> from lightning.pytorch import Trainer
+ >>> from lightning.pytorch.callbacks import EmailCallback
+ >>> your_password = os.getenv("EMAIL_PASSWORD") # strongly suggest not to hardcode password
+ >>> email_callback = EmailCallback(
+ ... sender_email = "[email protected]",
+ ... password = your_password,
+ ... receiver_emails = ["[email protected]"]
+ ... )
+ >>> trainer = Trainer(callbacks=[email_callback])
+
+ SMTP Providers:
+ Currently supported SMTP servers
+
+ - GMAIL: Gmail SMTP server address and port.
+
+ Attributes:
+ EMAIL_BODY_TEMPLATE (str): Template for the body of the email.
+
+ Methods:
+ on_train_end(trainer, pl_module): Called when training ends to send an email notification.
+
+ Raises:
+ Exception: If there is an error while sending the email.
+
+ """
+
+ EMAIL_BODY_TEMPLATE = textwrap.dedent(
+ """
+ Hello,
+
+ The training for the model {module} has been completed.
+
+ - Final Epoch: {final_epoch}
+ - Total Steps: {total_steps}
+
+ Logged Metrics:
+ """
+ )
+
+ def __init__(
+ self,
+ sender_email: str,
+ password: str,
+ receiver_emails: Optional[List[str]] = None,
+ smtp_provider: SMTPProvider = SMTPProvider.GMAIL,
+ metric_precision: int = 5,
+ ):
+ self.sender_email = sender_email
+ self.receiver_emails = receiver_emails if receiver_emails else [sender_email]
+ self.password = password
+ self.smtp_server, self.smtp_port = smtp_provider.value
+ self.metric_precision = metric_precision
+
+ @override
+ def on_train_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+ if trainer.fast_dev_run:
+ return
+ try:
+ # Create the email message
+ msg = MIMEMultipart()
+ msg["From"] = self.sender_email
+ msg["To"] = ", ".join(self.receiver_emails)
+ msg["Subject"] = f"Training for {pl_module.__class__.__name__} completed"
+
+ # Gather detailed training information
+ final_epoch = trainer.current_epoch
+ total_steps = trainer.global_step
+ metrics = trainer.callback_metrics
+
+ # Format the body of the email with named placeholders
+ body = self.EMAIL_BODY_TEMPLATE.format(
+ module=pl_module.__class__.__name__,
+ final_epoch=final_epoch,
+ total_steps=total_steps,
+ )
+
+ for key, value in metrics.items():
+ if isinstance(value, (float, int)): # Ensure value is numeric
+ value = round(value, self.metric_precision)
+ elif hasattr(value, "item"): # For tensors or numpy values
+ value = round(value.item(), self.metric_precision)
+ body += f"- {key}: {value}\n"
+
+ body += "\nBest regards,\nPytorch Lightning"
+
+ # Attach the body with the msg instance
+ msg.attach(MIMEText(body, "plain"))
+
+ # Set up the SMTP server
+ server = smtplib.SMTP(self.smtp_server, self.smtp_port)
+ server.starttls()
+ server.login(self.sender_email, self.password)
+
+ # Send the email to each recipient
+ for recipient in self.receiver_emails:
+ server.sendmail(self.sender_email, recipient, msg.as_string())
+
+ # Quit the server
+ server.quit()
+ log.info(f"Completion email successfully sent to: {', '.join(self.receiver_emails)}")
+ except Exception as e:
+ log.exception(f"An error occurred while sending an email to confirm training completion: {e}")