Add search classmethod for Record (#68)

* Add search classmethod for Record

* limit search result to prevent throttling

Author: l2yao

src/smexperiments/_base_types.py

@@ -131,6 +131,35 @@ def _list(
         except StopIteration:
             return
 
+    @classmethod
+    def _search(
+        cls,
+        search_resource,
+        search_item_factory,
+        boto_next_token_name="NextToken",
+        sagemaker_boto_client=None,
+        **kwargs
+    ):
+        sagemaker_boto_client = sagemaker_boto_client or _utils.sagemaker_client()
+        next_token = None
+        try:
+            while True:
+                search_request_kwargs = _boto_functions.to_boto(kwargs, cls._custom_boto_names, cls._custom_boto_types)
+                search_request_kwargs["Resource"] = search_resource
+                if next_token:
+                    search_request_kwargs[boto_next_token_name] = next_token
+                search_method = getattr(sagemaker_boto_client, "search")
+                search_method_response = search_method(**search_request_kwargs)
+                search_items = search_method_response.get("Results", [])
+                next_token = search_method_response.get(boto_next_token_name)
+                for item in search_items:
+                    if cls.__name__ in item:
+                        yield search_item_factory(item[cls.__name__])
+                if not next_token:
+                    break
+        except StopIteration:
+            return
+
     @classmethod
     def _construct(cls, boto_method_name, sagemaker_boto_client=None, **kwargs):
         sagemaker_boto_client = sagemaker_boto_client or _utils.sagemaker_client()

src/smexperiments/api_types.py

@@ -270,3 +270,157 @@ class BatchPutMetricsError(_base_types.ApiObject):
 
     def __init__(self, code=None, message=None, metric_index=None, **kwargs):
         super(BatchPutMetricsError, self).__init__(code=code, message=message, metric_index=metric_index, **kwargs)
+
+
+class ExperimentSearchResult(_base_types.ApiObject):
+    """Summary model of an Experiment search result.
+
+    Attributes:
+        experiment_arn (str): Arn of the experiment.
+        experiment_name (str): Name of the experiment.
+        display_name (str): Display name of the experiment.
+        source (dict): The source of the experiment
+        tags (list): The list of tags that are associated with the experiment.
+    """
+
+    experiment_arn = None
+    experiment_name = None
+    display_name = None
+    source = None
+    tags = None
+
+    def __init__(self, experiment_arn=None, experiment_name=None, display_name=None, source=None, tags=None, **kwargs):
+        super(ExperimentSearchResult, self).__init__(
+            experiment_arn=experiment_arn,
+            experiment_name=experiment_name,
+            display_name=display_name,
+            source=source,
+            tags=tags,
+            **kwargs
+        )
+
+
+class TrialSearchResult(_base_types.ApiObject):
+    """Summary model of an Trial search result.
+
+    Attributes:
+        trial_arn (str): Arn of the trial.
+        trial_name (str): Name of the trial.
+        display_name (str): Display name of the trial.
+        source (dict): The source of the trial.
+        tags (list): The list of tags that are associated with the trial.
+        trial_component_summaries (dict):
+    """
+
+    trial_arn = None
+    trial_name = None
+    display_name = None
+    source = None
+    tags = None
+    trial_component_summaries = None
+
+    def __init__(
+        self,
+        trial_arn=None,
+        trial_name=None,
+        display_name=None,
+        source=None,
+        tags=None,
+        trial_component_summaries=None,
+        **kwargs
+    ):
+        super(TrialSearchResult, self).__init__(
+            trial_arn=trial_arn,
+            trial_name=trial_name,
+            display_name=display_name,
+            source=source,
+            tags=tags,
+            trial_component_summaries=trial_component_summaries,
+            **kwargs
+        )
+
+
+class TrialComponentSearchResult(_base_types.ApiObject):
+    """Summary model of an Trial Component search result.
+
+    Attributes:
+        trial_component_arn (str): Arn of the trial component.
+        trial_component_name (str): Name of the trial component.
+        display_name (str): Display name of the trial component.
+        source (dict): The source of the trial component.
+        status (dict): The status of the trial component.
+        start_time (datetime): Start time.
+        end_time (datetime): End time.
+        creation_time (datetime): Creation time.
+        created_by (str): Created by.
+        last_modified_time (datetime): Date last modified.
+        last_modified_by (datetime): User last modified.
+        parameters (dict): The hyperparameters of the component.
+        input_artifacts (dict): The input artifacts of the component.
+        output_artifacts (dict): The output artifacts of the component.
+        metrics (list): The metrics for the component.
+        source_detail (dict): The source of the trial component.
+        tags (list): The list of tags that are associated with the trial component.
+        parents (dict): The parent of trial component
+    """
+
+    trial_component_arn = None
+    trial_component_name = None
+    display_name = None
+    source = None
+    status = None
+    start_time = None
+    end_time = None
+    creation_time = None
+    created_by = None
+    last_modified_time = None
+    last_modified_by = None
+    parameters = None
+    input_artifacts = None
+    output_artifacts = None
+    metrics = None
+    source_detail = None
+    tags = None
+    parents = None
+
+    def __init__(
+        self,
+        trial_component_arn=None,
+        trial_component_name=None,
+        start_time=None,
+        end_time=None,
+        display_name=None,
+        source=None,
+        status=None,
+        creation_time=None,
+        created_by=None,
+        last_modified_time=None,
+        last_modified_by=None,
+        parameters=None,
+        input_artifacts=None,
+        output_artifacts=None,
+        metrics=None,
+        source_detail=None,
+        tags=None,
+        parents=None,
+    ):
+        super(TrialComponentSearchResult, self).__init__(
+            trial_component_arn=trial_component_arn,
+            trial_component_name=trial_component_name,
+            display_name=display_name,
+            source=source,
+            status=status,
+            start_time=start_time,
+            end_time=end_time,
+            creation_time=creation_time,
+            created_by=created_by,
+            last_modified_by=last_modified_by,
+            last_modified_time=last_modified_time,
+            parameters=parameters,
+            input_artifacts=input_artifacts,
+            output_artifacts=output_artifacts,
+            metrics=metrics,
+            source_detail=source_detail,
+            tags=tags,
+            parents=parents,
+        )

src/smexperiments/experiment.py

@@ -145,6 +145,38 @@ def list(
             sagemaker_boto_client=sagemaker_boto_client,
         )
 
+    @classmethod
+    def search(
+        cls, search_expression=None, sort_by=None, sort_order=None, max_results=None, sagemaker_boto_client=None,
+    ):
+        """
+        Search experiments. Returns SearchResults in the account matching the search criteria.
+
+        Args:
+            search_expression: (dict, optional): A Boolean conditional statement. Resource objects
+                must satisfy this condition to be included in search results. You must provide at
+                least one subexpression, filter, or nested filter.
+            sort_by (str, optional): The name of the resource property used to sort the SearchResults.
+                The default is LastModifiedTime
+            sort_order (str, optional): How SearchResults are ordered. Valid values are Ascending or
+                Descending . The default is Descending .
+            max_results (int, optional): The maximum number of results to return in a SearchResponse.
+            sagemaker_boto_client (SageMaker.Client, optional): Boto3 client for SageMaker. If not
+                supplied, a default boto3 client will be used.
+
+        Returns:
+            collections.Iterator[SearchResult] : An iterator over search results matching the search criteria.
+        """
+        return super(Experiment, cls)._search(
+            search_resource="Experiment",
+            search_item_factory=api_types.ExperimentSearchResult.from_boto,
+            search_expression=search_expression,
+            sort_by=sort_by,
+            sort_order=sort_order,
+            max_results=max_results,
+            sagemaker_boto_client=sagemaker_boto_client,
+        )
+
     def list_trials(self, created_before=None, created_after=None, sort_by=None, sort_order=None):
         """List trials in this experiment matching the specified criteria.
 

src/smexperiments/trial.py

@@ -164,6 +164,38 @@ def list(
             sagemaker_boto_client=sagemaker_boto_client,
         )
 
+    @classmethod
+    def search(
+        cls, search_expression=None, sort_by=None, sort_order=None, max_results=None, sagemaker_boto_client=None,
+    ):
+        """
+        Search experiments. Returns SearchResults in the account matching the search criteria.
+
+        Args:
+            search_expression: (dict, optional): A Boolean conditional statement. Resource objects
+                must satisfy this condition to be included in search results. You must provide at
+                least one subexpression, filter, or nested filter.
+            sort_by (str, optional): The name of the resource property used to sort the SearchResults.
+                The default is LastModifiedTime
+            sort_order (str, optional): How SearchResults are ordered. Valid values are Ascending or
+                Descending . The default is Descending .
+            max_results (int, optional): The maximum number of results to return in a SearchResponse.
+            sagemaker_boto_client (SageMaker.Client, optional): Boto3 client for SageMaker. If not
+                supplied, a default boto3 client will be used.
+
+        Returns:
+            collections.Iterator[SearchResult] : An iterator over search results matching the search criteria.
+        """
+        return super(Trial, cls)._search(
+            search_resource="ExperimentTrial",
+            search_item_factory=api_types.TrialSearchResult.from_boto,
+            search_expression=search_expression,
+            sort_by=sort_by,
+            sort_order=sort_order,
+            max_results=max_results,
+            sagemaker_boto_client=sagemaker_boto_client,
+        )
+
     def add_trial_component(self, tc):
         """Add the specified trial component to this ``Trial``.
 

src/smexperiments/trial_component.py

@@ -179,3 +179,35 @@ def list(
             max_results=max_results,
             next_token=next_token,
         )
+
+    @classmethod
+    def search(
+        cls, search_expression=None, sort_by=None, sort_order=None, max_results=None, sagemaker_boto_client=None,
+    ):
+        """
+        Search experiments. Returns SearchResults in the account matching the search criteria.
+
+        Args:
+            search_expression: (dict, optional): A Boolean conditional statement. Resource objects
+                must satisfy this condition to be included in search results. You must provide at
+                least one subexpression, filter, or nested filter.
+            sort_by (str, optional): The name of the resource property used to sort the SearchResults.
+                The default is LastModifiedTime
+            sort_order (str, optional): How SearchResults are ordered. Valid values are Ascending or
+                Descending . The default is Descending .
+            max_results (int, optional): The maximum number of results to return in a SearchResponse.
+            sagemaker_boto_client (SageMaker.Client, optional): Boto3 client for SageMaker. If not
+                supplied, a default boto3 client will be used.
+
+        Returns:
+            collections.Iterator[SearchResult] : An iterator over search results matching the search criteria.
+        """
+        return super(TrialComponent, cls)._search(
+            search_resource="ExperimentTrialComponent",
+            search_item_factory=api_types.TrialComponentSearchResult.from_boto,
+            search_expression=search_expression,
+            sort_by=sort_by,
+            sort_order=sort_order,
+            max_results=max_results,
+            sagemaker_boto_client=sagemaker_boto_client,
+        )

tests/integ/test_experiment.py

@@ -85,6 +85,16 @@ def test_list_sort(sagemaker_boto_client, experiments):
     assert experiment_names  # sanity test
 
 
+def test_search(sagemaker_boto_client):
+    experiment_names_searched = []
+    for s in experiment.Experiment.search(max_results=10, sagemaker_boto_client=sagemaker_boto_client):
+        if "smexperiments-integ-" in s.experiment_name:
+            experiment_names_searched.append(s.experiment_name)
+
+    assert len(experiment_names_searched) > 0
+    assert experiment_names_searched  # sanity test
+
+
 def test_create_trial(experiment_obj, sagemaker_boto_client):
     trial_obj = experiment_obj.create_trial()
     try:

tests/integ/test_trial.py

@@ -76,6 +76,16 @@ def test_list_sort(trials, sagemaker_boto_client):
         assert trial_names_listed  # sanity test
 
 
+def test_search(sagemaker_boto_client):
+    trial_names_searched = []
+    for s in trial.Trial.search(max_results=10, sagemaker_boto_client=sagemaker_boto_client):
+        if "smexperiments-integ-" in s.trial_name:
+            trial_names_searched.append(s.trial_name)
+
+    assert len(trial_names_searched) > 0
+    assert trial_names_searched  # sanity test
+
+
 def test_add_remove_trial_component(trial_obj, trial_component_obj):
     trial_obj.add_trial_component(trial_component_obj)
     trial_components = list(trial_obj.list_trial_components())

tests/integ/test_trial_component.py

@@ -99,3 +99,13 @@ def test_list_trial_components_by_experiment(experiment_obj, trial_component_obj
     )
     assert 0 == len(trial_components)
     trial_obj.delete()
+
+
+def test_search(sagemaker_boto_client):
+    trial_component_names_searched = []
+    for s in trial_component.TrialComponent.search(max_results=10, sagemaker_boto_client=sagemaker_boto_client):
+        if "smexperiments-integ-" in s.trial_component_name:
+            trial_component_names_searched.append(s.trial_component_name)
+
+    assert len(trial_component_names_searched) > 0
+    assert trial_component_names_searched  # sanity test

tests/unit/test_experiment.py

@@ -129,6 +129,36 @@ def test_list_trials_call_args(sagemaker_boto_client):
     sagemaker_boto_client.list_trials.assert_called_with(CreatedBefore=created_before, CreatedAfter=created_after)
 
 
+def test_search(sagemaker_boto_client):
+    sagemaker_boto_client.search.return_value = {
+        "Results": [
+            {
+                "Experiment": {
+                    "ExperimentName": "experiment-1",
+                    "ExperimentArn": "arn::experiment-1",
+                    "DisplayName": "Experiment1",
+                }
+            },
+            {
+                "Experiment": {
+                    "ExperimentName": "experiment-2",
+                    "ExperimentArn": "arn::experiment-2",
+                    "DisplayName": "Experiment2",
+                }
+            },
+        ]
+    }
+    expected = [
+        api_types.ExperimentSearchResult(
+            experiment_name="experiment-1", experiment_arn="arn::experiment-1", display_name="Experiment1"
+        ),
+        api_types.ExperimentSearchResult(
+            experiment_name="experiment-2", experiment_arn="arn::experiment-2", display_name="Experiment2",
+        ),
+    ]
+    assert expected == list(experiment.Experiment.search(sagemaker_boto_client=sagemaker_boto_client))
+
+
 def test_experiment_create_trial_with_name(sagemaker_boto_client):
     experiment_obj = experiment.Experiment(sagemaker_boto_client=sagemaker_boto_client)
     experiment_obj.experiment_name = "someExperimentName"