Remove classmethod refactoring + make cell tracking optional

Author: agupta01

jupyter_scheduler/executors.py

@@ -29,7 +29,9 @@ class ExecutionManager(ABC):
     _model = None
     _db_session = None
 
-    def __init__(self, job_id: str, root_dir: str, db_url: str, staging_paths: Dict[str, str]):
+    def __init__(
+        self, job_id: str, root_dir: str, db_url: str, staging_paths: Dict[str, str]
+    ):
         self.job_id = job_id
         self.staging_paths = staging_paths
         self.root_dir = root_dir
@@ -74,18 +76,16 @@ def execute(self):
         """
         pass
 
-    @classmethod
     @abstractmethod
-    def supported_features(cls) -> Dict[JobFeature, bool]:
+    def supported_features(self) -> Dict[JobFeature, bool]:
         """Returns a configuration of supported features
         by the execution engine. Implementors are expected
         to override this to return a dictionary of supported
         job creation features.
         """
         pass
 
-    @classmethod
-    def validate(cls, input_path: str) -> bool:
+    def validate(self, input_path: str) -> bool:
         """Returns True if notebook has valid metadata to execute, False otherwise"""
         return True
 
@@ -134,7 +134,9 @@ def execute(self):
         staging_dir = os.path.dirname(self.staging_paths["input"])
 
         ep = ExecutePreprocessor(
-            kernel_name=nb.metadata.kernelspec["name"], store_widget_state=True, cwd=staging_dir
+            kernel_name=nb.metadata.kernelspec["name"],
+            store_widget_state=True,
+            cwd=staging_dir,
         )
 
         if self.supported_features().get(JobFeature.track_cell_execution, False):
@@ -173,10 +175,14 @@ def add_side_effects_files(self, staging_dir: str):
         if new_files_set:
             with self.db_session() as session:
                 current_packaged_files_set = set(
-                    session.query(Job.packaged_files).filter(Job.job_id == self.job_id).scalar()
+                    session.query(Job.packaged_files)
+                    .filter(Job.job_id == self.job_id)
+                    .scalar()
                     or []
                 )
-                updated_packaged_files = list(current_packaged_files_set.union(new_files_set))
+                updated_packaged_files = list(
+                    current_packaged_files_set.union(new_files_set)
+                )
                 session.query(Job).filter(Job.job_id == self.job_id).update(
                     {"packaged_files": updated_packaged_files}
                 )
@@ -186,11 +192,12 @@ def create_output_files(self, job: DescribeJob, notebook_node):
         for output_format in job.output_formats:
             cls = nbconvert.get_exporter(output_format)
             output, _ = cls().from_notebook_node(notebook_node)
-            with fsspec.open(self.staging_paths[output_format], "w", encoding="utf-8") as f:
+            with fsspec.open(
+                self.staging_paths[output_format], "w", encoding="utf-8"
+            ) as f:
                 f.write(output)
 
-    @classmethod
-    def supported_features(cls) -> Dict[JobFeature, bool]:
+    def supported_features(self) -> Dict[JobFeature, bool]:
         return {
             JobFeature.job_name: True,
             JobFeature.output_formats: True,
@@ -205,11 +212,10 @@ def supported_features(cls) -> Dict[JobFeature, bool]:
             JobFeature.output_filename_template: False,
             JobFeature.stop_job: True,
             JobFeature.delete_job: True,
-            JobFeature.track_cell_execution: True,
+            JobFeature.track_cell_execution: False,
         }
 
-    @classmethod
-    def validate(cls, input_path: str) -> bool:
+    def validate(self, input_path: str) -> bool:
         with open(input_path, encoding="utf-8") as f:
             nb = nbformat.read(f, as_version=4)
             try:

jupyter_scheduler/scheduler.py

@@ -442,7 +442,7 @@ def create_job(self, model: CreateJob) -> str:
             raise InputUriError(model.input_uri)
 
         input_path = os.path.join(self.root_dir, model.input_uri)
-        if not self.execution_manager_class.validate(input_path):
+        if not self.execution_manager_class.validate(self.execution_manager_class, input_path):
             raise SchedulerError(
                 """There is no kernel associated with the notebook. Please open
                     the notebook, select a kernel, and re-submit the job to execute.

jupyter_scheduler/tests/mocks.py

@@ -16,7 +16,7 @@ def execute(self):
     def process(self):
         pass
 
