trainml 0.5.6__py3-none-any.whl → 0.5.8__py3-none-any.whl

tests/integration/test_jobs_integration.py

@@ -6,10 +6,19 @@ import asyncio
 import aiohttp
 from pytest import mark, fixture, raises
 from trainml.exceptions import ApiError
+from urllib.parse import urlparse
 
 pytestmark = [mark.sdk, mark.integration, mark.jobs]
 
 
+def extract_domain_suffix(hostname):
+    parts = hostname.split(".")
+    if len(parts) >= 2:
+        return ".".join(parts[-2:])
+    else:
+        return None
+
+
 @fixture(scope="class")
 async def job(trainml):
     job = await trainml.jobs.create(
@@ -34,6 +43,8 @@ class JobLifeCycleTests:
         assert job.status != "running"
         job = await job.wait_for("running")
         assert job.status == "running"
+        assert job.url
+        assert extract_domain_suffix(urlparse(job.url).hostname) == "proximl.cloud"
 
     async def test_stop_job(self, job):
         assert job.status == "running"
@@ -518,6 +529,7 @@ class JobIOTests:
 @mark.asyncio
 class JobTypeTests:
     async def test_endpoint(self, trainml):
+
         job = await trainml.jobs.create(
             "CLI Automated Tests - Endpoint",
             type="endpoint",
@@ -544,6 +556,7 @@ class JobTypeTests:
         await job.wait_for("running")
         await job.refresh()
         assert job.url
+        assert extract_domain_suffix(urlparse(job.url).hostname) == "proximl.cloud"
         tries = 0
         await asyncio.sleep(30)
         async with aiohttp.ClientSession() as session:

tests/unit/cli/cloudbender/test_cli_service_unit.py

@@ -0,0 +1,34 @@
+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.services]
+
+from trainml.cli.cloudbender import service as specimen
+from trainml.cloudbender.services import Service
+
+
+def test_list(runner, mock_services):
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+        mock_trainml.cloudbender = AsyncMock()
+        mock_trainml.cloudbender.services = AsyncMock()
+        mock_trainml.cloudbender.services.list = AsyncMock(return_value=mock_services)
+        result = runner.invoke(
+            specimen,
+            args=["list", "--provider=prov-id-1", "--region=reg-id-1"],
+        )
+        assert result.exit_code == 0
+        mock_trainml.cloudbender.services.list.assert_called_once_with(
+            provider_uuid="prov-id-1", region_uuid="reg-id-1"
+        )
+
+
+def test_list_no_provider(runner, mock_services):
+    with patch("trainml.cli.TrainML", new=AsyncMock) as mock_trainml:
+        mock_trainml.cloudbender = AsyncMock()
+        mock_trainml.cloudbender.services = AsyncMock()
+        mock_trainml.cloudbender.services.list = AsyncMock(return_value=mock_services)
+        result = runner.invoke(specimen, ["list"])
+        assert result.exit_code != 0

tests/unit/cloudbender/test_data_connectors_unit.py

@@ -0,0 +1,176 @@
+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.data_connectors as specimen
+from trainml.exceptions import (
+    ApiError,
+    SpecificationError,
+    TrainMLException,
+)
+
+pytestmark = [mark.sdk, mark.unit, mark.cloudbender, mark.data_connectors]
+
+
+@fixture
+def data_connectors(mock_trainml):
+    yield specimen.DataConnectors(mock_trainml)
+
+
+@fixture
+def data_connector(mock_trainml):
+    yield specimen.DataConnector(
+        mock_trainml,
+        provider_uuid="1",
+        region_uuid="a",
+        connector_id="x",
+        name="On-Prem Data Connector",
+        type="custom",
+        cidr="192.168.0.50/32",
+        port="443",
+        protocol="tcp",
+    )
+
+
+class RegionsTests:
+    @mark.asyncio
+    async def test_get_data_connector(
+        self,
+        data_connectors,
+        mock_trainml,
+    ):
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await data_connectors.get("1234", "5687", "91011")
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1234/region/5687/data_connector/91011", "GET", {}
+        )
+
+    @mark.asyncio
+    async def test_list_data_connectors(
+        self,
+        data_connectors,
+        mock_trainml,
+    ):
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await data_connectors.list("1234", "5687")
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1234/region/5687/data_connector", "GET", {}
+        )
+
+    @mark.asyncio
+    async def test_remove_data_connector(
+        self,
+        data_connectors,
+        mock_trainml,
+    ):
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await data_connectors.remove("1234", "4567", "8910")
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1234/region/4567/data_connector/8910", "DELETE", {}
+        )
+
+    @mark.asyncio
+    async def test_create_data_connector(self, data_connectors, mock_trainml):
+        requested_config = dict(
+            provider_uuid="provider-id-1",
+            region_uuid="region-id-1",
+            name="On-Prem DataConnector",
+            type="custom",
+            cidr="192.168.0.50/32",
+            port="443",
+            protocol="tcp",
+        )
+        expected_payload = dict(
+            name="On-Prem DataConnector",
+            type="custom",
+            cidr="192.168.0.50/32",
+            port="443",
+            protocol="tcp",
+        )
+        api_response = {
+            "provider_uuid": "provider-id-1",
+            "region_uuid": "region-id-1",
+            "connector_id": "connector-id-1",
+            "name": "On-Prem DataConnector",
+            "type": "custom",
+            "cidr": "192.168.0.50/32",
+            "port": "443",
+            "protocol": "tcp",
+            "createdAt": "2020-12-31T23:59:59.000Z",
+        }
+
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        response = await data_connectors.create(**requested_config)
+        mock_trainml._query.assert_called_once_with(
+            "/provider/provider-id-1/region/region-id-1/data_connector",
+            "POST",
+            None,
+            expected_payload,
+        )
+        assert response.id == "connector-id-1"
+
+
+class DataConnectorTests:
+    def test_data_connector_properties(self, data_connector):
+        assert isinstance(data_connector.id, str)
+        assert isinstance(data_connector.provider_uuid, str)
+        assert isinstance(data_connector.region_uuid, str)
+        assert isinstance(data_connector.type, str)
+        assert isinstance(data_connector.name, str)
+
+    def test_data_connector_str(self, data_connector):
+        string = str(data_connector)
+        regex = r"^{.*\"connector_id\": \"" + data_connector.id + r"\".*}$"
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    def test_data_connector_repr(self, data_connector):
+        string = repr(data_connector)
+        regex = (
+            r"^DataConnector\( trainml , \*\*{.*'connector_id': '"
+            + data_connector.id
+            + r"'.*}\)$"
+        )
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    def test_data_connector_bool(self, data_connector, mock_trainml):
+        empty_data_connector = specimen.DataConnector(mock_trainml)
+        assert bool(data_connector)
+        assert not bool(empty_data_connector)
+
+    @mark.asyncio
+    async def test_data_connector_remove(self, data_connector, mock_trainml):
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await data_connector.remove()
+        mock_trainml._query.assert_called_once_with(
+            "/provider/1/region/a/data_connector/x", "DELETE"
+        )
+
+    @mark.asyncio
+    async def test_data_connector_refresh(self, data_connector, mock_trainml):
+        api_response = {
+            "provider_uuid": "provider-id-1",
+            "region_uuid": "region-id-1",
+            "connector_id": "connector-id-1",
+            "name": "On-Prem DataConnector",
+            "type": "custom",
+            "cidr": "192.168.0.50/32",
+            "port": "443",
+            "protocol": "tcp",
+            "createdAt": "2020-12-31T23:59:59.000Z",
+        }
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        response = await data_connector.refresh()
+        mock_trainml._query.assert_called_once_with(
+            f"/provider/1/region/a/data_connector/x", "GET"
+        )
+        assert data_connector.id == "connector-id-1"
+        assert response.id == "connector-id-1"

