genesis-flow 1.0.1__py3-none-any.whl → 1.0.4__py3-none-any.whl

mlflow/store/artifact/azure_blob_artifact_repo.py

@@ -41,7 +41,7 @@ class AzureBlobArtifactRepository(ArtifactRepository, MultipartUploadMixin):
     - DefaultAzureCredential is configured
     """
 
-    def __init__(self, artifact_uri: str, tracking_uri: Optional[str] = None, client=None) -> None:
+    def __init__(self, artifact_uri: str, client=None, tracking_uri: Optional[str] = None) -> None:
         super().__init__(artifact_uri, tracking_uri)
 
         _DEFAULT_TIMEOUT = 600  # 10 minutes

mlflow/store/artifact/azure_data_lake_artifact_repo.py

@@ -82,9 +82,9 @@ class AzureDataLakeArtifactRepository(CloudArtifactRepository):
     def __init__(
         self,
         artifact_uri: str,
-        tracking_uri: Optional[str] = None,
         credential=None,
         credential_refresh_def=None,
+        tracking_uri: Optional[str] = None,
     ) -> None:
         super().__init__(artifact_uri, tracking_uri)
         _DEFAULT_TIMEOUT = 600  # 10 minutes

mlflow/store/artifact/gcs_artifact_repo.py

@@ -43,9 +43,9 @@ class GCSArtifactRepository(ArtifactRepository, MultipartUploadMixin):
     def __init__(
         self,
         artifact_uri: str,
-        tracking_uri: Optional[str] = None,
         client=None,
         credential_refresh_def=None,
+        tracking_uri: Optional[str] = None,
     ) -> None:
         super().__init__(artifact_uri, tracking_uri)
         from google.auth.exceptions import DefaultCredentialsError

mlflow/store/artifact/local_artifact_repo.py

@@ -14,6 +14,7 @@ from mlflow.utils.file_utils import (
     local_file_uri_to_path,
     mkdir,
     relative_path_to_artifact_path,
+    shutil_copytree_without_file_permissions,
 )
 from mlflow.utils.uri import validate_path_is_safe
 
@@ -64,7 +65,7 @@ class LocalArtifactRepository(ArtifactRepository):
         )
         if not os.path.exists(artifact_dir):
             mkdir(artifact_dir)
-        shutil.copytree(src=local_dir, dst=artifact_dir, dirs_exist_ok=True)
+        shutil_copytree_without_file_permissions(local_dir, artifact_dir)
 
     def download_artifacts(self, artifact_path, dst_path=None):
         """

mlflow/store/artifact/s3_artifact_repo.py

@@ -122,16 +122,54 @@ def _get_s3_client(
 
 
 class S3ArtifactRepository(ArtifactRepository, MultipartUploadMixin):
-    """Stores artifacts on Amazon S3."""
+    """
+    Stores artifacts on Amazon S3.
+
+    This repository provides MLflow artifact storage using Amazon S3 as the backend.
+    It supports both single-file uploads and multipart uploads for large files,
+    with automatic content type detection and configurable upload parameters.
+
+    The repository uses boto3 for S3 operations and supports various authentication
+    methods including AWS credentials, IAM roles, and environment variables.
+
+    Environment Variables:
+        AWS_ACCESS_KEY_ID: AWS access key ID for authentication
+        AWS_SECRET_ACCESS_KEY: AWS secret access key for authentication
+        AWS_SESSION_TOKEN: AWS session token for temporary credentials
+        AWS_DEFAULT_REGION: Default AWS region for S3 operations
+        MLFLOW_S3_ENDPOINT_URL: Custom S3 endpoint URL (for S3-compatible storage)
+        MLFLOW_S3_IGNORE_TLS: Set to 'true' to disable TLS verification
+        MLFLOW_S3_UPLOAD_EXTRA_ARGS: JSON string of extra arguments for S3 uploads
+        MLFLOW_BOTO_CLIENT_ADDRESSING_STYLE: S3 addressing style ('path' or 'virtual')
+
+    Note:
+        This class inherits from both ArtifactRepository and MultipartUploadMixin,
+        providing full artifact management capabilities including efficient large file uploads.
+    """
 
     def __init__(
         self,
         artifact_uri: str,
-        tracking_uri: Optional[str] = None,
         access_key_id=None,
         secret_access_key=None,
         session_token=None,
+        tracking_uri: Optional[str] = None,
     ) -> None:
+        """
+        Initialize an S3 artifact repository.
+
+        Args:
+            artifact_uri: S3 URI in the format 's3://bucket-name/path/to/artifacts'.
+                The URI must be a valid S3 URI with a bucket that exists and is accessible.
+            access_key_id: Optional AWS access key ID. If None, uses default AWS credential
+                resolution (environment variables, IAM roles, etc.).
+            secret_access_key: Optional AWS secret access key. Must be provided if
+                access_key_id is provided.
+            session_token: Optional AWS session token for temporary credentials.
+                Used with STS tokens or IAM roles.
+            tracking_uri: Optional URI for the MLflow tracking server.
+                If None, uses the current tracking URI context.
+        """
         super().__init__(artifact_uri, tracking_uri)
         self._access_key_id = access_key_id
         self._secret_access_key = secret_access_key
@@ -145,7 +183,17 @@ class S3ArtifactRepository(ArtifactRepository, MultipartUploadMixin):
         )
 
     def parse_s3_compliant_uri(self, uri):
