Track Real-Time Progress Of Jupyter Notebook Execution With Cell-Level Granularity

When executing notebooks through Jupyter Scheduler, users often face the challenge of monitoring the real-time progress of long-running jobs. This lack of visibility can lead to user frustration and uncertainty about whether a job is actively progressing, stuck on a particular cell, or how much work remains. This article proposes a solution to address this issue by introducing a mechanism to track completed cells during Jupyter Notebook execution, providing users with valuable insights into job progress and enabling more informed decision-making.

Problem Statement: The Need for Real-Time Progress Tracking

Currently, when users execute notebooks using Jupyter Scheduler, they lack real-time visibility into the progress of these jobs. This absence of feedback creates several pain points:

• Uncertainty about Job Status: Users cannot determine if a job is actively progressing or stalled on a specific cell, leading to anxiety and uncertainty.
• Difficulty in Estimating Remaining Work: Without progress updates, it's challenging to estimate how much longer a job will take to complete, making it difficult to plan and manage workflows effectively.
• Unnecessary Job Termination: Due to the lack of feedback, users may prematurely stop jobs, assuming they are stuck or not progressing, leading to wasted resources and time.
• Debugging Challenges: When a job fails, it's difficult to pinpoint the exact cell where the failure occurred, making debugging a time-consuming process.

These challenges highlight the critical need for a solution that provides real-time progress tracking during Jupyter Notebook execution. By offering users clear insights into the number of cells completed, the system can empower them to make informed decisions, optimize workflows, and debug issues more efficiently.

Proposed Solution: Tracking Completed Cells

To address the challenges outlined above, we propose adding a completed_cells field to the Job model within Jupyter Scheduler. This field will track the number of cells executed during a job's execution. By updating this field after each cell execution, users can monitor the job's progress in real-time.

This approach leverages the code_cells_executed attribute in the nbclient NotebookClient, which provides the number of code cells executed so far. By adapting the existing ExecutePreprocessor used by Jupyter Scheduler, we can seamlessly integrate this functionality into the execution process.

Implementing a Custom TrackingExecutePreprocessor

To achieve real-time cell tracking, we propose creating a custom TrackingExecutePreprocessor that extends the existing ExecutePreprocessor. This custom preprocessor will override the preprocess_cell method to update the completed_cells field in the database after each cell is executed.

The following Python code demonstrates the implementation of the TrackingExecutePreprocessor:

class TrackingExecutePreprocessor(ExecutePreprocessor):
    """Custom ExecutePreprocessor that tracks completed cells and updates the database"""

    def __init__(self, db_session, job_id, **kwargs):
        super().__init__(**kwargs)
        self.db_session = db_session
        self.job_id = job_id

    def preprocess_cell(self, cell, resources, index):
        """
        Override to track completed cells in the database.
        Calls the superclass implementation and then updates the database.
        """
        # Call the superclass implementation
        cell, resources = super().preprocess_cell(cell, resources, index)

        # Update the database with the current count of completed cells
        with self.db_session() as session:
            session.query(Job).filter(Job.job_id == self.job_id).update(
                {"completed_cells": self.code_cells_executed}
            )
            session.commit()

        return cell, resources

This class inherits from ExecutePreprocessor and overrides the preprocess_cell method. This method is called for each cell in the notebook before it is executed. The overridden method first calls the superclass implementation to execute the cell and then updates the database with the current count of completed cells.

Key components of the TrackingExecutePreprocessor:

• __init__(self, db_session, job_id, **kwargs): The constructor initializes the preprocessor with the database session (db_session) and the job ID (job_id).
• preprocess_cell(self, cell, resources, index): This method is called for each cell in the notebook. It first calls the superclass implementation to execute the cell. Then, it updates the completed_cells field in the database with the current value of self.code_cells_executed, which is an attribute inherited from nbclient.NotebookClient.

Implementation Steps

To fully implement this solution, we propose the following steps:

