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

tests/unit/conftest.py

@@ -22,6 +22,7 @@ from trainml.cloudbender import Cloudbender
 from trainml.cloudbender.providers import Provider, Providers
 from trainml.cloudbender.regions import Region, Regions
 from trainml.cloudbender.nodes import Node, Nodes
+from trainml.cloudbender.devices import Device, Devices
 from trainml.cloudbender.datastores import Datastore, Datastores
 from trainml.cloudbender.reservations import Reservation, Reservations
 from trainml.cloudbender.device_configs import DeviceConfig, DeviceConfigs
@@ -484,19 +485,19 @@ def mock_jobs():
                     {
                         "rig_uuid": "rig-id-1",
                         "job_worker_uuid": "worker-id-11",
-                        "command": "PYTHONPATH=$PYTHONPATH:$TRAINML_MODEL_PATH python -m official.vision.image_classification.resnet_cifar_main --num_gpus=1 --data_dir=$TRAINML_DATA_PATH --model_dir=$TRAINML_OUTPUT_PATH --enable_checkpoint_and_export=True --train_epochs=1 --batch_size=1024",
+                        "command": "PYTHONPATH=$PYTHONPATH:$ML_MODEL_PATH python -m official.vision.image_classification.resnet_cifar_main --num_gpus=1 --data_dir=$ML_DATA_PATH --model_dir=$ML_OUTPUT_PATH --enable_checkpoint_and_export=True --train_epochs=1 --batch_size=1024",
                         "status": "stopped",
                     },
                     {
                         "rig_uuid": "rig-id-2",
                         "job_worker_uuid": "worker-id-12",
-                        "command": "PYTHONPATH=$PYTHONPATH:$TRAINML_MODEL_PATH python -m official.vision.image_classification.resnet_cifar_main --num_gpus=1 --data_dir=$TRAINML_DATA_PATH --model_dir=$TRAINML_OUTPUT_PATH --enable_checkpoint_and_export=True --train_epochs=1 --batch_size=1024",
+                        "command": "PYTHONPATH=$PYTHONPATH:$ML_MODEL_PATH python -m official.vision.image_classification.resnet_cifar_main --num_gpus=1 --data_dir=$ML_DATA_PATH --model_dir=$ML_OUTPUT_PATH --enable_checkpoint_and_export=True --train_epochs=1 --batch_size=1024",
                         "status": "stopped",
                     },
                     {
                         "rig_uuid": "rig-id-2",
                         "job_worker_uuid": "worker-id-13",
-                        "command": "PYTHONPATH=$PYTHONPATH:$TRAINML_MODEL_PATH python -m official.vision.image_classification.resnet_cifar_main --num_gpus=1 --data_dir=$TRAINML_DATA_PATH --model_dir=$TRAINML_OUTPUT_PATH --enable_checkpoint_and_export=True --train_epochs=1 --batch_size=1024",
+                        "command": "PYTHONPATH=$PYTHONPATH:$ML_MODEL_PATH python -m official.vision.image_classification.resnet_cifar_main --num_gpus=1 --data_dir=$ML_DATA_PATH --model_dir=$ML_OUTPUT_PATH --enable_checkpoint_and_export=True --train_epochs=1 --batch_size=1024",
                         "status": "stopped",
                     },
                 ],
@@ -706,6 +707,53 @@ def mock_nodes():
     ]
 
 
+@fixture(scope="session")
+def mock_devices():
+    trainml = Mock()
+    yield [
+        Device(
+            trainml,
+            **{
+                "provider_uuid": "prov-id-1",
+                "region_uuid": "reg-id-1",
+                "device_id": "dev-id-1",
+                "type": "device",
+                "service": "compute",
+                "friendly_name": "hq-orin-01",
+                "hostname": "hq-orin-01",
+                "status": "active",
+                "online": True,
+                "maintenance_mode": False,
+                "job_status": "running",
+                "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",
+            },
+        ),
+        Device(
+            trainml,
+            **{
+                "provider_uuid": "prov-id-1",
+                "region_uuid": "reg-id-1",
+                "device_id": "dev-id-2",
+                "type": "device",
+                "service": "compute",
+                "friendly_name": "hq-orin-02",
+                "hostname": "hq-orin-02",
+                "status": "active",
+                "online": True,
+                "maintenance_mode": False,
+                "job_status": "running",
+                "job_last_deployed": "2023-06-02T21:22:40.084Z",
+                "job_config_id": "job-id-2",
+                "job_config_revision": "1685740490096",
+                "device_config_id": "conf-id-1",
+            },
+        ),
+    ]
+
+
 @fixture(scope="session")
 def mock_datastores():
     trainml = Mock()