tests/unit/cloudbender/test_services_unit.py

@@ -28,6 +28,7 @@ def service(mock_trainml):
         region_uuid="a",
         service_id="x",
         name="On-Prem Service",
+        type="https",
         public=False,
         hostname="app1.proximl.cloud",
     )
@@ -79,10 +80,12 @@ class RegionsTests:
             provider_uuid="provider-id-1",
             region_uuid="region-id-1",
             name="On-Prem Service",
+            type="https",
             public=False,
         )
         expected_payload = dict(
             name="On-Prem Service",
+            type="https",
             public=False,
         )
         api_response = {
@@ -90,6 +93,7 @@ class RegionsTests:
             "region_uuid": "region-id-1",
             "service_id": "service-id-1",
             "name": "On-Prem Service",
+            "type": "https",
             "public": False,
             "hostname": "app1.proximl.cloud",
             "createdAt": "2020-12-31T23:59:59.000Z",
@@ -114,6 +118,7 @@ class serviceTests:
         assert isinstance(service.public, bool)
         assert isinstance(service.name, str)
         assert isinstance(service.hostname, str)
+        assert isinstance(service.type, str)
 
     def test_service_str(self, service):
         string = str(service)
@@ -148,6 +153,7 @@ class serviceTests:
             "region_uuid": "region-id-1",
             "service_id": "service-id-1",
             "name": "On-Prem Service",
+            "type": "https",
             "public": False,
             "hostname": "app1.proximl.cloud",
             "createdAt": "2020-12-31T23:59:59.000Z",

tests/unit/test_projects_unit.py

@@ -48,6 +48,18 @@ def project_datastore(mock_trainml):
     )
 
 
