frogml-core 0.0.1__py3-none-any.whl

frogml_services_mock/mocks/batch_job_manager_service.py

@@ -0,0 +1,337 @@
+import uuid
+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 (
+    AwsTemporaryCredentials,
+)
+from _frogml_proto.qwak.batch_job.v1.batch_job_service_pb2 import (
+    BatchJobDeploymentSize,
+    BatchJobDetails,
+    BatchJobMessage,
+    BatchJobStatusMessage,
+    BatchTaskStatusMessage,
+    CancelBatchJobRequest,
+    CancelBatchJobResponse,
+    CancelWarmupJobRequest,
+    CancelWarmupJobResponse,
+    ExecutionReportLine,
+    GetBatchJobDetailsRequest,
+    GetBatchJobDetailsResponse,
+    GetBatchJobDownloadDetailsRequest,
+    GetBatchJobDownloadDetailsResponse,
+    GetBatchJobPreSignedDownloadUrlRequest,
+    GetBatchJobPreSignedDownloadUrlResponse,
+    GetBatchJobPreSignedUploadUrlRequest,
+    GetBatchJobPreSignedUploadUrlResponse,
+    GetBatchJobReportRequest,
+    GetBatchJobReportResponse,
+    GetBatchJobStatusRequest,
+    GetBatchJobStatusResponse,
+    GetBatchJobUploadDetailsRequest,
+    GetBatchJobUploadDetailsResponse,
+    ListBatchJobsRequest,
+    ListBatchJobsResponse,
+    StartBatchJobRequest,
+    StartBatchJobResponse,
+    StartWarmupJobRequest,
+    StartWarmupJobResponse,
+    TaskExecutionDetails,
+)
+from _frogml_proto.qwak.batch_job.v1.batch_job_service_pb2_grpc import (
+    BatchJobManagementServiceServicer,
+)
+from frogml_services_mock.mocks.utils.exception_handlers import (
+    raise_internal_grpc_error,
+)
+
+BATCH_JOB_NOT_FOUND_ERROR = "Batch Job Not Found"
+
+
+class BatchJobManagerService(BatchJobManagementServiceServicer):
+    def __init__(self):
+        self.id_to_batch_job: Dict[str, MockBatchJob] = dict()
+        self.model_id_active_warmup: Set[str] = set()
+        self.models_to_fail: Set[str] = set()
+
+    def StartBatchJob(
+        self, request: StartBatchJobRequest, context
+    ) -> StartBatchJobResponse:
+        try:
+            model_id = (
+                request.model_id
+                if request.model_id
+                else request.batch_job_request.model_details.model_id
+            )
+
+            if model_id in self.models_to_fail:
+                return StartBatchJobResponse(
+                    success=False, failure_message="Model start failed"
+                )
+            mock_batch_job: MockBatchJob = MockBatchJob(request)
+            self.id_to_batch_job[mock_batch_job.id] = mock_batch_job
+            return StartBatchJobResponse(batch_id=mock_batch_job.id, success=True)
+        except Exception as e:
+            return StartBatchJobResponse(success=False, failure_message=str(e))
+
+    def CancelBatchJob(
+        self, request: CancelBatchJobRequest, context
+    ) -> 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
+                return CancelBatchJobResponse(success=True)
+            else:
+                return CancelBatchJobResponse(
+                    success=False, failure_message=BATCH_JOB_NOT_FOUND_ERROR
+                )
+        except Exception as e:
+            return CancelBatchJobResponse(success=False, failure_message=str(e))
+
+    def StartWarmupJob(
+        self, request: StartWarmupJobRequest, context
+    ) -> StartWarmupJobResponse:
+        try:
+            self.model_id_active_warmup.add(request.model_id)
+            return StartWarmupJobResponse(success=True)
+        except Exception as e:
+            return StartWarmupJobResponse(success=False, failure_message=str(e))
+
+    def CancelWarmupJob(
+        self, request: CancelWarmupJobRequest, context
+    ) -> CancelWarmupJobResponse:
+        try:
+            if request.model_id in self.model_id_active_warmup:
+                self.model_id_active_warmup.remove(request.model_id)
+            return CancelWarmupJobResponse(success=True)
+        except Exception as e:
+            return CancelWarmupJobResponse(success=False, failure_message=str(e))
+
+    def GetBatchJobStatus(
+        self, request: GetBatchJobStatusRequest, context
+    ) -> GetBatchJobStatusResponse:
+        try:
+            if request.batch_id in self.id_to_batch_job:
+                batch_job = self.id_to_batch_job[request.batch_id]
+                return GetBatchJobStatusResponse(
+                    success=True,
+                    job_status=batch_job.status,
+                    finished_files=batch_job.finished_files,
+                    total_files=batch_job.total_files,
+                )
+            else:
+                return GetBatchJobStatusResponse(
+                    success=False, failure_message=BATCH_JOB_NOT_FOUND_ERROR
+                )
+        except Exception as e:
+            return GetBatchJobStatusResponse(success=False, failure_message=str(e))
+
+    def GetBatchJobReport(
+        self, request: GetBatchJobReportRequest, context
+    ) -> GetBatchJobReportResponse:
+        try:
+            if request.batch_id in self.id_to_batch_job:
+                batch_job = self.id_to_batch_job[request.batch_id]
+                return GetBatchJobReportResponse(
+                    successful=True, report_messages=batch_job.report
+                )
+            else:
+                return GetBatchJobReportResponse(
+                    successful=False, failure_message=BATCH_JOB_NOT_FOUND_ERROR
+                )
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def ListBatchJobs(
+        self, request: ListBatchJobsRequest, context
+    ) -> ListBatchJobsResponse:
+        try:
+            if request.build_id:
+                batch_jobs = list(
+                    filter(
+                        lambda job: job.start_request.build_id == request.build_id,
+                        self.id_to_batch_job.values(),
+                    )
+                )
+            else:
+                batch_jobs = list(
+                    filter(
+                        lambda job: job.start_request.model_id == request.model_id,
+                        self.id_to_batch_job.values(),
+                    )
+                )
+
+            batch_jobs_details = list(map(batch_job_to_details, batch_jobs))
+            return ListBatchJobsResponse(batch_jobs=batch_jobs_details)
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def GetBatchJobDetails(
+        self, request: GetBatchJobDetailsRequest, context
+    ) -> GetBatchJobDetailsResponse:
+        try:
+            if request.job_id in self.id_to_batch_job:
+                return GetBatchJobDetailsResponse(
+                    success=True,
+                    batch_job=batch_job_to_message(
+                        self.id_to_batch_job[request.job_id]
+                    ),
+                )
+            else:
+                return GetBatchJobDetailsResponse(
+                    success=False, failure_message=BATCH_JOB_NOT_FOUND_ERROR
+                )
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def GetBatchJobPreSignedUploadUrl(
+        self, request: GetBatchJobPreSignedUploadUrlRequest, context
+    ) -> GetBatchJobPreSignedUploadUrlResponse:
+        try:
+            return GetBatchJobPreSignedUploadUrlResponse(
+                input_path="/input_path",
+                output_path="/output_path",
+                bucket="bucket",
+                urls=[
+                    f"{request.model_id}_{i}.{request.file_type}"
+                    for i in range(request.number_of_files)
+                ],
+                success=True,
+            )
+        except Exception as e:
+            return GetBatchJobPreSignedUploadUrlResponse(
+                success=False, failure_message=str(e)
+            )
+
+    def GetBatchJobPreSignedDownloadUrl(
+        self, request: GetBatchJobPreSignedDownloadUrlRequest, context
+    ) -> GetBatchJobPreSignedDownloadUrlResponse:
+        try:
+            if request.job_id in self.id_to_batch_job:
+                batch_job = self.id_to_batch_job[request.job_id]
+                return GetBatchJobPreSignedDownloadUrlResponse(
+                    urls=[
+                        f"bucket://output_path/{batch_job.id}/{i}.{batch_job.start_request.output_file_type}"
+                        for i in range(batch_job.total_files)
+                    ],
+                    success=True,
+                )
+            else:
+                return GetBatchJobPreSignedDownloadUrlResponse(
+                    success=False, failure_message=BATCH_JOB_NOT_FOUND_ERROR
+                )
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def GetBatchJobUploadDetails(
+        self, request: GetBatchJobUploadDetailsRequest, context
+    ) -> GetBatchJobUploadDetailsResponse:
+        try:
+            return GetBatchJobUploadDetailsResponse(
+                input_path="/input_path",
+                output_path="/output_path",
+                bucket="bucket",
+                credentials=AwsTemporaryCredentials(
+                    access_key_id="access_key_id",
+                    secret_access_key="secret_access_key",
+                    session_token="session_token",
+                    expiration_time=Timestamp(seconds=3600),
+                ),
+            )
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def GetBatchJobDownloadDetails(
+        self, request: GetBatchJobDownloadDetailsRequest, context
+    ) -> GetBatchJobDownloadDetailsResponse:
+        try:
+            if request.job_id in self.id_to_batch_job:
+                batch_job = self.id_to_batch_job[request.job_id]
+                return GetBatchJobDownloadDetailsResponse(
+                    bucket="bucket",
+                    keys=[
+                        f"bucket://output_path/{batch_job.id}/{i}.{batch_job.start_request.output_file_type}"
+                        for i in range(batch_job.total_files)
+                    ],
+                    credentials=AwsTemporaryCredentials(
+                        access_key_id="access_key_id",
+                        secret_access_key="secret_access_key",
+                        session_token="session_token",
+                        expiration_time=Timestamp(seconds=3600),
+                    ),
+                )
+            else:
+                raise Exception(BATCH_JOB_NOT_FOUND_ERROR)
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+
+class MockBatchJob:
+    def __init__(self, request: StartBatchJobRequest):
+        self.start_request: StartBatchJobRequest = request
+        self.status: BatchJobStatusMessage = (
+            BatchJobStatusMessage.BATCH_JOB_PENDING_STATUS
+        )
+        self.start_time: Timestamp = Timestamp()
+        self.end_time: Optional[Timestamp] = None
+        self.id: str = str(uuid.uuid4())
+        self.finished_files: int = 0
+        self.total_files: int = 0
+        self.report: List[str] = []
+        self.tasks: List[TaskExecutionDetails] = []
+
+    def add_report_messages(self, messages: List[str]):
+        self.report.extend(messages)
+
+    def add_task_executions(self):
+        self.tasks.append(
+            TaskExecutionDetails(
+                task_id=str(uuid.uuid4()),
+                status=BatchTaskStatusMessage.BATCH_TASK_RUNNING_STATUS,
+                start_time=Timestamp(),
+                end_time=Timestamp(),
+                filename="test",
+            )
+        )
+
+
+def batch_job_to_details(batch_job: MockBatchJob) -> BatchJobDetails:
+    return BatchJobDetails(
+        build_id=batch_job.start_request.build_id,
+        job_id=batch_job.id,
+        job_status=batch_job.status,
+        start_time=batch_job.start_time,
+        end_time=batch_job.end_time,
+        job_size=BatchJobDeploymentSize(
+            number_of_pods=batch_job.start_request.batch_job_deployment_size.number_of_pods,
+            cpu=batch_job.start_request.batch_job_deployment_size.cpu,
+            memory_amount=batch_job.start_request.batch_job_deployment_size.memory_amount,
+            memory_units=batch_job.start_request.batch_job_deployment_size.memory_units,
+        ),
+    )
+
+
+def batch_job_to_message(batch_job: MockBatchJob) -> BatchJobMessage:
+    return BatchJobMessage(
+        job_id=batch_job.id,
+        model_id=batch_job.start_request.model_id,
+        branch_id=batch_job.start_request.branch_id,
+        build_id=batch_job.start_request.build_id,
+        start_time=batch_job.start_time,
+        job_status=batch_job.status,
+        end_time=batch_job.end_time,
+        job_size=BatchJobDeploymentSize(
+            number_of_pods=batch_job.start_request.batch_job_deployment_size.number_of_pods,
+            cpu=batch_job.start_request.batch_job_deployment_size.cpu,
+            memory_amount=batch_job.start_request.batch_job_deployment_size.memory_amount,
+            memory_units=batch_job.start_request.batch_job_deployment_size.memory_units,
+        ),
+        report_messages=[
+            ExecutionReportLine(time=Timestamp(), line=line)
+            for line in batch_job.report
+        ],
+        task_executions=batch_job.tasks,
+    )

