frogml 1.2.50__py3-none-any.whl → 2.0.0__py3-none-any.whl

frogml/core/clients/batch_job_management/client.py

@@ -1,4 +1,5 @@
-import grpc
+from typing import Optional
+
 from dependency_injector.wiring import Provide
 
 from frogml._proto.qwak.batch_job.v1.batch_job_service_pb2 import (
@@ -36,7 +37,6 @@ from frogml._proto.qwak.batch_job.v1.batch_job_service_pb2 import (
     StartBatchJobResponse,
     StartWarmupJobRequest,
     StartWarmupJobResponse,
-    UpdateTasksDetailsResponse,
     UpdateTasksDetailsRequest,
     BatchTaskDetails,
     InputFileDetails,
@@ -56,6 +56,7 @@ from frogml.core.clients.logging_client import LoggingClient
 from frogml.core.clients.model_management import ModelsManagementClient
 from frogml.core.exceptions import FrogmlException
 from frogml.core.inner.di_configuration import FrogmlContainer
+from frogml.core.inner.tool.grpc.grpc_try_wrapping import grpc_try_catch_wrapper
 
 from .executions_config import (
     INPUT_FORMATTERS_MAP,
@@ -71,7 +72,6 @@ from .results import (
     GetExecutionReportResult,
     StartExecutionResult,
 )
-from typing import List, Optional
 
 CLIENT_TIMEOUT = 180  # Seconds
 
@@ -80,7 +80,7 @@ class BatchJobManagerClient:
     def __init__(
         self,
         grpc_channel=Provide[FrogmlContainer.core_grpc_channel],
-        logging_client: LoggingClient = None,
+        logging_client: Optional[LoggingClient] = None,
     ):
         self.batch_job_management = BatchJobManagementServiceStub(grpc_channel)
         self.logging_client = logging_client
@@ -98,80 +98,80 @@ class BatchJobManagerClient:
         Returns:
             The response received from the api. On successful start of execution a batch job id is returned
         """
-        job_size = self.get_batch_deployment_size_from_resources(execution_config)
+        job_size: BatchJobDeploymentSize = (
+            self.get_batch_deployment_size_from_resources(execution_config)
+        )
 
-        try:
-            user_raw_input_format = execution_config.execution.input_file_type.upper()
-            user_raw_output_format = execution_config.execution.output_file_type.upper()
+        user_raw_input_format: str = execution_config.execution.input_file_type.upper()
+        user_raw_output_format: str = (
+            execution_config.execution.output_file_type.upper()
+        )
 
-            if user_raw_input_format not in INPUT_FORMATTERS_MAP:
-                raise ValueError(
-                    f"Invalid input format - please choose one of {list(INPUT_FORMATTERS_MAP.keys())}"
-                )
+        if user_raw_input_format not in INPUT_FORMATTERS_MAP:
+            raise ValueError(
+                f"Invalid input format - please choose one of {list(INPUT_FORMATTERS_MAP.keys())}"
+            )
 
-            if user_raw_output_format not in OUTPUT_FORMATTERS_MAP:
+        if user_raw_output_format not in OUTPUT_FORMATTERS_MAP:
+            raise ValueError(
+                f"Invalid output format - please choose one of {list(OUTPUT_FORMATTERS_MAP.keys())}"
+            )
+
+        user_purchase_option: Optional[str] = (
+            execution_config.advanced_options.purchase_option
+        )
+        if user_purchase_option is not None:
+            if user_purchase_option not in PURCHASE_OPTION_SET:
                 raise ValueError(
-                    f"Invalid input format - please choose one of {list(OUTPUT_FORMATTERS_MAP.keys())}"
+                    f"Invalid purchase option - please choose one of {list(PURCHASE_OPTION_SET)}"
                 )
+            user_purchase_option = user_purchase_option.replace("-", "")
 
-            user_purchase_option = execution_config.advanced_options.purchase_option
-            if user_purchase_option is not None:
-                if user_purchase_option not in PURCHASE_OPTION_SET:
-                    raise ValueError(
-                        f"Invalid purchase option - please choose one of {list(PURCHASE_OPTION_SET)}"
-                    )
-                else:
-                    user_purchase_option = user_purchase_option.replace("-", "")
-
-            start_job_result: StartBatchJobResponse = self._start_batch_job(
-                batch_job_request=BatchJobRequest(
-                    model_details=BatchJobModelDetails(
-                        model_id=execution_config.execution.model_id,
-                        build_id=execution_config.execution.build_id,
+        start_job_result: StartBatchJobResponse = self._start_batch_job(
+            batch_job_request=BatchJobRequest(
+                model_details=BatchJobModelDetails(
+                    model_id=execution_config.execution.model_id,
+                    build_id=execution_config.execution.build_id,
+                ),
+                data_details=BatchJobDataDetails(
+                    source_path=BatchJobSourcePath(
+                        source_folder=execution_config.execution.source_folder,
+                        source_bucket=execution_config.execution.source_bucket,
+                        input_file_type=INPUT_FORMATTERS_MAP.get(user_raw_input_format),
                     ),
-                    data_details=BatchJobDataDetails(
-                        source_path=BatchJobSourcePath(
-                            source_folder=execution_config.execution.source_folder,
-                            source_bucket=execution_config.execution.source_bucket,
-                            input_file_type=INPUT_FORMATTERS_MAP.get(
-                                user_raw_input_format
-                            ),
-                        ),
-                        destination_path=BatchJobDestinationPath(
-                            destination_bucket=execution_config.execution.destination_bucket,
-                            destination_folder=execution_config.execution.destination_folder,
-                            output_file_type=OUTPUT_FORMATTERS_MAP.get(
-                                user_raw_output_format
-                            ),
+                    destination_path=BatchJobDestinationPath(
+                        destination_bucket=execution_config.execution.destination_bucket,
+                        destination_folder=execution_config.execution.destination_folder,
+                        output_file_type=OUTPUT_FORMATTERS_MAP.get(
+                            user_raw_output_format
                         ),
-                        token_secret=execution_config.execution.access_token_name,
-                        secret_secret=execution_config.execution.access_secret_name,
-                        session_token=execution_config.execution.session_token,
-                        service_account_json_key_secret=execution_config.execution.service_account_key_secret_name,
                     ),
-                    execution_details=BatchJobExecutionDetails(
-                        job_timeout=execution_config.execution.job_timeout,
-                        task_timeout=execution_config.execution.file_timeout,
-                        batch_job_deployment_size=job_size,
-                        advanced_deployment_options=AdvancedDeploymentOptions(
-                            custom_iam_role_arn=execution_config.advanced_options.custom_iam_role_arn,
-                            purchase_option=user_purchase_option,
-                            service_account_key_secret_name=execution_config.advanced_options.service_account_key_secret_name,
-                        ),
-                        parameters=BatchJobManagerClient._batch_job_parameters_as_list(
-                            execution_config.execution.parameters
-                        ),
+                    token_secret=execution_config.execution.access_token_name,
+                    secret_secret=execution_config.execution.access_secret_name,
+                    session_token=execution_config.execution.session_token,
+                    service_account_json_key_secret=execution_config.execution.service_account_key_secret_name,
+                ),
+                execution_details=BatchJobExecutionDetails(
+                    job_timeout=execution_config.execution.job_timeout,
+                    task_timeout=execution_config.execution.file_timeout,
+                    batch_job_deployment_size=job_size,
+                    advanced_deployment_options=AdvancedDeploymentOptions(
+                        custom_iam_role_arn=execution_config.advanced_options.custom_iam_role_arn,
+                        purchase_option=user_purchase_option,
+                        service_account_key_secret_name=execution_config.advanced_options.service_account_key_secret_name,
                     ),
-                )
+                    parameters=BatchJobManagerClient._batch_job_parameters_as_list(
+                        execution_config.execution.parameters
+                    ),
+                ),
             )
+        )
 
-            return StartExecutionResult(
-                success=start_job_result.success,
-                execution_id=start_job_result.batch_id,
-                failure_message=start_job_result.failure_message,
-            )
-        except grpc.RpcError as e:
-            raise FrogmlException(f"Failed to start execution, error is {e}")
+        return StartExecutionResult(
+            success=start_job_result.success,
+            execution_id=start_job_result.batch_id,
+            failure_message=start_job_result.failure_message,
+        )
 
     @staticmethod
     def get_batch_deployment_size_from_resources(execution_config):
@@ -207,6 +207,9 @@ class BatchJobManagerClient:
             ),
         )
 
+    @grpc_try_catch_wrapper(
+        error_message="Failed to start execution", operation="Start Batch Execution"
+    )
     def _start_batch_job(
         self,
         batch_job_request: BatchJobRequest,
@@ -244,28 +247,30 @@ class BatchJobManagerClient:
         """
         job_size = self.get_batch_deployment_size_from_resources(execution_config)
 
-        try:
-            start_warmup_job_result: StartWarmupJobRequest = (
-                self.batch_job_management.StartWarmupJob(
-                    StartWarmupJobRequest(
-                        model_id=execution_config.execution.model_id,
-                        build_id=execution_config.execution.build_id,
-                        branch_id=ModelsManagementClient().get_model_uuid(
-                            execution_config.execution.model_id,
-                        ),
-                        warmup_timeout=execution_config.warmup.timeout,
-                        batch_job_deployment_size=job_size,
-                    ),
-                    timeout=CLIENT_TIMEOUT,
-                )
+        start_warmup_job_result: StartWarmupJobResponse = self._start_warmup_job(
+            StartWarmupJobRequest(
+                model_id=execution_config.execution.model_id,
+                build_id=execution_config.execution.build_id,
+                branch_id=ModelsManagementClient().get_model_uuid(
+                    execution_config.execution.model_id,
+                ),
+                warmup_timeout=execution_config.warmup.timeout,
+                batch_job_deployment_size=job_size,
             )
+        )
 
-            return StartWarmupJobResponse(
-                success=start_warmup_job_result.success,
-                failure_message=start_warmup_job_result.failure_message,
-            )
-        except grpc.RpcError as e:
-            raise FrogmlException(f"Failed to start warmup, error is {e}")
+        return StartWarmupJobResponse(
+            success=start_warmup_job_result.success,
+            failure_message=start_warmup_job_result.failure_message,
+        )
+
+    @grpc_try_catch_wrapper(
+        error_message="Failed to start warmup", operation="Start Batch Warmup Job"
+    )
+    def _start_warmup_job(
+        self, request: StartWarmupJobRequest
+    ) -> StartWarmupJobResponse:
+        return self.batch_job_management.StartWarmupJob(request, timeout=CLIENT_TIMEOUT)
 
     def get_execution_status(self, execution_id: str) -> ExecutionStatusResult:
         """
@@ -277,25 +282,26 @@ class BatchJobManagerClient:
             the status of the execution
 
         """
-        try:
-            batch_job_status: GetBatchJobStatusResponse = (
-                self.batch_job_management.GetBatchJobStatus(
-                    GetBatchJobStatusRequest(
-                        batch_id=execution_id,
-                    )
-                )
-            )
-            return ExecutionStatusResult(
-                success=batch_job_status.success,
-                status=BatchJobStatusMessage.Name(batch_job_status.job_status),
-                finished_files=batch_job_status.finished_files,
-                total_files=batch_job_status.total_files,
-                failure_message=batch_job_status.failure_message,
-            )
-        except grpc.RpcError as e:
-            raise FrogmlException(
-                f"Failed to get execution status, error is {e.details()}"
-            )
+        batch_job_status: GetBatchJobStatusResponse = self._get_batch_job_status(
+            execution_id
+        )
+
+        return ExecutionStatusResult(
+            success=batch_job_status.success,
+            status=BatchJobStatusMessage.Name(batch_job_status.job_status),
+            finished_files=batch_job_status.finished_files,
+            total_files=batch_job_status.total_files,
+            failure_message=batch_job_status.failure_message,
+        )
+
+    @grpc_try_catch_wrapper(
+        error_message="Failed to get execution status for execution '{execution_id}'",
+        operation="Get Batch Execution Status",
+    )
+    def _get_batch_job_status(self, execution_id: str) -> GetBatchJobStatusResponse:
+        return self.batch_job_management.GetBatchJobStatus(
+            GetBatchJobStatusRequest(batch_id=execution_id)
+        )
 
     def cancel_warmup(
         self, execution_config: ExecutionConfig
@@ -309,27 +315,30 @@ class BatchJobManagerClient:
         Returns:
             The response received from the api. On successful start of execution a batch job id is returned
         """
-        try:
-            cancel_warmup_job_response: CancelWarmupJobResponse = (
-                self.batch_job_management.CancelWarmupJob(
-                    CancelWarmupJobRequest(
-                        model_id=execution_config.execution.model_id,
-                        build_id=execution_config.execution.build_id,
-                        branch_id=ModelsManagementClient().get_model_uuid(
-                            execution_config.execution.model_id,
-                        ),
-                    ),
-                    timeout=CLIENT_TIMEOUT,
-                )
+        cancel_warmup_job_response: CancelWarmupJobResponse = self._cancel_warmup_job(
+            CancelWarmupJobRequest(
+                model_id=execution_config.execution.model_id,
+                build_id=execution_config.execution.build_id,
+                branch_id=ModelsManagementClient().get_model_uuid(
+                    execution_config.execution.model_id,
+                ),
             )
+        )
 
-            return CancelWarmupJobResponse(
-                success=cancel_warmup_job_response.success,
-                failure_message=cancel_warmup_job_response.failure_message,
-            )
+        return CancelWarmupJobResponse(
+            success=cancel_warmup_job_response.success,
+            failure_message=cancel_warmup_job_response.failure_message,
+        )
 
-        except grpc.RpcError as e:
-            raise FrogmlException(f"Failed to cancel execution, error is {e}")
+    @grpc_try_catch_wrapper(
+        error_message="Failed to cancel warmup", operation="Cancel Batch Warmup"
+    )
+    def _cancel_warmup_job(
+        self, request: CancelWarmupJobRequest
+    ) -> CancelWarmupJobResponse:
+        return self.batch_job_management.CancelWarmupJob(
+            request, timeout=CLIENT_TIMEOUT
+        )
 
     def cancel_execution(self, execution_id: str) -> CancelExecutionResult:
         """
@@ -341,23 +350,23 @@ class BatchJobManagerClient:
             A successful response or failure of the cancel process
 
         """
-        try:
-            cancel_batch_job_response: CancelBatchJobResponse = (
-                self.batch_job_management.CancelBatchJob(
-                    CancelBatchJobRequest(
-                        batch_id=execution_id,
-                    ),
-                    timeout=CLIENT_TIMEOUT,
-                )
-            )
+        cancel_batch_job_response: CancelBatchJobResponse = self._cancel_batch_job(
+            execution_id
+        )
 
-            return CancelExecutionResult(
-                success=cancel_batch_job_response.success,
-                failure_message=cancel_batch_job_response.failure_message,
-            )
+        return CancelExecutionResult(
+            success=cancel_batch_job_response.success,
+            failure_message=cancel_batch_job_response.failure_message,
+        )
 
-        except grpc.RpcError as e:
-            raise FrogmlException(f"Failed to cancel execution, error is {e}")
+    @grpc_try_catch_wrapper(
+        error_message="Failed to cancel execution '{execution_id}'",
+        operation="Cancel Batch Execution",
+    )
+    def _cancel_batch_job(self, execution_id: str) -> CancelBatchJobResponse:
+        return self.batch_job_management.CancelBatchJob(
+            CancelBatchJobRequest(batch_id=execution_id), timeout=CLIENT_TIMEOUT
+        )
 
     def get_execution_report(
         self,
@@ -375,42 +384,45 @@ class BatchJobManagerClient:
         Returns:
             A full report of all the events that occurred as part of the execution job.
         """
-        try:
-            self.logging_client = (
-                self.logging_client if self.logging_client else LoggingClient()
-            )
-            batch_job_report: GetBatchJobReportResponse = (
-                self.batch_job_management.GetBatchJobReport(
-                    GetBatchJobReportRequest(batch_id=execution_id),
-                    timeout=CLIENT_TIMEOUT,
-                )
-            )
+        self.logging_client = (
+            self.logging_client if self.logging_client else LoggingClient()
+        )
 
-            try:
-                execution_log_response: ReadLogsResponse = (
-                    self.logging_client.read_execution_models_logs(
-                        execution_id=execution_id,
-                        model_id=model_id,
-                        model_group_name=model_group_name,
-                    )
+        batch_job_report: GetBatchJobReportResponse = self._get_batch_job_report(
+            execution_id
+        )
+
+        try:
+            execution_log_response: ReadLogsResponse = (
+                self.logging_client.read_execution_models_logs(
+                    execution_id=execution_id,
+                    model_id=model_id,
+                    model_group_name=model_group_name,
                 )
-                execution_logs = execution_log_response.log_line
-                execution_logs.sort(key=lambda line: line.ingested_iso_timestamp)
-                execution_logs_msgs = [line.text for line in execution_logs]
-            except FrogmlException as e:
-                execution_logs_msgs = [
-                    f"Error reading the execution model run logs due to: {e.message}"
-                ]
-
-            return GetExecutionReportResult(
-                success=batch_job_report.successful,
-                failure_message=batch_job_report.failure_message,
-                records=batch_job_report.report_messages,
-                model_logs=execution_logs_msgs,
             )
+            execution_logs = execution_log_response.log_line
+            execution_logs.sort(key=lambda line: line.ingested_iso_timestamp)
+            execution_logs_msgs = [line.text for line in execution_logs]
+        except FrogmlException as e:
+            execution_logs_msgs = [
+                f"Error reading the execution model run logs due to: {e.error_message}"
+            ]
+
+        return GetExecutionReportResult(
+            success=batch_job_report.successful,
+            failure_message=batch_job_report.failure_message,
+            records=batch_job_report.report_messages,
+            model_logs=execution_logs_msgs,
+        )
 
-        except grpc.RpcError as e:
-            raise FrogmlException(f"Failed to get report for execution, error is: {e}")
+    @grpc_try_catch_wrapper(
+        error_message="Failed to get report for execution '{execution_id}'",
+        operation="Get Batch Execution Report",
+    )
+    def _get_batch_job_report(self, execution_id: str) -> GetBatchJobReportResponse:
+        return self.batch_job_management.GetBatchJobReport(
+            GetBatchJobReportRequest(batch_id=execution_id), timeout=CLIENT_TIMEOUT
+        )
 
     def get_pre_signed_upload_urls_details(
         self, model_id: str, number_of_batches: int, file_type: str = "csv"
@@ -424,31 +436,35 @@ class BatchJobManagerClient:
 
         Returns: GetBatchJobPreSignedUploadUrlResult which contains the input/output path, the bucket, and the pre-signed urls
         """
-        try:
-            response: GetBatchJobPreSignedUploadUrlResponse = (
-                self.batch_job_management.GetBatchJobPreSignedUploadUrl(
-                    GetBatchJobPreSignedUploadUrlRequest(
-                        model_id=model_id,
-                        number_of_files=number_of_batches,
-                        file_type=file_type,
-                    ),
-                    timeout=CLIENT_TIMEOUT,
+        response: GetBatchJobPreSignedUploadUrlResponse = (
+            self._get_pre_signed_upload_url(
+                GetBatchJobPreSignedUploadUrlRequest(
+                    model_id=model_id,
+                    number_of_files=number_of_batches,
+                    file_type=file_type,
                 )
             )
+        )
 
-            return GetBatchJobPreSignedUploadUrlResult(
-                success=response.success,
-                failure_message=response.failure_message,
-                input_path=response.input_path,
-                output_path=response.output_path,
-                bucket=response.bucket,
-                urls=response.urls,
-            )
+        return GetBatchJobPreSignedUploadUrlResult(
+            success=response.success,
+            failure_message=response.failure_message,
+            input_path=response.input_path,
+            output_path=response.output_path,
+            bucket=response.bucket,
+            urls=response.urls,
+        )
 
-        except grpc.RpcError as e:
-            raise FrogmlException(
-                f"Failed to get pre signed urls for execution, error is: {e}"
-            )
+    @grpc_try_catch_wrapper(
+        error_message="Failed to get pre signed urls for execution",
+        operation="Get Batch Pre-Signed Upload URLs",
+    )
+    def _get_pre_signed_upload_url(
+        self, request: GetBatchJobPreSignedUploadUrlRequest
+    ) -> GetBatchJobPreSignedUploadUrlResponse:
+        return self.batch_job_management.GetBatchJobPreSignedUploadUrl(
+            request, timeout=CLIENT_TIMEOUT
+        )
 
     def get_pre_signed_download_urls_details(
         self, execution_id: str
@@ -461,27 +477,32 @@ class BatchJobManagerClient:
         Returns: GetBatchJobPreSignedDownloadUrlResult which contains the pre-signed urls of the output files
 
         """
-        try:
-            response: GetBatchJobPreSignedDownloadUrlResponse = (
-                self.batch_job_management.GetBatchJobPreSignedDownloadUrl(
-                    GetBatchJobPreSignedDownloadUrlRequest(
-                        job_id=execution_id,
-                    ),
-                    timeout=CLIENT_TIMEOUT,
-                )
-            )
+        response: GetBatchJobPreSignedDownloadUrlResponse = (
+            self._get_pre_signed_download_url(execution_id)
+        )
 
-            return GetBatchJobPreSignedDownloadUrlResult(
-                success=response.success,
-                failure_message=response.failure_message,
-                urls=response.urls,
-            )
+        return GetBatchJobPreSignedDownloadUrlResult(
+            success=response.success,
+            failure_message=response.failure_message,
+            urls=response.urls,
+        )
 
-        except grpc.RpcError as e:
-            raise FrogmlException(
-                f"Failed to get pre signed urls for execution, error is: {e}"
-            )
+    @grpc_try_catch_wrapper(
+        error_message="Failed to get pre signed download urls for execution '{execution_id}'",
+        operation="Get Batch Pre-Signed Download URLs",
+    )
+    def _get_pre_signed_download_url(
+        self, execution_id: str
+    ) -> GetBatchJobPreSignedDownloadUrlResponse:
+        return self.batch_job_management.GetBatchJobPreSignedDownloadUrl(
+            GetBatchJobPreSignedDownloadUrlRequest(job_id=execution_id),
+            timeout=CLIENT_TIMEOUT,
+        )
 
+    @grpc_try_catch_wrapper(
+        error_message="Failed to get upload details",
+        operation="Get Batch Upload Details",
+    )
     def get_upload_details(self, model_id: str) -> GetBatchJobUploadDetailsResponse:
         """
         Get upload details in order to start a job using FrogML cloud bucket
@@ -490,17 +511,17 @@ class BatchJobManagerClient:
 
         Returns: GetBatchJobPreSignedUploadUrlResult which contains the input/output path, the bucket, and temporary credentials
         """
-        try:
-            return self.batch_job_management.GetBatchJobUploadDetails(
-                GetBatchJobUploadDetailsRequest(
-                    model_id=model_id,
-                ),
-                timeout=CLIENT_TIMEOUT,
-            )
-
-        except grpc.RpcError as e:
-            raise FrogmlException(f"Failed to get upload details, error is: {e}")
+        return self.batch_job_management.GetBatchJobUploadDetails(
+            GetBatchJobUploadDetailsRequest(
+                model_id=model_id,
+            ),
+            timeout=CLIENT_TIMEOUT,
+        )
 
+    @grpc_try_catch_wrapper(
+        error_message="Failed to get download details",
+        operation="Get Batch Download Details",
+    )
     def get_download_details(
         self, execution_id: str
     ) -> GetBatchJobDownloadDetailsResponse:
@@ -511,17 +532,16 @@ class BatchJobManagerClient:
 
         Returns: GetBatchJobDownloadDetailsResponse which contains the keys the bucket, and temporary credentials
         """
-        try:
-            return self.batch_job_management.GetBatchJobDownloadDetails(
-                GetBatchJobDownloadDetailsRequest(
-                    job_id=execution_id,
-                ),
-                timeout=CLIENT_TIMEOUT,
-            )
-
-        except grpc.RpcError as e:
-            raise FrogmlException(f"Failed to get download details, error is: {e}")
+        return self.batch_job_management.GetBatchJobDownloadDetails(
+            GetBatchJobDownloadDetailsRequest(
+                job_id=execution_id,
+            ),
+            timeout=CLIENT_TIMEOUT,
+        )
 
+    @grpc_try_catch_wrapper(
+        error_message="Failed to list batch jobs", operation="List Batch Jobs"
+    )
     def list_batch_jobs(self, model_id: str, build_id: str) -> ListBatchJobsResponse:
         """
         List batch jobs by its model ID
@@ -531,15 +551,15 @@ class BatchJobManagerClient:
 
         Returns: ListBatchJobsResponse which contains list of batch jobs details
         """
-        try:
-            return self.batch_job_management.ListBatchJobs(
-                ListBatchJobsRequest(model_id=model_id, build_id=build_id),
-                timeout=CLIENT_TIMEOUT,
-            )
-
-        except grpc.RpcError as e:
-            raise FrogmlException(f"Failed to list batch jobs, error is: {e}")
+        return self.batch_job_management.ListBatchJobs(
+            ListBatchJobsRequest(model_id=model_id, build_id=build_id),
+            timeout=CLIENT_TIMEOUT,
+        )
 
+    @grpc_try_catch_wrapper(
+        error_message="Failed to get batch job details",
+        operation="Get Batch Job Details",
+    )
     def get_batch_job_details(self, job_id: str) -> GetBatchJobDetailsResponse:
         """
         Get batch jos by its job ID
@@ -548,16 +568,16 @@ class BatchJobManagerClient:
 
         Returns: GetBatchJobDetailsResponse which contains list of task execution details
         """
-        try:
-            return self.batch_job_management.GetBatchJobDetails(
-                GetBatchJobDetailsRequest(job_id=job_id),
-                timeout=CLIENT_TIMEOUT,
-            )
-
-        except grpc.RpcError as e:
-            raise FrogmlException(f"Failed to get batch job details, error is: {e}")
+        return self.batch_job_management.GetBatchJobDetails(
+            GetBatchJobDetailsRequest(job_id=job_id),
+            timeout=CLIENT_TIMEOUT,
+        )
 
-    def update_task_details(self, task_id: str, input_file_paths: List[str]):
+    @grpc_try_catch_wrapper(
+        error_message="Failed to update task details for task '{task_id}'",
+        operation="Update Batch Task Details",
+    )
+    def update_task_details(self, task_id: str, input_file_paths: list[str]):
         """
         Update task details for a specific task using the provided list of input file paths.
 
@@ -579,14 +599,6 @@ class BatchJobManagerClient:
             tasks_details=[batch_task_details]  # Wrap in a list
         )
 
-        try:
-            response: UpdateTasksDetailsResponse = (
-                self.batch_job_management.UpdateTasksDetails(
-                    update_request, timeout=CLIENT_TIMEOUT
-                )
-            )
-            return response
-        except grpc.RpcError as e:
-            raise FrogmlException(
-                f"Failed to update task details for task '{task_id}', error is: {e}"
-            )
+        return self.batch_job_management.UpdateTasksDetails(
+            update_request, timeout=CLIENT_TIMEOUT
+        )