trainml 0.5.9__py3-none-any.whl → 0.5.11__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (53) hide show
  1. tests/integration/projects/conftest.py +3 -1
  2. tests/integration/projects/test_projects_data_connectors_integration.py +44 -0
  3. tests/integration/projects/test_projects_datastores_integration.py +42 -0
  4. tests/integration/projects/test_projects_services_integration.py +44 -0
  5. tests/integration/test_checkpoints_integration.py +1 -2
  6. tests/integration/test_models_integration.py +0 -1
  7. tests/unit/cli/projects/__init__.py +0 -0
  8. tests/unit/cli/projects/test_cli_project_data_connector_unit.py +28 -0
  9. tests/unit/cli/projects/test_cli_project_datastore_unit.py +26 -0
  10. tests/unit/cli/projects/test_cli_project_key_unit.py +26 -0
  11. tests/unit/cli/projects/test_cli_project_secret_unit.py +26 -0
  12. tests/unit/cli/projects/test_cli_project_service_unit.py +26 -0
  13. tests/unit/cli/projects/test_cli_project_unit.py +19 -0
  14. tests/unit/cloudbender/test_datastores_unit.py +1 -5
  15. tests/unit/conftest.py +77 -4
  16. tests/unit/test_checkpoints_unit.py +15 -23
  17. tests/unit/test_datasets_unit.py +15 -20
  18. tests/unit/test_models_unit.py +13 -16
  19. tests/unit/test_volumes_unit.py +3 -0
  20. trainml/__init__.py +1 -1
  21. trainml/checkpoints.py +14 -3
  22. trainml/cli/cloudbender/datastore.py +2 -7
  23. trainml/cli/project/__init__.py +3 -72
  24. trainml/cli/project/data_connector.py +61 -0
  25. trainml/cli/project/datastore.py +61 -0
  26. trainml/cli/project/service.py +61 -0
  27. trainml/cloudbender/data_connectors.py +8 -0
  28. trainml/cloudbender/datastores.py +9 -19
  29. trainml/cloudbender/nodes.py +44 -1
  30. trainml/cloudbender/providers.py +53 -0
  31. trainml/cloudbender/regions.py +48 -0
  32. trainml/datasets.py +14 -3
  33. trainml/exceptions.py +51 -0
  34. trainml/jobs.py +2 -13
  35. trainml/models.py +14 -3
  36. trainml/volumes.py +15 -3
  37. {trainml-0.5.9.dist-info → trainml-0.5.11.dist-info}/METADATA +1 -1
  38. {trainml-0.5.9.dist-info → trainml-0.5.11.dist-info}/RECORD +42 -40
  39. tests/integration/test_projects_integration.py +0 -44
  40. tests/unit/cli/cloudbender/test_cli_reservation_unit.py +0 -34
  41. tests/unit/cli/test_cli_project_unit.py +0 -42
  42. tests/unit/cloudbender/test_reservations_unit.py +0 -173
  43. tests/unit/test_auth.py +0 -30
  44. tests/unit/test_projects_unit.py +0 -320
  45. tests/unit/test_trainml.py +0 -54
  46. trainml/cli/cloudbender/reservation.py +0 -159
  47. trainml/cli/project.py +0 -149
  48. trainml/cloudbender/reservations.py +0 -126
  49. trainml/projects.py +0 -228
  50. {trainml-0.5.9.dist-info → trainml-0.5.11.dist-info}/LICENSE +0 -0
  51. {trainml-0.5.9.dist-info → trainml-0.5.11.dist-info}/WHEEL +0 -0
  52. {trainml-0.5.9.dist-info → trainml-0.5.11.dist-info}/entry_points.txt +0 -0
  53. {trainml-0.5.9.dist-info → trainml-0.5.11.dist-info}/top_level.txt +0 -0
@@ -3,6 +3,8 @@ from pytest import fixture
3
3
 
4
4
  @fixture(scope="module")
5
5
  async def project(trainml):
6
- project = await trainml.projects.create(name="New Project", copy_keys=False)
6
+ project = await trainml.projects.create(
7
+ name="New Project", copy_keys=False, copy_secrets=False
8
+ )
7
9
  yield project
8
10
  await project.remove()
@@ -0,0 +1,44 @@
1
+ import re
2
+ import sys
3
+ import asyncio
4
+ from pytest import mark, fixture
5
+
6
+ pytestmark = [mark.sdk, mark.integration, mark.projects]
7
+
8
+
9
+ @mark.create
10
+ @mark.asyncio
11
+ class ProjectDataConnectorsTests:
12
+ @fixture(scope="class")
13
+ async def project_data_connector(self, project):
14
+ data_connectors = await project.data_connectors.list()
15
+ yield data_connectors[0]
16
+
17
+ async def test_list_project_data_connectors(self, project, project_data_connector):
18
+ data_connectors = await project.data_connectors.list()
19
+ assert len(data_connectors) > 0
20
+
21
+ async def test_project_data_connector_properties(
22
+ self, project, project_data_connector
23
+ ):
24
+ assert isinstance(project_data_connector.project_uuid, str)
25
+ assert isinstance(project_data_connector.type, str)
26
+ assert isinstance(project_data_connector.name, str)
27
+ assert isinstance(project_data_connector.region_uuid, str)
28
+ assert project.id == project_data_connector.project_uuid
29
+
30
+ async def test_project_data_connector_str(self, project_data_connector):
31
+ string = str(project_data_connector)
32
+ regex = r"^{.*\"name\": \"" + project_data_connector.name + r"\".*}$"
33
+ assert isinstance(string, str)
34
+ assert re.match(regex, string)
35
+
36
+ async def test_project_data_connector_repr(self, project_data_connector):
37
+ string = repr(project_data_connector)
38
+ regex = (
39
+ r"^ProjectDataConnector\( trainml , \*\*{.*'name': '"
40
+ + project_data_connector.name
41
+ + r"'.*}\)$"
42
+ )
43
+ assert isinstance(string, str)
44
+ assert re.match(regex, string)
@@ -0,0 +1,42 @@
1
+ import re
2
+ import sys
3
+ import asyncio
4
+ from pytest import mark, fixture
5
+
6
+ pytestmark = [mark.sdk, mark.integration, mark.projects]
7
+
8
+
9
+ @mark.create
10
+ @mark.asyncio
11
+ class ProjectDatastoresTests:
12
+ @fixture(scope="class")
13
+ async def project_datastore(self, project):
14
+ datastores = await project.datastores.list()
15
+ yield datastores[0]
16
+
17
+ async def test_list_project_datastores(self, project):
18
+ datastores = await project.datastores.list()
19
+ assert len(datastores) > 0
20
+
21
+ async def test_project_datastore_properties(self, project, project_datastore):
22
+ assert isinstance(project_datastore.project_uuid, str)
23
+ assert isinstance(project_datastore.type, str)
24
+ assert isinstance(project_datastore.name, str)
25
+ assert isinstance(project_datastore.region_uuid, str)
26
+ assert project.id == project_datastore.project_uuid
27
+
28
+ async def test_project_datastore_str(self, project_datastore):
29
+ string = str(project_datastore)
30
+ regex = r"^{.*\"name\": \"" + project_datastore.name + r"\".*}$"
31
+ assert isinstance(string, str)
32
+ assert re.match(regex, string)
33
+
34
+ async def test_project_datastore_repr(self, project_datastore):
35
+ string = repr(project_datastore)
36
+ regex = (
37
+ r"^ProjectDatastore\( trainml , \*\*{.*'name': '"
38
+ + project_datastore.name
39
+ + r"'.*}\)$"
40
+ )
41
+ assert isinstance(string, str)
42
+ assert re.match(regex, string)
@@ -0,0 +1,44 @@
1
+ import re
2
+ import sys
3
+ import asyncio
4
+ from pytest import mark, fixture
5
+
6
+ pytestmark = [mark.sdk, mark.integration, mark.projects]
7
+
8
+
9
+ @mark.create
10
+ @mark.asyncio
11
+ class ProjectServicesTests:
12
+ @fixture(scope="class")
13
+ async def project_service(self, project):
14
+ services = await project.services.list()
15
+ yield services[0]
16
+
17
+ async def test_list_project_services(self, project):
18
+ services = await project.services.list()
19
+ assert len(services) > 0
20
+
21
+ async def test_project_service_properties(self, project, project_service):
22
+ assert isinstance(project_service.project_uuid, str)
23
+ assert isinstance(project_service.name, str)
24
+ assert isinstance(project_service.id, str)
25
+ assert isinstance(project_service.hostname, str)
26
+ assert isinstance(project_service.region_uuid, str)
27
+ assert isinstance(project_service.public, bool)
28
+ assert project.id == project_service.project_uuid
29
+
30
+ async def test_project_service_str(self, project_service):
31
+ string = str(project_service)
32
+ regex = r"^{.*\"name\": \"" + project_service.name + r"\".*}$"
33
+ assert isinstance(string, str)
34
+ assert re.match(regex, string)
35
+
36
+ async def test_project_service_repr(self, project_service):
37
+ string = repr(project_service)
38
+ regex = (
39
+ r"^ProjectService\( trainml , \*\*{.*'name': '"
40
+ + project_service.name
41
+ + r"'.*}\)$"
42
+ )
43
+ assert isinstance(string, str)
44
+ assert re.match(regex, string)
@@ -16,7 +16,7 @@ class GetCheckpointTests:
16
16
  source_type="git",
17
17
  source_uri="git@github.com:trainML/environment-tests.git",
18
18
  )
19
- checkpoint = await checkpoint.wait_for("ready", 120)
19
+ checkpoint = await checkpoint.wait_for("ready", 180)
20
20
  yield checkpoint
21
21
  await checkpoint.remove()
22
22
  checkpoint = await checkpoint.wait_for("archived", 60)
@@ -59,7 +59,6 @@ async def test_checkpoint_wasabi(trainml, capsys):
59
59
  name="CLI Automated Wasabi",
60
60
  source_type="wasabi",
61
61
  source_uri="s3://trainml-example/models/trainml-examples",
62
- capacity="10G",
63
62
  source_options=dict(endpoint_url="https://s3.wasabisys.com"),
64
63
  )
65
64
  checkpoint = await checkpoint.wait_for("ready", 300)
@@ -55,7 +55,6 @@ async def test_model_wasabi(trainml, capsys):
55
55
  name="CLI Automated Wasabi",
56
56
  source_type="wasabi",
57
57
  source_uri="s3://trainml-example/models/trainml-examples",
58
- capacity="10G",
59
58
  source_options=dict(endpoint_url="https://s3.wasabisys.com"),
60
59
  )
61
60
  model = await model.wait_for("ready", 300)
