Additional prediction model run arguments (#144)

VERSION

@@ -1 +1 @@
-4.1.0
+4.2.0

exabel_data_sdk/client/api/data_classes/prediction_model_run.py

@@ -1,60 +1,121 @@
+from enum import Enum
 from typing import Optional
 
 from exabel_data_sdk.stubs.exabel.api.analytics.v1.all_pb2 import (
     PredictionModelRun as ProtoPredictionModelRun,
 )
+from exabel_data_sdk.stubs.exabel.api.analytics.v1.prediction_model_messages_pb2 import (
+    ModelConfiguration as ProtoModelConfiguration,
+)
+
+
+class ModelConfiguration(Enum):
+    """Specifies a model configuration."""
+
+    # Latest configuration.
+    LATEST = ProtoModelConfiguration.LATEST
+    # Configuration of the active run.
+    ACTIVE = ProtoModelConfiguration.ACTIVE
+    # Configuration of a specific run.
+    SPECIFIC_RUN = ProtoModelConfiguration.SPECIFIC_RUN
 
 
 class PredictionModelRun:
     """
     A prediction model run in the Analytics API.
 
     Attributes:
-        name (str):         The resource name of the model run, for example
-                            "predictionModels/123/runs/4".
-        description (str):  A description of the run.
+        name (str):                         The resource name of the model run, for example
+                                            "predictionModels/123/runs/4".
+        description (str):                  A description of the run.
+        configuration (ModelConfiguration): Which model configuration to use.
+        configuration_source (int):         When using ModelConfiguration.SPECIFIC_RUN,
+                                            this specifies the run from which to retrieve the
+                                            configuration. It is not used for other configurations.
+        auto_activate (bool):               Whether to automatically activate the run if it
+                                            completes successfully. The run will not be activated
+                                            if it fails for any of the entities in the model.
     """
 
     def __init__(
         self,
         name: Optional[str] = None,
         description: str = "",
+        configuration: ModelConfiguration = ModelConfiguration.LATEST,
+        configuration_source: Optional[int] = None,
+        auto_activate: bool = False,
     ):
         """
         Create a prediction model run in the Analytics API.
 
         When creating a new prediction model run in the API, the name should be set to None.
 
         Args:
-            name (str):         The resource name of the prediction model run, for example
-                                "predictionModels/123/runs/4".
-            description (str):  A description of the run.
+            name (str):                         The resource name of the model run, for example
+                                                "predictionModels/123/runs/4".
+            description (str):                  A description of the run.
+            configuration (ModelConfiguration): Which model configuration to use.
+            configuration_source (int):         When using ModelConfiguration.SPECIFIC_RUN,
+                                                this specifies the run from which to retrieve the
+                                                configuration. It is not used for other
+                                                configurations.
+            auto_activate (bool):               Whether to automatically activate the run if it
+                                                completes successfully. The run will not be
+                                                activated if it fails for any of the entities in
+                                                the model.
         """
         self.name = name
         self.description = description
+        self.configuration = configuration
+        self.configuration_source = configuration_source
+        self.auto_activate = auto_activate
+        if self.configuration == ModelConfiguration.SPECIFIC_RUN and not self.configuration_source:
+            raise ValueError(
+                "The argument 'configuration_source' must be specified when using "
+                "ModelConfiguration.SPECIFIC_RUN."
+            )
 
     @staticmethod
     def from_proto(model_run: ProtoPredictionModelRun) -> "PredictionModelRun":
         """Create a PredictionModelRun from the given protobuf PredictionModelRun."""
         return PredictionModelRun(
             name=model_run.name,
             description=model_run.description,
+            configuration=ModelConfiguration(model_run.configuration),
+            configuration_source=model_run.configuration_source
+            if model_run.HasField("configuration_source")
+            else None,
+            auto_activate=model_run.auto_activate,
         )
 
     def to_proto(self) -> ProtoPredictionModelRun:
         """Create a protobuf PredictionModelRun from this PredictionModelRun."""
         return ProtoPredictionModelRun(
             name=self.name,
             description=self.description,
+            configuration=self.configuration.value,
+            configuration_source=self.configuration_source,
+            auto_activate=self.auto_activate,
         )
 
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, PredictionModelRun):
             return False
