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

tests/unit/projects/test_project_services_unit.py

@@ -0,0 +1,102 @@
+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.projects.services as specimen
+from trainml.exceptions import (
+    ApiError,
+    SpecificationError,
+    TrainMLException,
+)
+
+pytestmark = [mark.sdk, mark.unit, mark.projects]
+
+
+@fixture
+def project_services(mock_trainml):
+    yield specimen.ProjectServices(mock_trainml, project_id="1")
+
+
+@fixture
+def project_service(mock_trainml):
+    yield specimen.ProjectService(
+        mock_trainml,
+        id="res-id-1",
+        name="service 1",
+        project_uuid="proj-id-1",
+        region_uuid="reg-id-1",
+        public=False,
+        hostname="asdf.proximl.cloud",
+    )
+
+
+class ProjectServicesTests:
+    @mark.asyncio
+    async def test_project_services_refresh(self, project_services, mock_trainml):
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await project_services.refresh()
+        mock_trainml._query.assert_called_once_with("/project/1/services", "PATCH")
+
+    @mark.asyncio
+    async def test_project_services_list(self, project_services, mock_trainml):
+        api_response = [
+            {
+                "project_uuid": "proj-id-1",
+                "region_uuid": "reg-id-1",
+                "id": "res-id-1",
+                "type": "port",
+                "name": "On-Prem Service A",
+                "resource": "8001",
+                "hostname": "service-a.local",
+            },
+            {
+                "project_uuid": "proj-id-1",
+                "region_uuid": "reg-id-2",
+                "id": "res-id-2",
+                "type": "port",
+                "name": "Cloud Service B",
+                "resource": "8001",
+                "hostname": "service-b.local",
+            },
+        ]
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        resp = await project_services.list()
+        mock_trainml._query.assert_called_once_with(
+            "/project/1/services", "GET", dict()
+        )
+        assert len(resp) == 2
+
+
+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.hostname, str)
+        assert isinstance(project_service.public, bool)
+        assert isinstance(project_service.region_uuid, str)
+
+    def test_project_service_str(self, project_service):
+        string = str(project_service)
+        regex = r"^{.*\"id\": \"" + project_service.id + r"\".*}$"
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    def test_project_service_repr(self, project_service):
+        string = repr(project_service)
+        regex = (
+            r"^ProjectService\( trainml , \*\*{.*'id': '"
+            + project_service.id
+            + r"'.*}\)$"
+        )
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    def test_project_service_bool(self, project_service, mock_trainml):
+        empty_project_service = specimen.ProjectService(mock_trainml)
+        assert bool(project_service)
+        assert not bool(empty_project_service)

tests/unit/projects/test_projects_unit.py

@@ -0,0 +1,128 @@
+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.projects as specimen
+from trainml.exceptions import (
+    ApiError,
+    SpecificationError,
+    TrainMLException,
+)
+
+pytestmark = [mark.sdk, mark.unit, mark.projects]
+
+
+@fixture
+def projects(mock_trainml):
+    yield specimen.Projects(mock_trainml)
+
+
+@fixture
+def project(mock_trainml):
+    yield specimen.Project(
+        mock_trainml,
+        id="1",
+        name="My Mock Project",
+        owner=True,
+        owner_name="Me",
+        created_name="Me",
+        job_all=True,
+        dataset_all=True,
+        model_all=True,
+        createdAt="2020-12-31T23:59:59.000Z",
+    )
+
+
+class ProjectsTests:
+    @mark.asyncio
+    async def test_get_project(
+        self,
+        projects,
+        mock_trainml,
+    ):
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await projects.get("1234")
+        mock_trainml._query.assert_called_once_with("/project/1234", "GET", dict())
+
+    @mark.asyncio
+    async def test_list_projects(
+        self,
+        projects,
+        mock_trainml,
+    ):
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await projects.list()
+        mock_trainml._query.assert_called_once_with("/project", "GET", dict())
+
+    @mark.asyncio
+    async def test_remove_project(
+        self,
+        projects,
+        mock_trainml,
+    ):
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await projects.remove("4567")
+        mock_trainml._query.assert_called_once_with("/project/4567", "DELETE", dict())
+
+    @mark.asyncio
+    async def test_create_project_simple(self, projects, mock_trainml):
+        requested_config = dict(
+            name="new project",
+        )
+        expected_payload = dict(name="new project", copy_keys=False, copy_secrets=False)
+        api_response = {
+            "id": "project-id-1",
+            "name": "new project",
+            "owner": True,
+            "owner_name": "Me",
+            "created_name": "Me",
+            "job_all": True,
+            "dataset_all": True,
+            "model_all": True,
+            "createdAt": "2020-12-31T23:59:59.000Z",
+        }
+
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        response = await projects.create(**requested_config)
+        mock_trainml._query.assert_called_once_with(
+            "/project", "POST", None, expected_payload
+        )
+        assert response.id == "project-id-1"
+
+
+class ProjectTests:
+    def test_project_properties(self, project):
+        assert isinstance(project.id, str)
+        assert isinstance(project.name, str)
+        assert isinstance(project.owner_name, str)
+        assert isinstance(project.is_owner, bool)
+
+    def test_project_str(self, project):
+        string = str(project)
+        regex = r"^{.*\"id\": \"" + project.id + r"\".*}$"
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    def test_project_repr(self, project):
+        string = repr(project)
+        regex = r"^Project\( trainml , \*\*{.*'id': '" + project.id + r"'.*}\)$"
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    def test_project_bool(self, project, mock_trainml):
+        empty_project = specimen.Project(mock_trainml)
+        assert bool(project)
+        assert not bool(empty_project)
+
+    @mark.asyncio
+    async def test_project_remove(self, project, mock_trainml):
+        api_response = dict()
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        await project.remove()
+        mock_trainml._query.assert_called_once_with("/project/1", "DELETE")