frogml_services_mock/mocks/build_management.py

@@ -0,0 +1,101 @@
+from collections import defaultdict
+
+import grpc
+
+from _frogml_proto.qwak.builds.builds_pb2 import (
+    SUCCESSFUL,
+    Build,
+    BuildSpec,
+    GetBuildResponse,
+    ListBuildsResponse,
+    RegisterExperimentTrackingValuesResponse,
+    RegisterModelSchemaResponse,
+    UpdateBuildStatusResponse,
+)
+from _frogml_proto.qwak.builds.builds_pb2_grpc import BuildsManagementServiceServicer
+
+
+class BuildsManagementServiceMock(BuildsManagementServiceServicer):
+    def __init__(self):
+        self.builds = defaultdict(lambda: {})
+        self.builds_to_models = defaultdict()
+        self.init_status = SUCCESSFUL
+
+    def UpdateBuildStatus(self, request, context):
+        model_uuid = self.builds_to_models[request.build_id]
+        if model_uuid in self.builds:
+            self.builds[model_uuid][request.build_id]["status"] = str(
+                request.build_status
+            )
+        else:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details("Received a non existing build id")
+
+        return UpdateBuildStatusResponse()
+
+    def RegisterModelSchema(self, request, context):
+        model_uuid = self.builds_to_models[request.build_id]
+        if model_uuid in self.builds:
+            self.builds[model_uuid][request.build_id]["schema"] = request.model_schema
+        return RegisterModelSchemaResponse()
+
+    def RegisterExperimentTrackingValues(self, request, context):
+        model_uuid = self.builds_to_models[request.build_id]
+        if model_uuid in self.builds:
+            spec = BuildSpec(
+                build_id=request.build_id,
+                experiment_tracking_values=request.experiment_tracking_values,
+            )
+            self.builds[model_uuid][request.build_id] = {
+                "spec": spec,
+                "status": 1,
+            }
+        else:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details("Received a non existing build id")
+        return RegisterExperimentTrackingValuesResponse()
+
+    def GetBuild(self, request, context):
+        model_uuid = self.builds_to_models[request.build_id]
+        if model_uuid in self.builds:
+            persistent_build = self.builds[model_uuid][request.build_id]
+
+            build = self._persistent_build_to_build_proto(persistent_build)
+
+            return GetBuildResponse(build=build)
+        else:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details("Received a non existing build id")
+            return GetBuildResponse()
+
+    def ListBuilds(self, request, context):
+        return ListBuildsResponse(
+            builds=[
+                self._persistent_build_to_build_proto(persistent_build)
+                for persistent_build in self.builds[request.model_uuid].values()
+            ]
+        )
+
+    @staticmethod
+    def _persistent_build_to_build_proto(persistent_build):
+        spec = persistent_build["spec"]
+        return Build(
+            build_spec=BuildSpec(
+                build_id=spec.build_id,
+                commit_id=spec.commit_id,
+                model_id=spec.model_id,
+                branch_name=spec.branch_name,
+                tags=spec.tags,
+                experiment_tracking_values=spec.experiment_tracking_values,
+            ),
+            build_status=SUCCESSFUL,
+            model_schema=(
+                persistent_build["schema"] if "schema" in persistent_build else None
+            ),
+        )
+
+    def given_build(self, build_id: str, model_uuid: str, build_spec: BuildSpec):
+        self.builds[model_uuid] = {
+            build_id: {"spec": build_spec, "status": self.init_status}
+        }
+        self.builds_to_models[build_id] = model_uuid

