Skip to content

Commit f5d39a9

Browse files
authored
Modified Watch in Model Deployment (#489)
2 parents 9bc802e + 0fb5a02 commit f5d39a9

File tree

2 files changed

+9
-6
lines changed

2 files changed

+9
-6
lines changed

ads/model/deployment/model_deployment.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +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]
6867

6968
MODEL_DEPLOYMENT_KIND = "deployment"
7069
MODEL_DEPLOYMENT_TYPE = "modelDeployment"
@@ -720,7 +719,7 @@ def update(
720719

721720
def watch(
722721
self,
723-
log_type: str = ModelDeploymentLogType.ACCESS,
722+
log_type: str = None,
724723
time_start: datetime = None,
725724
interval: int = LOG_INTERVAL,
726725
log_filter: str = None,
@@ -731,7 +730,7 @@ def watch(
731730
----------
732731
log_type: str, optional
733732
The log type. Can be `access`, `predict` or None.
734-
Defaults to access.
733+
Defaults to None.
735734
time_start : datetime.datetime, optional
736735
Starting time for the log query.
737736
Defaults to None.
@@ -757,7 +756,7 @@ def watch(
757756
count = self.logs(log_type).stream(
758757
source=self.model_deployment_id,
759758
interval=interval,
760-
stop_condition=self._stop_condition,
759+
stop_condition=self._stream_stop_condition,
761760
time_start=time_start,
762761
log_filter=log_filter,
763762
)
@@ -773,7 +772,11 @@ def watch(
773772

774773
def _stop_condition(self):
775774
"""Stops the sync once the model deployment is in a terminal state."""
776-
return self.state in TERMINAL_STATES
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]
777780

778781
def _check_and_print_status(self, prev_status) -> str:
779782
"""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)