trainml 0.5.0__py3-none-any.whl → 0.5.2__py3-none-any.whl

tests/integration/test_jobs_integration.py

@@ -100,7 +100,7 @@ class JobLifeCycleTests:
         training_job = await job.copy(
             name="CLI Automated Tests - Job Convert",
             type="training",
-            workers=["python $TRAINML_MODEL_PATH/tensorflow/main.py"],
+            workers=["python $ML_MODEL_PATH/tensorflow/main.py"],
             data=dict(
                 datasets=[
                     dict(
@@ -191,8 +191,7 @@ class JobAPIResourceValidationTests:
                 disk_size=10,
             )
         assert (
-            "Invalid Request - CPU Count must be a multiple of 4"
-            in error.value.message
+            "Invalid Request - CPU Count must be a multiple of 4" in error.value.message
         )
 
     async def test_invalid_gpu_count_for_cpu(self, trainml):
@@ -373,7 +372,7 @@ class JobIOTests:
             type="training",
             gpu_types=["gtx1060"],
             disk_size=10,
-            workers=["python $TRAINML_MODEL_PATH/tensorflow/main.py"],
+            workers=["python $ML_MODEL_PATH/tensorflow/main.py"],
             environment=dict(
                 type="DEEPLEARNING_PY310",
                 env=[
@@ -431,7 +430,7 @@ class JobIOTests:
             source_uri="git@github.com:trainML/environment-tests.git",
         )
         await model.wait_for("ready", 300)
-        assert model.size >= 500000
+        assert model.size >= 200000
 
         job = await trainml.jobs.create(
             "CLI Automated Tests - Training With trainML Model Output",
@@ -440,7 +439,7 @@ class JobIOTests:
             gpu_count=1,
             cpu_count=8,
             disk_size=10,
-            worker_commands=["python $TRAINML_MODEL_PATH/tensorflow/main.py"],
+            worker_commands=["python $ML_MODEL_PATH/tensorflow/main.py"],
             data=dict(
                 datasets=[
                     dict(
@@ -548,22 +547,21 @@ class JobTypeTests:
             ),
             environment=dict(
                 type="CUSTOM",
-                custom_image="tensorflow/tensorflow:2.10.0-gpu",
+                custom_image="tensorflow/tensorflow:2.13.0-gpu",
                 packages=dict(
                     pip=[
-                        "tensorflow_addons==0.18.0",
                         "matplotlib",
                         "scipy",
-                        "tensorflow_hub==0.12.0",
+                        "tensorflow_hub",
                         "keras_applications",
                         "keras_preprocessing",
-                        "protobuf==3.20.1",
-                        "typing-extensions==4.4.0",
+                        "protobuf",
+                        "typing-extensions",
                     ]
                 ),
             ),
             worker_commands=[
-                "python $TRAINML_MODEL_PATH/tensorflow/main.py",
+                "python $ML_MODEL_PATH/tensorflow/main.py",
             ],
             data=dict(
                 datasets=[
@@ -589,10 +587,7 @@ class JobTypeTests:
         sys.stderr.write(captured.err)
         assert "Epoch 1/2" in captured.out
         assert "Epoch 2/2" in captured.out
-        assert (
-            "Uploading s3://trainml-example/output/resnet_cifar10"
-            in captured.out
-        )
+        assert "Uploading s3://trainml-example/output/resnet_cifar10" in captured.out
         assert (
             "upload: ./model.ckpt-0002.data-00000-of-00001 to s3://trainml-example/output/resnet_cifar10/model.ckpt-0002.data-00000-of-00001"
             in captured.out
@@ -615,7 +610,7 @@ class JobFeatureTests:
                 source_uri="git@github.com:trainML/environment-tests.git",
             ),
             worker_commands=[
-                "python $TRAINML_MODEL_PATH/pytorch/main.py",
+                "python $ML_MODEL_PATH/pytorch/main.py",
             ],
             data=dict(
                 datasets=[dict(id="MNIST", public=True)],
@@ -641,7 +636,7 @@ class JobFeatureTests:
             gpu_count=1,
             disk_size=10,
             workers=[
-                "python $TRAINML_MODEL_PATH/tensorflow/main.py",
+                "python $ML_MODEL_PATH/tensorflow/main.py",
             ],
             data=dict(
                 input_type="wasabi",

tests/unit/cli/cloudbender/test_cli_device_unit.py

@@ -0,0 +1,38 @@
+import re
+import json
+import click
+from unittest.mock import AsyncMock, patch
+from pytest import mark, fixture, raises
+
+pytestmark = [mark.cli, mark.unit, mark.cloudbender, mark.devices]
+
+from trainml.cli.cloudbender import device as specimen
+from trainml.cloudbender.devices import Device
+
+
+def test_list(runner, mock_devices):
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+        mock_trainml.cloudbender = AsyncMock()
+        mock_trainml.cloudbender.devices = AsyncMock()
+        mock_trainml.cloudbender.devices.list = AsyncMock(
+            return_value=mock_devices
+        )
+        result = runner.invoke(
+            specimen,
+            args=["list", "--provider=prov-id-1", "--region=reg-id-1"],
+        )
+        assert result.exit_code == 0
+        mock_trainml.cloudbender.devices.list.assert_called_once_with(
+            provider_uuid="prov-id-1", region_uuid="reg-id-1"
+        )
+
+
+def test_list_no_provider(runner, mock_devices):
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+        mock_trainml.cloudbender = AsyncMock()
+        mock_trainml.cloudbender.devices = AsyncMock()
+        mock_trainml.cloudbender.devices.list = AsyncMock(
+            return_value=mock_devices
+        )
+        result = runner.invoke(specimen, ["list"])
+        assert result.exit_code != 0

tests/unit/cloudbender/test_datastores_unit.py

@@ -151,3 +151,23 @@ class datastoreTests:
         mock_trainml._query.assert_called_once_with(
             "/provider/1/region/a/datastore/x", "DELETE"
         )
+
+    @mark.asyncio
+    async def test_datastore_refresh(self, datastore, mock_trainml):
+        api_response = {
+            "provider_uuid": "provider-id-1",
+            "region_uuid": "region-id-1",
+            "store_id": "store-id-1",
+            "name": "On-Prem Datastore",
+            "type": "nfs",
+            "uri": "192.168.0.50",
+            "root": "/exports",
+            "createdAt": "2020-12-31T23:59:59.000Z",
+        }
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        response = await datastore.refresh()
+        mock_trainml._query.assert_called_once_with(
+            f"/provider/1/region/a/datastore/x", "GET"
+        )
+        assert datastore.id == "store-id-1"
+        assert response.id == "store-id-1"

tests/unit/cloudbender/test_device_configs_unit.py

@@ -150,3 +150,24 @@ class device_configTests:
         mock_trainml._query.assert_called_once_with(
             "/provider/1/region/a/device/config/x", "DELETE"
         )
+
+    @mark.asyncio
+    async def test_device_config_refresh(self, device_config, mock_trainml):
+        api_response = {
+            "provider_uuid": "provider-id-1",
+            "region_uuid": "region-id-1",
+            "config_id": "device_config-id-1",
+            "name": "IoT 1",
+            "model_uuid": "model-id-1",
+            "model_project_uuid": "proj-id-1",
+            "image": "nvidia/cuda",
+            "command": "python run.py",
+            "createdAt": "2020-12-31T23:59:59.000Z",
+        }
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        response = await device_config.refresh()
+        mock_trainml._query.assert_called_once_with(
+            f"/provider/1/region/a/device/config/x", "GET"
+        )
+        assert device_config.id == "device_config-id-1"
+        assert response.id == "device_config-id-1"

tests/unit/cloudbender/test_devices_unit.py

@@ -0,0 +1,270 @@
+import re
+import json
+import logging
+from unittest.mock import AsyncMock, patch
+from pytest import mark, fixture, raises
+from aiohttp import WSMessage, WSMsgType
+
+import trainml.cloudbender.devices as specimen
+from trainml.exceptions import (
+    ApiError,
+    SpecificationError,
+    TrainMLException,
+)
+
+pytestmark = [mark.sdk, mark.unit, mark.cloudbender, mark.devices]
+
+
+@fixture
+def devices(mock_trainml):
+    yield specimen.Devices(mock_trainml)
+
+
+@fixture
+def device(mock_trainml):
+    yield specimen.Device(
+        mock_trainml,
+        provider_uuid="1",
+        region_uuid="a",
+        device_id="x",
+        type="device",
+        service="compute",
+        friendly_name="hq-orin-01",
+        hostname="hq-orin-01",
+        status="active",
+        online=True,
+        maintenance_mode=False,
+        job_status="stopped",
+        job_last_deployed="2023-06-02T21:22:40.084Z",
+        job_config_id="job-id-1",
+        job_config_revision="1685740490096",
+        device_config_id="conf-id-2",
+    )
+
+
+class RegionsTests:
+    @mark.asyncio
+    async def test_get_device(
+        self,
+        devices,
+        mock_trainml,
+    ):
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await devices.get("1234", "5687", "91011")
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1234/region/5687/device/91011", "GET", {}
+        )
+
+    @mark.asyncio
+    async def test_list_devices(
+        self,
+        devices,
+        mock_trainml,
+    ):
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await devices.list("1234", "5687")
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1234/region/5687/device", "GET", {}
+        )
+
+    @mark.asyncio
+    async def test_remove_device(
+        self,
+        devices,
+        mock_trainml,
+    ):
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await devices.remove("1234", "4567", "8910")
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1234/region/4567/device/8910", "DELETE", {}
+        )
+
+    @mark.asyncio
+    async def test_create_device(self, devices, mock_trainml):
+        requested_config = dict(
+            provider_uuid="provider-id-1",
+            region_uuid="region-id-1",
+            friendly_name="phys-device",
+            hostname="phys-device",
+            minion_id="asdf",
+        )
+        expected_payload = dict(
+            friendly_name="phys-device",
+            hostname="phys-device",
+            minion_id="asdf",
+            type="device",
+            service="compute",
+        )
+        api_response = {
+            "provider_uuid": "provider-id-1",
+            "region_uuid": "region-id-1",
+            "device_id": "rig-id-1",
+            "name": "phys-device",
+            "type": "device",
+            "service": "compute",
+            "status": "new",
+            "online": False,
+            "maintenance_mode": True,
+            "job_status": "stopped",
+            "job_last_deployed": "2023-06-02T21:22:40.084Z",
+            "job_config_id": "job-id-1",
+            "job_config_revision": "1685740490096",
+            "device_config_id": "conf-id-1",
+            "createdAt": "2020-12-31T23:59:59.000Z",
+        }
+
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        response = await devices.create(**requested_config)
+        mock_trainml._query.assert_called_once_with(
+            "/provider/provider-id-1/region/region-id-1/device",
+            "POST",
+            None,
+            expected_payload,
+        )
+        assert response.id == "rig-id-1"
+
+
+class deviceTests:
+    def test_device_properties(self, device):
+        assert isinstance(device.id, str)
+        assert isinstance(device.provider_uuid, str)
+        assert isinstance(device.region_uuid, str)
+        assert isinstance(device.name, str)
+        assert isinstance(device.hostname, str)
+        assert isinstance(device.status, str)
+        assert isinstance(device.online, bool)
+        assert isinstance(device.maintenance_mode, bool)
+        assert isinstance(device.device_config_id, str)
+        assert isinstance(device.job_status, str)
+        assert isinstance(device.job_last_deployed, str)
+        assert isinstance(device.job_config_id, str)
+        assert isinstance(device.job_config_revision, str)
+
+    def test_device_str(self, device):
+        string = str(device)
+        regex = r"^{.*\"device_id\": \"" + device.id + r"\".*}$"
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    def test_device_repr(self, device):
+        string = repr(device)
+        regex = (
+            r"^Device\( trainml , \*\*{.*'device_id': '"
+            + device.id
+            + r"'.*}\)$"
+        )
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    def test_device_bool(self, device, mock_trainml):
+        empty_device = specimen.Device(mock_trainml)
+        assert bool(device)
+        assert not bool(empty_device)
+
+    @mark.asyncio
+    async def test_device_remove(self, device, mock_trainml):
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await device.remove()
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1/region/a/device/x", "DELETE"
+        )
+
+    @mark.asyncio
+    async def test_device_refresh(self, device, mock_trainml):
+        api_response = {
+            "provider_uuid": "provider-id-1",
+            "region_uuid": "region-id-1",
+            "device_id": "device-id-1",
+            "name": "phys-device",
+            "type": "device",
+            "service": "compute",
+            "status": "new",
+            "online": False,
+            "maintenance_mode": True,
+            "job_status": "stopped",
+            "job_last_deployed": "2023-06-02T21:22:40.084Z",
+            "job_config_id": "job-id-1",
+            "job_config_revision": "1685740490096",
+            "device_config_id": "conf-id-1",
+            "createdAt": "2020-12-31T23:59:59.000Z",
+        }
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        response = await device.refresh()
+        mock_trainml._query.assert_called_once_with(
+            f"/provider/1/region/a/device/x", "GET"
+        )
+        assert device.id == "device-id-1"
+        assert response.id == "device-id-1"
+
+    @mark.asyncio
+    async def test_device_toggle_maintenance(self, device, mock_trainml):
+        api_response = None
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await device.toggle_maintenance()
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1/region/a/device/x/maintenance", "PATCH"
+        )
+
+    @mark.asyncio
+    async def test_device_run_action(self, device, mock_trainml):
+        api_response = None
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await device.run_action(command="report")
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1/region/a/device/x/action",
+            "POST",
+            None,
+            dict(command="report"),
+        )
+
+    @mark.asyncio
+    async def test_device_set_config(self, device, mock_trainml):
+        api_response = {
+            "provider_uuid": "provider-id-1",
+            "region_uuid": "region-id-1",
+            "device_id": "device-id-1",
+            "name": "phys-device",
+            "type": "device",
+            "service": "compute",
+            "status": "new",
+            "online": False,
+            "maintenance_mode": True,
+            "job_status": "stopped",
+            "job_last_deployed": "2023-06-02T21:22:40.084Z",
+            "job_config_id": "job-id-1",
+            "job_config_revision": "1685740490096",
+            "device_config_id": "config-id-1",
+            "createdAt": "2020-12-31T23:59:59.000Z",
+        }
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        response = await device.set_config(device_config_id="config-id-1")
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1/region/a/device/x",
+            "PATCH",
+            None,
+            dict(device_config_id="config-id-1"),
+        )
+        assert device.id == "device-id-1"
+        assert response.id == "device-id-1"
+
+    @mark.asyncio
+    async def test_device_deploy_endpoint(self, device, mock_trainml):
+        api_response = None
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await device.deploy_endpoint()
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1/region/a/device/x/deploy", "PUT"
+        )
+
+    @mark.asyncio
+    async def test_device_stop_endpoint(self, device, mock_trainml):
+        api_response = None
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await device.stop_endpoint()
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1/region/a/device/x/stop", "PUT"
+        )

tests/unit/cloudbender/test_nodes_unit.py

@@ -157,3 +157,46 @@ class nodeTests:
         mock_trainml._query.assert_called_once_with(
             "/provider/1/region/a/node/x", "DELETE"
         )
+
+    @mark.asyncio
+    async def test_node_refresh(self, node, mock_trainml):
+        api_response = {
+            "provider_uuid": "provider-id-1",
+            "region_uuid": "region-id-1",
+            "rig_uuid": "rig-id-1",
+            "name": "phys-node",
+            "type": "permanent",
+            "service": "compute",
+            "status": "new",
+            "online": False,
+            "maintenance_mode": True,
+            "createdAt": "2020-12-31T23:59:59.000Z",
+        }
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        response = await node.refresh()
+        mock_trainml._query.assert_called_once_with(
+            f"/provider/1/region/a/node/x", "GET"
+        )
+        assert node.id == "rig-id-1"
+        assert response.id == "rig-id-1"
+
+    @mark.asyncio
+    async def test_node_toggle_maintenance(self, node, mock_trainml):
+        api_response = None
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await node.toggle_maintenance()
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1/region/a/node/x/maintenance", "PATCH"
+        )
+
+    @mark.asyncio
+    async def test_node_run_action(self, node, mock_trainml):
+        api_response = None
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await node.run_action(command="report")
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1/region/a/node/x/action",
+            "POST",
+            None,
+            dict(command="report"),
+        )