frogml_services_mock/mocks/build_orchestrator_build_api.py

@@ -0,0 +1,138 @@
+from typing import List
+
+from google.protobuf.timestamp_pb2 import Timestamp
+
+from _frogml_proto.qwak.build.v1.build_api_pb2 import (
+    GetBuildRequest,
+    GetBuildResponse,
+    ListBuildsRequest,
+    ListBuildsResponse,
+    LogPhaseStatusRequest,
+    LogPhaseStatusResponse,
+    RegisterBuildRequest,
+    RegisterBuildResponse,
+    RegisterExperimentTrackingRequest,
+    RegisterExperimentTrackingResponse,
+    RegisterModelSchemaRequest,
+    RegisterModelSchemaResponse,
+    RegisterTagsRequest,
+    RegisterTagsResponse,
+    SaveFrameworkModelsRequest,
+    SaveFrameworkModelsResponse,
+    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 (
+    FrameworkModelDataClass,
+    HuggingModelDataClass,
+)
+from frogml_services_mock.mocks.utils.exception_handlers import (
+    raise_internal_grpc_error,
+)
+
+
+class BuildOrchestratorBuildApiMock(BuildAPIServicer):
+    def __init__(self):
+        super(BuildOrchestratorBuildApiMock, self).__init__()
+        self._builds: dict[str:Build] = {}
+        self._framework_models: dict[str : List[FrameworkModelDataClass]] = {}
+
+    def given_build(self, build: Build):
+        self._builds[build.buildId] = build
+
+    def SaveFrameworkModels(self, request: SaveFrameworkModelsRequest, context):
+        self._framework_models[request.spec.build_id] = request.spec.framework_models
+        return SaveFrameworkModelsResponse()
+
+    def RegisterBuild(
+        self, request: RegisterBuildRequest, context
+    ) -> RegisterBuildResponse:
+        timestamp = Timestamp()
+        timestamp.GetCurrentTime()
+        now = timestamp
+
+        audit = Audit(
+            created_by="", created_at=now, last_modified_by="", last_modified_at=now
+        )
+        self._builds[request.buildId] = Build(
+            buildId=request.buildId,
+            commitId=request.commitId,
+            branchId="",
+            buildConfig=request.buildConfig,
+            build_status=BuildStatus.IN_PROGRESS,
+            tags=request.tags,
+            steps=request.steps,
+            audit=audit,
+        )
+
+        return RegisterBuildResponse()
+
+    def UpdateBuildStatus(
+        self, request: UpdateBuildStatusRequest, context
+    ) -> UpdateBuildStatusResponse:
+        self._builds[request.buildId].build_status = request.build_status
+        return UpdateBuildStatusResponse()
+
+    def RegisterTags(
+        self, request: RegisterTagsRequest, context
+    ) -> RegisterTagsResponse:
+        self._builds[request.build_id].tags.extend(request.tags)
+        return RegisterTagsResponse()
+
+    def RegisterExperimentTracking(
+        self, request: RegisterExperimentTrackingRequest, context
+    ) -> RegisterExperimentTrackingResponse:
+        self._builds[request.build_id].params.update(request.params)
+        self._builds[request.build_id].metrics.update(request.metrics)
+        return RegisterExperimentTrackingResponse()
+
+    def RegisterModelSchema(
+        self, request: RegisterModelSchemaRequest, context
+    ) -> RegisterModelSchemaResponse:
+        self._builds[request.build_id].model_schema.CopyFrom(request.model_schema)
+        return RegisterModelSchemaResponse()
+
+    def GetBuild(self, request: GetBuildRequest, context) -> GetBuildResponse:
+        return GetBuildResponse(build=self._builds[request.build_id])
+
+    def ListBuilds(self, request: ListBuildsRequest, context) -> ListBuildsResponse:
+        valid_builds = [
+            self._builds[buildId]
+            for buildId in self._builds
+            if self._builds[buildId].model_uuid == request.model_uuid
+        ]
+        builds_to_return = []
+
+        if len(request.filter.tags) > 0:
+            for build in valid_builds:
+                if request.filter.require_all_tags:
+                    should_be_included = all(
+                        tag in build.tags for tag in request.filter.tags
+                    )
+                else:
+                    should_be_included = any(
+                        tag in build.tags for tag in request.filter.tags
+                    )
+                if not request.filter.include_extra_tags:
+                    should_be_included = should_be_included and not any(
+                        tag not in request.filter.tags for tag in build.tags
+                    )
+                if should_be_included:
+                    builds_to_return.append(build)
+        else:
+            builds_to_return = valid_builds
+
+        return ListBuildsResponse(build=builds_to_return)
+
+    def LogPhaseStatus(
+        self, request: LogPhaseStatusRequest, context
+    ) -> LogPhaseStatusResponse:
+        try:
+            if request.build_id not in self._builds:
+                raise Exception(f"Build {request.build_id} doesn't exist")
+
+            return LogPhaseStatusResponse()
+        except Exception as e:
+            raise_internal_grpc_error(context, e)

frogml_services_mock/mocks/build_orchestrator_build_settings_api.py

@@ -0,0 +1,38 @@
+from google.protobuf.timestamp_pb2 import Timestamp
+
+from _frogml_proto.qwak.build_settings.build_settings_api_pb2 import (
+    DeleteBuildSettingsRequest,
+    DeleteBuildSettingsResponse,
+    GetBuildSettingsRequest,
+    GetBuildSettingsResponse,
+    SaveBuildSettingsRequest,
+    SaveBuildSettingsResponse,
+)
+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_services_mock.mocks.utils.exception_handlers import (
+    raise_internal_grpc_error,
+)
+
+
+class BuildOrchestratorBuildSettingsApiMock(BuildSettingsApiServicer):
+    def __init__(self):
+        self._build_settings = dict()
+
+    def given_build_settings(self, environment_id: str, build_settings: BuildSettings):
+        self._build_settings[environment_id] = build_settings
+
+    def GetBuildSettings(
+        self, request: GetBuildSettingsRequest, context
+    ) -> GetBuildSettingsResponse:
+        build_settings = self._build_settings.get(request.environment_id)
+        if not build_settings:
+            raise_internal_grpc_error(
+                context,
+                Exception(
+                    f"Build settings for environment id {request.environment_id} doesn't exist"
+                ),
+            )
+        return GetBuildSettingsResponse(build_settings=build_settings)