wehub-resource-sync
92 lines
2.1 KiB
Python
92 lines
2.1 KiB
Python
import random
|
|
import uuid
|
|
|
|
import pytest
|
|
|
|
from mlflow.entities import (
|
|
Dataset,
|
|
DatasetInput,
|
|
InputTag,
|
|
LifecycleStage,
|
|
Metric,
|
|
Param,
|
|
RunData,
|
|
RunInfo,
|
|
RunInputs,
|
|
RunStatus,
|
|
RunTag,
|
|
)
|
|
from mlflow.utils.time import get_current_time_millis
|
|
|
|
from tests.helper_functions import random_int, random_str
|
|
|
|
|
|
@pytest.fixture
|
|
def run_data():
|
|
metrics = [
|
|
Metric(
|
|
key=random_str(10),
|
|
value=random_int(0, 1000),
|
|
timestamp=get_current_time_millis() + random_int(-1e4, 1e4),
|
|
step=random_int(),
|
|
)
|
|
]
|
|
params = [Param(random_str(10), random_str(random_int(10, 35))) for _ in range(10)]
|
|
tags = [RunTag(random_str(10), random_str(random_int(10, 35))) for _ in range(10)]
|
|
|
|
rd = RunData(metrics=metrics, params=params, tags=tags)
|
|
|
|
return rd, metrics, params, tags
|
|
|
|
|
|
@pytest.fixture
|
|
def run_info():
|
|
run_id = str(uuid.uuid4())
|
|
experiment_id = str(random_int(10, 2000))
|
|
user_id = random_str(random_int(10, 25))
|
|
run_name = random_str(random_int(10, 25))
|
|
status = RunStatus.to_string(random.choice(RunStatus.all_status()))
|
|
start_time = random_int(1, 10)
|
|
end_time = start_time + random_int(1, 10)
|
|
lifecycle_stage = LifecycleStage.ACTIVE
|
|
artifact_uri = random_str(random_int(10, 40))
|
|
ri = RunInfo(
|
|
run_id=run_id,
|
|
run_name=run_name,
|
|
experiment_id=experiment_id,
|
|
user_id=user_id,
|
|
status=status,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
lifecycle_stage=lifecycle_stage,
|
|
artifact_uri=artifact_uri,
|
|
)
|
|
return (
|
|
ri,
|
|
run_id,
|
|
run_name,
|
|
experiment_id,
|
|
user_id,
|
|
status,
|
|
start_time,
|
|
end_time,
|
|
lifecycle_stage,
|
|
artifact_uri,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def run_inputs():
|
|
datasets = [
|
|
DatasetInput(
|
|
dataset=Dataset(
|
|
name="name1", digest="digest1", source_type="my_source_type", source="source"
|
|
),
|
|
tags=[InputTag(key="key", value="value")],
|
|
)
|
|
]
|
|
|
|
run_inputs = RunInputs(dataset_inputs=datasets)
|
|
|
|
return run_inputs, datasets
|