-        """Parse an S3 URI, returning (bucket, path)"""
+        """
+        Parse an S3 URI into bucket and path components.
+
+        Args:
+            uri: S3 URI in the format 's3://bucket-name/path/to/object'
+
+        Returns:
+            A tuple containing (bucket_name, object_path) where:
+            - bucket_name: The S3 bucket name
+            - object_path: The path within the bucket (without leading slash)
+        """
         parsed = urllib.parse.urlparse(uri)
         if parsed.scheme != "s3":
             raise Exception(f"Not an S3 URI: {uri}")
@@ -156,6 +204,17 @@ class S3ArtifactRepository(ArtifactRepository, MultipartUploadMixin):
 
     @staticmethod
     def get_s3_file_upload_extra_args():
+        """
+        Get additional S3 upload arguments from environment variables.
+
+        Returns:
+            Dictionary of extra arguments for S3 uploads, or None if not configured.
+            These arguments are passed to boto3's upload_file method.
+
+        Environment Variables:
+            MLFLOW_S3_UPLOAD_EXTRA_ARGS: JSON string containing extra arguments
+                for S3 uploads (e.g., '{"ServerSideEncryption": "AES256"}')
+        """
         s3_file_upload_extra_args = MLFLOW_S3_UPLOAD_EXTRA_ARGS.get()
         if s3_file_upload_extra_args:
             return json.loads(s3_file_upload_extra_args)
@@ -175,6 +234,19 @@ class S3ArtifactRepository(ArtifactRepository, MultipartUploadMixin):
         s3_client.upload_file(Filename=local_file, Bucket=bucket, Key=key, ExtraArgs=extra_args)
 
     def log_artifact(self, local_file, artifact_path=None):
+        """
+        Log a local file as an artifact to S3.
+
+        This method uploads a single file to S3 with automatic content type detection
+        and optional extra upload arguments from environment variables.
+
+        Args:
+            local_file: Absolute path to the local file to upload. The file must
+                exist and be readable.
+            artifact_path: Optional relative path within the S3 bucket where the
+                artifact should be stored. If None, the file is stored in the root
+                of the configured S3 path. Use forward slashes (/) for path separators.
+        """
         (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
         if artifact_path:
             dest_path = posixpath.join(dest_path, artifact_path)
@@ -184,6 +256,20 @@ class S3ArtifactRepository(ArtifactRepository, MultipartUploadMixin):
         )
 
     def log_artifacts(self, local_dir, artifact_path=None):
