proximl 0.5.4__py3-none-any.whl → 0.5.6__py3-none-any.whl

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

proximl/volumes.py

@@ -0,0 +1,255 @@
+import json
+import logging
+import math
+import asyncio
+from datetime import datetime
+
+from .exceptions import (
+    VolumeError,
+    ApiError,
+    SpecificationError,
+    ProxiMLException,
+)
+from .connections import Connection
+
+
+class Volumes(object):
+    def __init__(self, proximl):
+        self.proximl = proximl
+
+    async def get(self, id, **kwargs):
+        resp = await self.proximl._query(f"/volume/{id}", "GET", kwargs)
+        return Volume(self.proximl, **resp)
+
+    async def list(self, **kwargs):
+        resp = await self.proximl._query(f"/volume", "GET", kwargs)
+        volumes = [Volume(self.proximl, **volume) for volume in resp]
+        return volumes
+
+    async def create(self, name, source_type, source_uri, capacity, **kwargs):
+        data = dict(
+            name=name,
+            source_type=source_type,
+            source_uri=source_uri,
+            capacity=capacity,
+            source_options=kwargs.get("source_options"),
+            project_uuid=kwargs.get("project_uuid") or self.proximl.active_project,
+        )
+        payload = {k: v for k, v in data.items() if v is not None}
+        logging.info(f"Creating Volume {name}")
+        resp = await self.proximl._query("/volume", "POST", None, payload)
+        volume = Volume(self.proximl, **resp)
+        logging.info(f"Created Volume {name} with id {volume.id}")
+
+        return volume
+
+    async def remove(self, id, **kwargs):
+        await self.proximl._query(f"/volume/{id}", "DELETE", dict(**kwargs, force=True))
+
+
+class Volume:
+    def __init__(self, proximl, **kwargs):
+        self.proximl = proximl
+        self._volume = kwargs
+        self._id = self._volume.get("id", self._volume.get("id"))
+        self._status = self._volume.get("status")
+        self._name = self._volume.get("name")
+        self._capacity = self._volume.get("capacity")
+        self._used_size = self._volume.get("used_size")
+        self._billed_size = self._volume.get("billed_size")
+        self._project_uuid = self._volume.get("project_uuid")
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    @property
+    def status(self) -> str:
+        return self._status
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def capacity(self) -> str:
+        return self._capacity
+
+    @property
+    def used_size(self) -> int:
+        return self._used_size
+
+    @property
+    def billed_size(self) -> int:
+        return self._billed_size
+
+    def __str__(self):
+        return json.dumps({k: v for k, v in self._volume.items()})
+
+    def __repr__(self):
+        return f"Volume( proximl , **{self._volume.__repr__()})"
+
+    def __bool__(self):
+        return bool(self._id)
+
+    async def get_log_url(self):
+        resp = await self.proximl._query(
+            f"/volume/{self._id}/logs",
+            "GET",
+            dict(project_uuid=self._project_uuid),
+        )
+        return resp
+
+    async def get_details(self):
+        resp = await self.proximl._query(
+            f"/volume/{self._id}/details",
+            "GET",
+            dict(project_uuid=self._project_uuid),
+        )
+        return resp
+
+    async def get_connection_utility_url(self):
+        resp = await self.proximl._query(
+            f"/volume/{self._id}/download",
+            "GET",
+            dict(project_uuid=self._project_uuid),
+        )
+        return resp
+
+    def get_connection_details(self):
+        if self._volume.get("vpn"):
+            details = dict(
+                entity_type="volume",
+                project_uuid=self._volume.get("project_uuid"),
+                cidr=self._volume.get("vpn").get("cidr"),
+                ssh_port=self._volume.get("vpn").get("client").get("ssh_port"),
+                input_path=(
+                    self._volume.get("source_uri")
+                    if self.status in ["new", "downloading"]
+                    else None
+                ),
+                output_path=(
+                    self._volume.get("output_uri")
+                    if self.status == "exporting"
+                    else None
+                ),
+            )
+        else:
+            details = dict()
+        return details
+
+    async def connect(self):
+        if self.status in ["ready", "failed"]:
+            raise SpecificationError(
+                "status",
+                f"You can only connect to downloading or exporting volumes.",
+            )
+        if self.status == "new":
+            await self.wait_for("downloading")
+        connection = Connection(
+            self.proximl, entity_type="volume", id=self.id, entity=self
+        )
+        await connection.start()
+        return connection.status
+
+    async def disconnect(self):
+        connection = Connection(
+            self.proximl, entity_type="volume", id=self.id, entity=self
+        )
+        await connection.stop()
+        return connection.status
+
+    async def remove(self, force=False):
+        await self.proximl._query(
+            f"/volume/{self._id}",
+            "DELETE",
+            dict(project_uuid=self._project_uuid, force=force),
+        )
+
+    async def rename(self, name):
+        resp = await self.proximl._query(
+            f"/volume/{self._id}",
+            "PATCH",
+            dict(project_uuid=self._project_uuid),
+            dict(name=name),
+        )
+        self.__init__(self.proximl, **resp)
+        return self
+
+    async def export(self, output_type, output_uri, output_options=dict()):
+        resp = await self.proximl._query(
+            f"/volume/{self._id}/export",
+            "POST",
+            dict(project_uuid=self._project_uuid),
+            dict(
+                output_type=output_type,
+                output_uri=output_uri,
+                output_options=output_options,
+            ),
+        )
+        self.__init__(self.proximl, **resp)
+        return self
+
+    def _get_msg_handler(self, msg_handler):
+        def handler(data):
+            if data.get("type") == "subscription":
+                if msg_handler:
+                    msg_handler(data)
+                else:
+                    timestamp = datetime.fromtimestamp(int(data.get("time")) / 1000)
+                    print(
+                        f"{timestamp.strftime('%m/%d/%Y, %H:%M:%S')}: {data.get('msg').rstrip()}"
+                    )
+
+        return handler
+
+    async def attach(self, msg_handler=None):
+        await self.refresh()
+        if self.status not in ["ready", "failed"]:
+            await self.proximl._ws_subscribe(
+                "volume",
+                self._project_uuid,
+                self.id,
+                self._get_msg_handler(msg_handler),
+            )
+
+    async def refresh(self):
+        resp = await self.proximl._query(
+            f"/volume/{self.id}",
+            "GET",
+            dict(project_uuid=self._project_uuid),
+        )
+        self.__init__(self.proximl, **resp)
+        return self
+
+    async def wait_for(self, status, timeout=300):
+        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
+        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 == status:
+                return self
+            elif self.status == "failed":
+                raise VolumeError(self.status, self)
+            else:
+                count += 1
+                logging.debug(f"self: {self}, retry count {count}")
+
+        raise ProxiMLException(f"Timeout waiting for {status}")

