proximl 0.5.3__py3-none-any.whl → 0.5.5__py3-none-any.whl

proximl/__init__.py

@@ -13,5 +13,5 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 
-__version__ = "0.5.3"
+__version__ = "0.5.5"
 __all__ = "ProxiML"

proximl/cli/__init__.py

@@ -142,9 +142,7 @@ def configure(config):
         project for project in projects if project.id == active_project_id
     ]
 
-    active_project_name = (
-        active_project[0].name if len(active_project) else "UNSET"
-    )
+    active_project_name = active_project[0].name if len(active_project) else "UNSET"
 
     click.echo(f"Current Active Project: {active_project_name}")
 
@@ -154,9 +152,7 @@ def configure(config):
         show_choices=True,
         default=active_project_name,
     )
-    selected_project = [
-        project for project in projects if project.name == name
-    ]
+    selected_project = [project for project in projects if project.name == name]
     config.proximl.client.set_active_project(selected_project[0].id)
 
 
@@ -164,6 +160,7 @@ from proximl.cli.connection import connection
 from proximl.cli.dataset import dataset
 from proximl.cli.model import model
 from proximl.cli.checkpoint import checkpoint
+from proximl.cli.volume import volume
 from proximl.cli.environment import environment
 from proximl.cli.gpu import gpu
 from proximl.cli.job import job

proximl/cli/job/create.py

@@ -389,7 +389,7 @@ def notebook(
         ],
         case_sensitive=False,
     ),
-    default="rtx3090",
+    default=["rtx3090"],
     multiple=True,
     show_default=True,
     help="GPU type.",
@@ -732,7 +732,7 @@ def training(
         ],
         case_sensitive=False,
     ),
-    default="rtx3090",
+    default=["rtx3090"],
     show_default=True,
     multiple=True,
     help="GPU type.",
@@ -1099,7 +1099,7 @@ def from_json(config, attach, connect, file):
         ],
         case_sensitive=False,
     ),
-    default="rtx3090",
+    default=["rtx3090"],
     show_default=True,
     multiple=True,
     help="GPU type.",

proximl/cli/volume.py

