Skip to content

Commit cc51a72

Browse files
ctimlctiml
authored
Merge pull request #106 from InfuseAI/feature/register-models
[Feature] Support register model
2 parents 2acdb92 + 67fa9ea commit cc51a72

File tree

3 files changed

+299
-2
lines changed

3 files changed

+299
-2
lines changed

‎docs/CLI/models.md

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,10 @@ Available Commands:
1212
get Get the model
1313
get-version Get a version of the model
1414
list List models
15+
list-artifacts List artifacts of a run
16+
list-runs List runs of an experiment
1517
list-versions List versions of the model
18+
register Register a model
1619
1720
Options:
1821
-h, --help Show the help
@@ -88,6 +91,38 @@ primehub models list
8891

8992

9093

94+
### list-artifacts
95+
96+
List artifacts of a run
97+
98+
99+
```
100+
primehub models list-artifacts <run_id>
101+
```
102+
103+
* run_id: The run id
104+
105+
106+
* *(optional)* path: Add `--path <path>` to filter artifacts matching this path
107+
108+
109+
110+
111+
### list-runs
112+
113+
List runs of an experiment
114+
115+
116+
```
117+
primehub models list-runs <experiment_name>
118+
```
119+
120+
* experiment_name: The experiment name
121+
122+
123+
124+
125+
91126
### list-versions
92127

93128
List versions of the model
@@ -102,6 +137,23 @@ primehub models list-versions <model>
102137

103138

104139

140+
141+
### register
142+
143+
Register a model
144+
145+
146+
```
147+
primehub models register <name> <run_id> <path>
148+
```
149+
150+
* name: Model name
151+
* run_id: The run id
152+
* path: The location of the model artifacts
153+
154+
155+
156+
105157

106158

107159
## Examples
@@ -226,4 +278,66 @@ endpoint: https://primehub-python-sdk.primehub.io/deployment/deployme
226278
availableReplicas: None
227279
message: Deployment is being deployed and not available now
228280
pods: [{'name': 'deploy-deployment-foo-dhpqt-6bcd77b854-lwcz8'}]
281+
```
282+
283+
### Example: register a model
284+
285+
List runs by experiment name
286+
287+
```
288+
$ primehub models list-runs mnist
289+
runId experimentId status startTime endTime
290+
-------------------------------- -------------- -------- ------------------- -------------------
291+
b14a7d0b39d344af8cf9914d0616f27e 1 FINISHED 2022-12-27 18:25:58 2022-12-27 18:25:58
292+
bd1694678bf2489f8de55d71bf9585db 1 FINISHED 2022-12-27 18:24:35 2022-12-27 18:24:35
293+
3a3545ed3cb04c1dbe7b8e1392cbaf1a 1 FAILED 2022-12-27 18:24:23 2022-12-27 18:24:23
294+
2fbd744173a44678ad9ca74bb3879d62 1 FAILED 2022-12-27 16:47:12 2022-12-27 16:47:12
295+
```
296+
297+
List artifacts by run id
298+
299+
```
300+
$ primehub models list-artifacts b14a7d0b39d344af8cf9914d0616f27e
301+
path is_dir file_size
302+
--------- -------- -----------
303+
outputs True
304+
hello.txt False 256
305+
```
306+
307+
List artifacts by run id and specify search path
308+
309+
```
310+
$ primehub models list-artifacts b14a7d0b39d344af8cf9914d0616f27e --path outputs
311+
path is_dir file_size
312+
-------------------------- -------- -----------
313+
outputs/.ipynb_checkpoints True
314+
outputs/bar True
315+
outputs/foo True
316+
outputs/test.txt False 12
317+
```
318+
319+
Register model
320+
321+
```
322+
$ primehub models register model1 b14a7d0b39d344af8cf9914d0616f27e outputs/foo
323+
name: model1
324+
version: 1
325+
creation_timestamp: 1672161284579
326+
last_updated_timestamp: 1672161284579
327+
current_stage: None
328+
source: /project/foo/phapplications/mlflow-kvmdg/mlruns/1/b14a7d0b39d344af8cf9914d0616f27e/artifacts/outputs/foo
329+
run_id: b14a7d0b39d344af8cf9914d0616f27e
330+
status: READY
331+
```
332+
333+
List the created model
334+
335+
```
336+
$ primehub models list
337+
name creationTimestamp lastUpdatedTimestamp description latestVersion
338+
------------------ ------------------- ---------------------- ------------- ---------------
339+
footest 2022-12-27 18:29:26 2022-12-27 18:29:26 1
340+
created-by-graphql 2022-12-27 19:25:11 2022-12-27 19:43:08 2
341+
abc 2022-12-27 21:37:29 2022-12-28 00:45:09 5
342+
model1 2022-12-28 01:14:44 2022-12-28 01:16:30 1
229343
```

‎primehub/extras/templates/examples/models.md

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,3 +119,65 @@ availableReplicas: None
119119
message: Deployment is being deployed and not available now
120120
pods: [{'name': 'deploy-deployment-foo-dhpqt-6bcd77b854-lwcz8'}]
121121
```
122+
123+
### Example: register a model
124+
125+
List runs by experiment name
126+
127+
```
128+
$ primehub models list-runs mnist
129+
runId experimentId status startTime endTime
130+
-------------------------------- -------------- -------- ------------------- -------------------
131+
b14a7d0b39d344af8cf9914d0616f27e 1 FINISHED 2022-12-27 18:25:58 2022-12-27 18:25:58
132+
bd1694678bf2489f8de55d71bf9585db 1 FINISHED 2022-12-27 18:24:35 2022-12-27 18:24:35
133+
3a3545ed3cb04c1dbe7b8e1392cbaf1a 1 FAILED 2022-12-27 18:24:23 2022-12-27 18:24:23
134+
2fbd744173a44678ad9ca74bb3879d62 1 FAILED 2022-12-27 16:47:12 2022-12-27 16:47:12
135+
```
136+
137+
List artifacts by run id
138+
139+
```
140+
$ primehub models list-artifacts b14a7d0b39d344af8cf9914d0616f27e
141+
path is_dir file_size
142+
--------- -------- -----------
143+
outputs True
144+
hello.txt False 256
145+
```
146+
147+
List artifacts by run id and specify search path
148+
149+
```
150+
$ primehub models list-artifacts b14a7d0b39d344af8cf9914d0616f27e --path outputs
151+
path is_dir file_size
152+
-------------------------- -------- -----------
153+
outputs/.ipynb_checkpoints True
154+
outputs/bar True
155+
outputs/foo True
156+
outputs/test.txt False 12
157+
```
158+
159+
Register model
160+
161+
```
162+
$ primehub models register model1 b14a7d0b39d344af8cf9914d0616f27e outputs/foo
163+
name: model1
164+
version: 1
165+
creation_timestamp: 1672161284579
166+
last_updated_timestamp: 1672161284579
167+
current_stage: None
168+
source: /project/foo/phapplications/mlflow-kvmdg/mlruns/1/b14a7d0b39d344af8cf9914d0616f27e/artifacts/outputs/foo
169+
run_id: b14a7d0b39d344af8cf9914d0616f27e
170+
status: READY
171+
```
172+
173+
List the created model
174+
175+
```
176+
$ primehub models list
177+
name creationTimestamp lastUpdatedTimestamp description latestVersion
178+
------------------ ------------------- ---------------------- ------------- ---------------
179+
footest 2022-12-27 18:29:26 2022-12-27 18:29:26 1
180+
created-by-graphql 2022-12-27 19:25:11 2022-12-27 19:43:08 2
181+
abc 2022-12-27 21:37:29 2022-12-28 00:45:09 5
182+
model1 2022-12-28 01:14:44 2022-12-28 01:16:30 1
183+
```

‎primehub/models.py

Lines changed: 123 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import textwrap
2-
from typing import Iterator, Any
2+
from typing import Iterator, Any, Optional
33

44
from primehub import Helpful, cmd, Module
55
from primehub.utils.display import display_tree_like_format
@@ -54,7 +54,8 @@ def list(self) -> Iterator:
5454
m['creationTimestamp'] = timestamp_to_isoformat(m['creationTimestamp'])
5555
m['lastUpdatedTimestamp'] = timestamp_to_isoformat(m['lastUpdatedTimestamp'])
5656
versions = m.pop('latestVersions')
57-
m['latestVersion'] = versions[0]['version']
57+
if len(versions):
58+
m['latestVersion'] = versions[0]['version']
5859
yield m
5960
return results
6061