{proximl-0.5.4.dist-info → proximl-0.5.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: proximl
-Version: 0.5.4
+Version: 0.5.6
 Summary: proxiML client SDK and command line utilities
 Home-page: https://github.com/proxiML/python-sdk
 Author: proxiML

{proximl-0.5.4.dist-info → proximl-0.5.6.dist-info}/RECORD

@@ -2,20 +2,21 @@ examples/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 examples/create_dataset_and_training_job.py,sha256=Fqueoz2KD1MTxnU960iqHUdxvo68Xe64HT5XS55lI9w,1178
 examples/local_storage.py,sha256=6K6LMO7ZPI7N2KdBcgqXSvdsqJfjISzN4yRO9YrJqbA,1745
 examples/training_inference_pipeline.py,sha256=pxux0QUUtRXxKj2rX-6fPEKBIi43mhRd1zJ-Lf1ZJGI,2334
-proximl/__init__.py,sha256=CW5DS3tZElDWIqGtqV8IS_5QxZk8vO7KIZcfdZ8g0gU,432
+proximl/__init__.py,sha256=3gq9ZsAAIsysDGCFRY4rF-eaMA6i7uPpZDDQ39rjBKM,432
 proximl/__main__.py,sha256=JgErYkiskih8Y6oRwowALtR-rwQhAAdqOYWjQraRIPI,59
 proximl/auth.py,sha256=LacGBDAVel5HcJx7JXp1wVn3s0gZc7nf--vDfSFOXYU,26565
 proximl/checkpoints.py,sha256=JPSq9iTSb6dBu0JFKJZgMqDBt9kSGxfl2OaxKehk4QY,8274
 proximl/connections.py,sha256=0C8VJSkDturQVlzg_3yAkQz8W9NAmgzDCBRK_SrzeUI,20035
 proximl/datasets.py,sha256=RB7IiDkWZPuwFw1o0xKQn10HS_3ui_BZf7tVlAvCIHA,7935
 proximl/environments.py,sha256=L_cRmau1wJxpGvnJAqgms-GiNdDhiuOntrlBqsdoE3A,1507
-proximl/exceptions.py,sha256=amVbfciOptsivh_z5rORxKGSj1QdeZx09myfQxObFZ8,3724
+proximl/exceptions.py,sha256=3q8qj-sahwEEHeYyjY89GrytVAgHDApN8HSCB3dhnf4,4094
 proximl/gpu_types.py,sha256=V-EZzE-hDLi5eVQ2_9yGLTm8-Qk1AnnzctfSVC44yLY,1901
-proximl/jobs.py,sha256=F-CMcU1l0RTPFT9NI5t5RtCdmU_gP2N3c7QHTnPY9cY,17807
+proximl/jobs.py,sha256=kQbC9ojrMEvRAysFfftwLdVtNDjpiFBfji4bz1To82E,17838
 proximl/models.py,sha256=QLtkIoznww5cZE_Eqoy7IaEgCEfNbJOe_21hdoExO-k,7750
 proximl/projects.py,sha256=87Fsce8BkPJLN_cEsZ3FLequpBgmd8WtjY7Vs7uZkO4,5209
-proximl/proximl.py,sha256=FAfAq6P8UOspuIavAEdYcXkM12TARgt86V5uB-zygTo,11045
-proximl/cli/__init__.py,sha256=G408VR_7LNs2Jww_hiLhFEPs0IHYAVOYEJ9barZcdLM,4338
+proximl/proximl.py,sha256=yzKbJ7ak8YiVv6q7CEagPufECNU4YTNeul6N3OuyH1M,10864
+proximl/volumes.py,sha256=VnlgTeQwoDKGpZRhMrSKtxe4EWzb4qQ-OdDyaJohH5Y,8094
+proximl/cli/__init__.py,sha256=R_8ExAKQp67N2xtwGM00gtK3zSoWPGruHAP_XFLddSI,4346
 proximl/cli/checkpoint.py,sha256=Iv1i1EAt2LJey2wy2ioQ6-ZysqwRG4kFj0lnE6INZCM,7170
 proximl/cli/connection.py,sha256=YiWqRIB9ZfTl30DjDFaJEpXuDrA-Ldl9PEzFFdZ_hFI,1700
 proximl/cli/dataset.py,sha256=ueoeicBY8aMLpvpKUIBICnS9GsEnDOj7ZlFmOfjjY4c,6871
@@ -23,6 +24,7 @@ proximl/cli/environment.py,sha256=nh7oYbG5oOrZEpZkMkKgvzFXmQJWnFTMw1-YbuvkdFU,10
 proximl/cli/gpu.py,sha256=xL8eqM5ca_Ueaj8cWit1iKn34KhaR0StrubVeRU2YQY,883
 proximl/cli/model.py,sha256=xdjveIaRPK7MdfrnFygPEuwYRJRW9VqheZ-11XnXDcE,6111
 proximl/cli/project.py,sha256=Er1twSiWQSAKir-hBIT9fRo2fc_UGqFoIJOwwjQGmlo,3522
+proximl/cli/volume.py,sha256=uyIrKov4zwCjyLyZrEJYoEbIkS0zdU3xSyWZk2BM1kA,6246
 proximl/cli/cloudbender/__init__.py,sha256=vxj62MyM3sC9h8M4ii3szH4s9JvEhicOQ0D0m7eNwPA,534
 proximl/cli/cloudbender/datastore.py,sha256=_vQOj-NfrL_nj4HfxNJL63TJZjLgfDyztRLyaRU58v8,3478
 proximl/cli/cloudbender/device.py,sha256=FdQZPESP6YBfUSzXq1Byu7eNMKi59qSOICONK-TEljI,3453
@@ -30,6 +32,7 @@ proximl/cli/cloudbender/node.py,sha256=xxzj68YvpRey2vZQasgYTnwv3x7TnwpuPSSf8Ma5a
 proximl/cli/cloudbender/provider.py,sha256=qhWbDK1tWi00wQWEYqGw7yGoZx0nEjV40GLHRuuE86c,1726
 proximl/cli/cloudbender/region.py,sha256=WnSkY4dXKRJ-FNaoxMfmoh6iuUx5dXCNJmEFT34Xtao,2892
 proximl/cli/cloudbender/reservation.py,sha256=xzHs5l8BbmYgKUq6kfFU-jEtRQY0j_vYnmRVcL4wwDo,3569
+proximl/cli/cloudbender/service.py,sha256=YimwXQml82-PKKQIqu9Cvggo9wOKLVTJPxMvDf4pgn8,2869
 proximl/cli/job/__init__.py,sha256=s8mU2PvCWDcv4gGT3EmjHn8MIZlXBAoayoZKmnKpXnY,6545
 proximl/cli/job/create.py,sha256=sGvbenY0yxvxHo-FZVbdw8FaZx5D4ekTCjD7P4YHG4g,34288
 proximl/cloudbender/__init__.py,sha256=iE29obtC0_9f0IhRvHQcG5aY58fVhVYipTakpjAhdss,64
@@ -41,19 +44,21 @@ proximl/cloudbender/nodes.py,sha256=QeWUaWW1HNvCune1lhakcve6LJyMzOy7cjCtvyOiaTs,
 proximl/cloudbender/providers.py,sha256=cH5lCew5WCFpXYS93vuoGNWkZyx7T_mdMH6YNWp2QGs,2036
 proximl/cloudbender/regions.py,sha256=Nu1LT6nuLD8Nt-5-7_FLlxDNZoDDAY6QduTdEBqfxJA,3570
 proximl/cloudbender/reservations.py,sha256=14ImJRLWQGG7CXDYhDnOI2W8pnP6CVVG2aVpysQVN0E,3586
+proximl/cloudbender/services.py,sha256=KR0EQ-BUVk1ov9Q195NsnfiSiMb9vOYD8XyxD87RRvw,3222
 tests/integration/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/integration/conftest.py,sha256=zRWpherX-yfbpk7xqZk9bIZCyJ-dwVeszY_7kekn2M4,1134
-tests/integration/test_checkpoints_integration.py,sha256=QJq6C1PHhS3MtmFIH3dX0aGNw4OvqBpVjvIzwLqPf4E,3128
-tests/integration/test_datasets_integration.py,sha256=yLAM2FUlMOIYN4WZYN82uumXHcXJacmwSgzmy88IIkg,3424
+tests/integration/test_checkpoints_integration.py,sha256=DUA6ZlX0gsnKCEvFmWCKbKkfJG3dQG6uD8T3fgo7y0M,3230
+tests/integration/test_datasets_integration.py,sha256=Ndp8itnncDSjVH0t5BM5R_-_yL4qt6JrkQAVOTMi1R8,3499
 tests/integration/test_environments_integration.py,sha256=7P6pKSyxA7rTwyNCD9HEaM2ablMG8WcBesOzkG-BgsQ,1403
 tests/integration/test_gpu_types_integration.py,sha256=Zv-yrHcAgKau9BJQvzi92bdrRHhLPl_hbhhzNLEWJ9w,1256
-tests/integration/test_jobs_integration.py,sha256=23bUE9Q0gLu41CrhCqLvoNSNMXj7_dWpg2O1--YumxI,23341
-tests/integration/test_models_integration.py,sha256=qbSzhCbxvrZQEFOvHJVBSmLr3fxoGD8PYwTcPyTX7J0,2848
-tests/integration/test_projects_integration.py,sha256=2iJ_I_LbV2jsGziWdwJ0H9o83IS24EGmS48wQrQHkbw,1492
+tests/integration/test_jobs_integration.py,sha256=fMqYfdzHOxUxtRGe0XZlsFhjUGoeo0Prxbq5_yQvroc,24730
+tests/integration/test_models_integration.py,sha256=JFgX6V652iMobpNKfn0XdP9_TDg5pnVtG8lL6yrEOZk,2902
+tests/integration/test_projects_integration.py,sha256=XdGGqvOQby77kMfoD_zAukKvN6tN7v1QydZ4hFp7u_s,1446
+tests/integration/test_volumes_integration.py,sha256=Xo2Whw2U-7jyvESIkyex3f0SMXKlExe3kLmsbpXTHhQ,3270
 tests/integration/cloudbender/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/integration/cloudbender/test_providers_integration.py,sha256=gFqPQom-Cn1iZC50_ChQ2us2_f4tIPATQSAUcWdf7ss,1473
 tests/unit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/unit/conftest.py,sha256=Hy2_Wz73PQak2b6MCDyz75sGZhbndRquifAFELJlSjE,29059
+tests/unit/conftest.py,sha256=dilDp7ZSmyg9tJbkJoGcqUF07JBTdkNoULOR6qK7NsU,31197
 tests/unit/test_auth.py,sha256=IJZHT5CZLNfu3MybTdEWIsKlvxNfFYpZ6oWYzbS856g,858
 tests/unit/test_checkpoints_unit.py,sha256=WfFE6HvQHglBKfcJ3B0IYXdXpaNmcL2BDurRQmIbFAU,16008
 tests/unit/test_connections_unit.py,sha256=LFAZzlrvL9oM8rZJTiC1oA9quw1KA2vMUCc3LV6SjXs,5507
@@ -65,6 +70,7 @@ tests/unit/test_jobs_unit.py,sha256=kCrm2FdogqYgji0Coe9KGaODYEsXmvuV46FN2mUE1Ts,
 tests/unit/test_models_unit.py,sha256=wgwdvJLt6QZ_IKCvp6Kmt2Z4b_CvIN8KCuFqlBaxDe8,15064
 tests/unit/test_projects_unit.py,sha256=_SwPFhy2Xl6HRt6w5_HYrqtZJHAzFzoCZ3kqAS1Zn3c,9530
 tests/unit/test_proximl.py,sha256=E-er5V1O-4ZVfTO2R-2wGsKN80q0_BVG5ho9Ab9lnE4,1701
+tests/unit/test_volumes_unit.py,sha256=YcyrGPCq7XfAtCdKdp44723o7mLGIu5_hKOdhAFVYV4,15305
 tests/unit/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/unit/cli/conftest.py,sha256=w6p_2URicywJKUCtY79tSD_mx8cwJtxHbK_Lu3grOYs,236
 tests/unit/cli/test_cli_checkpoint_unit.py,sha256=Z4v80C2gqzRXdHons3_82lNC_VqL_YOt53HrUQnzReI,702
@@ -74,6 +80,7 @@ tests/unit/cli/test_cli_gpu_unit.py,sha256=xdwIZEZJcJlWSuLBEcLhZXXH9EogZoKRiJlMR
 tests/unit/cli/test_cli_job_unit.py,sha256=UioSx_ZRY0qh4AJf90rIepputVUPBZ0KSqNR_u592UY,611
 tests/unit/cli/test_cli_model_unit.py,sha256=AucngxvYjW6GidDGBPHnKyYOb82ff7xMX5mVVUcbrCA,629
 tests/unit/cli/test_cli_project_unit.py,sha256=ms9gJ8pgMNGeIMdFcvBcwSPmb0i2qo9-rk9CCF53-9M,1756
+tests/unit/cli/test_cli_volume_unit.py,sha256=VeYb9PHuD45E23JCybF-C7ZaHsuERw7me2OUmYnmAAg,644
 tests/unit/cli/cloudbender/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/unit/cli/cloudbender/test_cli_datastore_unit.py,sha256=sUZWnzzCCG7NdgzEHsQ2zHpEb-ZD6FDIazca6FqMOc0,1381
 tests/unit/cli/cloudbender/test_cli_device_unit.py,sha256=o1vrPlbaanYK1iJG5pE6tDwgmOXDuLUY0VH8DxtaPYI,1342
@@ -89,9 +96,10 @@ tests/unit/cloudbender/test_nodes_unit.py,sha256=ehOHkNroiLKNTR09SbnBPpwELE72GcG
 tests/unit/cloudbender/test_providers_unit.py,sha256=y63VCqHXb4Yu8sh0kW30-ojRvv9aUa5j1jNkmb46KTc,4373
 tests/unit/cloudbender/test_regions_unit.py,sha256=9bvP268gpNyygjh1IEpSSiUt2aP6okv7QOsV1XoaIS0,6299
 tests/unit/cloudbender/test_reservations_unit.py,sha256=ICuFT5sexnLvS7taoC18yQYuDZHpBRrNuCj3Uq_Arwo,5624
-proximl-0.5.4.dist-info/LICENSE,sha256=ADFxLEZDxKY0j4MdyUd5GNuhQ18rnWH5rOz1ZG7yiOA,1069
-proximl-0.5.4.dist-info/METADATA,sha256=z0_DE7Fbt70q0F1oEKF-br_S5ZqiUCL9j2Ieo-tJZgs,7344
-proximl-0.5.4.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
-proximl-0.5.4.dist-info/entry_points.txt,sha256=HmI311IIabkZReMCXu-nGbvIEW-KfaduAOyfiSqt5SY,63
-proximl-0.5.4.dist-info/top_level.txt,sha256=-TWqc9tAaxmWmW4c7uYsmzPEYUIoh6z02xxqPbv7Kys,23
-proximl-0.5.4.dist-info/RECORD,,
+tests/unit/cloudbender/test_services_unit.py,sha256=iYaQpyCXDg77GQEIhmgiVwKX83jyvIf-4-4oya5WA_o,5043
+proximl-0.5.6.dist-info/LICENSE,sha256=ADFxLEZDxKY0j4MdyUd5GNuhQ18rnWH5rOz1ZG7yiOA,1069
+proximl-0.5.6.dist-info/METADATA,sha256=wS2y66puGjq7ZkG2GsyFIJWCiSjbDiA5uUdA8PzlIqw,7344
+proximl-0.5.6.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
+proximl-0.5.6.dist-info/entry_points.txt,sha256=HmI311IIabkZReMCXu-nGbvIEW-KfaduAOyfiSqt5SY,63
+proximl-0.5.6.dist-info/top_level.txt,sha256=-TWqc9tAaxmWmW4c7uYsmzPEYUIoh6z02xxqPbv7Kys,23
+proximl-0.5.6.dist-info/RECORD,,

tests/integration/test_checkpoints_integration.py

@@ -54,18 +54,20 @@ class GetCheckpointTests:
 
 @mark.create
 @mark.asyncio
-async def test_checkpoint_aws(proximl, capsys):
+async def test_checkpoint_wasabi(proximl, capsys):
     checkpoint = await proximl.checkpoints.create(
-        name="CLI Automated AWS",
-        source_type="aws",
-        source_uri="s3://proximl-examples/models/mxnet-model.zip",
+        name="CLI Automated Wasabi",
+        source_type="wasabi",
+        source_uri="s3://proximl-example/models/proximl-examples",
+        capacity="10G",
+        source_options=dict(endpoint_url="https://s3.wasabisys.com"),
     )
     checkpoint = await checkpoint.wait_for("ready", 300)
     status = checkpoint.status
     size = checkpoint.size
     await checkpoint.remove()
     assert status == "ready"
-    assert size >= 1000000
+    assert size >= 500000
 
 
 @mark.create

tests/integration/test_datasets_integration.py

@@ -13,8 +13,9 @@ class GetDatasetTests:
     async def dataset(self, proximl):
         dataset = await proximl.datasets.create(
             name="CLI Automated",
-            source_type="aws",
-            source_uri="s3://proximl-examples/data/cifar10",
+            source_type="wasabi",
+            source_uri="s3://proximl-example/input/cifar-10/cifar-10-batches-bin",
+            source_options=dict(endpoint_url="https://s3.wasabisys.com"),
         )
         dataset = await dataset.wait_for("ready", 300)
         yield dataset
@@ -48,9 +49,7 @@ class GetDatasetTests:
     async def test_dataset_repr(self, dataset):
         string = repr(dataset)
         regex = (
-            r"^Dataset\( proximl , \*\*{.*'dataset_uuid': '"
-            + dataset.id
-            + r"'.*}\)$"
+            r"^Dataset\( proximl , \*\*{.*'dataset_uuid': '" + dataset.id + r"'.*}\)$"
         )
         assert isinstance(string, str)
         assert re.match(regex, string)