From 46c1fd1188a6c1b49f33b0c3edf60572d8f4a776 Mon Sep 17 00:00:00 2001 From: Aaron Wilson Date: Tue, 22 Aug 2023 14:53:50 -0500 Subject: [PATCH] SDK/python: Update dsort job info query and related tests Signed-off-by: Aaron Wilson --- python/aistore/sdk/dsort.py | 16 +-- python/aistore/sdk/dsort_types.py | 31 +++++- .../tests/integration/sdk/test_dsort_ops.py | 7 +- python/tests/unit/sdk/test_dsort.py | 104 +++++++++--------- 4 files changed, 91 insertions(+), 67 deletions(-) diff --git a/python/aistore/sdk/dsort.py b/python/aistore/sdk/dsort.py index 86f8e186ef..4307bc32ec 100644 --- a/python/aistore/sdk/dsort.py +++ b/python/aistore/sdk/dsort.py @@ -12,7 +12,7 @@ DSORT_ABORT, DSORT_UUID, ) -from aistore.sdk.dsort_types import DsortMetrics +from aistore.sdk.dsort_types import JobInfo from aistore.sdk.errors import Timeout from aistore.sdk.utils import validate_file, probing_frequency @@ -56,17 +56,17 @@ def abort(self): HTTP_METHOD_DELETE, path=f"{URL_PATH_DSORT}/{DSORT_ABORT}", params=qparam ) - def metrics(self) -> Dict[str, DsortMetrics]: + def get_job_info(self) -> Dict[str, JobInfo]: """ - Get metrics for a dsort job + Get info for a dsort job Returns: - Dictionary of metrics for jobs associated with this dsort job + Dictionary of job info for all jobs associated with this dsort """ qparam = {DSORT_UUID: [self._dsort_id]} return self._client.request_deserialize( HTTP_METHOD_GET, path=URL_PATH_DSORT, - res_model=Dict[str, DsortMetrics], + res_model=Dict[str, JobInfo], params=qparam, ) @@ -97,12 +97,12 @@ def wait( if passed > timeout: raise Timeout("dsort job to finish") finished = True - for metric in self.metrics().values(): - if metric.aborted: + for job_info in self.get_job_info().values(): + if job_info.metrics.aborted: logger.info("DSort job '%s' aborted", self._dsort_id) return # Shard creation is the last phase, so check if it's finished - finished = metric.shard_creation.finished and finished + finished = job_info.metrics.shard_creation.finished and finished if finished: logger.info("DSort job '%s' finished", self._dsort_id) return diff --git a/python/aistore/sdk/dsort_types.py b/python/aistore/sdk/dsort_types.py index bfe4ba17d5..28784710ab 100644 --- a/python/aistore/sdk/dsort_types.py +++ b/python/aistore/sdk/dsort_types.py @@ -1,9 +1,14 @@ from typing import List -from pydantic import BaseModel +from pydantic import BaseModel, Field + +from aistore.sdk.types import BucketModel + # See ext/dsort/metric.go for cluster-side type definitions +# pylint: disable=too-few-public-methods + class TimeStats(BaseModel): """ @@ -99,3 +104,27 @@ class DsortMetrics(BaseModel): warnings: List[str] = None errors: List[str] = None extended: bool = None + + +class JobInfo(BaseModel): + """ + Info about a dsort Job, including metrics + """ + + id: str + src_bck: BucketModel = Field(alias="src-bck") + dst_bck: BucketModel = Field(alias="dst-bck") + started_time: str = None + finish_time: str = None + extracted_duration: str = Field(alias="started_meta_sorting", default=None) + sorting_duration: str = Field(alias="started_shard_creation", default=None) + creation_duration: str = Field(alias="finished_shard_creation", default=None) + objects: int = Field(alias="loc-objs") + bytes: int = Field(alias="loc-bytes") + metrics: DsortMetrics = Field(alias="Metrics") + aborted: bool + archived: bool + + # pylint: disable=missing-class-docstring + class Config: + allow_population_by_field_name = True diff --git a/python/tests/integration/sdk/test_dsort_ops.py b/python/tests/integration/sdk/test_dsort_ops.py index 0900d6c7ce..8d9757ac2d 100644 --- a/python/tests/integration/sdk/test_dsort_ops.py +++ b/python/tests/integration/sdk/test_dsort_ops.py @@ -132,7 +132,6 @@ def test_abort(self): ) dsort.abort() dsort.wait(timeout=TEST_TIMEOUT) - metrics = dsort.metrics() - for metric in metrics.values(): - self.assertTrue(metric.aborted) - self.assertEqual(1, len(metric.errors)) + for job_info in dsort.get_job_info().values(): + self.assertTrue(job_info.metrics.aborted) + self.assertEqual(1, len(job_info.metrics.errors)) diff --git a/python/tests/unit/sdk/test_dsort.py b/python/tests/unit/sdk/test_dsort.py index 9802f7fe31..f48c7a4314 100644 --- a/python/tests/unit/sdk/test_dsort.py +++ b/python/tests/unit/sdk/test_dsort.py @@ -11,7 +11,7 @@ HTTP_METHOD_GET, ) from aistore.sdk.dsort import Dsort -from aistore.sdk.dsort_types import DsortMetrics +from aistore.sdk.dsort_types import DsortMetrics, JobInfo from aistore.sdk.errors import Timeout from aistore.sdk.utils import probing_frequency @@ -22,6 +22,15 @@ def setUp(self) -> None: self.dsort_id = "123" self.dsort = Dsort(client=self.mock_client, dsort_id=self.dsort_id) + @staticmethod + def _get_mock_job_info(finished, aborted=False): + mock_metrics = Mock(DsortMetrics) + mock_metrics.aborted = aborted + mock_metrics.shard_creation = Mock(finished=finished) + mock_job_info = Mock(JobInfo) + mock_job_info.metrics = mock_metrics + return mock_job_info + def test_properties(self): self.assertEqual(self.dsort_id, self.dsort.dsort_id) @@ -52,24 +61,24 @@ def test_abort(self): params={DSORT_UUID: [self.dsort_id]}, ) - def test_metrics(self): - metrics = {"id_1": Mock(DsortMetrics)} - self.mock_client.request_deserialize.return_value = metrics - res = self.dsort.metrics() - self.assertEqual(metrics, res) + def test_get_job_info(self): + mock_job_info = {"id_1": Mock(JobInfo)} + self.mock_client.request_deserialize.return_value = mock_job_info + res = self.dsort.get_job_info() + self.assertEqual(mock_job_info, res) self.mock_client.request_deserialize.assert_called_with( HTTP_METHOD_GET, path=URL_PATH_DSORT, - res_model=Dict[str, DsortMetrics], + res_model=Dict[str, JobInfo], params={DSORT_UUID: [self.dsort_id]}, ) @patch("aistore.sdk.dsort.time.sleep") - @patch("aistore.sdk.dsort.Dsort.metrics") - def test_wait_default_timeout(self, mock_metrics, mock_sleep): + @patch("aistore.sdk.dsort.Dsort.get_job_info") + def test_wait_default_timeout(self, mock_get_job_info, mock_sleep): timeout = 300 frequency = probing_frequency(timeout) - expected_metrics_calls = [ + expected_job_info_calls = [ call(), call(), call(), @@ -77,41 +86,40 @@ def test_wait_default_timeout(self, mock_metrics, mock_sleep): expected_sleep_calls = [call(frequency), call(frequency)] self._wait_test_helper( self.dsort, - mock_metrics, + mock_get_job_info, mock_sleep, - expected_metrics_calls, + expected_job_info_calls, expected_sleep_calls, ) @patch("aistore.sdk.dsort.time.sleep") - @patch("aistore.sdk.dsort.Dsort.metrics") - def test_wait(self, mock_status, mock_sleep): + @patch("aistore.sdk.dsort.Dsort.get_job_info") + def test_wait(self, mock_get_job_info, mock_sleep): timeout = 20 frequency = probing_frequency(timeout) - expected_metrics_calls = [call(), call(), call()] + expected_job_info_calls = [call(), call(), call()] expected_sleep_calls = [call(frequency), call(frequency)] self._wait_test_helper( self.dsort, - mock_status, + mock_get_job_info, mock_sleep, - expected_metrics_calls, + expected_job_info_calls, expected_sleep_calls, timeout=timeout, ) @patch("aistore.sdk.dsort.time.sleep") - @patch("aistore.sdk.dsort.Dsort.metrics") + @patch("aistore.sdk.dsort.Dsort.get_job_info") # pylint: disable=unused-argument - def test_wait_timeout(self, mock_metrics, mock_sleep): - mock_metric = Mock(DsortMetrics) - mock_metric.aborted = False - mock_metric.shard_creation = Mock(finished=False) - mock_metrics.return_value = {"key": mock_metric} + def test_wait_timeout(self, mock_get_job_info, mock_sleep): + mock_get_job_info.return_value = { + "key": self._get_mock_job_info(finished=False, aborted=False) + } self.assertRaises(Timeout, self.dsort.wait) @patch("aistore.sdk.dsort.time.sleep") - @patch("aistore.sdk.dsort.Dsort.metrics") - def test_wait_aborted(self, mock_metrics, mock_sleep): + @patch("aistore.sdk.dsort.Dsort.get_job_info") + def test_wait_aborted(self, mock_get_job_info, mock_sleep): timeout = 300 frequency = probing_frequency(timeout) expected_metrics_calls = [ @@ -119,21 +127,15 @@ def test_wait_aborted(self, mock_metrics, mock_sleep): call(), ] expected_sleep_calls = [call(frequency)] - unfinished_metric = Mock(DsortMetrics) - unfinished_metric.aborted = False - unfinished_metric.shard_creation = Mock(finished=False) - aborted_metric = Mock(DsortMetrics) - aborted_metric.aborted = True - aborted_metric.shard_creation = Mock(finished=False) - mock_metrics.side_effect = [ - {"key": unfinished_metric}, - {"key": aborted_metric}, - {"key": unfinished_metric}, + mock_get_job_info.side_effect = [ + {"key": self._get_mock_job_info(finished=False)}, + {"key": self._get_mock_job_info(finished=False, aborted=True)}, + {"key": self._get_mock_job_info(finished=False)}, ] self._wait_exec_assert( self.dsort, - mock_metrics, + mock_get_job_info, mock_sleep, expected_metrics_calls, expected_sleep_calls, @@ -143,28 +145,22 @@ def test_wait_aborted(self, mock_metrics, mock_sleep): def _wait_test_helper( self, dsort, - mock_metrics, + mock_get_job_info, mock_sleep, - expected_metrics_calls, + expected_job_info_calls, expected_sleep_calls, **kwargs, ): - unfinished_metric = Mock(DsortMetrics) - unfinished_metric.aborted = False - unfinished_metric.shard_creation = Mock(finished=False) - finished_metric = Mock(DsortMetrics) - finished_metric.aborted = False - finished_metric.shard_creation = Mock(finished=True) - mock_metrics.side_effect = [ - {"key": unfinished_metric}, - {"key": unfinished_metric}, - {"key": finished_metric}, + mock_get_job_info.side_effect = [ + {"job_id": self._get_mock_job_info(finished=False)}, + {"job_id": self._get_mock_job_info(finished=False)}, + {"job_id": self._get_mock_job_info(finished=True)}, ] self._wait_exec_assert( dsort, - mock_metrics, + mock_get_job_info, mock_sleep, - expected_metrics_calls, + expected_job_info_calls, expected_sleep_calls, **kwargs, ) @@ -172,15 +168,15 @@ def _wait_test_helper( def _wait_exec_assert( self, dsort, - mock_metrics, + mock_get_job_info, mock_sleep, - expected_metrics_calls, + expected_job_info_calls, expected_sleep_calls, **kwargs, ): dsort.wait(**kwargs) - mock_metrics.assert_has_calls(expected_metrics_calls) + mock_get_job_info.assert_has_calls(expected_job_info_calls) mock_sleep.assert_has_calls(expected_sleep_calls) - self.assertEqual(len(expected_metrics_calls), mock_metrics.call_count) + self.assertEqual(len(expected_job_info_calls), mock_get_job_info.call_count) self.assertEqual(len(expected_sleep_calls), mock_sleep.call_count)