-    def supported_features(cls) -> Dict[JobFeature, bool]:
+    def supported_features(self) -> Dict[JobFeature, bool]:
         return {
             JobFeature.job_name: True,
             JobFeature.output_formats: True,

jupyter_scheduler/tests/test_execution_manager.py

@@ -61,8 +61,8 @@ def test_add_side_effects_files(
     assert side_effect_file_name in job.packaged_files
 
 
-def test_default_execution_manager_cell_tracking_hook():
-    """Test that DefaultExecutionManager sets up on_cell_executed hook when track_cell_execution is supported"""
+def test_default_execution_manager_cell_tracking_hook_not_set_by_default():
+    """Test that DefaultExecutionManager does NOT set up on_cell_executed hook when track_cell_execution is disabled by default"""
     job_id = "test-job-id"
 
     with patch.object(DefaultExecutionManager, "model") as mock_model:
@@ -100,9 +100,8 @@ def test_default_execution_manager_cell_tracking_hook():
                             # Verify ExecutePreprocessor was created
                             mock_ep_class.assert_called_once()
 
-                            # Verify on_cell_executed hook was set
-                            assert hasattr(mock_ep, "on_cell_executed")
-                            assert mock_ep.on_cell_executed is not None
+                            # Verify patching method was never called
+                            mock_model.__update_completed_cells_hook.assert_not_called()
 
 
 def test_update_completed_cells_hook():
@@ -158,8 +157,8 @@ def test_update_completed_cells_hook_database_error():
     # Mock db_session with error
     mock_db_session = MagicMock()
     mock_session_context = MagicMock()
-    mock_session_context.query.return_value.filter.return_value.update.side_effect = Exception(
-        "DB Error"
+    mock_session_context.query.return_value.filter.return_value.update.side_effect = (
+        Exception("DB Error")
     )
     mock_db_session.return_value.__enter__.return_value = mock_session_context
     manager._db_session = mock_db_session
@@ -181,12 +180,18 @@ def test_update_completed_cells_hook_database_error():
 
 def test_supported_features_includes_track_cell_execution():
     """Test that DefaultExecutionManager supports track_cell_execution feature"""
-    features = DefaultExecutionManager.supported_features()
+    manager = DefaultExecutionManager(
+        job_id="test-job-id",
+        root_dir="/test",
+        db_url="sqlite:///:memory:",
+        staging_paths={"input": "/test/input.ipynb"},
+    )
+    features = manager.supported_features()
 
     from jupyter_scheduler.models import JobFeature
 
     assert JobFeature.track_cell_execution in features
-    assert features[JobFeature.track_cell_execution] is True
+    assert features[JobFeature.track_cell_execution] is False
 
 
 def test_hook_uses_correct_job_id():
@@ -233,8 +238,7 @@ def test_cell_tracking_disabled_when_feature_false():
 
     # Create a custom execution manager class with track_cell_execution = False
     class DisabledTrackingExecutionManager(DefaultExecutionManager):
-        @classmethod
-        def supported_features(cls):
+        def supported_features(self):
             features = super().supported_features()
             from jupyter_scheduler.models import JobFeature
 
@@ -256,8 +260,12 @@ def supported_features(cls):
         with patch.object(DisabledTrackingExecutionManager, "model") as mock_model:
             with patch("jupyter_scheduler.executors.open", mock=MagicMock()):
                 with patch("jupyter_scheduler.executors.nbformat.read") as mock_nb_read:
-                    with patch.object(DisabledTrackingExecutionManager, "add_side_effects_files"):
-                        with patch.object(DisabledTrackingExecutionManager, "create_output_files"):
+                    with patch.object(
+                        DisabledTrackingExecutionManager, "add_side_effects_files"
+                    ):
+                        with patch.object(
+                            DisabledTrackingExecutionManager, "create_output_files"
+                        ):
                             with patch(
                                 "jupyter_scheduler.executors.ExecutePreprocessor"
                             ) as mock_ep_class:
@@ -288,15 +296,20 @@ def test_disabled_tracking_feature_support():
 
     # Create a custom execution manager class with track_cell_execution = False
     class DisabledTrackingExecutionManager(DefaultExecutionManager):
-        @classmethod
-        def supported_features(cls):
+        def supported_features(self):
             features = super().supported_features()
             from jupyter_scheduler.models import JobFeature
 
             features[JobFeature.track_cell_execution] = False
             return features
 
-    features = DisabledTrackingExecutionManager.supported_features()
+    manager = DisabledTrackingExecutionManager(
+        job_id="test-job-id",
+        root_dir="/test",
+        db_url="sqlite:///:memory:",
+        staging_paths={"input": "/test/input.ipynb"},
+    )
+    features = manager.supported_features()
 
     from jupyter_scheduler.models import JobFeature