File without changes
@@ -0,0 +1,28 @@
1
+ import re
2
+ import json
3
+ import click
4
+ from unittest.mock import AsyncMock, patch, create_autospec
5
+ from pytest import mark, fixture, raises
6
+
7
+ pytestmark = [mark.cli, mark.unit, mark.projects]
8
+
9
+ from trainml.cli.project import data_connector as specimen
10
+ from trainml.projects import (
11
+ Project,
12
+ )
13
+
14
+
15
+ def test_list_data_connectors(runner, mock_project_data_connectors):
16
+ with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
17
+ project = create_autospec(Project)
18
+ mock_trainml.projects = AsyncMock()
19
+ mock_trainml.projects.get = AsyncMock(return_value=project)
20
+ mock_trainml.projects.get_current = AsyncMock(return_value=project)
21
+ project.data_connectors = AsyncMock()
22
+ project.data_connectors.list = AsyncMock(
23
+ return_value=mock_project_data_connectors
24
+ )
25
+ result = runner.invoke(specimen, ["list"])
26
+ print(result)
27
+ assert result.exit_code == 0
28
+ project.data_connectors.list.assert_called_once()
@@ -0,0 +1,26 @@
1
+ import re
2
+ import json
3
+ import click
4
+ from unittest.mock import AsyncMock, patch, create_autospec
5
+ from pytest import mark, fixture, raises
6
+
7
+ pytestmark = [mark.cli, mark.unit, mark.projects]
8
+
9
+ from trainml.cli.project import datastore as specimen
10
+ from trainml.projects import (
11
+ Project,
12
+ )
13
+
14
+
15
+ def test_list_datastores(runner, mock_project_datastores):
16
+ with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
17
+ project = create_autospec(Project)
18
+ mock_trainml.projects = AsyncMock()
19
+ mock_trainml.projects.get = AsyncMock(return_value=project)
20
+ mock_trainml.projects.get_current = AsyncMock(return_value=project)
21
+ project.datastores = AsyncMock()
22
+ project.datastores.list = AsyncMock(return_value=mock_project_datastores)
23
+ result = runner.invoke(specimen, ["list"])
24
+ print(result)
25
+ assert result.exit_code == 0
26
+ project.datastores.list.assert_called_once()
@@ -0,0 +1,26 @@
1
+ import re
2
+ import json
3
+ import click
4
+ from unittest.mock import AsyncMock, patch, create_autospec
5
+ from pytest import mark, fixture, raises
6
+
7
+ pytestmark = [mark.cli, mark.unit, mark.projects]
8
+
9
+ from trainml.cli.project import key as specimen
10
+ from trainml.projects import (
11
+ Project,
12
+ )
13
+
14
+
15
+ def test_list_keys(runner, mock_project_keys):
16
+ with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
17
+ project = create_autospec(Project)
18
+ mock_trainml.projects = AsyncMock()
19
+ mock_trainml.projects.get = AsyncMock(return_value=project)
20
+ mock_trainml.projects.get_current = AsyncMock(return_value=project)
21
+ project.keys = AsyncMock()
22
+ project.keys.list = AsyncMock(return_value=mock_project_keys)
23
+ result = runner.invoke(specimen, ["list"])
24
+ print(result)
25
+ assert result.exit_code == 0
26
+ project.keys.list.assert_called_once()
@@ -0,0 +1,26 @@
1
+ import re
2
+ import json
3
+ import click
4
+ from unittest.mock import AsyncMock, patch, create_autospec
5
+ from pytest import mark, fixture, raises
6
+
7
+ pytestmark = [mark.cli, mark.unit, mark.projects]
8
+
9
+ from trainml.cli.project import secret as specimen
10
+ from trainml.projects import (
11
+ Project,
12
+ )
13
+
14
+
15
+ def test_list_secrets(runner, mock_project_secrets):
16
+ with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
17
+ project = create_autospec(Project)
18
+ mock_trainml.projects = AsyncMock()
19
+ mock_trainml.projects.get = AsyncMock(return_value=project)
20
+ mock_trainml.projects.get_current = AsyncMock(return_value=project)
21
+ project.secrets = AsyncMock()
22
+ project.secrets.list = AsyncMock(return_value=mock_project_secrets)
23
+ result = runner.invoke(specimen, ["list"])
24
+ print(result)
25
+ assert result.exit_code == 0
26
+ project.secrets.list.assert_called_once()
@@ -0,0 +1,26 @@
1
+ import re
2
+ import json
3
+ import click
4
+ from unittest.mock import AsyncMock, patch, create_autospec
5
+ from pytest import mark, fixture, raises
6
+
7
+ pytestmark = [mark.cli, mark.unit, mark.projects]
8
+
9
+ from trainml.cli.project import service as specimen
10
+ from trainml.projects import (
11
+ Project,
12
+ )
13
+
14
+
15
+ def test_list_services(runner, mock_project_services):
16
+ with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
17
+ project = create_autospec(Project)
18
+ mock_trainml.projects = AsyncMock()
19
+ mock_trainml.projects.get = AsyncMock(return_value=project)
20
+ mock_trainml.projects.get_current = AsyncMock(return_value=project)
21
+ project.services = AsyncMock()
22
+ project.services.list = AsyncMock(return_value=mock_project_services)
23
+ result = runner.invoke(specimen, ["list"])
24
+ print(result)
25
+ assert result.exit_code == 0
26
+ project.services.list.assert_called_once()
@@ -0,0 +1,19 @@
1
+ import re
2
+ import json
3
+ import click
4
+ from unittest.mock import AsyncMock, patch, create_autospec
5
+ from pytest import mark, fixture, raises
6
+
7
+ pytestmark = [mark.cli, mark.unit, mark.projects]
8
+
9
+ from trainml.cli import project as specimen
10
+
11
+
12
+ def test_list(runner, mock_projects):
13
+ with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
14
+ mock_trainml.projects = AsyncMock()
15
+ mock_trainml.projects.list = AsyncMock(return_value=mock_projects)
16
+ result = runner.invoke(specimen, ["list"])
17
+ print(result)
18
+ assert result.exit_code == 0
19
+ mock_trainml.projects.list.assert_called_once()
@@ -119,8 +119,6 @@ class datastoreTests:
119
119
  assert isinstance(datastore.region_uuid, str)