@@ -863,6 +911,7 @@ def mock_trainml(
     mock_providers,
     mock_regions,
     mock_nodes,
+    mock_devices,
     mock_datastores,
     mock_reservations,
     mock_device_configs,
@@ -898,6 +947,8 @@ def mock_trainml(
     trainml.cloudbender.regions.list = AsyncMock(return_value=mock_regions)
     trainml.cloudbender.nodes = create_autospec(Nodes)
     trainml.cloudbender.nodes.list = AsyncMock(return_value=mock_nodes)
+    trainml.cloudbender.devices = create_autospec(Nodes)
+    trainml.cloudbender.devices.list = AsyncMock(return_value=mock_devices)
     trainml.cloudbender.datastores = create_autospec(Datastores)
     trainml.cloudbender.datastores.list = AsyncMock(
         return_value=mock_datastores

tests/unit/test_auth.py

@@ -22,7 +22,7 @@ pytestmark = [mark.sdk, mark.unit]
     },
 )
 def test_auth_from_envs():
-    auth = specimen.Auth()
+    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"

trainml/__init__.py

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

trainml/auth.py

@@ -222,10 +222,6 @@ n_hex = (
 g_hex = "2"
 info_bits = bytearray("Caldera Derived Key", "utf-8")
 
-CONFIG_DIR = os.path.expanduser(
-    os.environ.get("TRAINML_CONFIG_DIR") or "~/.trainml"
-)
-
 
 def hash_sha256(buf):
     """AuthenticationHelper.hash"""
@@ -512,9 +508,9 @@ class AWSSRP(object):
 
 
 class Auth(object):
-    def __init__(self, domain_suffix="trainml.ai", **kwargs):
+    def __init__(self, config_dir, domain_suffix="trainml.ai", **kwargs):
         try:
-            with open(f"{CONFIG_DIR}/environment.json", "r") as file:
+            with open(f"{config_dir}/environment.json", "r") as file:
                 env_str = file.read().replace("\n", "")
             env = json.loads(env_str)
         except:
@@ -544,7 +540,7 @@ class Auth(object):
         )
 
         try:
-            with open(f"{CONFIG_DIR}/credentials.json", "r") as file:
+            with open(f"{config_dir}/credentials.json", "r") as file:
                 key_str = file.read().replace("\n", "")
             keys = json.loads(key_str)
         except:

trainml/cli/cloudbender/__init__.py

@@ -13,5 +13,6 @@ def cloudbender(config):
 from trainml.cli.cloudbender.provider import provider
 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.reservation import reservation

trainml/cli/cloudbender/device.py

@@ -0,0 +1,157 @@
+import click
+from trainml.cli import cli, pass_config, search_by_id_name
+from trainml.cli.cloudbender import cloudbender
+
+
+@cloudbender.group()
+@pass_config
+def device(config):
+    """trainML CloudBender device commands."""
+    pass
+
+
+@device.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 devices for.",
+)
+@pass_config
+def list(config, provider, region):
+    """List devices."""
+    data = [
+        [
+            "ID",
+            "NAME",
+            "STATUS",
+            "JOB STATUS",
+            "ONLINE",
+            "MAINTENANCE",
+        ],
+        [
+            "-" * 80,
+            "-" * 80,
+            "-" * 80,
+            "-" * 80,
+            "-" * 80,
+            "-" * 80,
+            "-" * 80,
+        ],
+    ]
+
+    devices = config.trainml.run(
+        config.trainml.client.cloudbender.devices.list(
+            provider_uuid=provider, region_uuid=region
+        )
+    )
+
+    for device in devices:
+        data.append(
+            [
+                device.id,
+                device.name,
+                device.status,
+                device.job_status,
+                "X" if device.online else "",
+                "X" if device.maintenance_mode else "",
+            ]
+        )
+
+    for row in data:
+        click.echo(
+            "{: >37.36} {: >29.28} {: >9.8} {: >11.10} {: >7.6} {: >12.11}"
+            "".format(*row),
+            file=config.stdout,
+        )
+
+
+@device.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 region in.",
+)
+@click.option(
+    "--minion-id",
+    "-m",
+    type=click.STRING,
+    required=True,
+    help="The minion_id of the new node.",
+)
+@click.option(
+    "--hostname",
+    "-h",
+    type=click.STRING,
+    help="The hostname (if different from name)",
+)
+@click.argument("name", type=click.STRING, required=True)
+@pass_config
+def create(config, provider, region, minion_id, hostname, name):
+    """
+    Creates a node.
+    """
+    if not hostname:
+        hostname = name
+    return config.trainml.run(
+        config.trainml.client.cloudbender.devices.create(
+            provider_uuid=provider,
+            region_uuid=region,
+            friendly_name=name,
+            hostname=hostname,
+            minion_id=minion_id,
+        )
+    )
+
+
+@device.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 delete the node from.",
+)
+@click.argument("device", type=click.STRING)
+@pass_config
+def remove(config, provider, region, node):
+    """
+    Remove a device.
+
+    DEVICE may be specified by name or ID, but ID is preferred.
+    """
+    devices = config.trainml.run(
+        config.trainml.client.cloudbender.devices.list(
+            provider_uuid=provider, region_uuid=region
+        )
+    )
+
+    found = search_by_id_name(device, devices)
+    if None is found:
+        raise click.UsageError("Cannot find specified device.")
+
+    return config.trainml.run(found.remove())

trainml/cli/dataset.py

@@ -252,9 +252,7 @@ def rename(config, dataset, name):
     DATASET may be specified by name or ID, but ID is preferred.
     """
     try:
-        dataset = config.trainml.run(
-            config.trainml.client.datasets.get(dataset)
-        )
+        dataset = config.trainml.run(config.trainml.client.datasets.get(dataset))
         if dataset is None:
             raise click.UsageError("Cannot find specified dataset.")
     except:

trainml/cli/job/create.py

@@ -683,7 +683,7 @@ def training(
         config.trainml.client.jobs.create(
             name=name,
             type="training",
-            gpu_type=gpu_type,
+            gpu_types=gpu_type,
             cpu_count=cpu_count,
             disk_size=disk_size,
             workers=[command for command in commands],
@@ -1021,7 +1021,7 @@ def inference(
         config.trainml.client.jobs.create(
             name=name,
             type="inference",
-            gpu_type=gpu_type,
+            gpu_types=gpu_type,
             cpu_count=cpu_count,
             disk_size=disk_size,
             workers=[command],
@@ -1330,7 +1330,7 @@ def endpoint(
         config.trainml.client.jobs.create(
             name=name,
             type="endpoint",
-            gpu_type=gpu_type,
+            gpu_types=gpu_type,
             cpu_count=cpu_count,
             disk_size=disk_size,
             endpoint=dict(routes=routes),

trainml/cli/model.py

@@ -1,4 +1,5 @@
 import click
+import logging
 from trainml.cli import cli, pass_config, search_by_id_name
 
 
@@ -64,6 +65,7 @@ def connect(config, model, attach):
     models = config.trainml.run(config.trainml.client.models.list())
 
     found = search_by_id_name(model, models)
+    logging.debug(found)
     if None is found:
         raise click.UsageError("Cannot find specified model.")
 

trainml/cloudbender/cloudbender.py

@@ -1,6 +1,7 @@
 from .providers import Providers
 from .regions import Regions
 from .nodes import Nodes
+from .devices import Devices
 from .datastores import Datastores
 from .reservations import Reservations
 from .device_configs import DeviceConfigs
@@ -12,6 +13,7 @@ class Cloudbender(object):
         self.providers = Providers(trainml)
         self.regions = Regions(trainml)
         self.nodes = Nodes(trainml)
+        self.devices = Devices(trainml)
         self.datastores = Datastores(trainml)
         self.reservations = Reservations(trainml)
         self.device_configs = DeviceConfigs(trainml)

trainml/cloudbender/datastores.py

@@ -118,3 +118,11 @@ class Datastore:
             f"/provider/{self._provider_uuid}/region/{self._region_uuid}/datastore/{self._id}",
             "DELETE",
         )
+
+    async def refresh(self):
+        resp = await self.trainml._query(
+            f"/provider/{self._provider_uuid}/region/{self._region_uuid}/datastore/{self._id}",
+            "GET",
+        )
+        self.__init__(self.trainml, **resp)
+        return self

trainml/cloudbender/device_configs.py

@@ -107,3 +107,11 @@ class DeviceConfig:
             f"/provider/{self._provider_uuid}/region/{self._region_uuid}/device/config/{self._id}",
             "DELETE",
         )
+
+    async def refresh(self):
+        resp = await self.trainml._query(
+            f"/provider/{self._provider_uuid}/region/{self._region_uuid}/device/config/{self._id}",
+            "GET",
+        )
+        self.__init__(self.trainml, **resp)
+        return self

trainml/cloudbender/devices.py

@@ -0,0 +1,190 @@
+import json
+import logging
+
+
+class Devices(object):
+    def __init__(self, trainml):
+        self.trainml = trainml
+
+    async def get(self, provider_uuid, region_uuid, id, **kwargs):
+        resp = await self.trainml._query(
+            f"/provider/{provider_uuid}/region/{region_uuid}/device/{id}",
+            "GET",
+            kwargs,
+        )
+        return Device(self.trainml, **resp)
+
+    async def list(self, provider_uuid, region_uuid, **kwargs):
+        resp = await self.trainml._query(
+            f"/provider/{provider_uuid}/region/{region_uuid}/device",
+            "GET",
+            kwargs,
+        )
+        devices = [Device(self.trainml, **device) for device in resp]
+        return devices
+
+    async def create(
+        self,
+        provider_uuid,
+        region_uuid,
+        friendly_name,
+        hostname,
+        minion_id,
+        **kwargs,
+    ):
+        logging.info(f"Creating Device {friendly_name}")
+        data = dict(
+            friendly_name=friendly_name,
+            hostname=hostname,
+            minion_id=minion_id,
+            type="device",
+            service="compute",
+            **kwargs,
+        )
+        payload = {k: v for k, v in data.items() if v is not None}
+        resp = await self.trainml._query(
+            f"/provider/{provider_uuid}/region/{region_uuid}/device",
+            "POST",
+            None,
+            payload,
+        )
+        device = Device(self.trainml, **resp)
+        logging.info(f"Created Device {friendly_name} with id {device.id}")
+        return device
+
+    async def remove(self, provider_uuid, region_uuid, id, **kwargs):
+        await self.trainml._query(
+            f"/provider/{provider_uuid}/region/{region_uuid}/device/{id}",
+            "DELETE",
+            kwargs,
+        )
+
+
+class Device:
+    def __init__(self, trainml, **kwargs):
+        self.trainml = trainml
+        self._device = kwargs
+        self._id = self._device.get("device_id")
+        self._provider_uuid = self._device.get("provider_uuid")
+        self._region_uuid = self._device.get("region_uuid")
+        self._name = self._device.get("friendly_name")
+        self._hostname = self._device.get("hostname")
+        self._status = self._device.get("status")
+        self._online = self._device.get("online")
+        self._maintenance_mode = self._device.get("maintenance_mode")
+        self._device_config_id = self._device.get("device_config_id")
+        self._job_status = self._device.get("job_status")
+        self._job_last_deployed = self._device.get("job_last_deployed")
+        self._job_config_id = self._device.get("job_config_id")
+        self._job_config_revision = self._device.get("job_config_revision")
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    @property
+    def provider_uuid(self) -> str:
+        return self._provider_uuid
+
+    @property
+    def region_uuid(self) -> str:
+        return self._region_uuid
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def hostname(self) -> str:
+        return self._hostname
+
+    @property
+    def status(self) -> str:
+        return self._status
+
+    @property
+    def online(self) -> bool:
+        return self._online
+
+    @property
+    def maintenance_mode(self) -> bool:
+        return self._maintenance_mode
+
+    @property
+    def device_config_id(self) -> str:
+        return self._device_config_id
+
+    @property
+    def job_status(self) -> str:
+        return self._job_status
+
+    @property
+    def job_last_deployed(self) -> str:
+        return self._job_last_deployed
+
+    @property
+    def job_config_id(self) -> str:
+        return self._job_config_id
+
+    @property
+    def job_config_revision(self) -> str:
+        return self._job_config_revision
+
+    def __str__(self):
+        return json.dumps({k: v for k, v in self._device.items()})
+
+    def __repr__(self):
+        return f"Device( trainml , **{self._device.__repr__()})"
+
+    def __bool__(self):
+        return bool(self._id)
+
+    async def remove(self):
+        await self.trainml._query(
+            f"/provider/{self._provider_uuid}/region/{self._region_uuid}/device/{self._id}",
+            "DELETE",
+        )
+
+    async def refresh(self):
+        resp = await self.trainml._query(
+            f"/provider/{self._provider_uuid}/region/{self._region_uuid}/device/{self._id}",
+            "GET",
+        )
+        self.__init__(self.trainml, **resp)
+        return self
+
+    async def toggle_maintenance(self):
+        await self.trainml._query(
+            f"/provider/{self._provider_uuid}/region/{self._region_uuid}/device/{self._id}/maintenance",
+            "PATCH",
+        )
+
+    async def run_action(self, command):
+        await self.trainml._query(
+            f"/provider/{self._provider_uuid}/region/{self._region_uuid}/device/{self._id}/action",
+            "POST",
+            None,
+            dict(command=command),
+        )
+
+    async def set_config(self, device_config_id):
+        resp = await self.trainml._query(
+            f"/provider/{self._provider_uuid}/region/{self._region_uuid}/device/{self._id}",
+            "PATCH",
+            None,
+            dict(device_config_id=device_config_id),
+        )
+        self.__init__(self.trainml, **resp)
+        return self
+
+    async def deploy_endpoint(self):
+        await self.trainml._query(
+            f"/provider/{self._provider_uuid}/region/{self._region_uuid}/device/{self._id}/deploy",
+            "PUT",
+        )
+
+    async def stop_endpoint(self):
+        await self.trainml._query(
+            f"/provider/{self._provider_uuid}/region/{self._region_uuid}/device/{self._id}/stop",
+            "PUT",
+        )

trainml/cloudbender/nodes.py

@@ -131,3 +131,25 @@ class Node:
             f"/provider/{self._provider_uuid}/region/{self._region_uuid}/node/{self._id}",
             "DELETE",
         )
+
+    async def refresh(self):
+        resp = await self.trainml._query(
+            f"/provider/{self._provider_uuid}/region/{self._region_uuid}/node/{self._id}",
+            "GET",
+        )
+        self.__init__(self.trainml, **resp)
+        return self
+
+    async def toggle_maintenance(self):
+        await self.trainml._query(
+            f"/provider/{self._provider_uuid}/region/{self._region_uuid}/node/{self._id}/maintenance",
+            "PATCH",
+        )
+
+    async def run_action(self, command):
+        await self.trainml._query(
+            f"/provider/{self._provider_uuid}/region/{self._region_uuid}/node/{self._id}/action",
+            "POST",
+            None,
+            dict(command=command),
+        )

trainml/cloudbender/providers.py

@@ -61,3 +61,11 @@ class Provider:
 
     async def remove(self):
         await self.trainml._query(f"/provider/{self._id}", "DELETE")
+
+    async def refresh(self):
+        resp = await self.trainml._query(
+            f"/provider/{self._id}",
+            "GET",
+        )
+        self.__init__(self.trainml, **resp)
+        return self

trainml/cloudbender/regions.py

@@ -80,6 +80,14 @@ class Region:
             f"/provider/{self._provider_uuid}/region/{self._id}", "DELETE"
         )
 
+    async def refresh(self):
+        resp = await self.trainml._query(
+            f"/provider/{self._provider_uuid}/region/{self._id}",
+            "GET",
+        )
+        self.__init__(self.trainml, **resp)
+        return self
+
     async def add_dataset(self, project_uuid, dataset_uuid, **kwargs):
         await self.trainml._query(
             f"/provider/{self._provider_uuid}/region/{self._id}/dataset",

trainml/cloudbender/reservations.py

@@ -116,3 +116,11 @@ class Reservation:
             f"/provider/{self._provider_uuid}/region/{self._region_uuid}/reservation/{self._id}",
             "DELETE",
         )
+
+    async def refresh(self):
+        resp = await self.trainml._query(
+            f"/provider/{self._provider_uuid}/region/{self._region_uuid}/reservation/{self._id}",
+            "GET",
+        )
+        self.__init__(self.trainml, **resp)
+        return self