tests/unit/test_auth_unit.py

@@ -0,0 +1,30 @@
+import re
+import logging
+import json
+import os
+from unittest.mock import AsyncMock, patch, mock_open, MagicMock
+from pytest import mark, fixture, raises
+from aiohttp import WSMessage, WSMsgType
+
+import trainml.auth as specimen
+
+pytestmark = [mark.sdk, mark.unit]
+
+
+@patch.dict(
+    os.environ,
+    {
+        "TRAINML_USER": "user-id",
+        "TRAINML_KEY": "key",
+        "TRAINML_REGION": "ap-east-1",
+        "TRAINML_CLIENT_ID": "client_id",
+        "TRAINML_POOL_ID": "pool_id",
+    },
+)
+def test_auth_from_envs():
+    auth = specimen.Auth(config_dir=os.path.expanduser("~/.trainml"))
+    assert auth.__dict__.get("username") == "user-id"
+    assert auth.__dict__.get("password") == "key"
+    assert auth.__dict__.get("region") == "ap-east-1"
+    assert auth.__dict__.get("client_id") == "client_id"
+    assert auth.__dict__.get("pool_id") == "pool_id"

tests/unit/test_trainml_unit.py

@@ -0,0 +1,54 @@
+import re
+import logging
+import json
+import os
+from unittest.mock import AsyncMock, patch, mock_open
+from pytest import mark, fixture, raises
+from aiohttp import WSMessage, WSMsgType
+
+import trainml.trainml as specimen
+
+pytestmark = [mark.sdk, mark.unit]
+
+
+@patch.dict(
+    os.environ,
+    {
+        "TRAINML_USER": "user-id",
+        "TRAINML_KEY": "key",
+        "TRAINML_REGION": "region",
+        "TRAINML_CLIENT_ID": "client_id",
+        "TRAINML_POOL_ID": "pool_id",
+        "TRAINML_API_URL": "api.example.com",
+        "TRAINML_WS_URL": "api-ws.example.com",
+    },
+)
+def test_trainml_from_envs():
+    trainml = specimen.TrainML()
+    assert trainml.__dict__.get("api_url") == "api.example.com"
+    assert trainml.__dict__.get("ws_url") == "api-ws.example.com"
+    assert trainml.auth.__dict__.get("username") == "user-id"
+    assert trainml.auth.__dict__.get("password") == "key"
+    assert trainml.auth.__dict__.get("region") == "region"
+    assert trainml.auth.__dict__.get("client_id") == "client_id"
+    assert trainml.auth.__dict__.get("pool_id") == "pool_id"
+
+
+def test_trainml_env_from_files():
+    with patch(
+        "trainml.trainml.open",
+        mock_open(
+            read_data=json.dumps(
+                dict(
+                    region="region_file",
+                    client_id="client_id_file",
+                    pool_id="pool_id_file",
+                    api_url="api.example.com_file",
+                    ws_url="api-ws.example.com_file",
+                )
+            )
+        ),
+    ):
+        trainml = specimen.TrainML()
+    assert trainml.__dict__.get("api_url") == "api.example.com_file"
+    assert trainml.__dict__.get("ws_url") == "api-ws.example.com_file"