+        """
+        Log all files in a local directory as artifacts to S3.
+
+        This method recursively uploads all files in the specified directory,
+        preserving the directory structure in S3. Each file is uploaded with
+        automatic content type detection.
+
+        Args:
+            local_dir: Absolute path to the local directory containing files to upload.
+                The directory must exist and be readable.
+            artifact_path: Optional relative path within the S3 bucket where the
+                artifacts should be stored. If None, files are stored in the root
+                of the configured S3 path. Use forward slashes (/) for path separators.
+        """
         (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
         if artifact_path:
             dest_path = posixpath.join(dest_path, artifact_path)
@@ -205,6 +291,25 @@ class S3ArtifactRepository(ArtifactRepository, MultipartUploadMixin):
                 )
 
     def list_artifacts(self, path=None):
+        """
+        List all artifacts directly under the specified S3 path.
+
+        This method uses S3's list_objects_v2 API with pagination to efficiently
+        list artifacts. It treats S3 prefixes as directories and returns both
+        files and directories as FileInfo objects.
+
+        Args:
+            path: Optional relative path within the S3 bucket to list. If None,
+                lists artifacts in the root of the configured S3 path. If the path
+                refers to a single file, returns an empty list per MLflow convention.
+
+        Returns:
+            A list of FileInfo objects representing artifacts directly under the
+            specified path. Each FileInfo contains:
+            - path: Relative path of the artifact from the repository root
+            - is_dir: True if the artifact represents a directory (S3 prefix)
+            - file_size: Size in bytes for files, None for directories
+        """
         (bucket, artifact_path) = self.parse_s3_compliant_uri(self.artifact_uri)
         dest_path = artifact_path
         if path:
@@ -247,6 +352,18 @@ class S3ArtifactRepository(ArtifactRepository, MultipartUploadMixin):
             )
 
     def _download_file(self, remote_file_path, local_path):
+        """
+        Download a file from S3 to the local filesystem.
+
+        This method downloads a single file from S3 to the specified local path.
+        It's used internally by the download_artifacts method.
+
+        Args:
+            remote_file_path: Relative path of the file within the S3 bucket,
+                relative to the repository's root path.
+            local_path: Absolute path where the file should be saved locally.
+                The parent directory must exist.
+        """
         (bucket, s3_root_path) = self.parse_s3_compliant_uri(self.artifact_uri)
         s3_full_path = posixpath.join(s3_root_path, remote_file_path)
         s3_client = self._get_s3_client()
@@ -273,6 +390,28 @@ class S3ArtifactRepository(ArtifactRepository, MultipartUploadMixin):
                 s3_client.delete_objects(Bucket=bucket, Delete={"Objects": keys})
 
     def create_multipart_upload(self, local_file, num_parts=1, artifact_path=None):
+        """
+        Initiate a multipart upload for efficient large file uploads to S3.
+
+        This method creates a multipart upload session in S3 and generates
+        presigned URLs for uploading each part. This is more efficient than
+        single-part uploads for large files and provides better error recovery.
+
+        Args:
+            local_file: Absolute path to the local file to upload. The file must
+                exist and be readable.
+            num_parts: Number of parts to split the upload into. Must be between
+                1 and 10,000 (S3 limit). More parts allow greater parallelism
+                but increase overhead.
+            artifact_path: Optional relative path within the S3 bucket where the
+                artifact should be stored. If None, the file is stored in the root
+                of the configured S3 path.
+
+        Returns:
+            CreateMultipartUploadResponse containing:
+            - credentials: List of MultipartUploadCredential objects with presigned URLs
+            - upload_id: S3 upload ID for tracking this multipart upload
+        """
         (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
         if artifact_path:
             dest_path = posixpath.join(dest_path, artifact_path)
@@ -307,6 +446,23 @@ class S3ArtifactRepository(ArtifactRepository, MultipartUploadMixin):
         )
 
     def complete_multipart_upload(self, local_file, upload_id, parts=None, artifact_path=None):
