-
Notifications
You must be signed in to change notification settings - Fork 67
/
main.py
73 lines (52 loc) · 2.09 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
import mlflow
import os
import hydra
from omegaconf import DictConfig, OmegaConf
# This automatically reads in the configuration
@hydra.main(config_name='config')
def go(config: DictConfig):
# Setup the wandb experiment. All runs will be grouped under this name
os.environ["WANDB_PROJECT"] = config["main"]["project_name"]
os.environ["WANDB_RUN_GROUP"] = config["main"]["experiment_name"]
# You can get the path at the root of the MLflow project with this:
root_path = hydra.utils.get_original_cwd()
# Check which steps we need to execute
if isinstance(config["main"]["execute_steps"], str):
# This was passed on the command line as a comma-separated list of steps
steps_to_execute = config["main"]["execute_steps"].split(",")
else:
assert isinstance(config["main"]["execute_steps"], list)
steps_to_execute = config["main"]["execute_steps"]
# Download step
if "download" in steps_to_execute:
_ = mlflow.run(
os.path.join(root_path, "download"),
"main",
parameters={
"file_url": config["data"]["file_url"],
"artifact_name": "raw_data.parquet",
"artifact_type": "raw_data",
"artifact_description": "Data as downloaded"
},
)
if "preprocess" in steps_to_execute:
## YOUR CODE HERE: call the preprocess step
pass
if "check_data" in steps_to_execute:
## YOUR CODE HERE: call the check_data step
pass
if "segregate" in steps_to_execute:
## YOUR CODE HERE: call the segregate step
pass
if "random_forest" in steps_to_execute:
# Serialize decision tree configuration
model_config = os.path.abspath("random_forest_config.yml")
with open(model_config, "w+") as fp:
fp.write(OmegaConf.to_yaml(config["random_forest_pipeline"]))
## YOUR CODE HERE: call the random_forest step
pass
if "evaluate" in steps_to_execute:
## YOUR CODE HERE: call the evaluate step
pass
if __name__ == "__main__":
go()