-        return self.name == other.name and self.description == other.description
+        return (
+            self.name == other.name
+            and self.description == other.description
+            and self.configuration == other.configuration
+            and self.configuration_source == other.configuration_source
+            and self.auto_activate == other.auto_activate
+        )
 
     def __repr__(self) -> str:
-        return f"PredictionModelRun(name='{self.name}', description='{self.description}'"
+        return (
+            f"PredictionModelRun(name='{self.name}', description='{self.description}', "
+            f"configuration={self.configuration}, "
+            f"configuration_source='{self.configuration_source}', "
+            f"auto_activate={self.auto_activate}"
+        )
 
     def __lt__(self, other: object) -> bool:
         if not isinstance(other, PredictionModelRun):

exabel_data_sdk/client/api/entity_api.py

@@ -8,7 +8,7 @@
 from exabel_data_sdk.client.api.data_classes.entity_type import EntityType
 from exabel_data_sdk.client.api.data_classes.paging_result import PagingResult
 from exabel_data_sdk.client.api.data_classes.request_error import ErrorType, RequestError
-from exabel_data_sdk.client.api.pagable_resource import PagableResourceMixin
+from exabel_data_sdk.client.api.pageable_resource import PageableResourceMixin
 from exabel_data_sdk.client.api.resource_creation_result import (
     ResourceCreationResults,
     ResourceCreationStatus,
@@ -32,7 +32,7 @@
 )
 
 