1. Model Update: Modify the Job model to include a completed_cells field, which will store the number of cells executed.
2. Implement TrackingExecutePreprocessor: Create the TrackingExecutePreprocessor as described above, which will track the number of completed cells during execution.
3. API Updates:
   • Update the GET /jobs/{job_id} endpoint to expose the completed_cells value in the response body, allowing users to retrieve the current progress of a job.
   • Update the PATCH /jobs/{job_id} endpoint to allow manual patching of the completed_cells value if needed. This could be useful for correcting discrepancies or manually adjusting progress.

By implementing these steps, we can provide users with real-time progress updates and enhance the overall user experience with Jupyter Scheduler.

Benefits of Tracking Completed Cells

Implementing the proposed solution offers several significant benefits for users of Jupyter Scheduler:

• Real-Time Job Monitoring: Users can monitor the progress of their jobs in real-time through the API, gaining insights into the number of cells executed and the overall progress of the job. This real-time feedback allows users to stay informed and make timely decisions.
• Informed Decision-Making: With real-time progress information, users can make more informed decisions about whether to continue or stop long-running jobs. For example, if a job appears to be stuck on a particular cell, users can choose to stop the job and investigate the issue, saving time and resources.
• Improved Debugging: When a job fails, the completed_cells field indicates the last cell that was successfully executed. This information can be invaluable for debugging, as it helps users pinpoint the location of the error and focus their troubleshooting efforts.

Tracking completed cells provides users with a clear understanding of job execution progress, enabling them to manage their workflows more effectively and efficiently.

Additional Benefits: Identifying Failure Points

In addition to the benefits mentioned above, the completed_cells field can also serve as a valuable debugging tool. When a job is stopped, the completed_cells field will retain the count of the last cell that was successfully executed. This information can help users quickly identify the cell where the job failed, streamlining the debugging process.

By knowing the last completed cell, users can focus their attention on the subsequent cell, where the error likely occurred. This can significantly reduce the time and effort required to identify and resolve issues in notebooks.

Similar Implementations in Other Systems

The concept of tracking notebook execution progress is not new and has been implemented in other popular notebook execution systems, such as:

• Papermill: Papermill, a popular library for parameterizing and executing Jupyter Notebooks, tracks notebook execution progress with cell-level granularity. It also includes a progress tracker in stdout, providing users with real-time feedback during execution.
• Google Colab: Google Colab, a cloud-based notebook environment, shows real-time cell execution progress in its user interface, providing users with a visual representation of job progress.

These examples demonstrate the value and feasibility of tracking notebook execution progress. By implementing a similar solution in Jupyter Scheduler, we can align with industry best practices and provide users with a more robust and user-friendly experience.

Conclusion: Enhancing Jupyter Scheduler with Real-Time Progress Tracking

In conclusion, the proposed solution of tracking completed cells in Jupyter Notebook execution offers a significant improvement to the user experience. By adding a completed_cells field to the Job model and implementing the TrackingExecutePreprocessor, we can provide users with real-time progress updates, enabling them to monitor job execution, make informed decisions, and debug issues more effectively.

This feature addresses a critical need in notebook execution systems and aligns with similar implementations in other popular platforms like Papermill and Google Colab. By implementing this solution, Jupyter Scheduler can become an even more powerful and user-friendly tool for executing and managing Jupyter Notebooks.

This enhancement will empower users to manage their long-running Jupyter Notebook jobs with greater confidence and efficiency. By providing real-time feedback and valuable insights into job progress, we can significantly improve the overall user experience and make Jupyter Scheduler an even more valuable tool for data scientists, researchers, and developers.

By incorporating this functionality, Jupyter Scheduler can provide a more transparent and user-friendly experience for users who run complex computational tasks within Jupyter Notebooks. The ability to monitor progress in real-time is crucial for users who need to ensure that their jobs are running as expected, and that any potential issues are identified and addressed quickly. Implementing this feature will not only enhance the user experience but also improve the overall reliability and efficiency of Jupyter Scheduler.

This approach to tracking completed cells during Jupyter Notebook execution is a critical step towards providing users with the tools they need to manage their computational workflows effectively. By offering real-time feedback on job progress, we can empower users to make informed decisions, optimize their processes, and ultimately achieve their research and development goals more efficiently. The integration of this feature into Jupyter Scheduler will not only enhance the user experience but also solidify its position as a leading platform for data science and scientific computing.