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

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))

trainml/cli/project/service.py

@@ -0,0 +1,61 @@
+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 service(config):
+    """trainML project service commands."""
+    pass
+
+
+@service.command()
+@pass_config
+def list(config):
+    """List project services."""
+    data = [
+        ["ID", "NAME", "TYPE", "REGION_UUID"],
+        [
+            "-" * 80,
+            "-" * 80,
+            "-" * 80,
+            "-" * 80,
+        ],
+    ]
+    project = config.trainml.run(
+        config.trainml.client.projects.get(config.trainml.client.project)
+    )
+
+    services = config.trainml.run(project.services.list())
+
+    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} {: >15.13} {: >38.36}" "".format(*row),
+            file=config.stdout,
+        )
+
+
+@service.command()
+@pass_config
+def refresh(config):
+    """
+    Refresh project service list.
+    """
+    project = config.trainml.run(config.trainml.client.projects.get_current())
+
+    return config.trainml.run(project.services.refresh())

trainml/cloudbender/data_connectors.py

@@ -1,5 +1,13 @@
 import json
 import logging
+import asyncio
+import math
+
+from trainml.exceptions import (
+    ApiError,
+    SpecificationError,
+    TrainMLException,
+)
 
 
 class DataConnectors(object):

trainml/cloudbender/datastores.py

@@ -1,5 +1,13 @@
 import json
 import logging
+import asyncio
+import math
+
+from trainml.exceptions import (
+    ApiError,
+    SpecificationError,
+    TrainMLException,
+)
 
 
 class Datastores(object):
@@ -20,9 +28,7 @@ class Datastores(object):
             "GET",
             kwargs,
         )
-        datastores = [
-            Datastore(self.trainml, **datastore) for datastore in resp
-        ]
+        datastores = [Datastore(self.trainml, **datastore) for datastore in resp]
         return datastores
 
     async def create(
@@ -31,18 +37,12 @@ class Datastores(object):
         region_uuid,
         name,
         type,
-        uri,
-        root,
-        options=None,
         **kwargs,
     ):
         logging.info(f"Creating Datastore {name}")
         data = dict(
             name=name,
             type=type,
-            uri=uri,
-            root=root,
-            options=options,
             **kwargs,
         )
         payload = {k: v for k, v in data.items() if v is not None}
@@ -73,8 +73,6 @@ class Datastore:
         self._region_uuid = self._datastore.get("region_uuid")
         self._type = self._datastore.get("type")
         self._name = self._datastore.get("name")
-        self._uri = self._datastore.get("uri")
-        self._root = self._datastore.get("root")
 
     @property
     def id(self) -> str:
@@ -96,14 +94,6 @@ class Datastore:
     def name(self) -> str:
         return self._name
 
-    @property
-    def uri(self) -> str:
-        return self._uri
-
-    @property
-    def root(self) -> str:
-        return self._root
-
     def __str__(self):
         return json.dumps({k: v for k, v in self._datastore.items()})
 

trainml/cloudbender/nodes.py

@@ -1,5 +1,9 @@
 import json
 import logging
+import asyncio
+import math
+
+from trainml.exceptions import ApiError, SpecificationError, TrainMLException, NodeError
 
 
 class Nodes(object):
@@ -29,7 +33,7 @@ class Nodes(object):
         region_uuid,
         friendly_name,
         hostname,
-        minion_id,
+        minion_id=None,
         type="permanent",
         service="compute",
         **kwargs,
@@ -153,3 +157,42 @@ class Node:
             None,
             dict(command=command),
         )
+
+    async def wait_for(self, status, timeout=300):
+        if self.status == status:
+            return
+        valid_statuses = ["active", "maintenance", "offline", "stopped", "archived"]
+        if not status in valid_statuses:
+            raise SpecificationError(
+                "status",
+                f"Invalid wait_for status {status}.  Valid statuses are: {valid_statuses}",
+            )
+        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)
+        retry_count = math.ceil(timeout / POLL_INTERVAL)
+        count = 0
+        while count < retry_count:
+            await asyncio.sleep(POLL_INTERVAL)
+            try:
+                await self.refresh()
+            except ApiError as e:
+                if status == "archived" and e.status == 404:
+                    return
+                raise e
+            if self.status in ["errored", "failed"]:
+                raise NodeError(self.status, self)
+            if self.status == status:
+                return self
+            else:
+                count += 1
+                logging.debug(f"self: {self}, retry count {count}")
+
+        raise TrainMLException(f"Timeout waiting for {status}")

