Skip to content

Commit 4b3775c

Browse files
authored
Updates on Job Run API (#814)
2 parents 5013703 + b1a151c commit 4b3775c

File tree

2 files changed

+111
-6
lines changed

2 files changed

+111
-6
lines changed

ads/jobs/builders/infrastructure/dsc_job.py

Lines changed: 49 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import inspect
1010
import logging
1111
import os
12+
import re
1213
import time
1314
import traceback
1415
import uuid
@@ -375,12 +376,13 @@ def delete(self, force_delete: bool = False) -> DSCJob:
375376
"""
376377
runs = self.run_list()
377378
for run in runs:
378-
if run.lifecycle_state in [
379-
DataScienceJobRun.LIFECYCLE_STATE_ACCEPTED,
380-
DataScienceJobRun.LIFECYCLE_STATE_IN_PROGRESS,
381-
DataScienceJobRun.LIFECYCLE_STATE_NEEDS_ATTENTION,
382-
]:
383-
run.cancel(wait_for_completion=True)
379+
if force_delete:
380+
if run.lifecycle_state in [
381+
DataScienceJobRun.LIFECYCLE_STATE_ACCEPTED,
382+
DataScienceJobRun.LIFECYCLE_STATE_IN_PROGRESS,
383+
DataScienceJobRun.LIFECYCLE_STATE_NEEDS_ATTENTION,
384+
]:
385+
run.cancel(wait_for_completion=True)
384386
run.delete()
385387
self.client.delete_job(self.id)
386388
return self
@@ -582,6 +584,25 @@ def logging(self) -> OCILog:
582584
id=self.log_id, log_group_id=self.log_details.log_group_id, **auth
583585
)
584586

587+
@property
588+
def exit_code(self):
589+
"""The exit code of the job run from the lifecycle details.
590+
Note that,
591+
None will be returned if the job run is not finished or failed without exit code.
592+
0 will be returned if job run succeeded.
593+
"""
594+
if self.lifecycle_state == self.LIFECYCLE_STATE_SUCCEEDED:
595+
return 0
596+
if not self.lifecycle_details:
597+
return None
598+
match = re.search(r"exit code (\d+)", self.lifecycle_details)
599+
if not match:
600+
return None
601+
try:
602+
return int(match.group(1))
603+
except Exception:
604+
return None
605+
585606
@staticmethod
586607
def _format_log(message: str, date_time: datetime.datetime) -> dict:
587608
"""Formats a message as log record with datetime.
@@ -655,6 +676,22 @@ def _check_and_print_status(self, prev_status) -> str:
655676
print(f"{timestamp} - {status}")
656677
return status
657678

679+
def wait(self, interval: float = SLEEP_INTERVAL):
680+
"""Waits for the job run until if finishes.
681+
682+
Parameters
683+
----------
684+
interval : float
685+
Time interval in seconds between each request to update the logs.
686+
Defaults to 3 (seconds).
687+
688+
"""
689+
self.sync()
690+
while self.status not in self.TERMINAL_STATES:
691+
time.sleep(interval)
692+
self.sync()
693+
return self
694+
658695
def watch(
659696
self,
660697
interval: float = SLEEP_INTERVAL,
@@ -830,6 +867,12 @@ def download(self, to_dir):
830867
self.job.download(to_dir)
831868
return self
832869

870+
def delete(self, force_delete: bool = False):
871+
if force_delete:
872+
self.cancel(wait_for_completion=True)
873+
super().delete()
874+
return
875+
833876

834877
# This is for backward compatibility
835878
DSCJobRun = DataScienceJobRun
Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
#!/usr/bin/env python
2+
3+
# Copyright (c) 2024 Oracle and/or its affiliates.
4+
# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/up
5+
"""Contains tests for DataScienceJobRun class."""
6+
7+
from unittest import TestCase, mock
8+
9+
import oci
10+
from ads.jobs import DataScienceJobRun
11+
12+
13+
class JobRunTestCase(TestCase):
14+
"""Contains test cases for Job runs."""
15+
16+
@mock.patch("ads.jobs.builders.infrastructure.dsc_job.DataScienceJobRun.sync")
17+
def test_job_run_wait(self, *args):
18+
"""Test waiting for job run."""
19+
run = DataScienceJobRun()
20+
run.lifecycle_state = oci.data_science.models.JobRun.LIFECYCLE_STATE_IN_PROGRESS
21+
22+
# Mock time.sleep to change the job run lifecycle state.
23+
def mock_sleep(*args, **kwargs):
24+
run.lifecycle_state = (
25+
oci.data_science.models.JobRun.LIFECYCLE_STATE_SUCCEEDED
26+
)
27+
28+
with mock.patch("time.sleep", wraps=mock_sleep):
29+
run.wait()
30+
31+
self.assertEqual(
32+
run.lifecycle_state,
33+
oci.data_science.models.JobRun.LIFECYCLE_STATE_SUCCEEDED,
34+
)
35+
36+
def test_job_run_exit_code(self):
37+
"""Tests job run exit code."""
38+
run = DataScienceJobRun()
39+
self.assertEqual(run.exit_code, None)
40+
run.lifecycle_state = run.LIFECYCLE_STATE_SUCCEEDED
41+
self.assertEqual(run.exit_code, 0)
42+
run.lifecycle_state = run.LIFECYCLE_STATE_IN_PROGRESS
43+
run.lifecycle_details = "Job run in progress."
44+
self.assertEqual(run.exit_code, None)
45+
run.lifecycle_state = run.LIFECYCLE_STATE_FAILED
46+
run.lifecycle_details = "Job run artifact execution failed with exit code 21."
47+
self.assertEqual(run.exit_code, 21)
48+
49+
@mock.patch("ads.jobs.builders.infrastructure.dsc_job.DataScienceJobRun.cancel")
50+
@mock.patch("ads.common.oci_datascience.OCIDataScienceMixin.delete")
51+
def test_job_run_delete(self, mock_delete, mock_cancel):
52+
"""Tests deleting job run."""
53+
run = DataScienceJobRun()
54+
# Cancel will not be called if job run is succeeded.
55+
run.lifecycle_state = run.LIFECYCLE_STATE_SUCCEEDED
56+
run.delete()
57+
mock_delete.assert_called_once()
58+
mock_cancel.assert_not_called()
59+
# Cancel will be called if job run is in progress and force_delete is set.
60+
run.lifecycle_state = run.LIFECYCLE_STATE_IN_PROGRESS
61+
run.delete(force_delete=True)
62+
mock_cancel.assert_called_once()

0 commit comments

Comments
 (0)