frogml-core 0.0.115__py3-none-any.whl → 0.0.117__py3-none-any.whl

frogml_services_mock/mocks/batch_job_manager_service.py

@@ -1,360 +0,0 @@
-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,
-    UpdateTasksDetailsRequest,
-    UpdateTasksDetailsResponse,
-    BatchTaskDetails,
-)
-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()
-        self.task_id_to_update_task_details: Dict[str, BatchTaskDetails] = dict()
-
-    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)
-
-    def UpdateTasksDetails(
-            self, request: UpdateTasksDetailsRequest, context
-    ) -> UpdateTasksDetailsResponse:
-        for task in request.tasks_details:
-            task_id = task.task_id
-            if not task_id :
-                raise ValueError("Task ID cannot be empty")
-            input_files = task.input_files_details
-
-            print(f"Updating task with ID: {task_id} and input files: {[file.path for file in input_files]}")
-            self.task_id_to_update_task_details[task_id] = task
-        return UpdateTasksDetailsResponse()
-
-    def get_task_details(self,task_id: str
-        ) -> Optional[BatchTaskDetails]:
-            """
-            Get task details by task ID
-            """
-            return self.task_id_to_update_task_details[task_id]
-
-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

@@ -1,101 +0,0 @@
-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

@@ -1,138 +0,0 @@
-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_core.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

@@ -1,38 +0,0 @@
-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)