-class EntityApi(PagableResourceMixin):
+class EntityApi(PageableResourceMixin):
     """
     API class for CRUD operations on entities and entity types.
 

exabel_data_sdk/client/api/pagable_resource.py → exabel_data_sdk/client/api/pageable_resource.py

@@ -9,8 +9,8 @@
 from exabel_data_sdk.client.api.data_classes.signal import Signal
 from exabel_data_sdk.client.api.data_classes.tag import Tag
 
-PagableResourceT = TypeVar(
-    "PagableResourceT",
+PageableResourceT = TypeVar(
+    "PageableResourceT",
     DataSet,
     Entity,
     EntityType,
@@ -22,22 +22,22 @@
 )
 
 
-class PagableResourceMixin:
-    """Mixin class for APIs that contain methods with pagable resources."""
+class PageableResourceMixin:
+    """Mixin class for APIs that contain methods with pageable resources."""
 
     @staticmethod
     def _get_resource_iterator(
-        pagable_func: Callable[..., PagingResult[PagableResourceT]],
+        pageable_func: Callable[..., PagingResult[PageableResourceT]],
         **kwargs: str,
-    ) -> Iterator[PagableResourceT]:
+    ) -> Iterator[PageableResourceT]:
         """
         Return an iterator with all the resources returnable by function, paging through the
         results.
         """
         page_token: Optional[str] = None
         resource_count = 0
         while True:
-            result = pagable_func(**kwargs, page_token=page_token)
+            result = pageable_func(**kwargs, page_token=page_token)
             for resource in result.results:
                 yield resource
             page_token = result.next_page_token

exabel_data_sdk/client/api/relationship_api.py

@@ -11,7 +11,7 @@
 from exabel_data_sdk.client.api.data_classes.relationship import Relationship
 from exabel_data_sdk.client.api.data_classes.relationship_type import RelationshipType
 from exabel_data_sdk.client.api.data_classes.request_error import ErrorType, RequestError
-from exabel_data_sdk.client.api.pagable_resource import PagableResourceMixin
+from exabel_data_sdk.client.api.pageable_resource import PageableResourceMixin
 from exabel_data_sdk.client.api.resource_creation_result import (
     ResourceCreationResults,
     ResourceCreationStatus,
@@ -31,7 +31,7 @@
 )
 
 
-class RelationshipApi(PagableResourceMixin):
+class RelationshipApi(PageableResourceMixin):
     """
     API class for entity relationship CRUD operations.
     """

exabel_data_sdk/client/api/signal_api.py

@@ -6,7 +6,7 @@
 from exabel_data_sdk.client.api.data_classes.paging_result import PagingResult
 from exabel_data_sdk.client.api.data_classes.request_error import ErrorType, RequestError
 from exabel_data_sdk.client.api.data_classes.signal import Signal
-from exabel_data_sdk.client.api.pagable_resource import PagableResourceMixin
+from exabel_data_sdk.client.api.pageable_resource import PageableResourceMixin
 from exabel_data_sdk.client.client_config import ClientConfig
 from exabel_data_sdk.stubs.exabel.api.data.v1.all_pb2 import (
     CreateSignalRequest,
@@ -17,7 +17,7 @@
 )
 
 
-class SignalApi(PagableResourceMixin):
+class SignalApi(PageableResourceMixin):
     """
     API class for Signal CRUD operations.
     """

exabel_data_sdk/client/api/tag_api.py

@@ -4,7 +4,7 @@
 from exabel_data_sdk.client.api.data_classes.paging_result import PagingResult
 from exabel_data_sdk.client.api.data_classes.request_error import ErrorType, RequestError
 from exabel_data_sdk.client.api.data_classes.tag import Tag
-from exabel_data_sdk.client.api.pagable_resource import PagableResourceMixin
+from exabel_data_sdk.client.api.pageable_resource import PageableResourceMixin
 from exabel_data_sdk.client.client_config import ClientConfig
 from exabel_data_sdk.stubs.exabel.api.analytics.v1.all_pb2 import (
     AddEntitiesRequest,
@@ -18,7 +18,7 @@
 )
 
 
-class TagApi(PagableResourceMixin):
+class TagApi(PageableResourceMixin):
     """
     API class for tag operations.
     """

exabel_data_sdk/client/api/time_series_api.py

@@ -12,7 +12,7 @@
 from exabel_data_sdk.client.api.data_classes.paging_result import PagingResult
 from exabel_data_sdk.client.api.data_classes.request_error import ErrorType, RequestError
 from exabel_data_sdk.client.api.error_handler import grpc_status_to_error_type
-from exabel_data_sdk.client.api.pagable_resource import PagableResourceMixin
+from exabel_data_sdk.client.api.pageable_resource import PageableResourceMixin
 from exabel_data_sdk.client.api.resource_creation_result import (
     ResourceCreationResult,
     ResourceCreationResults,
@@ -42,7 +42,7 @@
 from exabel_data_sdk.util.deprecate_arguments import deprecate_arguments
 
 
-class TimeSeriesApi(PagableResourceMixin):
+class TimeSeriesApi(PageableResourceMixin):
     """
     API class for time series CRUD operations.
     """

exabel_data_sdk/stubs/exabel/api/__init__.pyi

@@ -1,5 +1,5 @@
-from . import data
-from . import math
-from . import analytics
 from . import management
+from . import math
 from . import time
+from . import data
+from . import analytics

exabel_data_sdk/stubs/exabel/api/analytics/v1/__init__.pyi

@@ -1,8 +1,8 @@
 from . import item_messages_pb2
-from . import tag_service_pb2
 from . import prediction_model_messages_pb2
-from . import tag_messages_pb2
-from . import derived_signal_messages_pb2
+from . import tag_service_pb2
 from . import service_pb2
-from . import derived_signal_service_pb2
 from . import prediction_model_service_pb2
+from . import tag_messages_pb2
+from . import derived_signal_service_pb2
+from . import derived_signal_messages_pb2

exabel_data_sdk/stubs/exabel/api/analytics/v1/prediction_model_messages_pb2.py

@@ -6,7 +6,7 @@
 _sym_db = _symbol_database.Default()
 from google.api import field_behavior_pb2 as google_dot_api_dot_field__behavior__pb2
 from .....protoc_gen_openapiv2.options import annotations_pb2 as protoc__gen__openapiv2_dot_options_dot_annotations__pb2
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n7exabel/api/analytics/v1/prediction_model_messages.proto\x12\x17exabel.api.analytics.v1\x1a\x1fgoogle/api/field_behavior.proto\x1a.protoc_gen_openapiv2/options/annotations.proto"~\n\x12PredictionModelRun\x123\n\x04name\x18\x01 \x01(\tB%\xe0A\x03\x92A\x1fJ\x1d"predictionModels/123/runs/3"\x123\n\x0bdescription\x18\x02 \x01(\tB\x1e\x92A\x1bJ\x19"Initated by API request"BZ\n\x1bcom.exabel.api.analytics.v1B\x1cPredictionModelMessagesProtoP\x01Z\x1bexabel.com/api/analytics/v1b\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n7exabel/api/analytics/v1/prediction_model_messages.proto\x12\x17exabel.api.analytics.v1\x1a\x1fgoogle/api/field_behavior.proto\x1a.protoc_gen_openapiv2/options/annotations.proto"\x96\x02\n\x12PredictionModelRun\x123\n\x04name\x18\x01 \x01(\tB%\xe0A\x03\x92A\x1fJ\x1d"predictionModels/123/runs/3"\x124\n\x0bdescription\x18\x02 \x01(\tB\x1f\x92A\x1cJ\x1a"Initiated by API request"\x12B\n\rconfiguration\x18\x03 \x01(\x0e2+.exabel.api.analytics.v1.ModelConfiguration\x12!\n\x14configuration_source\x18\x04 \x01(\x05H\x00\x88\x01\x01\x12\x15\n\rauto_activate\x18\x05 \x01(\x08B\x17\n\x15_configuration_source*e\n\x12ModelConfiguration\x12%\n!MODEL_CONFIGURATION_NOT_SPECIFIED\x10\x00\x12\n\n\x06LATEST\x10\x01\x12\n\n\x06ACTIVE\x10\x02\x12\x10\n\x0cSPECIFIC_RUN\x10\x03BZ\n\x1bcom.exabel.api.analytics.v1B\x1cPredictionModelMessagesProtoP\x01Z\x1bexabel.com/api/analytics/v1b\x06proto3')
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
 _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'exabel.api.analytics.v1.prediction_model_messages_pb2', globals())
 if _descriptor._USE_C_DESCRIPTORS == False:
@@ -15,6 +15,8 @@
     _PREDICTIONMODELRUN.fields_by_name['name']._options = None
     _PREDICTIONMODELRUN.fields_by_name['name']._serialized_options = b'\xe0A\x03\x92A\x1fJ\x1d"predictionModels/123/runs/3"'
     _PREDICTIONMODELRUN.fields_by_name['description']._options = None
-    _PREDICTIONMODELRUN.fields_by_name['description']._serialized_options = b'\x92A\x1bJ\x19"Initated by API request"'
-    _PREDICTIONMODELRUN._serialized_start = 165
-    _PREDICTIONMODELRUN._serialized_end = 291
+    _PREDICTIONMODELRUN.fields_by_name['description']._serialized_options = b'\x92A\x1cJ\x1a"Initiated by API request"'
+    _MODELCONFIGURATION._serialized_start = 446
+    _MODELCONFIGURATION._serialized_end = 547
+    _PREDICTIONMODELRUN._serialized_start = 166
+    _PREDICTIONMODELRUN._serialized_end = 444

exabel_data_sdk/stubs/exabel/api/analytics/v1/prediction_model_messages_pb2.pyi

@@ -4,28 +4,72 @@ isort:skip_file
 Copyright (c) 2022 Exabel AS. All rights reserved."""
 import builtins
 import google.protobuf.descriptor