+@fixture
+def project_data_connector(mock_trainml):
+    yield specimen.ProjectDataConnector(
+        mock_trainml,
+        id="ds-id-1",
+        name="connector 1",
+        project_uuid="proj-id-1",
+        type="custom",
+        region_uuid="reg-id-1",
+    )
+
+
 @fixture
 def project_service(mock_trainml):
     yield specimen.ProjectService(
@@ -56,9 +68,8 @@ def project_service(mock_trainml):
         name="service 1",
         project_uuid="proj-id-1",
         region_uuid="reg-id-1",
-        type="port",
-        resource="8001",
-        hostname="service.local",
+        public=False,
+        hostname="asdf.proximl.cloud",
     )
 
 
@@ -152,14 +163,43 @@ class ProjectDatastoreTests:
         assert not bool(empty_project_datastore)
 
 
+class ProjectDataConnectorTests:
+    def test_project_data_connector_properties(self, project_data_connector):
+        assert isinstance(project_data_connector.id, str)
+        assert isinstance(project_data_connector.name, str)
+        assert isinstance(project_data_connector.project_uuid, str)
+        assert isinstance(project_data_connector.type, str)
+        assert isinstance(project_data_connector.region_uuid, str)
+
+    def test_project_data_connector_str(self, project_data_connector):
+        string = str(project_data_connector)
+        regex = r"^{.*\"id\": \"" + project_data_connector.id + r"\".*}$"
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    def test_project_data_connector_repr(self, project_data_connector):
+        string = repr(project_data_connector)
+        regex = (
+            r"^ProjectDataConnector\( trainml , \*\*{.*'id': '"
+            + project_data_connector.id
+            + r"'.*}\)$"
+        )
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    def test_project_data_connector_bool(self, project_data_connector, mock_trainml):
+        empty_project_data_connector = specimen.ProjectDataConnector(mock_trainml)
+        assert bool(project_data_connector)
+        assert not bool(empty_project_data_connector)
+
+
 class ProjectServiceTests:
     def test_project_service_properties(self, project_service):
         assert isinstance(project_service.id, str)
         assert isinstance(project_service.name, str)
         assert isinstance(project_service.project_uuid, str)
-        assert isinstance(project_service.type, str)
         assert isinstance(project_service.hostname, str)
-        assert isinstance(project_service.resource, str)
+        assert isinstance(project_service.public, bool)
         assert isinstance(project_service.region_uuid, str)
 
     def test_project_service_str(self, project_service):

trainml/__init__.py

@@ -13,5 +13,5 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 
-__version__ = "0.5.6"
+__version__ = "0.5.8"
 __all__ = "TrainML"

trainml/checkpoints.py

@@ -23,9 +23,7 @@ class Checkpoints(object):
 
     async def list(self, **kwargs):
         resp = await self.trainml._query(f"/checkpoint", "GET", kwargs)
-        checkpoints = [
-            Checkpoint(self.trainml, **checkpoint) for checkpoint in resp
-        ]
+        checkpoints = [Checkpoint(self.trainml, **checkpoint) for checkpoint in resp]
         return checkpoints
 
     async def list_public(self, **kwargs):
