wehub-resource-sync
142 lines
4.3 KiB
Python
142 lines
4.3 KiB
Python
"""
|
|
This simple example shows how you could use MLflow REST API to create new
|
|
runs inside an experiment to log parameters/metrics. Using MLflow REST API
|
|
instead of MLflow library might be useful to embed in an application where
|
|
you don't want to depend on the whole MLflow library, or to make
|
|
your own HTTP requests in another programming language (not Python).
|
|
For more details on MLflow REST API endpoints check the following page:
|
|
|
|
https://www.mlflow.org/docs/latest/rest-api.html
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import pwd
|
|
|
|
import requests
|
|
|
|
from mlflow.utils.time import get_current_time_millis
|
|
|
|
_DEFAULT_USER_ID = "unknown"
|
|
|
|
|
|
class MlflowTrackingRestApi:
|
|
def __init__(self, hostname, port, experiment_id):
|
|
self.base_url = "http://" + hostname + ":" + str(port) + "/api/2.0/mlflow"
|
|
self.experiment_id = experiment_id
|
|
self.run_id = self.create_run()
|
|
|
|
def create_run(self):
|
|
"""Create a new run for tracking."""
|
|
url = self.base_url + "/runs/create"
|
|
# user_id is deprecated and will be removed from the API in a future release
|
|
payload = {
|
|
"experiment_id": self.experiment_id,
|
|
"start_time": get_current_time_millis(),
|
|
"user_id": _get_user_id(),
|
|
}
|
|
r = requests.post(url, json=payload)
|
|
run_id = None
|
|
if r.status_code == 200:
|
|
run_id = r.json()["run"]["info"]["run_uuid"]
|
|
else:
|
|
print("Creating run failed!")
|
|
return run_id
|
|
|
|
def search_experiments(self):
|
|
"""Get all experiments."""
|
|
url = self.base_url + "/experiments/search"
|
|
r = requests.get(url)
|
|
experiments = None
|
|
if r.status_code == 200:
|
|
experiments = r.json()["experiments"]
|
|
return experiments
|
|
|
|
def log_param(self, param):
|
|
"""Log a parameter dict for the given run."""
|
|
url = self.base_url + "/runs/log-parameter"
|
|
payload = {"run_id": self.run_id, "key": param["key"], "value": param["value"]}
|
|
r = requests.post(url, json=payload)
|
|
return r.status_code
|
|
|
|
def log_metric(self, metric):
|
|
"""Log a metric dict for the given run."""
|
|
url = self.base_url + "/runs/log-metric"
|
|
payload = {
|
|
"run_id": self.run_id,
|
|
"key": metric["key"],
|
|
"value": metric["value"],
|
|
"timestamp": metric["timestamp"],
|
|
"step": metric["step"],
|
|
}
|
|
r = requests.post(url, json=payload)
|
|
return r.status_code
|
|
|
|
|
|
def _get_user_id():
|
|
"""Get the ID of the user for the current run."""
|
|
try:
|
|
return pwd.getpwuid(os.getuid())[0]
|
|
except ImportError:
|
|
return _DEFAULT_USER_ID
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Command-line arguments
|
|
parser = argparse.ArgumentParser(description="MLflow REST API Example")
|
|
|
|
parser.add_argument(
|
|
"--hostname",
|
|
type=str,
|
|
default="localhost",
|
|
dest="hostname",
|
|
help="MLflow server hostname/ip (default: localhost)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--port",
|
|
type=int,
|
|
default=5000,
|
|
dest="port",
|
|
help="MLflow server port number (default: 5000)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--experiment-id",
|
|
type=int,
|
|
default=0,
|
|
dest="experiment_id",
|
|
help="Experiment ID (default: 0)",
|
|
)
|
|
|
|
print("Running mlflow_tracking_rest_api.py")
|
|
|
|
args = parser.parse_args()
|
|
|
|
mlflow_rest = MlflowTrackingRestApi(args.hostname, args.port, args.experiment_id)
|
|
# Parameter is a key/val pair (str types)
|
|
param = {"key": "alpha", "value": "0.1980"}
|
|
status_code = mlflow_rest.log_param(param)
|
|
if status_code == 200:
|
|
print(
|
|
"Successfully logged parameter: {} with value: {}".format(param["key"], param["value"])
|
|
)
|
|
else:
|
|
print("Logging parameter failed!")
|
|
# Metric is a key/val pair (key/val have str/float types)
|
|
metric = {
|
|
"key": "precision",
|
|
"value": 0.769,
|
|
"timestamp": get_current_time_millis(),
|
|
"step": 1,
|
|
}
|
|
status_code = mlflow_rest.log_metric(metric)
|
|
if status_code == 200:
|
|
print(
|
|
"Successfully logged parameter: {} with value: {}".format(
|
|
metric["key"], metric["value"]
|
|
)
|
|
)
|
|
else:
|
|
print("Logging metric failed!")
|