tests/unit/cloudbender/test_providers_unit.py

@@ -123,3 +123,19 @@ class providerTests:
         mock_trainml._query = AsyncMock(return_value=api_response)
         await provider.remove()
         mock_trainml._query.assert_called_once_with("/provider/1", "DELETE")
+
+    @mark.asyncio
+    async def test_provider_refresh(self, provider, mock_trainml):
+        api_response = {
+            "customer_uuid": "cust-id-1",
+            "provider_uuid": "provider-id-1",
+            "type": "new provider",
+            "credits": 0.0,
+            "payment_mode": "credits",
+            "createdAt": "2020-12-31T23:59:59.000Z",
+        }
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        response = await provider.refresh()
+        mock_trainml._query.assert_called_once_with(f"/provider/1", "GET")
+        assert provider.id == "provider-id-1"
+        assert response.id == "provider-id-1"

tests/unit/cloudbender/test_regions_unit.py

@@ -139,6 +139,24 @@ class regionTests:
             "/provider/1/region/a", "DELETE"
         )
 
+    @mark.asyncio
+    async def test_region_refresh(self, region, mock_trainml):
+        api_response = {
+            "provider_uuid": "provider-id-1",
+            "region_uuid": "region-id-1",
+            "provider_type": "physical",
+            "name": "phys-region",
+            "status": "new",
+            "createdAt": "2020-12-31T23:59:59.000Z",
+        }
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        response = await region.refresh()
+        mock_trainml._query.assert_called_once_with(
+            f"/provider/1/region/a", "GET"
+        )
+        assert region.id == "region-id-1"
+        assert response.id == "region-id-1"
+
     @mark.asyncio
     async def test_region_stage_dataset(self, region, mock_trainml):
         api_response = dict()

tests/unit/cloudbender/test_reservations_unit.py

@@ -151,3 +151,23 @@ class reservationTests:
         mock_trainml._query.assert_called_once_with(
             "/provider/1/region/a/reservation/x", "DELETE"
         )
+
+    @mark.asyncio
+    async def test_reservation_refresh(self, reservation, mock_trainml):
+        api_response = {
+            "provider_uuid": "provider-id-1",
+            "region_uuid": "region-id-1",
+            "reservation_id": "reservation-id-1",
+            "name": "On-Prem Reservation",
+            "type": "port",
+            "resource": "8001",
+            "hostname": "service.local",
+            "createdAt": "2020-12-31T23:59:59.000Z",
+        }
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        response = await reservation.refresh()
+        mock_trainml._query.assert_called_once_with(
+            f"/provider/1/region/a/reservation/x", "GET"
+        )
+        assert reservation.id == "reservation-id-1"
+        assert response.id == "reservation-id-1"