@@ -39,8 +37,7 @@ class Checkpoints(object):
             source_type=source_type,
             source_uri=source_uri,
             source_options=kwargs.get("source_options"),
-            project_uuid=kwargs.get("project_uuid")
-            or self.trainml.active_project,
+            project_uuid=kwargs.get("project_uuid") or self.trainml.active_project,
         )
         payload = {k: v for k, v in data.items() if v is not None}
         logging.info(f"Creating Checkpoint {name}")
@@ -60,12 +57,13 @@ class Checkpoint:
     def __init__(self, trainml, **kwargs):
         self.trainml = trainml
         self._checkpoint = kwargs
-        self._id = self._checkpoint.get(
-            "id", self._checkpoint.get("checkpoint_uuid")
-        )
+        self._id = self._checkpoint.get("id", self._checkpoint.get("checkpoint_uuid"))
         self._status = self._checkpoint.get("status")
         self._name = self._checkpoint.get("name")
-        self._size = self._checkpoint.get("size")
+        self._size = self._checkpoint.get("size") or self._checkpoint.get("used_size")
+        self._billed_size = self._checkpoint.get("billed_size") or self._checkpoint.get(
+            "size"
+        )
         self._project_uuid = self._checkpoint.get("project_uuid")
 
     @property
@@ -84,6 +82,10 @@ class Checkpoint:
     def size(self) -> int:
         return self._size
 
+    @property
+    def billed_size(self) -> int:
+        return self._billed_size
+
     def __str__(self):
         return json.dumps({k: v for k, v in self._checkpoint.items()})
 
@@ -123,15 +125,17 @@ class Checkpoint:
                 entity_type="checkpoint",
                 project_uuid=self._checkpoint.get("project_uuid"),
                 cidr=self._checkpoint.get("vpn").get("cidr"),
-                ssh_port=self._checkpoint.get("vpn")
-                .get("client")
-                .get("ssh_port"),
-                input_path=self._checkpoint.get("source_uri")
-                if self.status in ["new", "downloading"]
-                else None,
-                output_path=self._checkpoint.get("output_uri")
-                if self.status == "exporting"
-                else None,
+                ssh_port=self._checkpoint.get("vpn").get("client").get("ssh_port"),
+                input_path=(
+                    self._checkpoint.get("source_uri")
+                    if self.status in ["new", "downloading"]
+                    else None
+                ),
+                output_path=(
+                    self._checkpoint.get("output_uri")
+                    if self.status == "exporting"
+                    else None
+                ),
             )
         else:
             details = dict()
@@ -195,9 +199,7 @@ class Checkpoint:
                 if msg_handler:
                     msg_handler(data)
                 else:
-                    timestamp = datetime.fromtimestamp(
-                        int(data.get("time")) / 1000
-                    )
+                    timestamp = datetime.fromtimestamp(int(data.get("time")) / 1000)
                     print(
                         f"{timestamp.strftime('%m/%d/%Y, %H:%M:%S')}: {data.get('msg').rstrip()}"
                     )
@@ -224,19 +226,24 @@ class Checkpoint:
         return self
 
     async def wait_for(self, status, timeout=300):
+        if self.status == status:
+            return
         valid_statuses = ["downloading", "ready", "archived"]
         if not status in valid_statuses:
             raise SpecificationError(
                 "status",
                 f"Invalid wait_for status {status}.  Valid statuses are: {valid_statuses}",
             )
-        if self.status == status:
-            return
+
+        MAX_TIMEOUT = 24 * 60 * 60
+        if timeout > MAX_TIMEOUT:
+            raise SpecificationError(
+                "timeout",
+                f"timeout must be less than {MAX_TIMEOUT} seconds.",
+            )
         POLL_INTERVAL_MIN = 5
         POLL_INTERVAL_MAX = 60
-        POLL_INTERVAL = max(
-            min(timeout / 60, POLL_INTERVAL_MAX), POLL_INTERVAL_MIN
-        )
+        POLL_INTERVAL = max(min(timeout / 60, POLL_INTERVAL_MAX), POLL_INTERVAL_MIN)
         retry_count = math.ceil(timeout / POLL_INTERVAL)
         count = 0
         while count < retry_count:

trainml/cli/cloudbender/__init__.py