+import google.protobuf.internal.enum_type_wrapper
 import google.protobuf.message
 import sys
-if sys.version_info >= (3, 8):
+import typing
+if sys.version_info >= (3, 10):
     import typing as typing_extensions
 else:
     import typing_extensions
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
+class _ModelConfiguration:
+    ValueType = typing.NewType('ValueType', builtins.int)
+    V: typing_extensions.TypeAlias = ValueType
+
+class _ModelConfigurationEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_ModelConfiguration.ValueType], builtins.type):
+    DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+    MODEL_CONFIGURATION_NOT_SPECIFIED: _ModelConfiguration.ValueType
+    'Not specified.'
+    LATEST: _ModelConfiguration.ValueType
+    'Latest configuration.'
+    ACTIVE: _ModelConfiguration.ValueType
+    'Configuration of the active run.'
+    SPECIFIC_RUN: _ModelConfiguration.ValueType
+    'Configuration of a specific run.'
+
+class ModelConfiguration(_ModelConfiguration, metaclass=_ModelConfigurationEnumTypeWrapper):
+    """Specifies which model configuration to use."""
+MODEL_CONFIGURATION_NOT_SPECIFIED: ModelConfiguration.ValueType
+'Not specified.'
+LATEST: ModelConfiguration.ValueType
+'Latest configuration.'
+ACTIVE: ModelConfiguration.ValueType
+'Configuration of the active run.'
+SPECIFIC_RUN: ModelConfiguration.ValueType
+'Configuration of a specific run.'
+global___ModelConfiguration = ModelConfiguration
+
 @typing_extensions.final
 class PredictionModelRun(google.protobuf.message.Message):
     """A prediction model run."""
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
     NAME_FIELD_NUMBER: builtins.int
     DESCRIPTION_FIELD_NUMBER: builtins.int