@@ -134,6 +135,7 @@ def list_versions(self, model: str) -> Iterator:
134135
version
135136
creationTimestamp
136137
lastUpdatedTimestamp
138+
source
137139
deployedBy
138140
}
139141
"""
@@ -182,6 +184,7 @@ def get_version(self, model: str, version: str) -> dict:
182184
version
183185
creationTimestamp
184186
lastUpdatedTimestamp
187+
source
185188
deployedBy
186189
}
187190
"""
@@ -192,6 +195,124 @@ def get_version(self, model: str, version: str) -> dict:
192195
results = results['data']['modelVersion']
193196
return results
194197

198+
@cmd(name='list-runs', description='List runs of an experiment', return_required=True)
199+
def list_runs(self, experiment_name: str):
200+
"""
201+
List runs of an experiment
202+
203+
:type experiment_name: str
204+
:param experiment_name: The experiment name
205+
206+
:rtype: list
207+
:return: The runs of an experiment
208+
"""
209+
210+
query = """
211+
query MLflowRuns($where: MLflowExperimentWhereUniqueInput!) {
212+
mlflowRuns(where: $where)
213+
}
214+
"""
215+
where = {
216+
'experimentName': experiment_name,
217+
'group': self.group_name,
218+
}
219+
220+
results = self.request({'where': where}, query)
221+
if 'data' in results:
222+
results = results['data']
223+
for r in results['mlflowRuns']:
224+
m = {
225+
'runId': r['info']['runId'],
226+
'experimentId': r['info']['experimentId'],
227+
'status': r['info']['status'],
228+
'startTime': timestamp_to_isoformat(r['info']['startTime']),
229+
'endTime': timestamp_to_isoformat(r['info']['endTime']),
230+
}
231+
yield m
232+
return results
233+
234+
@cmd(name='list-artifacts', description='List artifacts of a run', return_required=True, optionals=[('path', str)])
235+
def _list_artifacts(self, run_id: str, **kwargs):
236+
"""
237+
List artifacts of a run
238+
239+
:type run_id: str
240+
:param run_id: The run id
241+
242+
:type path: str
243+
:param path: Add `--path <path>` to filter artifacts matching this path
244+
245+
:rtype: list
246+
:return: The artifacts of a run
247+
"""
248+
path = kwargs.get('path')
249+
return self.list_artifacts(run_id, path)
250+
251+
def list_artifacts(self, run_id: str, path: Optional[str] = None):
252+
"""
253+
List artifacts of a run
254+
255+
:type run_id: str
256+
:param run_id: The run id
257+
258+
:type path: str
259+
:param path: Filter artifacts matching this path, defaults to None
260+
261+
:rtype: list
262+
:return: The artifacts of a run
263+
"""
264+
265+
query = """
266+
query MLflowArtifact($where: MLflowRunWhereUniqueInput!) {
267+
mlflowArtifact(where: $where)
268+
}
269+
"""
270+
where = {
271+
'runId': run_id,
272+
'group': self.group_name,
273+
}
274+
if path:
275+
where['path'] = path
276+
277+
results = self.request({'where': where}, query)
278+
if 'data' not in results:
279+
return results
280+
if 'files' in results['data']['mlflowArtifact']:
281+
return results['data']['mlflowArtifact']['files']
282+
return []
283+
284+
@cmd(name='register', description='Register a model')
285+
def register(self, name: str, run_id: str, path: str):
286+
"""
287+
Register a model
288+
289+
:type name: str
290+
:param name: Model name
291+
292+
:type run_id: str
293+
:param run_id: The run id
294+
295+
:type path: str
296+
:param path: The location of the model artifacts
297+
"""
298+
299+
query = """
300+
mutation RegisterModel($where: ModelRegisterInput!) {
301+
registerModel(where: $where)
302+
}
303+
"""
304+
where = {
305+
'name': name,
306+
'runId': run_id,
307+
'path': path,
308+
'group': self.group_name,
309+
}
310+
311+
results = self.request({'where': where}, query)
312+
if 'data' not in results:
313+
return results
314+
return results['data']['registerModel']
315+
195316
@cmd(name='deploy', description='Deploy the model version to the speific deployment', return_required=True)
196317
def deploy(self, model: str, version: str, deploy_id: str) -> dict:
197318
"""

0 commit comments

Comments
 (0)