120
120
  assert isinstance(datastore.type, str)
121
121
  assert isinstance(datastore.name, str)
122
- assert isinstance(datastore.uri, str)
123
- assert isinstance(datastore.root, str)
124
122
 
125
123
  def test_datastore_str(self, datastore):
126
124
  string = str(datastore)
@@ -131,9 +129,7 @@ class datastoreTests:
131
129
  def test_datastore_repr(self, datastore):
132
130
  string = repr(datastore)
133
131
  regex = (
134
- r"^Datastore\( trainml , \*\*{.*'store_id': '"
135
- + datastore.id
136
- + r"'.*}\)$"
132
+ r"^Datastore\( trainml , \*\*{.*'store_id': '" + datastore.id + r"'.*}\)$"
137
133
  )
138
134
  assert isinstance(string, str)
139
135
  assert re.match(regex, string)
tests/unit/conftest.py CHANGED
@@ -17,9 +17,11 @@ from trainml.projects import (
17
17
  Projects,
18
18
  Project,
19
19
  )
20
- from trainml.projects.datastores import ProjectDatastore
21
- from trainml.projects.services import ProjectService
22
- from trainml.projects.data_connectors import ProjectDataConnector
20
+ from trainml.projects.datastores import ProjectDatastores, ProjectDatastore
21
+ from trainml.projects.services import ProjectServices, ProjectService
22
+ from trainml.projects.data_connectors import ProjectDataConnectors, ProjectDataConnector
23
+ from trainml.projects.keys import ProjectKeys, ProjectKey
24
+ from trainml.projects.secrets import ProjectSecrets, ProjectSecret
23
25
 
24
26
  from trainml.cloudbender import Cloudbender
25
27
  from trainml.cloudbender.providers import Provider, Providers
@@ -1038,6 +1040,60 @@ def mock_device_configs():
1038
1040
  ]
1039
1041
 
1040
1042
 
1043
+ @fixture(
1044
+ scope="session",
1045
+ )
1046
+ def mock_project_keys():
1047
+ trainml = Mock()
1048
+ yield [
1049
+ ProjectKey(
1050
+ trainml,
1051
+ **{
1052
+ "project_uuid": "proj-id-1",
1053
+ "type": "aws",
1054
+ "key_id": "AKSHFKHDFS",
1055
+ "updatedAt": "2023-06-02T21:22:40.084Z",
1056
+ },
1057
+ ),
1058
+ ProjectKey(
1059
+ trainml,
1060
+ **{
1061
+ "project_uuid": "proj-id-1",
1062
+ "type": "gcp",
1063
+ "key_id": "credentials.json",
1064
+ "updatedAt": "2023-06-02T21:22:40.084Z",
1065
+ },
1066
+ ),
1067
+ ]
1068
+
1069
+
1070
+ @fixture(
1071
+ scope="session",
1072
+ )
1073
+ def mock_project_secrets():
1074
+ trainml = Mock()
1075
+ yield [
1076
+ ProjectSecret(
1077
+ trainml,
1078
+ **{
1079
+ "project_uuid": "proj-id-1",
1080
+ "name": "super_secret",
1081
+ "created_by": "User",
1082
+ "updatedAt": "2023-06-02T21:22:40.084Z",
1083
+ },
1084
+ ),
1085
+ ProjectSecret(
1086
+ trainml,
1087
+ **{
1088
+ "project_uuid": "proj-id-1",
1089
+ "name": "super_secret_2",
1090
+ "created_by": "User",
1091
+ "updatedAt": "2023-06-02T21:22:40.084Z",
1092
+ },
1093
+ ),
1094
+ ]
1095
+
1096
+
1041
1097
  @fixture(scope="function")