trainml/__init__.py

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

trainml/checkpoints.py

@@ -60,7 +60,10 @@ class Checkpoint:
         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
@@ -79,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()})
 

trainml/cli/project/__init__.py

@@ -0,0 +1,153 @@
+import click
+from trainml.cli import cli, pass_config, search_by_id_name
+
+
+@cli.group()
+@pass_config
+def project(config):
+    """trainML project commands."""
+    pass
+
+
+@project.command()
+@pass_config
+def list(config):
+    """List projects."""
+    data = [
+        ["ID", "NAME", "OWNER", "MINE"],
+        [
+            "-" * 80,
+            "-" * 80,
+            "-" * 80,
+            "-" * 80,
+        ],
+    ]
+
+    projects = config.trainml.run(config.trainml.client.projects.list())
+
+    for project in projects:
+        data.append(
+            [
+                project.id,
+                project.name,
+                project.owner_name,
+                "X" if project.is_owner else "",
+            ]
+        )
+
+    for row in data:
+        click.echo(
+            "{: >38.36} {: >30.28} {: >15.13} {: >4.4}" "".format(*row),
+            file=config.stdout,
+        )
+
+
+@project.command()
+@click.argument("name", type=click.STRING)
+@pass_config
+def create(config, name):
+    """
+    Create a project.
+
+    Project is created with the specified NAME.
+    """
+
+    return config.trainml.run(
+        config.trainml.client.projects.create(
+            name=name,
+        )
+    )
+
+
+@project.command()
+@click.argument("project", type=click.STRING)
+@pass_config
+def remove(config, project):
+    """
+    Remove a project.
+
+    PROJECT may be specified by name or ID, but ID is preferred.
+    """
+    projects = config.trainml.run(config.trainml.client.projects.list())
+
+    found = search_by_id_name(project, projects)
+    if None is found:
+        raise click.UsageError("Cannot find specified project.")
+
+    return config.trainml.run(found.remove())
+
+
+@project.command()
+@pass_config
+def list_datastores(config):
+    """List project datastores."""
+    data = [
+        ["ID", "NAME", "TYPE", "REGION_UUID"],
+        [
+            "-" * 80,
+            "-" * 80,
+            "-" * 80,
+            "-" * 80,
+        ],
+    ]
+    project = config.trainml.run(
+        config.trainml.client.projects.get(config.trainml.client.project)
+    )
+
+    datastores = config.trainml.run(project.list_datastores())
+
+    for datastore in datastores:
+        data.append(
+            [
+                datastore.id,
+                datastore.name,
+                datastore.type,
+                datastore.region_uuid,
+            ]
+        )
+
+    for row in data:
+        click.echo(
+            "{: >38.36} {: >30.28} {: >15.13} {: >38.36}" "".format(*row),
+            file=config.stdout,
+        )
+
+
+@project.command()
+@pass_config
+def list_services(config):
+    """List project services."""
+    data = [
+        ["ID", "NAME", "HOSTNAME", "REGION_UUID"],
+        [
+            "-" * 80,
+            "-" * 80,
+            "-" * 80,
+            "-" * 80,
+        ],
+    ]
+    project = config.trainml.run(
+        config.trainml.client.projects.get(config.trainml.client.project)
+    )
+
+    services = config.trainml.run(project.list_services())
+
+    for service in services:
+        data.append(
+            [
+                service.id,
+                service.name,
+                service.hostname,
+                service.region_uuid,
+            ]
+        )
+
+    for row in data:
+        click.echo(
+            "{: >38.36} {: >30.28} {: >30.28} {: >38.36}" "".format(*row),
+            file=config.stdout,
+        )
+
+
+from trainml.cli.project.secret import secret
+from trainml.cli.project.key import key

trainml/cli/project/key.py

