Added support to cancel all job runs (#346)

Author: lu-ohai

ads/jobs/ads_job.py

@@ -3,14 +3,20 @@
 
 # Copyright (c) 2021, 2023 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+import time
 from typing import List, Union, Dict
 from urllib.parse import urlparse
 
 import fsspec
+import oci
 from ads.common.auth import default_signer
 from ads.jobs.builders.base import Builder
 from ads.jobs.builders.infrastructure.dataflow import DataFlow, DataFlowRun
-from ads.jobs.builders.infrastructure.dsc_job import DataScienceJob, DataScienceJobRun
+from ads.jobs.builders.infrastructure.dsc_job import (
+    DataScienceJob, 
+    DataScienceJobRun, 
+    SLEEP_INTERVAL
+)
 from ads.jobs.builders.runtimes.pytorch_runtime import PyTorchDistributedRuntime
 from ads.jobs.builders.runtimes.container_runtime import ContainerRuntime
 from ads.jobs.builders.runtimes.python_runtime import (
@@ -460,7 +466,29 @@ def run_list(self, **kwargs) -> list:
             A list of job run instances, the actual object type depends on the infrastructure.
         """
         return self.infrastructure.run_list(**kwargs)
+    
+    def cancel(self, wait_for_completion: bool = True) -> None:
+        """Cancels the runs of the job.
 
+        Parameters
+        ----------
+        wait_for_completion: bool
+            Whether to wait for run to be cancelled before proceeding.
+            Defaults to True.
+        """
+        runs = self.run_list()
+        for run in runs:
+            run.cancel(wait_for_completion=False)
+        
+        if wait_for_completion:
+            for run in runs:
+                while (
+                    run.lifecycle_state !=
+                    oci.data_science.models.JobRun.LIFECYCLE_STATE_CANCELED
+                ):
+                    run.sync()
+                    time.sleep(SLEEP_INTERVAL)
+        
     def delete(self) -> None:
         """Deletes the job from the infrastructure."""
         self.infrastructure.delete()

ads/jobs/builders/infrastructure/dsc_job.py

@@ -725,19 +725,28 @@ def stop_condition():
 
         return self
 
-    def cancel(self) -> DataScienceJobRun:
+    def cancel(self, wait_for_completion: bool = True) -> DataScienceJobRun:
         """Cancels a job run
-        This method will wait for the job run to be canceled before returning.
+
+        Parameters
+        ----------
+        wait_for_completion: bool
+            Whether to wait for job run to be cancelled before proceeding.
+            Defaults to True.
 
         Returns
         -------
         self
             The job run instance.
         """
         self.client.cancel_job_run(self.id)
-        while self.lifecycle_state != "CANCELED":
-            self.sync()
-            time.sleep(SLEEP_INTERVAL)
+        if wait_for_completion:
+            while (
+                self.lifecycle_state != 
+                oci.data_science.models.JobRun.LIFECYCLE_STATE_CANCELED
+            ):
+                self.sync()
+                time.sleep(SLEEP_INTERVAL)
         return self
 
     def __repr__(self) -> str:

ads/opctl/backend/ads_ml_job.py

@@ -217,10 +217,22 @@ def cancel(self):
         """
         Cancel Job Run from OCID.
         """
-        run_id = self.config["execution"]["run_id"]
         with AuthContext(auth=self.auth_type, profile=self.profile):
-            DataScienceJobRun.from_ocid(run_id).cancel()
-            print(f"Job run {run_id} has been cancelled.")
+            wait_for_completion = self.config["execution"].get("wait_for_completion")
+            if self.config["execution"].get("id"):
+                id = self.config["execution"]["id"]
+                Job.from_datascience_job(id).cancel(
+                    wait_for_completion=wait_for_completion
+                )
+                if wait_for_completion:
+                    print(f"All job runs under {id} have been cancelled.")
+            elif self.config["execution"].get("run_id"):
+                run_id = self.config["execution"]["run_id"]
+                DataScienceJobRun.from_ocid(run_id).cancel(
+                    wait_for_completion=wait_for_completion
+                )
+                if wait_for_completion:
+                    print(f"Job run {run_id} has been cancelled.")
 
     def watch(self):
         """

ads/opctl/cmds.py

@@ -374,7 +374,7 @@ def delete(**kwargs) -> None:
     ):
         kwargs["id"] = kwargs.pop("ocid")
     else:
-        raise ValueError(f"{kwargs['ocid']} is valid or supported.")
+        raise ValueError(f"{kwargs['ocid']} is invalid or not supported.")
 
     p = ConfigProcessor().step(ConfigMerger, **kwargs)
     return _BackendFactory(p.config).backend.delete()
@@ -388,13 +388,24 @@ def cancel(**kwargs) -> None:
     ----------
     kwargs: dict
         keyword argument, stores command line args
+    
     Returns
     -------
     None
     """
-    kwargs["run_id"] = kwargs.pop("ocid")
-    if not kwargs.get("backend"):
-        kwargs["backend"] = _get_backend_from_run_id(kwargs["run_id"])
+    kwargs["backend"] = _get_backend_from_ocid(kwargs["ocid"])
+    if (
+        DataScienceResourceRun.JOB_RUN in kwargs["ocid"]
+        or DataScienceResourceRun.DATAFLOW_RUN in kwargs["ocid"]
+        or DataScienceResourceRun.PIPELINE_RUN in kwargs["ocid"]
+    ):
+        kwargs["run_id"] = kwargs.pop("ocid")
+    elif (
+        DataScienceResource.JOB in kwargs["ocid"]
+    ):
+        kwargs["id"] = kwargs.pop("ocid")
+    else:
+        raise ValueError(f"{kwargs['ocid']} is invalid or not supported.")
     p = ConfigProcessor().step(ConfigMerger, **kwargs)
     return _BackendFactory(p.config).backend.cancel()
 

tests/unitary/default_setup/jobs/test_jobs_base.py

@@ -20,6 +20,7 @@
     ScriptRuntime,
     NotebookRuntime,
 )
+from ads.jobs.builders.infrastructure.dsc_job import DataScienceJobRun
 from ads.jobs.builders.infrastructure.dsc_job_runtime import (
     CondaRuntimeHandler,
     ScriptRuntimeHandler,
@@ -603,3 +604,25 @@ def test_run_details_link_fail(self, mock_extract_region):
         test_run_instance = RunInstance()
         test_result = test_run_instance.run_details_link
         assert test_result == ""
+
+
+class DataScienceJobMethodTest(DataScienceJobPayloadTest):
+
+    @patch("ads.jobs.builders.infrastructure.dsc_job.DataScienceJobRun.cancel")
+    @patch("ads.jobs.ads_job.Job.run_list")
+    def test_job_cancel(self, mock_run_list, mock_cancel):
+        mock_run_list.return_value = [
+            DataScienceJobRun(
+                lifecycle_state="CANCELED"
+            )
+        ] * 3
+
+        job = (
+            Job(name="test")
+            .with_infrastructure(infrastructure.DataScienceJob())
+            .with_runtime(ScriptRuntime().with_script(self.SCRIPT_URI))
+        )
+
+        job.cancel()
+        mock_run_list.assert_called()
+        mock_cancel.assert_called_with(wait_for_completion=False)

tests/unitary/with_extras/opctl/test_opctl_cmds.py

@@ -181,8 +181,9 @@ def test_cancel(self, job_cancel_func, pipeline_cancel_func, monkeypatch):
         monkeypatch.delenv("NB_SESSION_OCID", raising=False)
         cancel(ocid="...datasciencejobrun...")
         job_cancel_func.assert_called()
-        with pytest.raises(ValueError):
-            cancel(ocid="....datasciencejob....")
+
+        cancel(ocid="....datasciencejob....")
+        job_cancel_func.assert_called()
 
         cancel(ocid="...datasciencepipelinerun...")
         pipeline_cancel_func.assert_called()