github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-08-28 14:05:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added Databricks support to MLflow Callback (#7906)

Browse Source 
Added a quick check to make integration easier with Databricks; another
option would be to make a new class, but this seemed more
straightfoward.

cc: @liangz1 Can this be done in a more straightfoward way?
This commit is contained in:
Rithwik Ediga Lakhamsani 2023-07-25 18:23:54 -07:00 committed by GitHub
parent 479cc086ba
commit d1d691caa4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 21 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
libs/langchain/langchain/callbacks/mlflow_callback.py
View File

@ -1,3 +1,4 @@
import os
import random
import string
import tempfile
@ -127,6 +128,11 @@ class MlflowLogger:
def __init__(self, **kwargs: Any):
self.mlflow = import_mlflow()
if "DATABRICKS_RUNTIME_VERSION" in os.environ:
self.mlflow.set_tracking_uri("databricks")
self.mlf_expid = self.mlflow.tracking.fluent._get_experiment_id()
self.mlf_exp = self.mlflow.get_experiment(self.mlf_expid)
else:
tracking_uri = get_from_dict_or_env(
kwargs, "tracking_uri", "MLFLOW_TRACKING_URI", ""
)