@@ -0,0 +1,124 @@
+import click
+import os
+import json
+import base64
+from pathlib import Path
+from trainml.cli import pass_config
+from trainml.cli.project import project
+
+
+@project.group()
+@pass_config
+def key(config):
+    """trainML project key commands."""
+    pass
+
+
+@key.command()
+@pass_config
+def list(config):
+    """List keys."""
+    data = [
+        ["TYPE", "KEY ID", "UPDATED AT"],
+        [
+            "-" * 80,
+            "-" * 80,
+            "-" * 80,
+        ],
+    ]
+    project = config.trainml.run(config.trainml.client.projects.get_current())
+    keys = config.trainml.run(project.keys.list())
+
+    for key in keys:
+        data.append(
+            [
+                key.type,
+                key.key_id,
+                key.updated_at.isoformat(timespec="seconds"),
+            ]
+        )
+
+    for row in data:
+        click.echo(
+            "{: >13.11} {: >37.35} {: >28.26}" "".format(*row),
+            file=config.stdout,
+        )
+
+
+@key.command()
+@click.argument(
+    "type",
+    type=click.Choice(
+        [
+            "aws",
+            "azure",
+            "docker",
+            "gcp",
+            "huggingface",
+            "kaggle",
+            "ngc",
+            "wasabi",
+        ],
+        case_sensitive=False,
+    ),
+)
+@pass_config
+def put(config, type):
+    """
+    Set a key.
+
+    A key is uploaded.
+    """
+    project = config.trainml.run(config.trainml.client.projects.get_current())
+
+    tenant = None
+
+    if type in ["aws", "wasabi"]:
+        key_id = click.prompt("Enter the key ID", type=str, hide_input=False)
+        secret = click.prompt("Enter the secret key", type=str, hide_input=True)
+    elif type == "azure":
+        key_id = click.prompt(
+            "Enter the Application (client) ID", type=str, hide_input=False
+        )
+        tenant = click.prompt(
+            "Enter the Directory (tenant) ley", type=str, hide_input=False
+        )
+        secret = click.prompt("Enter the client secret", type=str, hide_input=True)
+    elif type in ["docker", "huggingface"]:
+        key_id = click.prompt("Enter the username", type=str, hide_input=False)
+        secret = click.prompt("Enter the access token", type=str, hide_input=True)
+    elif type in ["gcp", "kaggle"]:
+        file_name = click.prompt(
+            "Enter the path of the credentials file",
+            type=click.Path(
+                exists=True, file_okay=True, dir_okay=False, resolve_path=True
+            ),
+            hide_input=False,
+        )
+        key_id = os.path.basename(file_name)
+        with open(file_name) as f:
+            secret = json.load(f)
+        secret = json.dumps(secret)
+    elif type == "ngc":
+        key_id = "$oauthtoken"
+        secret = click.prompt("Enter the access token", type=str, hide_input=True)
+    else:
+        raise click.UsageError("Unsupported key type")
+
+    return config.trainml.run(
+        project.keys.put(type=type, key_id=key_id, secret=secret, tenant=tenant)
+    )
+
+
+@key.command()
+@click.argument("name", type=click.STRING)
+@pass_config
+def remove(config, name):
+    """
+    Remove a key.
+
+
+    """
+    project = config.trainml.run(config.trainml.client.projects.get_current())
+
+    return config.trainml.run(project.key.remove(name))

trainml/cli/project/secret.py

@@ -0,0 +1,71 @@
+import click
+from trainml.cli import pass_config
+from trainml.cli.project import project
+
+
+@project.group()
+@pass_config
+def secret(config):
+    """trainML project secret commands."""
+    pass
+
+
+@secret.command()
+@pass_config
+def list(config):
+    """List secrets."""
+    data = [
+        ["NAME", "CREATED BY", "UPDATED AT"],
+        [
+            "-" * 80,
+            "-" * 80,
+            "-" * 80,
+        ],
+    ]
+    project = config.trainml.run(config.trainml.client.projects.get_current())
+    secrets = config.trainml.run(project.secrets.list())
+
+    for secret in secrets:
+        data.append(
+            [
+                secret.name,
+                secret.created_by,
+                secret.updated_at.isoformat(timespec="seconds"),
+            ]
+        )
+
+    for row in data:
+        click.echo(
+            "{: >38.36} {: >30.28} {: >28.26}" "".format(*row),
+            file=config.stdout,
+        )
+
+
+@secret.command()
+@click.argument("name", type=click.STRING)
+@pass_config
+def put(config, name):
+    """
+    Set a secret value.
+
+    Secret is created with the specified NAME.
+    """
+    project = config.trainml.run(config.trainml.client.projects.get_current())
+
+    value = click.prompt("Enter the secret value", type=str, hide_input=True)
+
+    return config.trainml.run(project.secrets.put(name=name, value=value))
+
+
+@secret.command()
+@click.argument("name", type=click.STRING)
+@pass_config
+def remove(config, name):
+    """
+    Remove a secret.
+
+
+    """
+    project = config.trainml.run(config.trainml.client.projects.get_current())
+
+    return config.trainml.run(project.secret.remove(name))