@@ -0,0 +1,235 @@
+import click
+from proximl.cli import cli, pass_config, search_by_id_name
+
+
+def pretty_size(num):
+    if not num:
+        num = 0.0
+    s = ("  B", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB")
+    n = 0
+    while num > 1023:
+        num = num / 1024
+        n += 1
+    return f"{num:.2f} {s[n]}"
+
+
+@cli.group()
+@pass_config
+def volume(config):
+    """proxiML volume commands."""
+    pass
+
+
+@volume.command()
+@click.argument("volume", type=click.STRING)
+@pass_config
+def attach(config, volume):
+    """
+    Attach to volume and show creation logs.
+
+    VOLUME may be specified by name or ID, but ID is preferred.
+    """
+    volumes = config.proximl.run(config.proximl.client.volumes.list())
+
+    found = search_by_id_name(volume, volumes)
+    if None is found:
+        raise click.UsageError("Cannot find specified volume.")
+
+    try:
+        config.proximl.run(found.attach())
+        return config.proximl.run(found.disconnect())
+    except:
+        try:
+            config.proximl.run(found.disconnect())
+        except:
+            pass
+        raise
+
+
+@volume.command()
+@click.option(
+    "--attach/--no-attach",
+    default=True,
+    show_default=True,
+    help="Auto attach to volume and show creation logs.",
+)
+@click.argument("volume", type=click.STRING)
+@pass_config
+def connect(config, volume, attach):
+    """
+    Connect local source to volume and begin upload.
+
+    VOLUME may be specified by name or ID, but ID is preferred.
+    """
+    volumes = config.proximl.run(config.proximl.client.volumes.list())
+
+    found = search_by_id_name(volume, volumes)
+    if None is found:
+        raise click.UsageError("Cannot find specified volume.")
+
+    try:
+        if attach:
+            config.proximl.run(found.connect(), found.attach())
+            return config.proximl.run(found.disconnect())
+        else:
+            return config.proximl.run(found.connect())
+    except:
+        try:
+            config.proximl.run(found.disconnect())
+        except:
+            pass
+        raise
+
+
+@volume.command()
+@click.option(
+    "--attach/--no-attach",
+    default=True,
+    show_default=True,
+    help="Auto attach to volume and show creation logs.",
+)
+@click.option(
+    "--connect/--no-connect",
+    default=True,
+    show_default=True,
+    help="Auto connect source and start volume creation.",
+)
+@click.option(
+    "--source",
+    "-s",
+    type=click.Choice(["local"], case_sensitive=False),
+    default="local",
+    show_default=True,
+    help="Dataset source type.",
+)
+@click.argument("name", type=click.STRING)
+@click.argument("capacity", type=click.INT)
+@click.argument(
+    "path", type=click.Path(exists=True, file_okay=False, resolve_path=True)
+)
+@pass_config
+def create(config, attach, connect, source, name, capacity, path):
+    """
+    Create a volume.
+
+    A volume with maximum size CAPACITY is created with the specified NAME using a local source at the PATH
+    specified. PATH should be a local directory containing the source data for
+    a local source or a URI for all other source types.
+    """
+
+    if source == "local":
+        volume = config.proximl.run(
+            config.proximl.client.volumes.create(
+                name=name, source_type="local", source_uri=path, capacity=capacity
+            )
+        )
+
+        try:
+            if connect and attach:
+                config.proximl.run(volume.attach(), volume.connect())
+                return config.proximl.run(volume.disconnect())
+            elif connect:
+                return config.proximl.run(volume.connect())
+            else:
+                raise click.UsageError(
+                    "Abort!\n"
+                    "No logs to show for local sourced volume without connect."
+                )
+        except:
+            try:
+                config.proximl.run(volume.disconnect())
+            except:
+                pass
+            raise
+
+
+@volume.command()
+@click.argument("volume", type=click.STRING)
+@pass_config
+def disconnect(config, volume):
+    """
+    Disconnect and clean-up volume upload.
+
+    VOLUME may be specified by name or ID, but ID is preferred.
+    """
+    volumes = config.proximl.run(config.proximl.client.volumes.list())
+
+    found = search_by_id_name(volume, volumes)
+    if None is found:
+        raise click.UsageError("Cannot find specified volume.")
+
+    return config.proximl.run(found.disconnect())
+
+
+@volume.command()
+@pass_config
+def list(config):
+    """List volumes."""
+    data = [
+        ["ID", "STATUS", "NAME", "CAPACITY"],
+        ["-" * 80, "-" * 80, "-" * 80, "-" * 80],
+    ]
+
+    volumes = config.proximl.run(config.proximl.client.volumes.list())
+
+    for volume in volumes:
+        data.append(
+            [
+                volume.id,
+                volume.status,
+                volume.name,
+                volume.capacity,
+            ]
+        )
+    for row in data:
+        click.echo(
+            "{: >38.36} {: >13.11} {: >40.38} {: >14.12}" "".format(*row),
+            file=config.stdout,
+        )
+
+
+@volume.command()
+@click.option(
+    "--force/--no-force",
+    default=False,
+    show_default=True,
+    help="Force removal.",
+)
+@click.argument("volume", type=click.STRING)
+@pass_config
+def remove(config, volume, force):
+    """
+    Remove a volume.
+
+    VOLUME may be specified by name or ID, but ID is preferred.
+    """
+    volumes = config.proximl.run(config.proximl.client.volumes.list())
+
+    found = search_by_id_name(volume, volumes)
+    if None is found:
+        if force:
+            config.proximl.run(found.client.volumes.remove(volume))
+        else:
+            raise click.UsageError("Cannot find specified volume.")
+
+    return config.proximl.run(found.remove(force=force))
+
+
+@volume.command()
+@click.argument("volume", type=click.STRING)
+@click.argument("name", type=click.STRING)
+@pass_config
+def rename(config, volume, name):
+    """
+    Renames a volume.
+
+    VOLUME may be specified by name or ID, but ID is preferred.
+    """
+    try:
+        volume = config.proximl.run(config.proximl.client.volumes.get(volume))
+        if volume is None:
+            raise click.UsageError("Cannot find specified volume.")
+    except:
+        raise click.UsageError("Cannot find specified volume.")
+
+    return config.proximl.run(volume.rename(name=name))

proximl/exceptions.py

@@ -97,14 +97,27 @@ class CheckpointError(ProxiMLException):
         return self._status
 
     def __repr__(self):
-        return "CheckpointError({self.status}, {self.message})".format(
-            self=self
-        )
+        return "CheckpointError({self.status}, {self.message})".format(self=self)
 
     def __str__(self):
-        return "CheckpointError({self.status}, {self.message})".format(
-            self=self
-        )
+        return "CheckpointError({self.status}, {self.message})".format(self=self)
+
+
+class VolumeError(ProxiMLException):
+    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 "VolumeError({self.status}, {self.message})".format(self=self)
+
+    def __str__(self):
+        return "VolumeError({self.status}, {self.message})".format(self=self)
 
 
 class ConnectionError(ProxiMLException):
@@ -130,11 +143,7 @@ class SpecificationError(ProxiMLException):
         return self._attribute
 
     def __repr__(self):
-        return "SpecificationError({self.attribute}, {self.message})".format(
-            self=self
-        )
+        return "SpecificationError({self.attribute}, {self.message})".format(self=self)
 
     def __str__(self):
-        return "SpecificationError({self.attribute}, {self.message})".format(
-            self=self
-        )
+        return "SpecificationError({self.attribute}, {self.message})".format(self=self)

proximl/jobs.py

@@ -77,8 +77,7 @@ class Jobs(object):
             model=model,
             endpoint=endpoint,
             source_job_uuid=kwargs.get("source_job_uuid"),
-            project_uuid=kwargs.get("project_uuid")
-            or self.proximl.active_project,
+            project_uuid=kwargs.get("project_uuid") or self.proximl.active_project,
         )
         payload = {
             k: v
@@ -103,9 +102,7 @@ class Jobs(object):
         return job
 
     async def remove(self, id, **kwargs):
-        await self.proximl._query(
-            f"/job/{id}", "DELETE", dict(**kwargs, force=True)
-        )
+        await self.proximl._query(f"/job/{id}", "DELETE", dict(**kwargs, force=True))
 
 
 class Job:
@@ -308,18 +305,26 @@ class Job:
             entity_type="job",
             project_uuid=self._job.get("project_uuid"),
             cidr=self.dict.get("vpn").get("cidr"),
-            ssh_port=self._job.get("vpn").get("client").get("ssh_port")
-            if self._job.get("vpn").get("client")
-            else None,
-            model_path=self._job.get("model").get("source_uri")
-            if self._job.get("model").get("source_type") == "local"
-            else None,
-            input_path=self._job.get("data").get("input_uri")
-            if self._job.get("data").get("input_type") == "local"
-            else None,
-            output_path=self._job.get("data").get("output_uri")
-            if self._job.get("data").get("output_type") == "local"
-            else None,
+            ssh_port=(
+                self._job.get("vpn").get("client").get("ssh_port")
+                if self._job.get("vpn").get("client")
+                else None
+            ),
+            model_path=(
+                self._job.get("model").get("source_uri")
+                if self._job.get("model").get("source_type") == "local"
+                else None
+            ),
+            input_path=(
+                self._job.get("data").get("input_uri")
+                if self._job.get("data").get("input_type") == "local"
+                else None
+            ),
+            output_path=(
+                self._job.get("data").get("output_uri")
+                if self._job.get("data").get("output_type") == "local"
+                else None
+            ),
         )
         return details
 
