qwak-core 0.4.273__py3-none-any.whl → 0.4.275__py3-none-any.whl

_qwak_proto/qwak/model_group/model_group_pb2.py

@@ -16,12 +16,14 @@ from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__
 from _qwak_proto.qwak.model_group import model_group_repository_details_pb2 as qwak_dot_model__group_dot_model__group__repository__details__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\"qwak/model_group/model_group.proto\x12\x1cqwak.model_groups.management\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x35qwak/model_group/model_group_repository_details.proto\"Y\n\"CreateIfNotExistsModelGroupRequest\x12\x18\n\x10model_group_name\x18\x01 \x01(\t\x12\x19\n\x11jfrog_project_key\x18\x02 \x01(\t\"\xdf\x01\n\x1bModelGroupBriefInfoResponse\x12\x16\n\x0emodel_group_id\x18\x01 \x01(\t\x12\x1f\n\x17model_group_description\x18\x02 \x01(\t\x12:\n\x16model_group_created_at\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12K\n\x12repository_details\x18\x04 \x03(\x0b\x32/.qwak.model_groups.management.RepositoryDetails2\xba\x01\n\x1bModelGroupManagementService\x12\x9a\x01\n\x1b\x43reateIfNotExistsModelGroup\x12@.qwak.model_groups.management.CreateIfNotExistsModelGroupRequest\x1a\x39.qwak.model_groups.management.ModelGroupBriefInfoResponseBL\n&com.qwak.ai.management.model.group.apiB ModelGroupManagementServiceProtoP\x01\x62\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\"qwak/model_group/model_group.proto\x12\x1cqwak.model_groups.management\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x35qwak/model_group/model_group_repository_details.proto\"Y\n\"CreateIfNotExistsModelGroupRequest\x12\x18\n\x10model_group_name\x18\x01 \x01(\t\x12\x19\n\x11jfrog_project_key\x18\x02 \x01(\t\"\xdf\x01\n\x1bModelGroupBriefInfoResponse\x12\x16\n\x0emodel_group_id\x18\x01 \x01(\t\x12\x1f\n\x17model_group_description\x18\x02 \x01(\t\x12:\n\x16model_group_created_at\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12K\n\x12repository_details\x18\x04 \x03(\x0b\x32/.qwak.model_groups.management.RepositoryDetails\"\x95\x01\n&CreateRepositoriesForModelGroupRequest\x12\x18\n\x10model_group_uuid\x18\x01 \x01(\t\x12Q\n\x12repository_details\x18\x02 \x03(\x0b\x32\x35.qwak.model_groups.management.RemoteRepositoryDetails\"\x82\x01\n\'CreateRepositoriesForModelGroupResponse\x12W\n\x1e\x61rtifactory_repository_details\x18\x01 \x03(\x0b\x32/.qwak.model_groups.management.RepositoryDetails2\xeb\x02\n\x1bModelGroupManagementService\x12\x9a\x01\n\x1b\x43reateIfNotExistsModelGroup\x12@.qwak.model_groups.management.CreateIfNotExistsModelGroupRequest\x1a\x39.qwak.model_groups.management.ModelGroupBriefInfoResponse\x12\xae\x01\n\x1f\x43reateRepositoriesForModelGroup\x12\x44.qwak.model_groups.management.CreateRepositoriesForModelGroupRequest\x1a\x45.qwak.model_groups.management.CreateRepositoriesForModelGroupResponseBL\n&com.qwak.ai.management.model.group.apiB ModelGroupManagementServiceProtoP\x01\x62\x06proto3')
 
 
 
 _CREATEIFNOTEXISTSMODELGROUPREQUEST = DESCRIPTOR.message_types_by_name['CreateIfNotExistsModelGroupRequest']
 _MODELGROUPBRIEFINFORESPONSE = DESCRIPTOR.message_types_by_name['ModelGroupBriefInfoResponse']