+    CONFIGURATION_FIELD_NUMBER: builtins.int
+    CONFIGURATION_SOURCE_FIELD_NUMBER: builtins.int
+    AUTO_ACTIVATE_FIELD_NUMBER: builtins.int
     name: builtins.str
     'Unique resource name of the run, e.g. `predictionModels/123/runs/3`.'
     description: builtins.str
     'You may use this to record some notes about the run. This is shown in the prediction model\n    interface when viewing all runs, and when viewing the results of a single run.\n    '
+    configuration: global___ModelConfiguration.ValueType
+    'Which model configuration to use.\n    If not specified, the latest model configuration is used.\n    '
+    configuration_source: builtins.int
+    'Specifies a prediction model run from which model configuration should be retrieved.\n    Only relevant when `configuration` is set to `ModelConfiguration.SPECIFIC_RUN`.\n    '
+    auto_activate: builtins.bool
+    'Whether to automatically set this run as active once it completes.\n    The run will not be activated if it fails for any of the entities in the model.\n    '
+
+    def __init__(self, *, name: builtins.str | None=..., description: builtins.str | None=..., configuration: global___ModelConfiguration.ValueType | None=..., configuration_source: builtins.int | None=..., auto_activate: builtins.bool | None=...) -> None:
+        ...
+
+    def HasField(self, field_name: typing_extensions.Literal['_configuration_source', b'_configuration_source', 'configuration_source', b'configuration_source']) -> builtins.bool:
+        ...
 
-    def __init__(self, *, name: builtins.str | None=..., description: builtins.str | None=...) -> None:
+    def ClearField(self, field_name: typing_extensions.Literal['_configuration_source', b'_configuration_source', 'auto_activate', b'auto_activate', 'configuration', b'configuration', 'configuration_source', b'configuration_source', 'description', b'description', 'name', b'name']) -> None:
         ...
 
-    def ClearField(self, field_name: typing_extensions.Literal['description', b'description', 'name', b'name']) -> None:
+    def WhichOneof(self, oneof_group: typing_extensions.Literal['_configuration_source', b'_configuration_source']) -> typing_extensions.Literal['configuration_source'] | None:
         ...
 global___PredictionModelRun = PredictionModelRun