@@ -396,8 +401,7 @@ class Job:
 
     def _get_msg_handler(self, msg_handler):
         worker_numbers = {
-            w.get("job_worker_uuid"): ind + 1
-            for ind, w in enumerate(self._workers)
+            w.get("job_worker_uuid"): ind + 1 for ind, w in enumerate(self._workers)
         }
         worker_numbers["data_worker"] = 0
 
@@ -407,9 +411,7 @@ class Job:
                 if msg_handler:
                     msg_handler(data)
                 else:
-                    timestamp = datetime.fromtimestamp(
-                        int(data.get("time")) / 1000
-                    )
+                    timestamp = datetime.fromtimestamp(int(data.get("time")) / 1000)
                     if len(self._workers) > 1:
                         print(
                             f"{timestamp.strftime('%m/%d/%Y, %H:%M:%S')}: Worker {data.get('worker_number')} - {data.get('msg').rstrip()}"
@@ -422,10 +424,7 @@ class Job:
         return handler
 
     async def attach(self, msg_handler=None):
-        if (
-            self.type == "notebook"
-            and self.status != "waiting for data/model download"
-        ):
+        if self.type == "notebook" and self.status != "waiting for data/model download":
             raise SpecificationError(
                 "type",
                 "Notebooks cannot be attached to after model download is complete.  Use open() instead.",
@@ -442,9 +441,7 @@ class Job:
     async def copy(self, name, **kwargs):
         logging.debug(f"copy request - name: {name} ; kwargs: {kwargs}")
         if self.type != "notebook":
-            raise SpecificationError(
-                "job", "Only notebook job types can be copied"
-            )
+            raise SpecificationError("job", "Only notebook job types can be copied")
 
         job = await self.proximl.jobs.create(
             name,
@@ -504,9 +501,7 @@ class Job:
 
         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:
@@ -519,23 +514,25 @@ class Job:
                 raise e
             if (
                 self.status == status
-                or (
-                    self.type == "training"
-                    and status == "finished"
-                    and self.status == "stopped"
-                )
                 or (
                     status
                     in [
                         "waiting for GPUs",
                         "waiting for resources",
                     ]  ## this status could be very short and the polling could miss it
-                    and self.status in ["starting", "provisioning", "running"]
+                    and self.status
+                    not in ["new", "waiting for GPUs", "waiting for resources"]
                 )
                 or (
                     status
                     == "waiting for data/model download"  ## this status could be very short and the polling could miss it
-                    and self.status in ["starting", "provisioning", "running"]
+                    and self.status
+                    not in [
+                        "new",
+                        "waiting for GPUs",
+                        "waiting for resources",
+                        "waiting for data/model download",
+                    ]
                 )
             ):
                 return self

proximl/proximl.py

@@ -10,6 +10,7 @@ from proximl.auth import Auth
 from proximl.datasets import Datasets
 from proximl.models import Models
 from proximl.checkpoints import Checkpoints
+from proximl.volumes import Volumes
 from proximl.jobs import Jobs
 from proximl.gpu_types import GpuTypes
 from proximl.environments import Environments
@@ -66,6 +67,7 @@ class ProxiML(object):
         self.datasets = Datasets(self)
         self.models = Models(self)
         self.checkpoints = Checkpoints(self)
+        self.volumes = Volumes(self)
         self.jobs = Jobs(self)
         self.gpu_types = GpuTypes(self)
         self.environments = Environments(self)
@@ -117,9 +119,7 @@ class ProxiML(object):
         )
         if params:
             if not isinstance(params, dict):
-                raise ProxiMLException(
-                    "Query parameters must be a valid dictionary"
-                )
+                raise ProxiMLException("Query parameters must be a valid dictionary")
             params = {
                 k: (str(v).lower() if isinstance(v, bool) else v)
                 for k, v in params.items()
@@ -155,13 +155,9 @@ class ProxiML(object):
                     content_type = resp.headers.get("content-type", "")
                     resp.close()
                     if content_type == "application/json":
-                        raise ApiError(
-                            resp.status, json.loads(what.decode("utf8"))
-                        )
+                        raise ApiError(resp.status, json.loads(what.decode("utf8")))
                     else:
-                        raise ApiError(
-                            resp.status, {"message": what.decode("utf8")}
-                        )
+                        raise ApiError(resp.status, {"message": what.decode("utf8")})
                 results = await resp.json()
                 return results
 
@@ -273,15 +269,11 @@ class ProxiML(object):
                     logging.debug(f"Websocket Disconnected.  Done? {done}")
                 except Exception as e:
                     connection_tries += 1
-                    logging.debug(
-                        f"Connection error: {traceback.format_exc()}"
-                    )
+                    logging.debug(f"Connection error: {traceback.format_exc()}")
                     if connection_tries == 5:
                         raise ApiError(
                             500,
-                            {
-                                "message": f"Connection error: {traceback.format_exc()}"
-                            },
+                            {"message": f"Connection error: {traceback.format_exc()}"},
                         )
 
     def set_active_project(self, project_uuid):