@@ -15,4 +15,5 @@ from trainml.cli.cloudbender.region import region
 from trainml.cli.cloudbender.node import node
 from trainml.cli.cloudbender.device import device
 from trainml.cli.cloudbender.datastore import datastore
+from trainml.cli.cloudbender.data_connector import data_connector
 from trainml.cli.cloudbender.service import service

trainml/cli/cloudbender/data_connector.py

@@ -0,0 +1,159 @@
+import click
+from trainml.cli import cli, pass_config, search_by_id_name
+from trainml.cli.cloudbender import cloudbender
+
+
+@cloudbender.group()
+@pass_config
+def data_connector(config):
+    """trainML CloudBender data connector commands."""
+    pass
+
+
+@data_connector.command()
+@click.option(
+    "--provider",
+    "-p",
+    type=click.STRING,
+    required=True,
+    help="The provider ID of the region.",
+)
+@click.option(
+    "--region",
+    "-r",
+    type=click.STRING,
+    required=True,
+    help="The region ID to list data connectors for.",
+)
+@pass_config
+def list(config, provider, region):
+    """List data connectors."""
+    data = [
+        ["ID", "NAME", "TYPE"],
+        [
+            "-" * 80,
+            "-" * 80,
+            "-" * 80,
+        ],
+    ]
+
+    data_connectors = config.trainml.run(
+        config.trainml.client.cloudbender.data_connectors.list(
+            provider_uuid=provider, region_uuid=region
+        )
+    )
+
+    for data_connector in data_connectors:
+        data.append(
+            [
+                data_connector.id,
+                data_connector.name,
+                data_connector.type,
+            ]
+        )
+
+    for row in data:
+        click.echo(
+            "{: >37.36} {: >29.28} {: >9.8}" "".format(*row),
+            file=config.stdout,
+        )
+
+
+@data_connector.command()
+@click.option(
+    "--provider",
+    "-p",
+    type=click.STRING,
+    required=True,
+    help="The provider ID of the region.",
+)
+@click.option(
+    "--region",
+    "-r",
+    type=click.STRING,
+    required=True,
+    help="The region ID to create the data_connector in.",
+)
+@click.option(
+    "--type",
+    "-t",
+    type=click.Choice(
+        [
+            "custom",
+        ],
+        case_sensitive=False,
+    ),
+    required=True,
+    help="The type of data connector to create.",
+)
+@click.option(
+    "--protocol",
+    "-r",
+    type=click.STRING,
+    help="The transport protocol of the data connector",
+)
+@click.option(
+    "--port-range",
+    "-p",
+    type=click.STRING,
+    help="The port range of the data connector",
+)
+@click.option(
+    "--cidr",
+    "-i",
+    type=click.STRING,
+    help="The IP range to allow in CIDR notation",
+)
+@click.argument("name", type=click.STRING, required=True)
+@pass_config
+def create(config, provider, region, type, protocol, port_range, cidr, name):
+    """
+    Creates a data_connector.
+    """
+    return config.trainml.run(
+        config.trainml.client.cloudbender.data_connectors.create(
+            provider_uuid=provider,
+            region_uuid=region,
+            name=name,
+            type=type,
+            protocol=protocol,
+            port_range=port_range,
+            cidr=cidr,
+        )
+    )
+
+
+@data_connector.command()
+@click.option(
+    "--provider",
+    "-p",
+    type=click.STRING,
+    required=True,
+    help="The provider ID of the region.",
+)
+@click.option(
+    "--region",
+    "-r",
+    type=click.STRING,
+    required=True,
+    help="The region ID to remove the data_connector from.",
+)
+@click.argument("data_connector", type=click.STRING)
+@pass_config
+def remove(config, provider, region, data_connector):
+    """
+    Remove a data_connector.
+
+    DATASTORE may be specified by name or ID, but ID is preferred.
+    """
+    data_connectors = config.trainml.run(
+        config.trainml.client.cloudbender.data_connectors.list(
+            provider_uuid=provider, region_uuid=region
+        )
+    )
+
+    found = search_by_id_name(data_connector, data_connectors)
+    if None is found:
+        raise click.UsageError("Cannot find specified data_connector.")
+
+    return config.trainml.run(found.remove())