trainml/cloudbender/providers.py

@@ -1,7 +1,16 @@
 import json
 import logging
+import asyncio
+import math
 from datetime import datetime
 
+from trainml.exceptions import (
+    ApiError,
+    SpecificationError,
+    TrainMLException,
+    ProviderError,
+)
+
 
 class Providers(object):
     def __init__(self, trainml):
@@ -36,6 +45,7 @@ class Provider:
         self._provider = kwargs
         self._id = self._provider.get("provider_uuid")
         self._type = self._provider.get("type")
+        self._status = self._provider.get("status")
         self._credits = self._provider.get("credits")
 
     @property
@@ -46,6 +56,10 @@ class Provider:
     def type(self) -> str:
         return self._type
 
+    @property
+    def status(self) -> str:
+        return self._status
+
     @property
     def credits(self) -> float:
         return self._credits
@@ -69,3 +83,42 @@ class Provider:
         )
         self.__init__(self.trainml, **resp)
         return self
+
+    async def wait_for(self, status, timeout=300):
+        if self.status == status:
+            return
+        valid_statuses = ["ready", "archived"]
+        if not status in valid_statuses:
+            raise SpecificationError(
+                "status",
+                f"Invalid wait_for status {status}.  Valid statuses are: {valid_statuses}",
+            )
+        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)
+        retry_count = math.ceil(timeout / POLL_INTERVAL)
+        count = 0
+        while count < retry_count:
+            await asyncio.sleep(POLL_INTERVAL)
+            try:
+                await self.refresh()
+            except ApiError as e:
+                if status == "archived" and e.status == 404:
+                    return
+                raise e
+            if self.status in ["errored", "failed"]:
+                raise ProviderError(self.status, self)
+            if self.status == status:
+                return self
+            else:
+                count += 1
+                logging.debug(f"self: {self}, retry count {count}")
+
+        raise TrainMLException(f"Timeout waiting for {status}")

trainml/cloudbender/regions.py

@@ -1,5 +1,14 @@
 import json
 import logging
+import asyncio
+import math
+
+from trainml.exceptions import (
+    ApiError,
+    SpecificationError,
+    TrainMLException,
+    RegionError,
+)
 
 
 class Regions(object):
@@ -111,3 +120,42 @@ class Region:
             None,
             dict(project_uuid=project_uuid, checkpoint_uuid=checkpoint_uuid),
         )
+
+    async def wait_for(self, status, timeout=300):
+        if self.status == status:
+            return
+        valid_statuses = ["healthy", "offline", "archived"]
+        if not status in valid_statuses:
+            raise SpecificationError(
+                "status",
+                f"Invalid wait_for status {status}.  Valid statuses are: {valid_statuses}",
+            )
+        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)
+        retry_count = math.ceil(timeout / POLL_INTERVAL)
+        count = 0
+        while count < retry_count:
+            await asyncio.sleep(POLL_INTERVAL)
+            try:
+                await self.refresh()
+            except ApiError as e:
+                if status == "archived" and e.status == 404:
+                    return
+                raise e
+            if self.status in ["errored", "failed"]:
+                raise RegionError(self.status, self)
+            if self.status == status:
+                return self
+            else:
+                count += 1
+                logging.debug(f"self: {self}, retry count {count}")
+
+        raise TrainMLException(f"Timeout waiting for {status}")

trainml/datasets.py

@@ -31,13 +31,24 @@ class Datasets(object):
         datasets = [Dataset(self.trainml, **dataset) for dataset in resp]
         return datasets
 
