Add completed_cells tracking during notebook execution

agupta01 · agupta01 · commit 1a46b8dae918 · 2025-07-16T14:54:37.000-04:00
diff --git a/jupyter_scheduler/executors.py b/jupyter_scheduler/executors.py
@@ -11,12 +11,38 @@
 import nbformat
 from nbconvert.preprocessors import CellExecutionError, ExecutePreprocessor
 
-from jupyter_scheduler.models import DescribeJob, JobFeature, Status
+from jupyter_scheduler.models import DescribeJob, JobFeature, Status, UpdateJob
 from jupyter_scheduler.orm import Job, create_session
 from jupyter_scheduler.parameterize import add_parameters
 from jupyter_scheduler.utils import get_utc_timestamp
 
 
+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
+
+
 class ExecutionManager(ABC):
     """Base execution manager.
     Clients are expected to override this class
@@ -132,8 +158,12 @@ def execute(self):
             nb = add_parameters(nb, job.parameters)
 
         staging_dir = os.path.dirname(self.staging_paths["input"])
-        ep = ExecutePreprocessor(
-            kernel_name=nb.metadata.kernelspec["name"], store_widget_state=True, cwd=staging_dir
+        ep = TrackingExecutePreprocessor(
+            db_session=self.db_session,
+            job_id=self.job_id,
+            kernel_name=nb.metadata.kernelspec["name"], 
+            store_widget_state=True, 
+            cwd=staging_dir
         )
 
         try:
diff --git a/jupyter_scheduler/models.py b/jupyter_scheduler/models.py
@@ -148,6 +148,7 @@ class DescribeJob(BaseModel):
     downloaded: bool = False
     package_input_folder: Optional[bool] = None
     packaged_files: Optional[List[str]] = []
+    completed_cells: Optional[int] = None
 
     class Config:
         orm_mode = True
@@ -193,6 +194,7 @@ class UpdateJob(BaseModel):
     status: Optional[Status] = None
     name: Optional[str] = None
     compute_type: Optional[str] = None
+    completed_cells: Optional[int] = None
 
 
 class DeleteJob(BaseModel):
diff --git a/jupyter_scheduler/orm.py b/jupyter_scheduler/orm.py
@@ -103,6 +103,7 @@ class Job(CommonColumns, Base):
     url = Column(String(256), default=generate_jobs_url)
     pid = Column(Integer)
     idempotency_token = Column(String(256))
+    completed_cells = Column(Integer)
     # All new columns added to this table must be nullable to ensure compatibility during database migrations.
     # Any default values specified for new columns will be ignored during the migration process.
 
diff --git a/jupyter_scheduler/tests/test_execution_manager.py b/jupyter_scheduler/tests/test_execution_manager.py
@@ -1,10 +1,12 @@
 import shutil
 from pathlib import Path
 from typing import Tuple
+from unittest.mock import MagicMock, patch
 
 import pytest
+import nbformat
 
-from jupyter_scheduler.executors import DefaultExecutionManager
+from jupyter_scheduler.executors import DefaultExecutionManager, TrackingExecutePreprocessor
 from jupyter_scheduler.orm import Job
 
 
@@ -58,3 +60,128 @@ def test_add_side_effects_files(
 
     job = jp_scheduler_db.query(Job).filter(Job.job_id == job_id).one()
     assert side_effect_file_name in job.packaged_files
+
+
+@pytest.fixture
+def mock_cell():
+    """Create a mock notebook cell for testing"""
+    cell = nbformat.v4.new_code_cell(source="print('test')")
+    return cell
+
+
+@pytest.fixture
+def mock_resources():
+    """Create mock resources for testing"""
+    return {"metadata": {"path": "/test/path"}}
+
+
+def test_tracking_execute_preprocessor_initialization():
+    """Test TrackingExecutePreprocessor initialization"""
+    mock_db_session = MagicMock()
+    job_id = "test-job-id"
+    
+    preprocessor = TrackingExecutePreprocessor(
+        db_session=mock_db_session,
+        job_id=job_id,
+        kernel_name="python3"
+    )
+    
+    assert preprocessor.db_session == mock_db_session
+    assert preprocessor.job_id == job_id
+    assert preprocessor.kernel_name == "python3"
+
+
+def test_tracking_execute_preprocessor_updates_database(mock_cell, mock_resources):
+    """Test that TrackingExecutePreprocessor updates the database after cell execution"""
+    mock_db_session = MagicMock()
+    mock_session_context = MagicMock()
+    mock_db_session.return_value.__enter__.return_value = mock_session_context
+    
+    job_id = "test-job-id"
+    
+    with patch.object(TrackingExecutePreprocessor, 'execute_cell') as mock_execute:
+        with patch.object(TrackingExecutePreprocessor, '_check_assign_resources'):
+            preprocessor = TrackingExecutePreprocessor(
+                db_session=mock_db_session,
+                job_id=job_id,
+                kernel_name="python3"
+            )
+            
+            # Mock the code_cells_executed attribute
+            preprocessor.code_cells_executed = 3
+            preprocessor.resources = mock_resources
+            
+            # Mock the execute_cell method to return the cell
+            mock_execute.return_value = mock_cell
+            
+            # Call preprocess_cell
+            result_cell, result_resources = preprocessor.preprocess_cell(mock_cell, mock_resources, 0)
+            
+            # Verify the superclass method was called
+            mock_execute.assert_called_once_with(mock_cell, 0, store_history=True)
+            
+            # Verify database update was called
+            mock_session_context.query.assert_called_once_with(Job)
+            mock_session_context.query.return_value.filter.return_value.update.assert_called_once_with(
+                {"completed_cells": 3}
+            )
+            mock_session_context.commit.assert_called_once()
+            
+            # Verify return values
+            assert result_cell == mock_cell
+            assert result_resources == mock_resources
+
+
+def test_tracking_execute_preprocessor_handles_database_errors(mock_cell, mock_resources):
+    """Test that TrackingExecutePreprocessor handles database errors gracefully"""
+    mock_db_session = MagicMock()
+    mock_session_context = MagicMock()
+    mock_db_session.return_value.__enter__.return_value = mock_session_context
+    
+    # Make the database update raise an exception
+    mock_session_context.query.return_value.filter.return_value.update.side_effect = Exception("DB Error")
+    
+    job_id = "test-job-id"
+    
+    with patch.object(TrackingExecutePreprocessor, 'execute_cell') as mock_execute:
+        with patch.object(TrackingExecutePreprocessor, '_check_assign_resources'):
+            preprocessor = TrackingExecutePreprocessor(
+                db_session=mock_db_session,
+                job_id=job_id,
+                kernel_name="python3"
+            )
+            
+            preprocessor.code_cells_executed = 1
+            preprocessor.resources = mock_resources
+            mock_execute.return_value = mock_cell
+            
+            # The database error should propagate
+            with pytest.raises(Exception, match="DB Error"):
+                preprocessor.preprocess_cell(mock_cell, mock_resources, 0)
+
+
+def test_tracking_execute_preprocessor_uses_correct_job_id(mock_cell, mock_resources):
+    """Test that TrackingExecutePreprocessor uses the correct job_id in database queries"""
+    mock_db_session = MagicMock()
+    mock_session_context = MagicMock()
+    mock_db_session.return_value.__enter__.return_value = mock_session_context
+    
+    job_id = "specific-job-id-123"
+    
+    with patch.object(TrackingExecutePreprocessor, 'execute_cell') as mock_execute:
+        with patch.object(TrackingExecutePreprocessor, '_check_assign_resources'):
+            preprocessor = TrackingExecutePreprocessor(
+                db_session=mock_db_session,
+                job_id=job_id,
+                kernel_name="python3"
+            )
+            
+            preprocessor.code_cells_executed = 2
+            preprocessor.resources = mock_resources
+            mock_execute.return_value = mock_cell
+            
+            preprocessor.preprocess_cell(mock_cell, mock_resources, 0)
+            
+            # Verify the correct job_id is used in the filter
+            filter_call = mock_session_context.query.return_value.filter.call_args[0][0]
+            assert str(filter_call).find(job_id) != -1 or filter_call.right.value == job_id
diff --git a/jupyter_scheduler/tests/test_handlers.py b/jupyter_scheduler/tests/test_handlers.py
@@ -131,6 +131,7 @@ async def test_get_jobs_for_single_job(jp_fetch):
             url="url_a",
             create_time=1664305872620,
             update_time=1664305872620,
+            completed_cells=5,
         )
         response = await jp_fetch("scheduler", "jobs", job_id, method="GET")
 
@@ -140,6 +141,7 @@ async def test_get_jobs_for_single_job(jp_fetch):
         assert body["job_id"] == job_id
         assert body["input_filename"]
         assert body["job_files"]
+        assert body["completed_cells"] == 5
 
 
 @pytest.mark.parametrize(
@@ -320,6 +322,28 @@ async def test_patch_jobs(jp_fetch):
         mock_update_job.assert_called_once_with(job_id, UpdateJob(**body))
 
 
+async def test_patch_jobs_with_completed_cells(jp_fetch):
+    with patch("jupyter_scheduler.scheduler.Scheduler.update_job") as mock_update_job:
+        job_id = "542e0fac-1274-4a78-8340-a850bdb559c8"
+        body = {"name": "updated job", "completed_cells": 10}
+        response = await jp_fetch(
+            "scheduler", "jobs", job_id, method="PATCH", body=json.dumps(body)
+        )
+        assert response.code == 204
+        mock_update_job.assert_called_once_with(job_id, UpdateJob(**body))
+
+
+async def test_patch_jobs_completed_cells_only(jp_fetch):
+    with patch("jupyter_scheduler.scheduler.Scheduler.update_job") as mock_update_job:
+        job_id = "542e0fac-1274-4a78-8340-a850bdb559c8"
+        body = {"completed_cells": 15}
+        response = await jp_fetch(
+            "scheduler", "jobs", job_id, method="PATCH", body=json.dumps(body)
+        )
+        assert response.code == 204
+        mock_update_job.assert_called_once_with(job_id, UpdateJob(**body))
+
+
 async def test_patch_jobs_for_stop_job(jp_fetch):
     with patch("jupyter_scheduler.scheduler.Scheduler.stop_job") as mock_stop_job:
         job_id = "542e0fac-1274-4a78-8340-a850bdb559c8"
@@ -677,3 +701,73 @@ async def test_delete_job_definition_for_unexpected_error(jp_fetch):
         assert expected_http_error(
             e, 500, "Unexpected error occurred while deleting the job definition."
         )
+
+
+# Model validation tests for completed_cells field
+def test_describe_job_completed_cells_validation():
+    """Test DescribeJob model validation for completed_cells field"""
+    # Test valid integer values
+    job_data = {
+        "name": "test_job",
+        "input_filename": "test.ipynb",
+        "runtime_environment_name": "test_env",
+        "job_id": "test-job-id",
+        "url": "http://test.com/jobs/test-job-id",
+        "create_time": 1234567890,
+        "update_time": 1234567890,
+        "completed_cells": 5
+    }
+    job = DescribeJob(**job_data)
+    assert job.completed_cells == 5
+    
+    # Test None value
+    job_data["completed_cells"] = None
+    job = DescribeJob(**job_data)
+    assert job.completed_cells is None
+    
+    # Test zero value
+    job_data["completed_cells"] = 0
+    job = DescribeJob(**job_data)
+    assert job.completed_cells == 0
+    
+    # Test invalid type
+    job_data["completed_cells"] = "invalid"
+    with pytest.raises(ValidationError):
+        DescribeJob(**job_data)
+
+
+def test_update_job_completed_cells_validation():
+    """Test UpdateJob model validation for completed_cells field"""
+    # Test valid integer values
+    update_data = {"completed_cells": 10}
+    update_job = UpdateJob(**update_data)
+    assert update_job.completed_cells == 10
+    
+    # Test None value
+    update_data = {"completed_cells": None}
+    update_job = UpdateJob(**update_data)
+    assert update_job.completed_cells is None
+    
+    # Test zero value
+    update_data = {"completed_cells": 0}
+    update_job = UpdateJob(**update_data)
+    assert update_job.completed_cells == 0
+    
+    # Test invalid type
+    update_data = {"completed_cells": "invalid"}
+    with pytest.raises(ValidationError):
+        UpdateJob(**update_data)
+    
+    # Test exclude_none behavior
+    update_data = {"name": "test", "completed_cells": None}
+    update_job = UpdateJob(**update_data)
+    job_dict = update_job.dict(exclude_none=True)
+    assert "completed_cells" not in job_dict
+    assert job_dict["name"] == "test"
+    
+    # Test include completed_cells when not None
+    update_data = {"name": "test", "completed_cells": 5}
+    update_job = UpdateJob(**update_data)
+    job_dict = update_job.dict(exclude_none=True)
+    assert job_dict["completed_cells"] == 5
+    assert job_dict["name"] == "test"
diff --git a/jupyter_scheduler/tests/test_orm.py b/jupyter_scheduler/tests/test_orm.py
