Skip to content

Commit 42c6cfa

Browse files
committed
ch.06 mlflow 관련 업데이트
1 parent 1e7e86c commit 42c6cfa

File tree

1 file changed

+5
-4
lines changed

1 file changed

+5
-4
lines changed

source/train.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,15 +10,16 @@
1010
import bentoml
1111
import mlflow
1212

13-
mlflow.set_tracking_uri(uri="http://mlflow-server:5000")
14-
mlflow.set_experiment("fraud_detection")
13+
if(Variable.get("isProd") == "true"):
14+
mlflow.set_tracking_uri(uri="http://mlflow-server:5000")
15+
mlflow.set_experiment("fraud_detection")
1516

1617
def build_model(**kwargs):
17-
mlflow.autolog()
1818
with mlflow.start_run():
1919
if(Variable.get("isProd") == "true"):
20+
mlflow.autolog()
2021
engine = create_engine("postgresql+psycopg2://" + Variable.get("db_username") + ":" + Variable.get("db_password") + "@" + Variable.get("db_host") + ":" + Variable.get("db_port") + "/" + Variable.get("db_name"))
21-
data = pd.read_sql("select * from public.sample_data limit 6362120", engine)
22+
data = pd.read_sql("select * from public.sample_data", engine)
2223
else:
2324
data = pd.read_csv('/opt/airflow/dags/sample_data/Fraud.csv')
2425

0 commit comments

Comments
 (0)