wehub-resource-sync
67 lines
1.8 KiB
Python
67 lines
1.8 KiB
Python
import os
|
|
import uuid
|
|
|
|
import mlflow.server
|
|
|
|
|
|
class User:
|
|
MLFLOW_TRACKING_USERNAME = "MLFLOW_TRACKING_USERNAME"
|
|
MLFLOW_TRACKING_PASSWORD = "MLFLOW_TRACKING_PASSWORD"
|
|
|
|
def __init__(self, username, password) -> None:
|
|
self.username = username
|
|
self.password = password
|
|
self.env = {}
|
|
|
|
def _record_env_var(self, key):
|
|
if key := os.environ.get(key):
|
|
self.env[key] = key
|
|
|
|
def _restore_env_var(self, key):
|
|
if value := self.env.get(key):
|
|
os.environ[key] = value
|
|
else:
|
|
del os.environ[key]
|
|
|
|
def __enter__(self):
|
|
self._record_env_var(User.MLFLOW_TRACKING_USERNAME)
|
|
self._record_env_var(User.MLFLOW_TRACKING_PASSWORD)
|
|
os.environ[User.MLFLOW_TRACKING_USERNAME] = self.username
|
|
os.environ[User.MLFLOW_TRACKING_PASSWORD] = self.password
|
|
return self
|
|
|
|
def __exit__(self, *_exc):
|
|
self._restore_env_var(User.MLFLOW_TRACKING_USERNAME)
|
|
self._restore_env_var(User.MLFLOW_TRACKING_PASSWORD)
|
|
self.env.clear()
|
|
|
|
|
|
tracking_uri = "http://localhost:5000"
|
|
mlflow.set_tracking_uri(tracking_uri)
|
|
client = mlflow.server.get_app_client("basic-auth", tracking_uri)
|
|
A = User("user_a", "password_a")
|
|
B = User("user_b", "password_b")
|
|
|
|
with A:
|
|
exp_a = mlflow.set_experiment(uuid.uuid4().hex)
|
|
with mlflow.start_run():
|
|
mlflow.log_metric("a", 1)
|
|
|
|
with B:
|
|
mlflow.set_experiment(exp_a.name)
|
|
try:
|
|
with mlflow.start_run(): # not allowed
|
|
mlflow.log_metric("b", 2)
|
|
except Exception as e:
|
|
print(str(e))
|
|
|
|
# Grant B permission to edit A's experiment
|
|
with A:
|
|
client.create_experiment_permission(str(exp_a.experiment_id), B.username, "EDIT")
|
|
|
|
# B can edit now, should be able to log a metric
|
|
with B:
|
|
mlflow.set_experiment(exp_a.name)
|
|
with mlflow.start_run():
|
|
mlflow.log_metric("b", 2)
|