1042
1098
  def mock_trainml(
1043
1099
  mock_my_datasets,
@@ -1058,6 +1114,11 @@ def mock_trainml(
1058
1114
  mock_services,
1059
1115
  mock_data_connectors,
1060
1116
  mock_device_configs,
1117
+ mock_project_datastores,
1118
+ mock_project_services,
1119
+ mock_project_data_connectors,
1120
+ mock_project_keys,
1121
+ mock_project_secrets,
1061
1122
  ):
1062
1123
  trainml = create_autospec(TrainML)
1063
1124
  trainml.active_project = "proj-id-1"
@@ -1081,6 +1142,18 @@ def mock_trainml(
1081
1142
  trainml.environments.list = AsyncMock(return_value=mock_environments)
1082
1143
  trainml.jobs.list = AsyncMock(return_value=mock_jobs)
1083
1144
  trainml.projects.list = AsyncMock(return_value=mock_projects)
1145
+ trainml.projects.datastores = create_autospec(ProjectDatastores)
1146
+ trainml.projects.datastores.list = AsyncMock(return_value=mock_project_datastores)
1147
+ trainml.projects.services = create_autospec(ProjectServices)
1148
+ trainml.projects.services.list = AsyncMock(return_value=mock_project_services)
1149
+ trainml.projects.data_connectors = create_autospec(ProjectDataConnectors)
1150
+ trainml.projects.data_connectors.list = AsyncMock(
1151
+ return_value=mock_project_data_connectors
1152
+ )
1153
+ trainml.projects.keys = create_autospec(ProjectKeys)
1154
+ trainml.projects.keys.list = AsyncMock(return_value=mock_project_keys)
1155
+ trainml.projects.secrets = create_autospec(ProjectSecrets)
1156
+ trainml.projects.secrets.list = AsyncMock(return_value=mock_project_secrets)
1084
1157
 
1085
1158
  trainml.cloudbender = create_autospec(Cloudbender)
1086
1159
 
@@ -1090,7 +1163,7 @@ def mock_trainml(
1090
1163
  trainml.cloudbender.regions.list = AsyncMock(return_value=mock_regions)
1091
1164
  trainml.cloudbender.nodes = create_autospec(Nodes)
1092
1165
  trainml.cloudbender.nodes.list = AsyncMock(return_value=mock_nodes)
1093
- trainml.cloudbender.devices = create_autospec(Nodes)
1166
+ trainml.cloudbender.devices = create_autospec(Devices)
1094
1167
  trainml.cloudbender.devices.list = AsyncMock(return_value=mock_devices)
1095
1168
  trainml.cloudbender.datastores = create_autospec(Datastores)
1096
1169
  trainml.cloudbender.datastores.list = AsyncMock(return_value=mock_datastores)
@@ -44,9 +44,7 @@ class CheckpointsTests:
44
44
  api_response = dict()
45
45
  mock_trainml._query = AsyncMock(return_value=api_response)
46
46
  await checkpoints.get("1234")
47
- mock_trainml._query.assert_called_once_with(
48
- "/checkpoint/1234", "GET", dict()
49
- )
47
+ mock_trainml._query.assert_called_once_with("/checkpoint/1234", "GET", dict())
50
48
 
51
49
  @mark.asyncio
52
50
  async def test_list_checkpoints(
@@ -57,9 +55,7 @@ class CheckpointsTests:
57
55
  api_response = dict()
58
56
  mock_trainml._query = AsyncMock(return_value=api_response)
59
57
  await checkpoints.list()
60
- mock_trainml._query.assert_called_once_with(
61
- "/checkpoint", "GET", dict()
62
- )
58
+ mock_trainml._query.assert_called_once_with("/checkpoint", "GET", dict())
63
59
 
64
60
  @mark.asyncio
65
61
  async def test_remove_checkpoint(
@@ -86,11 +82,13 @@ class CheckpointsTests:
86
82
  name="new checkpoint",
87
83
  source_type="aws",
88
84
  source_uri="s3://trainml-examples/checkpoints/resnet50",
85
+ type="evefs",
89
86
  )
90
87
  api_response = {
91
88
  "project_uuid": "cus-id-1",
92
89
  "checkpoint_uuid": "checkpoint-id-1",
93
90
  "name": "new checkpoint",
91
+ "type": "evefs",
94
92
  "status": "new",
95
93
  "source_type": "aws",
96
94
  "source_uri": "s3://trainml-examples/checkpoints/resnet50",
@@ -135,7 +133,9 @@ class CheckpointTests:
135
133
 
136
134
  @mark.asyncio
137
135
  async def test_checkpoint_get_log_url(self, checkpoint, mock_trainml):
138
- api_response = "https://trainml-jobs-dev.s3.us-east-2.amazonaws.com/1/logs/first_one.zip"
136
+ api_response = (
137
+ "https://trainml-jobs-dev.s3.us-east-2.amazonaws.com/1/logs/first_one.zip"
138
+ )
139
139
  mock_trainml._query = AsyncMock(return_value=api_response)
140
140
  response = await checkpoint.get_log_url()
141
141
  mock_trainml._query.assert_called_once_with(
@@ -163,7 +163,9 @@ class CheckpointTests:
163
163
  async def test_checkpoint_get_connection_utility_url(
164
164
  self, checkpoint, mock_trainml
165
165
  ):
166
- api_response = "https://trainml-jobs-dev.s3.us-east-2.amazonaws.com/1/vpn/first_one.zip"
166
+ api_response = (
167
+ "https://trainml-jobs-dev.s3.us-east-2.amazonaws.com/1/vpn/first_one.zip"
168
+ )
167
169
  mock_trainml._query = AsyncMock(return_value=api_response)
168
170
  response = await checkpoint.get_connection_utility_url()
169
171
  mock_trainml._query.assert_called_once_with(
@@ -338,9 +340,7 @@ class CheckpointTests:
338
340
  assert response.id == "data-id-1"
339
341
 
340
342
  @mark.asyncio
341
- async def test_checkpoint_wait_for_successful(
342
- self, checkpoint, mock_trainml
343
- ):
343
+ async def test_checkpoint_wait_for_successful(self, checkpoint, mock_trainml):
344
344
  api_response = {
345
345
  "customer_uuid": "cus-id-1",
346
346
  "checkpoint_uuid": "data-id-1",
@@ -373,9 +373,7 @@ class CheckpointTests:
373
373
  mock_trainml._query.assert_not_called()
374
374
 
375
375
  @mark.asyncio
376
- async def test_checkpoint_wait_for_incorrect_status(
377
- self, checkpoint, mock_trainml
378
- ):
376
+ async def test_checkpoint_wait_for_incorrect_status(self, checkpoint, mock_trainml):
379
377
  api_response = None
380
378
  mock_trainml._query = AsyncMock(return_value=api_response)
381
379
  with raises(SpecificationError):
@@ -383,9 +381,7 @@ class CheckpointTests:
383
381
  mock_trainml._query.assert_not_called()
384
382
 
385
383
  @mark.asyncio
386
- async def test_checkpoint_wait_for_with_delay(
387
- self, checkpoint, mock_trainml
388
- ):
384
+ async def test_checkpoint_wait_for_with_delay(self, checkpoint, mock_trainml):
389
385
  api_response_initial = dict(
390
386
  checkpoint_uuid="1",
391
387
  name="first one",
@@ -441,9 +437,7 @@ class CheckpointTests:
441
437
  self, checkpoint, mock_trainml
442
438
  ):
443
439
  mock_trainml._query = AsyncMock(
444
- side_effect=ApiError(
445
- 404, dict(errorMessage="Checkpoint Not Found")
446
- )
440
+ side_effect=ApiError(404, dict(errorMessage="Checkpoint Not Found"))
447
441
  )
448
442
  await checkpoint.wait_for("archived")
449
443
  mock_trainml._query.assert_called()
@@ -453,9 +447,7 @@ class CheckpointTests:
453
447
  self, checkpoint, mock_trainml
454
448
  ):
455
449
  mock_trainml._query = AsyncMock(
456
- side_effect=ApiError(
457
- 404, dict(errorMessage="Checkpoint Not Found")
458
- )
450
+ side_effect=ApiError(404, dict(errorMessage="Checkpoint Not Found"))
459
451
  )
460
452
  with raises(ApiError):
461
453
  await checkpoint.wait_for("ready")