+_CREATEREPOSITORIESFORMODELGROUPREQUEST = DESCRIPTOR.message_types_by_name['CreateRepositoriesForModelGroupRequest']
+_CREATEREPOSITORIESFORMODELGROUPRESPONSE = DESCRIPTOR.message_types_by_name['CreateRepositoriesForModelGroupResponse']
 CreateIfNotExistsModelGroupRequest = _reflection.GeneratedProtocolMessageType('CreateIfNotExistsModelGroupRequest', (_message.Message,), {
   'DESCRIPTOR' : _CREATEIFNOTEXISTSMODELGROUPREQUEST,
   '__module__' : 'qwak.model_group.model_group_pb2'
@@ -36,6 +38,20 @@ ModelGroupBriefInfoResponse = _reflection.GeneratedProtocolMessageType('ModelGro
   })
 _sym_db.RegisterMessage(ModelGroupBriefInfoResponse)
 
+CreateRepositoriesForModelGroupRequest = _reflection.GeneratedProtocolMessageType('CreateRepositoriesForModelGroupRequest', (_message.Message,), {
+  'DESCRIPTOR' : _CREATEREPOSITORIESFORMODELGROUPREQUEST,
+  '__module__' : 'qwak.model_group.model_group_pb2'
+  # @@protoc_insertion_point(class_scope:qwak.model_groups.management.CreateRepositoriesForModelGroupRequest)
+  })
+_sym_db.RegisterMessage(CreateRepositoriesForModelGroupRequest)
+
+CreateRepositoriesForModelGroupResponse = _reflection.GeneratedProtocolMessageType('CreateRepositoriesForModelGroupResponse', (_message.Message,), {
+  'DESCRIPTOR' : _CREATEREPOSITORIESFORMODELGROUPRESPONSE,
+  '__module__' : 'qwak.model_group.model_group_pb2'
+  # @@protoc_insertion_point(class_scope:qwak.model_groups.management.CreateRepositoriesForModelGroupResponse)
+  })
+_sym_db.RegisterMessage(CreateRepositoriesForModelGroupResponse)
+
 _MODELGROUPMANAGEMENTSERVICE = DESCRIPTOR.services_by_name['ModelGroupManagementService']
 if _descriptor._USE_C_DESCRIPTORS == False:
 
@@ -45,6 +61,10 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _CREATEIFNOTEXISTSMODELGROUPREQUEST._serialized_end=245
   _MODELGROUPBRIEFINFORESPONSE._serialized_start=248
   _MODELGROUPBRIEFINFORESPONSE._serialized_end=471
-  _MODELGROUPMANAGEMENTSERVICE._serialized_start=474
-  _MODELGROUPMANAGEMENTSERVICE._serialized_end=660
+  _CREATEREPOSITORIESFORMODELGROUPREQUEST._serialized_start=474
+  _CREATEREPOSITORIESFORMODELGROUPREQUEST._serialized_end=623
+  _CREATEREPOSITORIESFORMODELGROUPRESPONSE._serialized_start=626
+  _CREATEREPOSITORIESFORMODELGROUPRESPONSE._serialized_end=756
+  _MODELGROUPMANAGEMENTSERVICE._serialized_start=759
+  _MODELGROUPMANAGEMENTSERVICE._serialized_end=1122
 # @@protoc_insertion_point(module_scope)

_qwak_proto/qwak/model_group/model_group_pb2.pyi

@@ -66,3 +66,39 @@ class ModelGroupBriefInfoResponse(google.protobuf.message.Message):
     def ClearField(self, field_name: typing_extensions.Literal["model_group_created_at", b"model_group_created_at", "model_group_description", b"model_group_description", "model_group_id", b"model_group_id", "repository_details", b"repository_details"]) -> None: ...
 
 global___ModelGroupBriefInfoResponse = ModelGroupBriefInfoResponse
+
+class CreateRepositoriesForModelGroupRequest(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    MODEL_GROUP_UUID_FIELD_NUMBER: builtins.int
+    REPOSITORY_DETAILS_FIELD_NUMBER: builtins.int
+    model_group_uuid: builtins.str
+    """the model group UUID for which to create repositories"""
+    @property
+    def repository_details(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[qwak.model_group.model_group_repository_details_pb2.RemoteRepositoryDetails]:
+        """List of remote repository details to create for the model group"""
+    def __init__(
+        self,
+        *,
+        model_group_uuid: builtins.str = ...,
+        repository_details: collections.abc.Iterable[qwak.model_group.model_group_repository_details_pb2.RemoteRepositoryDetails] | None = ...,
+    ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["model_group_uuid", b"model_group_uuid", "repository_details", b"repository_details"]) -> None: ...
+
+global___CreateRepositoriesForModelGroupRequest = CreateRepositoriesForModelGroupRequest
+
+class CreateRepositoriesForModelGroupResponse(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    ARTIFACTORY_REPOSITORY_DETAILS_FIELD_NUMBER: builtins.int
+    @property
+    def artifactory_repository_details(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[qwak.model_group.model_group_repository_details_pb2.RepositoryDetails]:
+        """List of repository details associated with this model group"""
+    def __init__(
+        self,
+        *,
+        artifactory_repository_details: collections.abc.Iterable[qwak.model_group.model_group_repository_details_pb2.RepositoryDetails] | None = ...,
+    ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["artifactory_repository_details", b"artifactory_repository_details"]) -> None: ...
+
+global___CreateRepositoriesForModelGroupResponse = CreateRepositoriesForModelGroupResponse

_qwak_proto/qwak/model_group/model_group_pb2_grpc.py

@@ -19,6 +19,11 @@ class ModelGroupManagementServiceStub(object):
                 request_serializer=qwak_dot_model__group_dot_model__group__pb2.CreateIfNotExistsModelGroupRequest.SerializeToString,
                 response_deserializer=qwak_dot_model__group_dot_model__group__pb2.ModelGroupBriefInfoResponse.FromString,
                 )
+        self.CreateRepositoriesForModelGroup = channel.unary_unary(
+                '/qwak.model_groups.management.ModelGroupManagementService/CreateRepositoriesForModelGroup',
+                request_serializer=qwak_dot_model__group_dot_model__group__pb2.CreateRepositoriesForModelGroupRequest.SerializeToString,
+                response_deserializer=qwak_dot_model__group_dot_model__group__pb2.CreateRepositoriesForModelGroupResponse.FromString,
+                )
 
 
 class ModelGroupManagementServiceServicer(object):
@@ -30,6 +35,12 @@ class ModelGroupManagementServiceServicer(object):
         context.set_details('Method not implemented!')
         raise NotImplementedError('Method not implemented!')
 
+    def CreateRepositoriesForModelGroup(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
 
 def add_ModelGroupManagementServiceServicer_to_server(servicer, server):
     rpc_method_handlers = {
@@ -38,6 +49,11 @@ def add_ModelGroupManagementServiceServicer_to_server(servicer, server):
                     request_deserializer=qwak_dot_model__group_dot_model__group__pb2.CreateIfNotExistsModelGroupRequest.FromString,
                     response_serializer=qwak_dot_model__group_dot_model__group__pb2.ModelGroupBriefInfoResponse.SerializeToString,
             ),
+            'CreateRepositoriesForModelGroup': grpc.unary_unary_rpc_method_handler(
+                    servicer.CreateRepositoriesForModelGroup,
+                    request_deserializer=qwak_dot_model__group_dot_model__group__pb2.CreateRepositoriesForModelGroupRequest.FromString,
+                    response_serializer=qwak_dot_model__group_dot_model__group__pb2.CreateRepositoriesForModelGroupResponse.SerializeToString,
+            ),
     }
     generic_handler = grpc.method_handlers_generic_handler(
             'qwak.model_groups.management.ModelGroupManagementService', rpc_method_handlers)
@@ -64,3 +80,20 @@ class ModelGroupManagementService(object):
             qwak_dot_model__group_dot_model__group__pb2.ModelGroupBriefInfoResponse.FromString,
             options, channel_credentials,
             insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def CreateRepositoriesForModelGroup(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/qwak.model_groups.management.ModelGroupManagementService/CreateRepositoriesForModelGroup',
+            qwak_dot_model__group_dot_model__group__pb2.CreateRepositoriesForModelGroupRequest.SerializeToString,
+            qwak_dot_model__group_dot_model__group__pb2.CreateRepositoriesForModelGroupResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)

_qwak_proto/qwak/model_group/model_group_repository_details_pb2.py

@@ -12,14 +12,14 @@ from google.protobuf import symbol_database as _symbol_database
 _sym_db = _symbol_database.Default()
 
 
-from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n5qwak/model_group/model_group_repository_details.proto\x12\x1cqwak.model_groups.management\x1a\x1fgoogle/protobuf/timestamp.proto\"r\n\x11RepositoryDetails\x12\x45\n\x0frepository_type\x18\x01 \x01(\x0b\x32,.qwak.model_groups.management.RepositoryType\x12\x16\n\x0erepository_key\x18\x02 \x01(\t\"\x12\n\x10\x44ockerRepository\"\x13\n\x11\x44\x61tasetRepository\"\x14\n\x12\x41rtifactRepository\"\x17\n\x15HuggingFaceRepository\"\xdd\x02\n\x0eRepositoryType\x12K\n\x11\x64ocker_repository\x18\x01 \x01(\x0b\x32..qwak.model_groups.management.DockerRepositoryH\x00\x12M\n\x12\x64\x61taset_repository\x18\x02 \x01(\x0b\x32/.qwak.model_groups.management.DatasetRepositoryH\x00\x12O\n\x13\x61rtifact_repository\x18\x03 \x01(\x0b\x32\x30.qwak.model_groups.management.ArtifactRepositoryH\x00\x12V\n\x17hugging_face_repository\x18\x04 \x01(\x0b\x32\x33.qwak.model_groups.management.HuggingFaceRepositoryH\x00\x42\x06\n\x04typeBN\n&com.qwak.ai.management.model.group.apiB\"ModelGroupRepositoriesDetailsProtoP\x01\x62\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n5qwak/model_group/model_group_repository_details.proto\x12\x1cqwak.model_groups.management\"r\n\x11RepositoryDetails\x12\x45\n\x0frepository_type\x18\x01 \x01(\x0b\x32,.qwak.model_groups.management.RepositoryType\x12\x16\n\x0erepository_key\x18\x02 \x01(\t\"\xf0\x01\n\x17RemoteRepositoryDetails\x12K\n\x11\x64ocker_repository\x18\x01 \x01(\x0b\x32..qwak.model_groups.management.DockerRepositoryH\x00\x12V\n\x17hugging_face_repository\x18\x02 \x01(\x0b\x32\x33.qwak.model_groups.management.HuggingFaceRepositoryH\x00\x12\x1d\n\x15repository_remote_url\x18\x03 \x01(\tB\x11\n\x0frepository_type\"\x12\n\x10\x44ockerRepository\"\x13\n\x11\x44\x61tasetRepository\"\x14\n\x12\x41rtifactRepository\"\x17\n\x15HuggingFaceRepository\"\xdd\x02\n\x0eRepositoryType\x12K\n\x11\x64ocker_repository\x18\x01 \x01(\x0b\x32..qwak.model_groups.management.DockerRepositoryH\x00\x12M\n\x12\x64\x61taset_repository\x18\x02 \x01(\x0b\x32/.qwak.model_groups.management.DatasetRepositoryH\x00\x12O\n\x13\x61rtifact_repository\x18\x03 \x01(\x0b\x32\x30.qwak.model_groups.management.ArtifactRepositoryH\x00\x12V\n\x17hugging_face_repository\x18\x04 \x01(\x0b\x32\x33.qwak.model_groups.management.HuggingFaceRepositoryH\x00\x42\x06\n\x04typeBN\n&com.qwak.ai.management.model.group.apiB\"ModelGroupRepositoriesDetailsProtoP\x01\x62\x06proto3')
 
 
 
 _REPOSITORYDETAILS = DESCRIPTOR.message_types_by_name['RepositoryDetails']
+_REMOTEREPOSITORYDETAILS = DESCRIPTOR.message_types_by_name['RemoteRepositoryDetails']
 _DOCKERREPOSITORY = DESCRIPTOR.message_types_by_name['DockerRepository']
 _DATASETREPOSITORY = DESCRIPTOR.message_types_by_name['DatasetRepository']
 _ARTIFACTREPOSITORY = DESCRIPTOR.message_types_by_name['ArtifactRepository']
@@ -32,6 +32,13 @@ RepositoryDetails = _reflection.GeneratedProtocolMessageType('RepositoryDetails'
   })
 _sym_db.RegisterMessage(RepositoryDetails)
 
+RemoteRepositoryDetails = _reflection.GeneratedProtocolMessageType('RemoteRepositoryDetails', (_message.Message,), {
+  'DESCRIPTOR' : _REMOTEREPOSITORYDETAILS,
+  '__module__' : 'qwak.model_group.model_group_repository_details_pb2'
+  # @@protoc_insertion_point(class_scope:qwak.model_groups.management.RemoteRepositoryDetails)
+  })
+_sym_db.RegisterMessage(RemoteRepositoryDetails)
+
 DockerRepository = _reflection.GeneratedProtocolMessageType('DockerRepository', (_message.Message,), {
   'DESCRIPTOR' : _DOCKERREPOSITORY,
   '__module__' : 'qwak.model_group.model_group_repository_details_pb2'
@@ -71,16 +78,18 @@ if _descriptor._USE_C_DESCRIPTORS == False:
 
   DESCRIPTOR._options = None
   DESCRIPTOR._serialized_options = b'\n&com.qwak.ai.management.model.group.apiB\"ModelGroupRepositoriesDetailsProtoP\001'
-  _REPOSITORYDETAILS._serialized_start=120
-  _REPOSITORYDETAILS._serialized_end=234
-  _DOCKERREPOSITORY._serialized_start=236
-  _DOCKERREPOSITORY._serialized_end=254
-  _DATASETREPOSITORY._serialized_start=256
-  _DATASETREPOSITORY._serialized_end=275
-  _ARTIFACTREPOSITORY._serialized_start=277
-  _ARTIFACTREPOSITORY._serialized_end=297
-  _HUGGINGFACEREPOSITORY._serialized_start=299
-  _HUGGINGFACEREPOSITORY._serialized_end=322
-  _REPOSITORYTYPE._serialized_start=325
-  _REPOSITORYTYPE._serialized_end=674
+  _REPOSITORYDETAILS._serialized_start=87
+  _REPOSITORYDETAILS._serialized_end=201
+  _REMOTEREPOSITORYDETAILS._serialized_start=204
+  _REMOTEREPOSITORYDETAILS._serialized_end=444
+  _DOCKERREPOSITORY._serialized_start=446
+  _DOCKERREPOSITORY._serialized_end=464
+  _DATASETREPOSITORY._serialized_start=466
+  _DATASETREPOSITORY._serialized_end=485
+  _ARTIFACTREPOSITORY._serialized_start=487
+  _ARTIFACTREPOSITORY._serialized_end=507
+  _HUGGINGFACEREPOSITORY._serialized_start=509
+  _HUGGINGFACEREPOSITORY._serialized_end=532
+  _REPOSITORYTYPE._serialized_start=535
+  _REPOSITORYTYPE._serialized_end=884
 # @@protoc_insertion_point(module_scope)

_qwak_proto/qwak/model_group/model_group_repository_details_pb2.pyi

@@ -35,6 +35,33 @@ class RepositoryDetails(google.protobuf.message.Message):
 
 global___RepositoryDetails = RepositoryDetails
 
+class RemoteRepositoryDetails(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    DOCKER_REPOSITORY_FIELD_NUMBER: builtins.int
+    HUGGING_FACE_REPOSITORY_FIELD_NUMBER: builtins.int
+    REPOSITORY_REMOTE_URL_FIELD_NUMBER: builtins.int
+    @property
+    def docker_repository(self) -> global___DockerRepository:
+        """Docker repository"""
+    @property
+    def hugging_face_repository(self) -> global___HuggingFaceRepository:
+        """HuggingFace repository"""
+    repository_remote_url: builtins.str
+    """The remote repository URL"""
+    def __init__(
+        self,
+        *,
+        docker_repository: global___DockerRepository | None = ...,
+        hugging_face_repository: global___HuggingFaceRepository | None = ...,
+        repository_remote_url: builtins.str = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["docker_repository", b"docker_repository", "hugging_face_repository", b"hugging_face_repository", "repository_type", b"repository_type"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["docker_repository", b"docker_repository", "hugging_face_repository", b"hugging_face_repository", "repository_remote_url", b"repository_remote_url", "repository_type", b"repository_type"]) -> None: ...
+    def WhichOneof(self, oneof_group: typing_extensions.Literal["repository_type", b"repository_type"]) -> typing_extensions.Literal["docker_repository", "hugging_face_repository"] | None: ...
+
+global___RemoteRepositoryDetails = RemoteRepositoryDetails
+
 class DockerRepository(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
 

frogml_storage/__init__.py

@@ -1 +1 @@
-__version__ = "0.4.273"
+__version__ = "0.4.275"

qwak/__init__.py

@@ -1,7 +1,7 @@
 """Top-level package for qwak-core."""
 
 __author__ = "Qwak.ai"
-__version__ = "0.4.273"
+__version__ = "0.4.275"
 
 from qwak.inner.di_configuration import wire_dependencies
 from qwak.model.experiment_tracking import log_metric, log_param

qwak/clients/batch_job_management/client.py

@@ -34,6 +34,10 @@ from _qwak_proto.qwak.batch_job.v1.batch_job_service_pb2 import (
     StartBatchJobResponse,
     StartWarmupJobRequest,
     StartWarmupJobResponse,
+    UpdateTasksDetailsResponse,
+    UpdateTasksDetailsRequest,
+    BatchTaskDetails,
+    InputFileDetails,
 )
 from _qwak_proto.qwak.batch_job.v1.batch_job_service_pb2_grpc import (
     BatchJobManagementServiceStub,
@@ -67,6 +71,8 @@ from .results import (
     StartExecutionResult,
 )
 
+from typing import List
+
 CLIENT_TIMEOUT = 180  # Seconds
 
 
@@ -541,3 +547,36 @@ class BatchJobManagerClient:
 
         except grpc.RpcError as e:
             raise QwakException(f"Failed to get batch job details, error is: {e}")
+
+    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.
+
+        :param task_id: The ID of the task to update.
+        :param input_file_paths: A list of paths for the input files associated with the task.
+        """
+        # Create InputFileDetails messages for each input file path
+        input_files_details = [
+            InputFileDetails(path=file_path) for file_path in input_file_paths
+        ]
+
+        # Create BatchTaskDetails message for the single task
+        batch_task_details = BatchTaskDetails(
+            task_id=task_id, input_files_details=input_files_details
+        )
+
+        update_request = UpdateTasksDetailsRequest(
+            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 QwakException(
+                f"Failed to update task details for task '{task_id}', error is: {e}"
+            )

{qwak_core-0.4.273.dist-info → qwak_core-0.4.275.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: qwak-core
-Version: 0.4.273
+Version: 0.4.275
 Summary: Qwak Core contains the necessary objects and communication tools for using the Qwak Platform
 License: Apache-2.0
 Keywords: mlops,ml,deployment,serving,model

{qwak_core-0.4.273.dist-info → qwak_core-0.4.275.dist-info}/RECORD

@@ -487,11 +487,11 @@ _qwak_proto/qwak/logging/log_source_pb2_grpc.py,sha256=1oboBPFxaTEXt9Aw7EAj8gXHD
 _qwak_proto/qwak/model_descriptor/open_ai_descriptor_pb2.py,sha256=HCM2-Ah66CkK5Io1jCouxgDHT5HmSgAYcHKvsfTZw48,6751
 _qwak_proto/qwak/model_descriptor/open_ai_descriptor_pb2.pyi,sha256=G5CaQGczrZ9kGsMbmJXMFWAwnqMZ4jPSl9KQpSq0yMo,12345
 _qwak_proto/qwak/model_descriptor/open_ai_descriptor_pb2_grpc.py,sha256=1oboBPFxaTEXt9Aw7EAj8gXHDCNMhZD2VXqocC9l_gk,159
-_qwak_proto/qwak/model_group/model_group_pb2.py,sha256=fP198ECc5Ca3WccDvRlcaNavZgWeXKrzKyFKZXL52So,3427
-_qwak_proto/qwak/model_group/model_group_pb2.pyi,sha256=6xP_BE9OccBWV784jfKS5BcBabUcDL22XtrtWq2_hlU,3038
-_qwak_proto/qwak/model_group/model_group_pb2_grpc.py,sha256=UyRBABlYr-W3uUKUnStENcfnc0mXOtShh8_mjpzpro0,3073
-_qwak_proto/qwak/model_group/model_group_repository_details_pb2.py,sha256=t7oUXXLQvgxNuLyrA0qzuyfazRXK-uCIArXD_4s_1J4,5022
-_qwak_proto/qwak/model_group/model_group_repository_details_pb2.pyi,sha256=h5ciIby9peqjZqHosajfrHdmX5kosmzDJiadoCVpTmc,4119
+_qwak_proto/qwak/model_group/model_group_pb2.py,sha256=Gh4oDESijyWcSudhDkBKPkS-KTj46mFVy6Be9rmCL4g,5374
+_qwak_proto/qwak/model_group/model_group_pb2.pyi,sha256=VxmZR2nKY6OjnM8yhTLSIp9eho-Uk-IkZBu6py2D_t4,5073
+_qwak_proto/qwak/model_group/model_group_pb2_grpc.py,sha256=RAvxDibL9olZBfQRIoqr3lIuVGRyneTntWZS6bhx_zo,5125
+_qwak_proto/qwak/model_group/model_group_repository_details_pb2.py,sha256=Kax2nuUQKQ5d0X2A3b_hw29rls8FX1JYAWoemEWmmDc,5799
+_qwak_proto/qwak/model_group/model_group_repository_details_pb2.pyi,sha256=GgYhrSXu8L2RdEV0clooKgB9Gye3_no4q0YrMy8QrZw,5701
 _qwak_proto/qwak/model_group/model_group_repository_details_pb2_grpc.py,sha256=1oboBPFxaTEXt9Aw7EAj8gXHDCNMhZD2VXqocC9l_gk,159
 _qwak_proto/qwak/models/models_pb2.py,sha256=dypKmBnjqzAmSLkXRPx0QMXiB1D-mUNxRu5pGqxVj68,32093
 _qwak_proto/qwak/models/models_pb2.pyi,sha256=6h6PCqMOwZgSf77CUEx_SiRWX0uTxjiFMysVG_Z0Qvc,39027
@@ -601,7 +601,7 @@ _qwak_proto/qwak/workspace/workspace_pb2_grpc.py,sha256=1oboBPFxaTEXt9Aw7EAj8gXH
 _qwak_proto/qwak/workspace/workspace_service_pb2.py,sha256=AB3C9S_AbOD7Nx1Ni4j1rW6PNtYTV1zjiqFQk-goQ74,21429
 _qwak_proto/qwak/workspace/workspace_service_pb2.pyi,sha256=nKKCHwnovZhsy8TSVmdz-Vtl0nviOOoX56HD-41Xo08,13726
 _qwak_proto/qwak/workspace/workspace_service_pb2_grpc.py,sha256=yKGuexxTBza99Ihe0DSTniV2ZSd_AG47inHenqfi890,27193
-frogml_storage/__init__.py,sha256=Wq_pUPT9jRznmLWLcQuuLFhFJPxX4mt7avCmVvx1DJI,24
+frogml_storage/__init__.py,sha256=5Cbsl32LEshQ8ZQ7F7gzUh2VGMGdkb06EdCOlAiWvGc,24
 frogml_storage/artifactory/__init__.py,sha256=C02rcm7kqsZBVA6c6Gztxamj96hn8Aj6BuzYWFRmWbQ,71
 frogml_storage/artifactory/_artifactory_api.py,sha256=z8YX90bCy82BWqMNUdpbbig1qF5CarkIOfP6V96uYYg,11188
 frogml_storage/authentication/login/__init__.py,sha256=ch8UhQwh3o5ddyoJykQ0Bnog5-8qKHmeDpwqny8xMzM,46
@@ -639,7 +639,7 @@ frogml_storage/utils/_environment.py,sha256=NEnRxaFdRDi3UA33IVTSfQxqc9ZB6nV9YlHD
 frogml_storage/utils/_input_checks_utility.py,sha256=ZUvkxhcTcbSKtpSo5ePWha3Ca5xohxIpiAlZlpcQ3xk,3221
 frogml_storage/utils/_storage_utils.py,sha256=HB2g7uY5A3b33yIcAUM1OjHb5jWsnpESsiDrEviQwrI,366
 frogml_storage/utils/_url_utils.py,sha256=NUEfz9Fp1iE8b676-A5wrMlSTsJVRKrUhcUItOFAJD8,821
-qwak/__init__.py,sha256=2Xg-Qy0BfDX-xLBTGxnWF9JZ2QFnpPKvTGk6wE82nKY,587
+qwak/__init__.py,sha256=DRLZ7BRqdDrMQ_UVJKLmWIYj1PoXgd3S4anrAb26WJ8,587
 qwak/automations/__init__.py,sha256=qFZRvCxUUn8gcxkJR0v19ulHW2oJ0x6-Rif7HiheDP4,1522
 qwak/automations/automation_executions.py,sha256=5MeH_epYYWb8NKXgAozwT_jPyyUDednBHG7izloi7RY,3228
 qwak/automations/automations.py,sha256=3yx8e2v0uSKDnXbqyknasyEoQ5vxGni6K40Hbi1_zkk,12599
@@ -675,7 +675,7 @@ qwak/clients/automation_management/client.py,sha256=oO3CihRxPPPlIWN764Dmf4hmWChH
 qwak/clients/autoscaling/__init__.py,sha256=A-zuZOWaZf16NjnpX38204m730c5Lde491KaXneGirE,38
 qwak/clients/autoscaling/client.py,sha256=55Pc2yc8dLl2C2MSdO2g5E25vYgkmvJSxXnKZriP_dk,1240
 qwak/clients/batch_job_management/__init__.py,sha256=zywxYf2JupkIWnGVB2C6ugZs5CrOdxTPrksD5P6uQu4,211
-qwak/clients/batch_job_management/client.py,sha256=L0vq9n_z0pUBIvv97B3tQt4iIEZHEZTvLdLmCWMCHeo,21054
+qwak/clients/batch_job_management/client.py,sha256=4fwZtljRW93Yz6hqNEpR4YP9T7s7rYXLn2jR7hUE9B8,22453
 qwak/clients/batch_job_management/executions_config.py,sha256=LYf9j8T5W1M1EK6cjXPwAmiyGNR7wfdjNCmeKNhiREU,6882
 qwak/clients/batch_job_management/results.py,sha256=nrFGZ3y538FBhej94CYV4cwqQssK1aUQOLGubYdQVrc,1846
 qwak/clients/build_management/__init__.py,sha256=oXurHX6Kho5X-fOuuNHB3B9oMX6ejYUqh6GWv5qJyjw,43
@@ -1072,7 +1072,7 @@ qwak_services_mock/mocks/audience_service_api.py,sha256=2PvXpzKBR-Dnu9Fmb70UJn-h
 qwak_services_mock/mocks/authentication_service.py,sha256=79A8hmYUsg4i7EwyGG54CidbANlp-ub4f716BbLCaJQ,1166
 qwak_services_mock/mocks/automation_management_service.py,sha256=C9rYKMy0hzOTaKTCrTkE5jTnPSenVkLV8DCptBKr02U,8211
 qwak_services_mock/mocks/autoscaling_service_api.py,sha256=YgucRD3ZjhxqttNVSF_HL7u3aDo55LdRSEE7q0fxp3M,1019
-qwak_services_mock/mocks/batch_job_manager_service.py,sha256=HeUZulDtfIlSm2h4cLdByR0voJjOy97m4_YWldg7pqQ,13129
+qwak_services_mock/mocks/batch_job_manager_service.py,sha256=zVBBBkYdvWFwEXP_VHBH8SBIvJosPt0jHnHlJjuf6AI,14080
 qwak_services_mock/mocks/build_management.py,sha256=LYtrEvDt-lJbNJEP-pFpvAKcS82mAzb2gyoCu1-XCOs,3686
 qwak_services_mock/mocks/build_orchestrator_build_api.py,sha256=mtb6wEZug9zCiMHZVrfsU0YOP59rgBrpvZDjEInG7B8,5004
 qwak_services_mock/mocks/build_orchestrator_build_settings_api.py,sha256=oj_JTxY_3xj8FkgWkxQER_EoBFDqL2g21NPrU31fD6o,1379
@@ -1113,6 +1113,6 @@ qwak_services_mock/mocks/workspace_manager_service_mock.py,sha256=O9ZSwln4T4kHVk
 qwak_services_mock/services_mock.py,sha256=zXtHcX8a_acz7ynxuCBxxVpHpde7aAGjIn6Uw52LY1s,19593
 qwak_services_mock/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 qwak_services_mock/utils/service_utils.py,sha256=ZlB0CnB1J6oBn6_m7fQO2U8tKoboHdUa6ljjkRMYNXU,265
-qwak_core-0.4.273.dist-info/METADATA,sha256=dg-g1DVAuLRRPWYRI2dG7fUCqg4bNMsWuMKqLY83xM0,15142
-qwak_core-0.4.273.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-qwak_core-0.4.273.dist-info/RECORD,,
+qwak_core-0.4.275.dist-info/METADATA,sha256=nhP_L33Xeef7gSOW0aW9KWIAxSMn-fNRmYHezoD-S88,15142
+qwak_core-0.4.275.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+qwak_core-0.4.275.dist-info/RECORD,,

qwak_services_mock/mocks/batch_job_manager_service.py

@@ -36,6 +36,9 @@ from _qwak_proto.qwak.batch_job.v1.batch_job_service_pb2 import (
     StartWarmupJobRequest,
     StartWarmupJobResponse,
     TaskExecutionDetails,
+    UpdateTasksDetailsRequest,
+    UpdateTasksDetailsResponse,
+    BatchTaskDetails,
 )
 from _qwak_proto.qwak.batch_job.v1.batch_job_service_pb2_grpc import (
     BatchJobManagementServiceServicer,
@@ -51,6 +54,7 @@ class BatchJobManagerService(BatchJobManagementServiceServicer):
         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
@@ -265,6 +269,25 @@ class BatchJobManagerService(BatchJobManagementServiceServicer):
         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):
@@ -331,4 +354,4 @@ def batch_job_to_message(batch_job: MockBatchJob) -> BatchJobMessage:
             for line in batch_job.report
         ],
         task_executions=batch_job.tasks,
-    )
+    )