-    async def create(self, name, source_type, source_uri, **kwargs):
+    async def create(
+        self,
+        name,
+        source_type,
+        source_uri,
+        type="evefs",
+        project_uuid=None,
+        **kwargs,
+    ):
+        if not project_uuid:
+            project_uuid = self.trainml.active_project
         data = dict(
             name=name,
             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=project_uuid,
+            type=type,
+            **kwargs,
         )
         payload = {k: v for k, v in data.items() if v is not None}
         logging.info(f"Creating Dataset {name}")

trainml/exceptions.py

@@ -147,3 +147,54 @@ class SpecificationError(TrainMLException):
 
     def __str__(self):
         return "SpecificationError({self.attribute}, {self.message})".format(self=self)
+
+
+class NodeError(TrainMLException):
+    def __init__(self, status, data, *args):
+        super().__init__(data, *args)
+        self._status = status
+        self._message = data
+
+    @property
+    def status(self) -> str:
+        return self._status
+
+    def __repr__(self):
+        return "NodeError({self.status}, {self.message})".format(self=self)
+
+    def __str__(self):
+        return "NodeError({self.status}, {self.message})".format(self=self)
+
+
+class ProviderError(TrainMLException):
+    def __init__(self, status, data, *args):
+        super().__init__(data, *args)
+        self._status = status
+        self._message = data
+
+    @property
+    def status(self) -> str:
+        return self._status
+
+    def __repr__(self):
+        return "ProviderError({self.status}, {self.message})".format(self=self)
+
+    def __str__(self):
+        return "ProviderError({self.status}, {self.message})".format(self=self)
+
+
+class RegionError(TrainMLException):
+    def __init__(self, status, data, *args):
+        super().__init__(data, *args)
+        self._status = status
+        self._message = data
+
+    @property
+    def status(self) -> str:
+        return self._status
+
+    def __repr__(self):
+        return "RegionError({self.status}, {self.message})".format(self=self)
+
+    def __str__(self):
+        return "RegionError({self.status}, {self.message})".format(self=self)

trainml/jobs.py

@@ -476,7 +476,6 @@ class Job:
             return
         valid_statuses = [
             "waiting for data/model download",
-            "waiting for GPUs",
             "waiting for resources",
             "running",
             "stopped",
@@ -488,11 +487,6 @@ class Job:
                 "status",
                 f"Invalid wait_for status {status}.  Valid statuses are: {valid_statuses}",
             )
-        if status == "waiting for GPUs":
-            warnings.warn(
-                "'waiting for GPUs' status is deprecated, use 'waiting for resources' instead.",
-                DeprecationWarning,
-            )
         if (self.type == "training") and status == "stopped":
             warnings.warn(
                 "'stopped' status is deprecated for training jobs, use 'finished' instead.",
@@ -523,12 +517,8 @@ class Job:
                 self.status == status
                 or (
                     status
-                    in [
-                        "waiting for GPUs",
-                        "waiting for resources",
-                    ]  ## this status could be very short and the polling could miss it
-                    and self.status
-                    not in ["new", "waiting for GPUs", "waiting for resources"]
+                    == "waiting for resources"  ## this status could be very short and the polling could miss it
+                    and self.status not in ["new", "waiting for resources"]
                 )
                 or (
                     status
@@ -536,7 +526,6 @@ class Job:
                     and self.status
                     not in [
                         "new",
-                        "waiting for GPUs",
                         "waiting for resources",
                         "waiting for data/model download",
                     ]

trainml/models.py

@@ -26,13 +26,24 @@ class Models(object):
         models = [Model(self.trainml, **model) for model in resp]
         return models
 
-    async def create(self, name, source_type, source_uri, **kwargs):
+    async def create(
+        self,
+        name,
+        source_type,
+        source_uri,
+        type="evefs",
+        project_uuid=None,
+        **kwargs,
+    ):
+        if not project_uuid:
+            project_uuid = self.trainml.active_project
         data = dict(
             name=name,
             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=project_uuid,
+            type=type,
+            **kwargs,
         )
         payload = {k: v for k, v in data.items() if v is not None}
         logging.info(f"Creating Model {name}")

trainml/projects/__init__.py

@@ -0,0 +1,3 @@
+from .projects import Projects, Project
+
+__all__ = ["Projects", "Project"]