Skip to content

Commit 0fb5a02

Browse files
committed
Updated pr.
1 parent 94cff58 commit 0fb5a02

File tree

2 files changed

+9
-7
lines changed

2 files changed

+9
-7
lines changed

ads/model/deployment/model_deployment.py

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -64,8 +64,6 @@
6464
DELETE_WORKFLOW_STEPS = 2
6565
DEACTIVATE_WORKFLOW_STEPS = 2
6666
DEFAULT_RETRYING_REQUEST_ATTEMPTS = 3
67-
TERMINAL_STATES = [State.ACTIVE, State.FAILED, State.DELETED, State.INACTIVE]
68-
WATCH_TERMINAL_STATES = [State.FAILED, State.DELETED, State.INACTIVE]
6967

7068
MODEL_DEPLOYMENT_KIND = "deployment"
7169
MODEL_DEPLOYMENT_TYPE = "modelDeployment"
@@ -749,7 +747,7 @@ def watch(
749747
The instance of ModelDeployment.
750748
"""
751749
status = ""
752-
while self.state not in TERMINAL_STATES:
750+
while not self._stop_condition():
753751
status = self._check_and_print_status(status)
754752
time.sleep(interval)
755753

@@ -758,7 +756,7 @@ def watch(
758756
count = self.logs(log_type).stream(
759757
source=self.model_deployment_id,
760758
interval=interval,
761-
stop_condition=self._stop_condition,
759+
stop_condition=self._stream_stop_condition,
762760
time_start=time_start,
763761
log_filter=log_filter,
764762
)
@@ -773,8 +771,12 @@ def watch(
773771
pass
774772

775773
def _stop_condition(self):
776-
"""Stops the watch sync once the model deployment is in a terminal state."""
777-
return self.state in WATCH_TERMINAL_STATES
774+
"""Stops the sync once the model deployment is in a terminal state."""
775+
return self.state in [State.ACTIVE, State.FAILED, State.DELETED, State.INACTIVE]
776+
777+
def _stream_stop_condition(self):
778+
"""Stops the stream sync once the model deployment is in a terminal state."""
779+
return self.state in [State.FAILED, State.DELETED, State.INACTIVE]
778780

779781
def _check_and_print_status(self, prev_status) -> str:
780782
"""Check and print the next status.

tests/unitary/default_setup/model_deployment/test_model_deployment_v2.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1286,7 +1286,7 @@ def test_watch(self, mock_from_ocid, mock_stream):
12861286
mock_stream.assert_called_with(
12871287
source=model_deployment.model_deployment_id,
12881288
time_start=time_start,
1289-
stop_condition=model_deployment._stop_condition,
1289+
stop_condition=model_deployment._stream_stop_condition,
12901290
interval=10,
12911291
log_filter="test_filter",
12921292
)

0 commit comments

Comments
 (0)