frogml-core 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

frogml_services_mock/mocks/batch_job_manager_service.py

@@ -3,10 +3,10 @@ from typing import Dict, List, Optional, Set
 
 from google.protobuf.timestamp_pb2 import Timestamp
 
-from _frogml_proto.qwak.administration.authenticated_user.v1.credentials_pb2 import (
+from frogml_proto.qwak.administration.authenticated_user.v1.credentials_pb2 import (
     AwsTemporaryCredentials,
 )
-from _frogml_proto.qwak.batch_job.v1.batch_job_service_pb2 import (
+from frogml_proto.qwak.batch_job.v1.batch_job_service_pb2 import (
     BatchJobDeploymentSize,
     BatchJobDetails,
     BatchJobMessage,
@@ -39,7 +39,7 @@ from _frogml_proto.qwak.batch_job.v1.batch_job_service_pb2 import (
     StartWarmupJobResponse,
     TaskExecutionDetails,
 )
-from _frogml_proto.qwak.batch_job.v1.batch_job_service_pb2_grpc import (
+from frogml_proto.qwak.batch_job.v1.batch_job_service_pb2_grpc import (
     BatchJobManagementServiceServicer,
 )
 from frogml_services_mock.mocks.utils.exception_handlers import (
@@ -80,9 +80,9 @@ class BatchJobManagerService(BatchJobManagementServiceServicer):
     ) -> CancelBatchJobResponse:
         try:
             if request.batch_id in self.id_to_batch_job.keys():
-                self.id_to_batch_job[
-                    request.batch_id
-                ].status = BatchJobStatusMessage.BATCH_JOB_CANCELLED_STATUS
+                self.id_to_batch_job[request.batch_id].status = (
+                    BatchJobStatusMessage.BATCH_JOB_CANCELLED_STATUS
+                )
                 return CancelBatchJobResponse(success=True)
             else:
                 return CancelBatchJobResponse(

frogml_services_mock/mocks/build_management.py

@@ -2,7 +2,7 @@ from collections import defaultdict
 
 import grpc
 
-from _frogml_proto.qwak.builds.builds_pb2 import (
+from frogml_proto.qwak.builds.builds_pb2 import (
     SUCCESSFUL,
     Build,
     BuildSpec,
@@ -12,7 +12,7 @@ from _frogml_proto.qwak.builds.builds_pb2 import (
     RegisterModelSchemaResponse,
     UpdateBuildStatusResponse,
 )
-from _frogml_proto.qwak.builds.builds_pb2_grpc import BuildsManagementServiceServicer
+from frogml_proto.qwak.builds.builds_pb2_grpc import BuildsManagementServiceServicer
 
 
 class BuildsManagementServiceMock(BuildsManagementServiceServicer):

frogml_services_mock/mocks/build_orchestrator_build_api.py

@@ -2,7 +2,7 @@ from typing import List
 
 from google.protobuf.timestamp_pb2 import Timestamp
 
-from _frogml_proto.qwak.build.v1.build_api_pb2 import (
+from frogml_proto.qwak.build.v1.build_api_pb2 import (
     GetBuildRequest,
     GetBuildResponse,
     ListBuildsRequest,
@@ -22,9 +22,9 @@ from _frogml_proto.qwak.build.v1.build_api_pb2 import (
     UpdateBuildStatusRequest,
     UpdateBuildStatusResponse,
 )
-from _frogml_proto.qwak.build.v1.build_api_pb2_grpc import BuildAPIServicer
-from _frogml_proto.qwak.build.v1.build_pb2 import Audit, Build, BuildStatus
-from frogml.clients.build_orchestrator.client import (
+from frogml_proto.qwak.build.v1.build_api_pb2_grpc import BuildAPIServicer
+from frogml_proto.qwak.build.v1.build_pb2 import Audit, Build, BuildStatus
+from frogml_core.clients.build_orchestrator.client import (
     FrameworkModelDataClass,
     HuggingModelDataClass,
 )

frogml_services_mock/mocks/build_orchestrator_build_settings_api.py

@@ -1,6 +1,6 @@
 from google.protobuf.timestamp_pb2 import Timestamp
 
-from _frogml_proto.qwak.build_settings.build_settings_api_pb2 import (
+from frogml_proto.qwak.build_settings.build_settings_api_pb2 import (
     DeleteBuildSettingsRequest,
     DeleteBuildSettingsResponse,
     GetBuildSettingsRequest,
@@ -8,10 +8,10 @@ from _frogml_proto.qwak.build_settings.build_settings_api_pb2 import (
     SaveBuildSettingsRequest,
     SaveBuildSettingsResponse,
 )
-from _frogml_proto.qwak.build_settings.build_settings_api_pb2_grpc import (
+from frogml_proto.qwak.build_settings.build_settings_api_pb2_grpc import (
     BuildSettingsApiServicer,
 )
-from _frogml_proto.qwak.build_settings.build_settings_pb2 import BuildSettings
+from frogml_proto.qwak.build_settings.build_settings_pb2 import BuildSettings
 from frogml_services_mock.mocks.utils.exception_handlers import (
     raise_internal_grpc_error,
 )

frogml_services_mock/mocks/build_orchestrator_service_api.py

@@ -1,12 +1,12 @@
 import grpc
 
-from _frogml_proto.qwak.builds.build_pb2 import BaseDockerImageType
-from _frogml_proto.qwak.builds.build_url_pb2 import (
+from frogml_proto.qwak.builds.build_pb2 import BaseDockerImageType
+from frogml_proto.qwak.builds.build_url_pb2 import (
     BuildVersioningTagsDefinition,
     BuildVersioningTagsProperties,
     BuildVersioningTagsType,
 )
-from _frogml_proto.qwak.builds.builds_orchestrator_service_pb2 import (
+from frogml_proto.qwak.builds.builds_orchestrator_service_pb2 import (
     BuildModelResponse,
     CancelBuildModelResponse,
     CreateDataTableResponse,
@@ -15,7 +15,7 @@ from _frogml_proto.qwak.builds.builds_orchestrator_service_pb2 import (
     GetBuildVersioningUploadURLResponse,
     ListBuildVersioningTagsResponse,
 )
-from _frogml_proto.qwak.builds.builds_orchestrator_service_pb2_grpc import (
+from frogml_proto.qwak.builds.builds_orchestrator_service_pb2_grpc import (
     BuildsOrchestratorServiceServicer,
 )
 from frogml_services_mock.mocks.utils.exception_handlers import (

frogml_services_mock/mocks/data_versioning_service.py

@@ -1,14 +1,14 @@
 from collections import defaultdict
 from typing import Dict, List
 
-from _frogml_proto.qwak.data_versioning.data_versioning_pb2 import DataTagSpec
-from _frogml_proto.qwak.data_versioning.data_versioning_service_pb2 import (
+from frogml_proto.qwak.data_versioning.data_versioning_pb2 import DataTagSpec
+from frogml_proto.qwak.data_versioning.data_versioning_service_pb2 import (
     GetModelDataTagsRequest,
     GetModelDataTagsResponse,
     RegisterDataTagRequest,
     RegisterDataTagResponse,
 )
-from _frogml_proto.qwak.data_versioning.data_versioning_service_pb2_grpc import (
+from frogml_proto.qwak.data_versioning.data_versioning_service_pb2_grpc import (
     DataVersioningManagementServiceServicer,
 )
 from frogml_services_mock.mocks.utils.exception_handlers import (

frogml_services_mock/mocks/deployment_management_service.py

@@ -4,7 +4,7 @@ from typing import Dict, Optional
 
 from google.protobuf.timestamp_pb2 import Timestamp
 
-from _frogml_proto.qwak.deployment.deployment_pb2 import (
+from frogml_proto.qwak.deployment.deployment_pb2 import (
     DeploymentBrief,
     DeploymentDetails,
     EnvironmentApplyModelTrafficResponse,
@@ -19,7 +19,7 @@ from _frogml_proto.qwak.deployment.deployment_pb2 import (
     ModelDeploymentStatus,
     Variation,
 )
-from _frogml_proto.qwak.deployment.deployment_service_pb2 import (
+from frogml_proto.qwak.deployment.deployment_service_pb2 import (
     ApplyModelTrafficRequest,
     ApplyModelTrafficResponse,
     DeployModelRequest,
@@ -37,7 +37,7 @@ from _frogml_proto.qwak.deployment.deployment_service_pb2 import (
     UpdateDeploymentRuntimeSettingsRequest,
     UpdateDeploymentRuntimeSettingsResponse,
 )
-from _frogml_proto.qwak.deployment.deployment_service_pb2_grpc import (
+from frogml_proto.qwak.deployment.deployment_service_pb2_grpc import (
     DeploymentManagementServiceServicer,
 )
 from frogml_services_mock.mocks.utils.exception_handlers import (
@@ -101,9 +101,9 @@ class DeploymentManagementServiceMock(DeploymentManagementServiceServicer):
                     )
                     if deployment_key in self.active_deployments:
                         prev_deployment = self.active_deployments[deployment_key]
-                        prev_deployment[
-                            "status"
-                        ] = ModelDeploymentStatus.SUCCESSFUL_UNDEPLOYMENT
+                        prev_deployment["status"] = (
+                            ModelDeploymentStatus.SUCCESSFUL_UNDEPLOYMENT
+                        )
                         self.deployments_history[env_request_details.model_uuid].insert(
                             0, prev_deployment
                         )
@@ -130,9 +130,7 @@ class DeploymentManagementServiceMock(DeploymentManagementServiceServicer):
                         env_request_details.hosting_service.kube_deployment.kube_deployment_type
                         == KubeDeploymentType.ONLINE
                     ):
-                        self.model_traffic[
-                            (env_request_details.model_id, env)
-                        ] = (
+                        self.model_traffic[(env_request_details.model_id, env)] = (
                             env_request_details.hosting_service.kube_deployment.serving_strategy.realtime_config.traffic_config
                         )
                     else:
@@ -170,9 +168,9 @@ class DeploymentManagementServiceMock(DeploymentManagementServiceServicer):
                     )
                 else:
                     deployment_details = self.active_deployments[deployment_key]
-                    deployment_details[
-                        "status"
-                    ] = ModelDeploymentStatus.SUCCESSFUL_UNDEPLOYMENT
+                    deployment_details["status"] = (
+                        ModelDeploymentStatus.SUCCESSFUL_UNDEPLOYMENT
+                    )
                     self.deployments_history[details.model_uuid].insert(
                         0, deployment_details
                     )
@@ -181,9 +179,9 @@ class DeploymentManagementServiceMock(DeploymentManagementServiceServicer):
                     )
                     del self.active_deployments[deployment_key]
                     if details.traffic_config:
-                        self.model_traffic[
-                            (details.model_id, env)
-                        ] = details.traffic_config
+                        self.model_traffic[(details.model_id, env)] = (
+                            details.traffic_config
+                        )
                     else:
                         try:
                             del self.model_traffic[(details.model_id, env)]
@@ -269,13 +267,13 @@ class DeploymentManagementServiceMock(DeploymentManagementServiceServicer):
                         self.__get_variation_name(deployment_details["hosting_service"])
                         == request.variation_name
                     ):
-                        results[
-                            deployment_details["env"]
-                        ] = EnvironmentRuntimeDeploymentSettingsResultMessage()
+                        results[deployment_details["env"]] = (
+                            EnvironmentRuntimeDeploymentSettingsResultMessage()
+                        )
                 else:
-                    results[
-                        deployment_details["env"]
-                    ] = EnvironmentRuntimeDeploymentSettingsResultMessage()
+                    results[deployment_details["env"]] = (
+                        EnvironmentRuntimeDeploymentSettingsResultMessage()
+                    )
             return UpdateDeploymentRuntimeSettingsResponse(
                 environment_to_runtime_deployment_settings_response=results
             )

frogml_services_mock/mocks/ecosystem_service_api.py

@@ -1,12 +1,12 @@
-from _frogml_proto.qwak.ecosystem.v0.credentials_pb2 import (
+from frogml_proto.qwak.ecosystem.v0.credentials_pb2 import (
     AwsTemporaryCredentials,
     CloudCredentials,
 )
-from _frogml_proto.qwak.ecosystem.v0.ecosystem_runtime_service_pb2 import (
+from frogml_proto.qwak.ecosystem.v0.ecosystem_runtime_service_pb2 import (
     GetAuthenticatedUserContextResponse,
     GetCloudCredentialsResponse,
 )
-from _frogml_proto.qwak.ecosystem.v0.ecosystem_runtime_service_pb2_grpc import (
+from frogml_proto.qwak.ecosystem.v0.ecosystem_runtime_service_pb2_grpc import (
     QwakEcosystemRuntimeServicer,
 )
 from frogml_services_mock.mocks.utils.exception_handlers import (

frogml_services_mock/mocks/execution_management_service.py

@@ -1,7 +1,7 @@
 from grpc import RpcError
 
-from _frogml_proto.qwak.execution.v1.execution_service_pb2 import TriggerBackfillResponse
-from _frogml_proto.qwak.execution.v1.execution_service_pb2_grpc import (
+from frogml_proto.qwak.execution.v1.execution_service_pb2 import TriggerBackfillResponse
+from frogml_proto.qwak.execution.v1.execution_service_pb2_grpc import (
     FeatureStoreExecutionServiceServicer,
 )
 

frogml_services_mock/mocks/feature_store_data_sources_manager_api.py

@@ -3,11 +3,11 @@ import uuid
 
 import grpc
 
-from _frogml_proto.qwak.feature_store.sources.data_source_pb2 import (
+from frogml_proto.qwak.feature_store.sources.data_source_pb2 import (
     DataSource,
     DataSourceDefinition,
 )
-from _frogml_proto.qwak.feature_store.sources.data_source_service_pb2 import (
+from frogml_proto.qwak.feature_store.sources.data_source_service_pb2 import (
     CreateDataSourceResponse,
     CreateDataSourceUploadURLRequest,
     CreateDataSourceUploadURLResponse,
@@ -15,7 +15,7 @@ from _frogml_proto.qwak.feature_store.sources.data_source_service_pb2 import (
     GetDataSourceByIdResponse,
     GetDataSourceByNameResponse,
 )
-from _frogml_proto.qwak.feature_store.sources.data_source_service_pb2_grpc import (
+from frogml_proto.qwak.feature_store.sources.data_source_service_pb2_grpc import (
     DataSourceServiceServicer,
 )
 

frogml_services_mock/mocks/feature_store_entities_manager_api.py

@@ -2,13 +2,13 @@ import uuid
 
 import grpc
 
-from _frogml_proto.qwak.feature_store.entities.entity_pb2 import (
+from frogml_proto.qwak.feature_store.entities.entity_pb2 import (
     Entity,
     EntityDefinition,
     EntityMetadata,
     FeatureSetBrief,
 )
-from _frogml_proto.qwak.feature_store.entities.entity_service_pb2 import (
+from frogml_proto.qwak.feature_store.entities.entity_service_pb2 import (
     CreateEntityResponse,
     DeleteEntityRequest,
     DeleteEntityResponse,
@@ -16,7 +16,7 @@ from _frogml_proto.qwak.feature_store.entities.entity_service_pb2 import (
     GetEntityByNameResponse,
     ListEntitiesResponse,
 )
-from _frogml_proto.qwak.feature_store.entities.entity_service_pb2_grpc import (
+from frogml_proto.qwak.feature_store.entities.entity_service_pb2_grpc import (
     EntityServiceServicer,
 )
 

frogml_services_mock/mocks/feature_store_feature_set_manager_api.py

@@ -3,7 +3,7 @@ from typing import Dict, Optional
 
 import grpc
 
-from _frogml_proto.qwak.feature_store.features.feature_set_pb2 import (
+from frogml_proto.qwak.feature_store.features.feature_set_pb2 import (
     DeployedModelInUseLink,
     FeatureSet,
     FeatureSetDefinition,
@@ -14,7 +14,7 @@ from _frogml_proto.qwak.feature_store.features.feature_set_pb2 import (
     FeatureSetSpec,
     FeatureStatus,
 )
-from _frogml_proto.qwak.feature_store.features.feature_set_service_pb2 import (
+from frogml_proto.qwak.feature_store.features.feature_set_service_pb2 import (
     DeleteFeatureSetResponse,
     DeleteFeaturesetVersionResponse,
     GetEnvToFeatureSetsMappingResponse,
@@ -26,7 +26,7 @@ from _frogml_proto.qwak.feature_store.features.feature_set_service_pb2 import (
     SetActiveFeaturesetVersionResponse,
     UpdateFeatureSetResponse,
 )
-from _frogml_proto.qwak.feature_store.features.feature_set_service_pb2_grpc import (
+from frogml_proto.qwak.feature_store.features.feature_set_service_pb2_grpc import (
     FeatureSetServiceServicer,
 )
 
@@ -110,9 +110,9 @@ class FeatureSetServiceMock(FeatureSetServiceServicer):
 
             curr_versions = self._fs_versions_by_name[request.feature_set_spec.name]
             curr_versions[curr_version] = featureset_definition
-            self._fs_versions_by_name[
-                featureset_definition.feature_set_spec.name
-            ] = curr_versions
+            self._fs_versions_by_name[featureset_definition.feature_set_spec.name] = (
+                curr_versions
+            )
 
             return UpdateFeatureSetResponse(
                 feature_set=FeatureSet(feature_set_definition=featureset_definition)
@@ -194,9 +194,9 @@ class FeatureSetServiceMock(FeatureSetServiceServicer):
 
     def DeleteFeaturesetVersion(self, request, context):
         version_number = request.featureset_version_number
-        version_definitions: Optional[
-            Dict[int, FeatureSetDefinition]
-        ] = self._fs_versions_by_name.get(request.featureset_name, None)
+        version_definitions: Optional[Dict[int, FeatureSetDefinition]] = (
+            self._fs_versions_by_name.get(request.featureset_name, None)
+        )
         if version_definitions:
             version_definition = version_definitions.get(version_number, None)
             if version_definition:
@@ -223,17 +223,17 @@ class FeatureSetServiceMock(FeatureSetServiceServicer):
             context.set_code(grpc.StatusCode.NOT_FOUND)
 
     def SetActiveFeaturesetVersion(self, request, context):
-        version_definitions: Optional[
-            Dict[int, FeatureSetDefinition]
-        ] = self._fs_versions_by_name.get(request.featureset_name, None)
+        version_definitions: Optional[Dict[int, FeatureSetDefinition]] = (
+            self._fs_versions_by_name.get(request.featureset_name, None)
+        )
         if version_definitions:
             version_definition = version_definitions.get(
                 request.featureset_version_number, None
             )
             if version_definition:
-                self._fs_definition_by_id[
-                    version_definition.feature_set_id
-                ] = version_definition
+                self._fs_definition_by_id[version_definition.feature_set_id] = (
+                    version_definition
+                )
                 return SetActiveFeaturesetVersionResponse()
             else:
                 context.set_details(
@@ -247,9 +247,9 @@ class FeatureSetServiceMock(FeatureSetServiceServicer):
             context.set_code(grpc.StatusCode.NOT_FOUND)
 
     def ListFeaturesetVersionsByName(self, request, context):
-        versions: Optional[
-            Dict[int, FeatureSetDefinition]
-        ] = self._fs_versions_by_name.get(request.featureset_name, None)
+        versions: Optional[Dict[int, FeatureSetDefinition]] = (
+            self._fs_versions_by_name.get(request.featureset_name, None)
+        )
         if versions:
             return ListFeaturesetVersionsByNameResponse(
                 featuresets=[

frogml_services_mock/mocks/features_online_serving_api.py

@@ -4,13 +4,13 @@ from datetime import datetime
 import grpc
 import pandas as pd
 
-from _frogml_proto.qwak.feature_store.serving.serving_pb2 import (
+from frogml_proto.qwak.feature_store.serving.serving_pb2 import (
     MultiFeatureValuesResponse,
 )
-from _frogml_proto.qwak.feature_store.serving.serving_pb2_grpc import (
+from frogml_proto.qwak.feature_store.serving.serving_pb2_grpc import (
     ServingServiceServicer,
 )
-from frogml.exceptions import FrogmlException
+from frogml_core.exceptions import FrogmlException
 
 
 class NpEncoder(json.JSONEncoder):

frogml_services_mock/mocks/features_operator_v3_service.py

@@ -1,10 +1,10 @@
 import uuid
 from typing import Optional
 
-from _frogml_proto.qwak.features_operator.v3.features_operator_async_service_pb2 import (
+from frogml_proto.qwak.features_operator.v3.features_operator_async_service_pb2 import (
     ValidationResponse,
 )
-from _frogml_proto.qwak.features_operator.v3.features_operator_async_service_pb2_grpc import (
+from frogml_proto.qwak.features_operator.v3.features_operator_async_service_pb2_grpc import (
     FeaturesOperatorAsyncServiceServicer,
 )
 

frogml_services_mock/mocks/features_set_state_service_api.py

@@ -1,13 +1,13 @@
 import grpc
 
-from _frogml_proto.qwak.feature_store.features.feature_set_state_pb2 import (
+from frogml_proto.qwak.feature_store.features.feature_set_state_pb2 import (
     FeatureSetState,
 )
-from _frogml_proto.qwak.feature_store.features.feature_set_state_service_pb2 import (
+from frogml_proto.qwak.feature_store.features.feature_set_state_service_pb2 import (
     GetFeatureSetStateByIdV1Response,
     GetFeatureSetStateByNameV1Response,
 )
-from _frogml_proto.qwak.feature_store.features.feature_set_state_service_pb2_grpc import (
+from frogml_proto.qwak.feature_store.features.feature_set_state_service_pb2_grpc import (
     FeatureSetStateServiceServicer,
 )
 from frogml_services_mock.mocks.utils.exception_handlers import (

frogml_services_mock/mocks/feedback_service.py

@@ -1,10 +1,10 @@
-from _frogml_proto.qwak.inference.feedback.feedback_pb2 import (
+from frogml_proto.qwak.inference.feedback.feedback_pb2 import (
     ActualValuesRequest,
     ActualValuesResponse,
     ConfigureFeedbackRequest,
     ConfigureFeedbackResponse,
 )
-from _frogml_proto.qwak.inference.feedback.feedback_pb2_grpc import (
+from frogml_proto.qwak.inference.feedback.feedback_pb2_grpc import (
     FeedbackServiceServicer,
 )
 from frogml_services_mock.mocks.utils.exception_handlers import (

frogml_services_mock/mocks/file_versioning_service.py

@@ -1,14 +1,14 @@
 from collections import defaultdict
 from typing import Dict, List
 
-from _frogml_proto.qwak.file_versioning.file_versioning_pb2 import FileTagSpec
-from _frogml_proto.qwak.file_versioning.file_versioning_service_pb2 import (
+from frogml_proto.qwak.file_versioning.file_versioning_pb2 import FileTagSpec
+from frogml_proto.qwak.file_versioning.file_versioning_service_pb2 import (
     GetModelFileTagsRequest,
     GetModelFileTagsResponse,
     RegisterFileTagRequest,
     RegisterFileTagResponse,
 )
-from _frogml_proto.qwak.file_versioning.file_versioning_service_pb2_grpc import (
+from frogml_proto.qwak.file_versioning.file_versioning_service_pb2_grpc import (
     FileVersioningManagementServiceServicer,
 )
 from frogml_services_mock.mocks.utils.exception_handlers import (

frogml_services_mock/mocks/fs_offline_serving_service.py

@@ -1,7 +1,7 @@
 import uuid
 from typing import Optional
 
-from _frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2 import (
+from frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2 import (
     GetFeatureValuesInRangeRequest,
     GetFeatureValuesInRangeResponse,
     GetFeatureValuesRequest,
@@ -9,10 +9,10 @@ from _frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2 imp
     GetFeatureValuesResultResponse,
     GetFileUploadUrlResponse,
 )
-from _frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2_grpc import (
+from frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2_grpc import (
     FeatureStoreOfflineServingAsyncServiceServicer,
 )
-from _frogml_proto.qwak.offline.serving.v1.options_pb2 import (
+from frogml_proto.qwak.offline.serving.v1.options_pb2 import (
     OfflineServingQueryOptions as ProtoOfflineServingQueryOptions,
 )
 

frogml_services_mock/mocks/instance_template_management_service.py

@@ -1,17 +1,17 @@
-from _frogml_proto.qwak.instance_template.instance_template_pb2 import (
+from frogml_proto.qwak.instance_template.instance_template_pb2 import (
     InstanceTemplateSpec,
     InstanceType,
 )
-from _frogml_proto.qwak.instance_template.instance_template_service_pb2 import (
+from frogml_proto.qwak.instance_template.instance_template_service_pb2 import (
     GetInstanceTemplateRequest,
     GetInstanceTemplateResponse,
     ListInstanceTemplatesRequest,
     ListInstanceTemplatesResponse,
 )
-from _frogml_proto.qwak.instance_template.instance_template_service_pb2_grpc import (
+from frogml_proto.qwak.instance_template.instance_template_service_pb2_grpc import (
     InstanceTemplateManagementServiceServicer,
 )
-from _frogml_proto.qwak.user_application.common.v0.resources_pb2 import (
+from frogml_proto.qwak.user_application.common.v0.resources_pb2 import (
     CpuResources,
     GpuResources,
     GpuType,
@@ -19,7 +19,7 @@ from _frogml_proto.qwak.user_application.common.v0.resources_pb2 import (
     NodeOptimizationType,
     PodComputeResources,
 )
-from frogml.exceptions import FrogmlException
+from frogml_core.exceptions import FrogmlException
 from frogml_services_mock.mocks.utils.exception_handlers import (
     raise_internal_grpc_error,
 )

frogml_services_mock/mocks/integration_management_service.py

@@ -3,8 +3,8 @@ from typing import Dict
 
 import grpc
 
-from _frogml_proto.qwak.integration.integration_pb2 import Integration, IntegrationSpec
-from _frogml_proto.qwak.integration.integration_service_pb2 import (
+from frogml_proto.qwak.integration.integration_pb2 import Integration, IntegrationSpec
+from frogml_proto.qwak.integration.integration_service_pb2 import (
     CreateIntegrationRequest,
     CreateIntegrationResponse,
     DeleteIntegrationRequest,
@@ -14,10 +14,10 @@ from _frogml_proto.qwak.integration.integration_service_pb2 import (
     ListIntegrationRequest,
     ListIntegrationsResponse,
 )
-from _frogml_proto.qwak.integration.integration_service_pb2_grpc import (
+from frogml_proto.qwak.integration.integration_service_pb2_grpc import (
     IntegrationManagementServiceServicer,
 )
-from _frogml_proto.qwak.integration.open_a_i_integration_pb2 import (
+from frogml_proto.qwak.integration.open_a_i_integration_pb2 import (
     OpenAIApiKeySystemSecretDescriptor,
     OpenAIIntegration,
 )

frogml_services_mock/mocks/internal_build_orchestrator_service.py

@@ -1,9 +1,9 @@
-from _frogml_proto.qwak.builds.internal_builds_orchestrator_service_pb2 import (
+from frogml_proto.qwak.builds.internal_builds_orchestrator_service_pb2 import (
     BuildInitDetails,
     InternalBuildModelRequest,
     InternalBuildModelResponse,
 )
-from _frogml_proto.qwak.builds.internal_builds_orchestrator_service_pb2_grpc import (
+from frogml_proto.qwak.builds.internal_builds_orchestrator_service_pb2_grpc import (
     InternalBuildsOrchestratorServiceServicer,
 )
 from frogml_services_mock.mocks.utils.exception_handlers import (

frogml_services_mock/mocks/job_registry_service_api.py

@@ -1,7 +1,7 @@
 import grpc
 
-from _frogml_proto.qwak.feature_store.jobs.v1.job_pb2 import JobRecord, JobState
-from _frogml_proto.qwak.feature_store.jobs.v1.job_service_pb2 import (
+from frogml_proto.qwak.feature_store.jobs.v1.job_pb2 import JobRecord, JobState
+from frogml_proto.qwak.feature_store.jobs.v1.job_service_pb2 import (
     ApplyJobRequest,
     ApplyJobResponse,
     GetLatestJobRequest,
@@ -9,7 +9,7 @@ from _frogml_proto.qwak.feature_store.jobs.v1.job_service_pb2 import (
     GetLatestSuccessfulJobRequest,
     GetLatestSuccessfulJobResponse,
 )
-from _frogml_proto.qwak.feature_store.jobs.v1.job_service_pb2_grpc import (
+from frogml_proto.qwak.feature_store.jobs.v1.job_service_pb2_grpc import (
     JobServiceServicer,
 )
 from frogml_services_mock.mocks.utils.exception_handlers import (

frogml_services_mock/mocks/kube_captain_service_api.py

@@ -1,8 +1,8 @@
-from _frogml_proto.qwak.kube_deployment_captain.kube_deployment_captain_service_pb2 import (
+from frogml_proto.qwak.kube_deployment_captain.kube_deployment_captain_service_pb2 import (
     DeployFeatureSetResponse,
     DeployStreamingAggregationCompactionResponse,
 )
-from _frogml_proto.qwak.kube_deployment_captain.kube_deployment_captain_service_pb2_grpc import (
+from frogml_proto.qwak.kube_deployment_captain.kube_deployment_captain_service_pb2_grpc import (
     KubeDeploymentCaptainServicer,
 )
 from frogml_services_mock.mocks.utils.exception_handlers import (

frogml_services_mock/mocks/logging_service.py

@@ -2,15 +2,15 @@ from collections import defaultdict
 from enum import Enum
 from typing import Dict, List
 
-from _frogml_proto.qwak.logging.log_line_pb2 import LogLine
-from _frogml_proto.qwak.logging.log_reader_service_pb2 import (
+from frogml_proto.qwak.logging.log_line_pb2 import LogLine
+from frogml_proto.qwak.logging.log_reader_service_pb2 import (
     ReadLogsRequest,
     ReadLogsResponse,
 )
-from _frogml_proto.qwak.logging.log_reader_service_pb2_grpc import (
+from frogml_proto.qwak.logging.log_reader_service_pb2_grpc import (
     LogReaderServiceServicer,
 )
-from _frogml_proto.qwak.logging.log_source_pb2 import (
+from frogml_proto.qwak.logging.log_source_pb2 import (
     InferenceExecutionSource,
     LogSource,
     ModelRuntimeSource,

frogml_services_mock/mocks/model_management_service.py

@@ -3,7 +3,7 @@ from typing import Dict
 
 import grpc
 
-from _frogml_proto.qwak.models.models_pb2 import (
+from frogml_proto.qwak.models.models_pb2 import (
     CreateModelRequest,
     CreateModelResponse,
     DeleteModelRequest,
@@ -18,7 +18,7 @@ from _frogml_proto.qwak.models.models_pb2 import (
     ModelSpec,
     ModelStatus,
 )
-from _frogml_proto.qwak.models.models_pb2_grpc import ModelsManagementServiceServicer
+from frogml_proto.qwak.models.models_pb2_grpc import ModelsManagementServiceServicer
 from frogml_services_mock.mocks.utils.exception_handlers import (
     raise_internal_grpc_error,
 )