+        """
+        Complete a multipart upload by combining all parts into a single S3 object.
+
+        This method should be called after all parts have been successfully uploaded
+        using the presigned URLs from create_multipart_upload. It tells S3 to combine
+        all the parts into the final object.
+
+        Args:
+            local_file: Absolute path to the local file that was uploaded. Must match
+                the local_file used in create_multipart_upload.
+            upload_id: The S3 upload ID returned by create_multipart_upload.
+            parts: List of MultipartUploadPart objects containing metadata for each
+                successfully uploaded part. Must include part_number and etag for each part.
+                Parts must be provided in order (part 1, part 2, etc.).
+            artifact_path: Optional relative path where the artifact should be stored.
+                Must match the artifact_path used in create_multipart_upload.
+        """
         (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
         if artifact_path:
             dest_path = posixpath.join(dest_path, artifact_path)
@@ -318,6 +474,20 @@ class S3ArtifactRepository(ArtifactRepository, MultipartUploadMixin):
         )
 
     def abort_multipart_upload(self, local_file, upload_id, artifact_path=None):
+        """
+        Abort a multipart upload and clean up any uploaded parts.
+
+        This method should be called if a multipart upload fails or is cancelled.
+        It cleans up any parts that were successfully uploaded and cancels the
+        multipart upload session in S3.
+
+        Args:
+            local_file: Absolute path to the local file that was being uploaded.
+                Must match the local_file used in create_multipart_upload.
+            upload_id: The S3 upload ID returned by create_multipart_upload.
+            artifact_path: Optional relative path where the artifact would have been stored.
+                Must match the artifact_path used in create_multipart_upload.
+        """
         (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
         if artifact_path:
             dest_path = posixpath.join(dest_path, artifact_path)

mlflow/tracing/client.py

@@ -2,15 +2,19 @@ import json
 import logging
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import nullcontext
-from typing import Optional, Sequence
+from typing import Optional, Sequence, Union
 
 import mlflow
-from mlflow.entities.assessment import Assessment
+from mlflow.entities.assessment import (
+    Assessment,
+)
 from mlflow.entities.model_registry import PromptVersion
 from mlflow.entities.span import NO_OP_SPAN_TRACE_ID
 from mlflow.entities.trace import Trace
 from mlflow.entities.trace_data import TraceData
 from mlflow.entities.trace_info import TraceInfo
+from mlflow.entities.trace_info_v2 import TraceInfoV2
+from mlflow.entities.trace_status import TraceStatus
 from mlflow.environment_variables import MLFLOW_SEARCH_TRACES_MAX_THREADS
 from mlflow.exceptions import (
     MlflowException,
@@ -59,65 +63,128 @@ class TracingClient:
     def store(self):
         return _get_store(self.tracking_uri)
 
-    def start_trace(self, trace_info: TraceInfo) -> TraceInfo:
+    def start_trace(
+        self,
+        experiment_id: str,
+        timestamp_ms: int,
+        request_metadata: dict[str, str],
+        tags: dict[str, str],
+    ):
         """
-        Create a new trace in the backend.
+        Start an initial TraceInfo object in the backend store.
 
         Args:
-            trace_info: The TraceInfo object to record in the backend.
+            experiment_id: String id of the experiment for this run.
+            timestamp_ms: Start time of the trace, in milliseconds since the UNIX epoch.
+            request_metadata: Metadata of the trace.
+            tags: Tags of the trace.
+
+        Returns:
+            The created TraceInfo object.
+        """
+        tags = exclude_immutable_tags(tags or {})
+        return self.store.start_trace(
+            experiment_id=experiment_id,
+            timestamp_ms=timestamp_ms,
+            request_metadata=request_metadata,
+            tags=tags,
+        )
+
+    def start_trace_v3(self, trace: Trace) -> TraceInfo:
+        """
+        Start a trace using the V3 API format.
+        NB: This method is named "Start" for internal reason in the backend, but actually
+        should be called at the end of the trace. We will migrate this to "CreateTrace"
+        API in the future to avoid confusion.
+
+        Args:
+            trace: The Trace object to create.
 
         Returns:
             The returned TraceInfoV3 object from the backend.
         """
-        return self.store.start_trace(trace_info=trace_info)
+        return self.store.start_trace_v3(trace=trace)
+
+    def end_trace(
+        self,
+        request_id: str,
+        timestamp_ms: int,
+        status: TraceStatus,
+        request_metadata: dict[str, str],
+        tags: dict[str, str],
+    ) -> TraceInfoV2:
+        """
+        Update the TraceInfo object in the backend store with the completed trace info.
+
+        Args:
+            request_id: Unique string identifier of the trace.
+            timestamp_ms: End time of the trace, in milliseconds. The execution time field
+                in the TraceInfo will be calculated by subtracting the start time from this.
+            status: Status of the trace.
+            request_metadata: Metadata of the trace. This will be merged with the existing
+                metadata logged during the start_trace call.
+            tags: Tags of the trace. This will be merged with the existing tags logged
+                during the start_trace or set_trace_tag calls.
+
+        Returns:
+            The updated TraceInfo object.
+        """
+        tags = exclude_immutable_tags(tags or {})
+        return self.store.end_trace(
+            request_id=request_id,
+            timestamp_ms=timestamp_ms,
+            status=status,
+            request_metadata=request_metadata,
+            tags=tags,
+        )
 
     def delete_traces(
         self,
         experiment_id: str,
         max_timestamp_millis: Optional[int] = None,
         max_traces: Optional[int] = None,
-        trace_ids: Optional[list[str]] = None,
+        request_ids: Optional[list[str]] = None,
     ) -> int:
         return self.store.delete_traces(
             experiment_id=experiment_id,
             max_timestamp_millis=max_timestamp_millis,
             max_traces=max_traces,
-            trace_ids=trace_ids,
+            request_ids=request_ids,
         )
 
-    def get_trace_info(self, trace_id: str) -> TraceInfo:
+    def get_trace_info(self, request_id, should_query_v3: bool = False) -> TraceInfoV2:
         """
-        Get the trace info matching the ``trace_id``.
+        Get the trace info matching the ``request_id``.
 
         Args:
-            trace_id: String id of the trace to fetch.
+            request_id: String id of the trace to fetch.
+            should_query_v3: If True, the backend store will query the V3 API for the trace info.
+                TODO: Remove this flag once the V3 API is the default in OSS.
 
         Returns:
             TraceInfo object, of type ``mlflow.entities.trace_info.TraceInfo``.
         """
-        with InMemoryTraceManager.get_instance().get_trace(trace_id) as trace:
-            if trace is not None:
-                return trace.info
+        return self.store.get_trace_info(request_id, should_query_v3=should_query_v3)
 
-        return self.store.get_trace_info(trace_id)
-
-    def get_trace(self, trace_id: str) -> Trace:
+    def get_trace(self, request_id) -> Trace:
         """
-        Get the trace matching the ``trace_id``.
+        Get the trace matching the ``request_id``.
 
         Args:
-            trace_id: String id of the trace to fetch.
+            request_id: String id of the trace to fetch.
 
         Returns:
             The fetched Trace object, of type ``mlflow.entities.Trace``.
         """
-        trace_info = self.get_trace_info(trace_id)
+        trace_info = self.get_trace_info(
+            request_id=request_id, should_query_v3=is_databricks_uri(self.tracking_uri)
+        )
         try:
             trace_data = self._download_trace_data(trace_info)
         except MlflowTraceDataNotFound:
             raise MlflowException(
                 message=(
-                    f"Trace with ID {trace_id} cannot be loaded because it is missing span data."
+                    f"Trace with ID {request_id} cannot be loaded because it is missing span data."
                     " Please try creating or loading another trace."
                 ),
                 error_code=BAD_REQUEST,
@@ -125,7 +192,7 @@ class TracingClient:
         except MlflowTraceDataCorrupted:
             raise MlflowException(
                 message=(
-                    f"Trace with ID {trace_id} cannot be loaded because its span data"
+                    f"Trace with ID {request_id} cannot be loaded because its span data"
                     " is corrupted. Please try creating or loading another trace."
                 ),
                 error_code=BAD_REQUEST,
@@ -253,23 +320,29 @@ class TracingClient:
             else:
                 filter_string = additional_filter
 
-        is_databricks = is_databricks_uri(self.tracking_uri)
-
-        def download_trace_extra_fields(trace_info: TraceInfo) -> Optional[Trace]:
+        def download_trace_extra_fields(
+            trace_info: Union[TraceInfoV2, TraceInfo],
+        ) -> Optional[Trace]:
             """
             Download trace data and assessments for the given trace_info and returns a Trace object.
             If the download fails (e.g., the trace data is missing or corrupted), returns None.
 
             The trace_info parameter can be either TraceInfo or TraceInfoV3 object.
             """
-            is_online_trace = is_uuid(trace_info.trace_id)
+            from mlflow.entities.trace_info import TraceInfo
+
+            # Determine if this is TraceInfo or TraceInfoV3
+            # Helps while transitioning to V3 traces for offline & online
+            is_v3 = isinstance(trace_info, TraceInfo)
+            trace_id = trace_info.trace_id if is_v3 else trace_info.request_id
+            is_online_trace = is_uuid(trace_id)
 
             # For online traces in Databricks, we need to get trace data from a different endpoint
             try:
                 if is_databricks and is_online_trace:
                     # For online traces, get data from the online API
                     trace_data = self.get_online_trace_details(
-                        trace_id=trace_info.trace_id,
+                        trace_id=trace_id,
                         sql_warehouse_id=sql_warehouse_id,
                         source_inference_table=trace_info.request_metadata.get(
                             "mlflow.sourceTable"
@@ -285,7 +358,7 @@ class TracingClient:
             except MlflowTraceDataException as e:
                 _logger.warning(
                     (
-                        f"Failed to download trace data for trace {trace_info.trace_id!r} "
+                        f"Failed to download trace data for trace {trace_id!r} "
                         f"with {e.ctx}. For full traceback, set logging level to DEBUG."
                     ),
                     exc_info=_logger.isEnabledFor(logging.DEBUG),
@@ -299,11 +372,7 @@ class TracingClient:
         next_token = page_token
 
         max_workers = MLFLOW_SEARCH_TRACES_MAX_THREADS.get()
-        executor = (
-            ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="MlflowTracingSearch")
-            if include_spans
-            else nullcontext()
-        )
+        executor = ThreadPoolExecutor(max_workers=max_workers) if include_spans else nullcontext()
         with executor:
             while len(traces) < max_results:
                 trace_infos, next_token = self._search_traces(
@@ -332,24 +401,24 @@ class TracingClient:
 
         return PagedList(traces, next_token)
 
-    def set_trace_tags(self, trace_id: str, tags: dict[str, str]):
+    def set_trace_tags(self, request_id, tags):
         """
-        Set tags on the trace with the given trace_id.
+        Set tags on the trace with the given request_id.
 
         Args:
-            trace_id: The ID of the trace.
+            request_id: The ID of the trace.
             tags: A dictionary of key-value pairs.
         """
         tags = exclude_immutable_tags(tags)
         for k, v in tags.items():
-            self.set_trace_tag(trace_id, k, v)
+            self.set_trace_tag(request_id, k, v)
 
-    def set_trace_tag(self, trace_id: str, key: str, value: str):
+    def set_trace_tag(self, request_id, key, value):
         """
         Set a tag on the trace with the given trace ID.
 
         Args:
-            trace_id: The ID of the trace to set the tag on.
+            request_id: The ID of the trace to set the tag on.
             key: The string key of the tag. Must be at most 250 characters long, otherwise
                 it will be truncated when stored.
             value: The string value of the tag. Must be at most 250 characters long, otherwise
@@ -362,7 +431,7 @@ class TracingClient:
             )
 
         # Trying to set the tag on the active trace first
-        with InMemoryTraceManager.get_instance().get_trace(trace_id) as trace:
+        with InMemoryTraceManager.get_instance().get_trace(request_id) as trace:
             if trace:
                 trace.info.tags[key] = str(value)
                 return
@@ -370,33 +439,33 @@ class TracingClient:
         if key in IMMUTABLE_TAGS:
             _logger.warning(f"Tag '{key}' is immutable and cannot be set on a trace.")
         else:
-            self.store.set_trace_tag(trace_id, key, str(value))
+            self.store.set_trace_tag(request_id, key, str(value))
 
-    def delete_trace_tag(self, trace_id: str, key: str):
+    def delete_trace_tag(self, request_id, key):
         """
         Delete a tag on the trace with the given trace ID.
 
         Args:
-            trace_id: The ID of the trace to delete the tag from.
+            request_id: The ID of the trace to delete the tag from.
             key: The string key of the tag. Must be at most 250 characters long, otherwise
                 it will be truncated when stored.
         """
         # Trying to delete the tag on the active trace first
-        with InMemoryTraceManager.get_instance().get_trace(trace_id) as trace:
+        with InMemoryTraceManager.get_instance().get_trace(request_id) as trace:
             if trace:
                 if key in trace.info.tags:
                     trace.info.tags.pop(key)
                     return
                 else:
                     raise MlflowException(
-                        f"Tag with key {key} not found in trace with ID {trace_id}.",
+                        f"Tag with key {key} not found in trace with ID {request_id}.",
                         error_code=RESOURCE_DOES_NOT_EXIST,
                     )
 
         if key in IMMUTABLE_TAGS:
             _logger.warning(f"Tag '{key}' is immutable and cannot be deleted on a trace.")
         else:
-            self.store.delete_trace_tag(trace_id, key)
+            self.store.delete_trace_tag(request_id, key)
 
     def get_assessment(self, trace_id: str, assessment_id: str) -> Assessment:
         """
@@ -492,12 +561,12 @@ class TracingClient:
 
         self.store.delete_assessment(trace_id=trace_id, assessment_id=assessment_id)
 
-    def _get_artifact_repo_for_trace(self, trace_info: TraceInfo):
+    def _get_artifact_repo_for_trace(self, trace_info: TraceInfoV2):
         artifact_uri = get_artifact_uri_for_trace(trace_info)
         artifact_uri = add_databricks_profile_info_to_artifact_uri(artifact_uri, self.tracking_uri)
         return get_artifact_repository(artifact_uri)
 
-    def _download_trace_data(self, trace_info: TraceInfo) -> TraceData:
+    def _download_trace_data(self, trace_info: Union[TraceInfoV2, TraceInfo]) -> TraceData:
         """
         Download trace data from artifact repository.
 
@@ -510,11 +579,32 @@ class TracingClient:
         artifact_repo = self._get_artifact_repo_for_trace(trace_info)
         return TraceData.from_dict(artifact_repo.download_trace_data())
 
-    def _upload_trace_data(self, trace_info: TraceInfo, trace_data: TraceData) -> None:
+    def _upload_trace_data(self, trace_info: TraceInfoV2, trace_data: TraceData) -> None:
         artifact_repo = self._get_artifact_repo_for_trace(trace_info)
         trace_data_json = json.dumps(trace_data.to_dict(), cls=TraceJSONEncoder, ensure_ascii=False)
         return artifact_repo.upload_trace_data(trace_data_json)
 
+    def _upload_ended_trace_info(
+        self,
+        trace_info: TraceInfoV2,
+    ) -> TraceInfoV2:
+        """
+        Update the TraceInfo object in the backend store with the completed trace info.
+
+        Args:
+            trace_info: Updated TraceInfo object to be stored in the backend store.
+
+        Returns:
+            The updated TraceInfo object.
+        """
+        return self.end_trace(
+            request_id=trace_info.request_id,
+            timestamp_ms=trace_info.timestamp_ms + trace_info.execution_time_ms,
+            status=trace_info.status,
+            request_metadata=trace_info.request_metadata,
+            tags=trace_info.tags or {},
+        )
+
     def link_prompt_versions_to_trace(
         self, trace_id: str, prompts: Sequence[PromptVersion]
     ) -> None: