genesis-flow 1.0.0__py3-none-any.whl → 1.0.3__py3-none-any.whl

mlflow/tracing/export/mlflow_v3.py

@@ -18,7 +18,6 @@ from mlflow.tracing.fluent import _EVAL_REQUEST_ID_TO_TRACE_ID, _set_last_active
 from mlflow.tracing.trace_manager import InMemoryTraceManager
 from mlflow.tracing.utils import add_size_stats_to_trace_metadata, maybe_get_request_id
 from mlflow.utils.databricks_utils import is_in_databricks_notebook
-from mlflow.utils.uri import is_databricks_uri
 
 _logger = logging.getLogger(__name__)
 
@@ -30,13 +29,15 @@ class MlflowV3SpanExporter(SpanExporter):
     """
 
     def __init__(self, tracking_uri: Optional[str] = None):
-        self._client = TracingClient(tracking_uri)
         self._is_async_enabled = self._should_enable_async_logging()
         if self._is_async_enabled:
             self._async_queue = AsyncTraceExportQueue()
+        self._client = TracingClient(tracking_uri)
 
-        # Display handler is no-op when running outside of notebooks.
-        self._display_handler = get_display_handler()
+        # Only display traces inline in Databricks notebooks
+        self._should_display_trace = is_in_databricks_notebook()
+        if self._should_display_trace:
+            self._display_handler = get_display_handler()
 
     def export(self, spans: Sequence[ReadableSpan]):
         """
@@ -64,7 +65,7 @@ class MlflowV3SpanExporter(SpanExporter):
             if eval_request_id := trace.info.tags.get(TraceTagKey.EVAL_REQUEST_ID):
                 _EVAL_REQUEST_ID_TO_TRACE_ID[eval_request_id] = trace.info.trace_id
 
-            if not maybe_get_request_id(is_evaluate=True):
+            if self._should_display_trace and not maybe_get_request_id(is_evaluate=True):
                 self._display_handler.display_traces([trace])
 
             if self._should_log_async():
@@ -88,7 +89,7 @@ class MlflowV3SpanExporter(SpanExporter):
         try:
             if trace:
                 add_size_stats_to_trace_metadata(trace)
-                returned_trace_info = self._client.start_trace(trace.info)
+                returned_trace_info = self._client.start_trace_v3(trace)
                 self._client._upload_trace_data(returned_trace_info, trace.data)
             else:
                 _logger.warning("No trace or trace info provided, unable to export")
@@ -110,11 +111,7 @@ class MlflowV3SpanExporter(SpanExporter):
             _logger.warning(f"Failed to link prompts to trace: {e}")
 
     def _should_enable_async_logging(self):
-        if (
-            is_in_databricks_notebook()
-            # NB: Not defaulting OSS backend to async logging for now to reduce blast radius.
-            or not is_databricks_uri(self._client.tracking_uri)
-        ):
+        if is_in_databricks_notebook():
             # NB: We don't turn on async logging in Databricks notebook by default
             # until we are confident that the async logging is working on the
             # offline workload on Databricks, to derisk the inclusion to the

mlflow/tracing/provider.py

@@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Optional
 
 from opentelemetry import context as context_api
 from opentelemetry import trace
+from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider
 
 import mlflow
@@ -277,7 +278,10 @@ def _setup_tracer_provider(disabled=False):
         # Default to MLflow Tracking Server
         processor = _get_mlflow_span_processor(tracking_uri=mlflow.get_tracking_uri())
 
-    tracer_provider = TracerProvider()
+    # Setting an empty resource to avoid triggering resource aggregation, which causes
+    # an issue in LiteLLM tracing: https://github.com/mlflow/mlflow/issues/16296
+    # MLflow tracing does not use resource right now.
+    tracer_provider = TracerProvider(resource=Resource.get_empty())
     tracer_provider.add_span_processor(processor)
     _MLFLOW_TRACER_PROVIDER = tracer_provider
 

mlflow/tracking/_model_registry/client.py

@@ -141,7 +141,11 @@ class ModelRegistryClient:
             obtained via the ``token`` attribute of the object.
 
         """
-        if is_prompt_supported_registry(self.registry_uri):
+        # Add prompt filter for prompt-supported registries that also support filter_string
+        # Unity Catalog supports prompts but not filter_string parameter
+        if is_prompt_supported_registry(self.registry_uri) and not (
+            self.registry_uri or ""
+        ).startswith("databricks-uc"):
             # Adjust filter string to include or exclude prompts
             filter_string = add_prompt_filter_string(filter_string, False)
 

mlflow/utils/file_utils.py

@@ -826,7 +826,8 @@ def shutil_copytree_without_file_permissions(src_dir, dst_dir):
             # For each directory <dirname> immediately under <dirpath>, create an equivalently-named
             # directory under the destination directory
             abs_dir_path = os.path.join(dst_dir, relative_dir_path)
-            os.mkdir(abs_dir_path)
+            if not os.path.exists(abs_dir_path):
+                os.mkdir(abs_dir_path)
         for filename in filenames:
             # For each file with name <filename> immediately under <dirpath>, copy that file to
             # the appropriate location in the destination directory

mlflow/utils/rest_utils.py

@@ -119,6 +119,7 @@ def http_request(
             host_creds.token,
             host_creds.databricks_auth_profile,
             retry_timeout_seconds=retry_timeout_seconds,
+            timeout=timeout,
         )
 
         def make_sdk_call():
@@ -259,6 +260,7 @@ def get_workspace_client(
     token,
     databricks_auth_profile,
     retry_timeout_seconds=None,
+    timeout=None,
 ):
     from databricks.sdk import WorkspaceClient
     from databricks.sdk.config import Config
@@ -267,6 +269,8 @@ def get_workspace_client(
         kwargs = {"host": host, "token": token}
     else:
         kwargs = {"profile": databricks_auth_profile}
+    if timeout is not None:
+        kwargs["http_timeout_seconds"] = timeout
     config = Config(
         **kwargs,
         retry_timeout_seconds=retry_timeout_seconds

mlflow/utils/secure_loading.py

@@ -22,14 +22,14 @@ SAFE_PICKLE_CLASSES = {
     'numpy.dtype',
     'numpy.int32', 'numpy.int64', 'numpy.float32', 'numpy.float64',
     'numpy.bool_', 'numpy.str_',
-    
+
     # Pandas types
     'pandas.core.frame.DataFrame',
     'pandas.core.series.Series',
     'pandas.core.index.Index',
     'pandas.core.dtypes.dtypes.CategoricalDtype',
-    
-    # Scikit-learn estimators (core models only)
+
+    # Scikit-learn estimators
     'sklearn.linear_model._base.LinearRegression',
     'sklearn.linear_model._logistic.LogisticRegression',
     'sklearn.ensemble._forest.RandomForestClassifier',
@@ -38,20 +38,74 @@ SAFE_PICKLE_CLASSES = {
     'sklearn.tree._classes.DecisionTreeRegressor',
     'sklearn.svm._classes.SVC',
     'sklearn.svm._classes.SVR',
-    
-    # Standard Python types
+
+    # Built-in types
     'builtins.dict', 'builtins.list', 'builtins.tuple', 'builtins.set',
     'builtins.str', 'builtins.int', 'builtins.float', 'builtins.bool',
-    
+    'builtins.type',
+
     # Collections
     'collections.OrderedDict',
     'collections.defaultdict',
-    
-    # Genesis-Flow internal types
+
+    # MLflow types
     'mlflow.models.signature.ModelSignature',
+    'mlflow.models.signature._TypeHints',
     'mlflow.types.schema.Schema',
+    'mlflow.pyfunc.model.PythonModel',
+
+    # Cloudpickle internals
+    'cloudpickle.cloudpickle._make_skeleton_class',
+    'cloudpickle.cloudpickle._class_setstate',
+    'cloudpickle.cloudpickle._make_function',
+    'cloudpickle.cloudpickle._builtin_type',
+    'cloudpickle.cloudpickle._function_setstate',
+    'cloudpickle.cloudpickle._make_empty_cell',
+    'cloudpickle.cloudpickle._make_cell',
+
+    # Sentence Transformers
+    'sentence_transformers.SentenceTransformer.SentenceTransformer',
+    'sentence_transformers.model_card.SentenceTransformerModelCardData',
+    'sentence_transformers.models.Transformer.Transformer',
+    'sentence_transformers.models.Pooling.Pooling',
+    'sentence_transformers.models.Normalize.Normalize',
+
+    # Torch
+    'torch.torch_version.TorchVersion',
+    'torch._utils._rebuild_tensor_v2',
+    'torch.storage._load_from_bytes',
+    'torch.nn.modules.sparse.Embedding',
+    'torch._utils._rebuild_parameter',
+    'torch.nn.modules.normalization.LayerNorm',
+    'torch.nn.modules.dropout.Dropout',
+    'torch.nn.modules.container.ModuleList',
+    'torch.nn.modules.linear.Linear',
+    'torch.nn.modules.activation.Tanh',
+    'torch.float32',
+    'torch._C._nn.gelu',
+
+    # Transformers
+    'transformers.models.bert.modeling_bert.BertModel',
+    'transformers.models.bert.modeling_bert.BertEmbeddings',
+    'transformers.models.bert.modeling_bert.BertEncoder',
+    'transformers.models.bert.modeling_bert.BertLayer',
+    'transformers.models.bert.modeling_bert.BertAttention',
+    'transformers.models.bert.modeling_bert.BertSdpaSelfAttention',
+    'transformers.models.bert.modeling_bert.BertSelfOutput',
+    'transformers.models.bert.modeling_bert.BertIntermediate',
+    'transformers.models.bert.modeling_bert.BertOutput',
+    'transformers.models.bert.modeling_bert.BertPooler',
+    'transformers.models.bert.configuration_bert.BertConfig',
+    'transformers.models.bert.tokenization_bert_fast.BertTokenizerFast',
+    'transformers.activations.GELUActivation',
+
+    # Tokenizers
+    'tokenizers.Tokenizer',
+    'tokenizers.models.Model',
+    'tokenizers.AddedToken',
 }
 
+
 class RestrictedUnpickler(pickle.Unpickler):
     """
     Secure unpickler that only allows safe, whitelisted classes.

mlflow/version.py

@@ -2,7 +2,7 @@
 import importlib.metadata
 import re
 
-VERSION = "3.1.2.dev0"
+VERSION = "3.1.4"
 
 
 def is_release_version():
@@ -21,4 +21,4 @@ def _is_package_installed(package_name: str) -> bool:
 # installed, or includes the full MLflow or mlflow-skinny package.
 # This is used to determine whether to import modules that require
 # dependencies that are not included in the tracing SDK.
-IS_TRACING_SDK_ONLY = not any(_is_package_installed(pkg) for pkg in ["mlflow", "mlflow-skinny"])
+IS_TRACING_SDK_ONLY = not any(_is_package_installed(pkg) for pkg in ["mlflow", "mlflow-skinny", "genesis-flow"])