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

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.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: proximl
-Version: 0.5.4
+Version: 0.5.5
 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.5.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
@@ -43,17 +45,18 @@ proximl/cloudbender/regions.py,sha256=Nu1LT6nuLD8Nt-5-7_FLlxDNZoDDAY6QduTdEBqfxJ
 proximl/cloudbender/reservations.py,sha256=14ImJRLWQGG7CXDYhDnOI2W8pnP6CVVG2aVpysQVN0E,3586
 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 +68,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 +78,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 +94,9 @@ 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,,
+proximl-0.5.5.dist-info/LICENSE,sha256=ADFxLEZDxKY0j4MdyUd5GNuhQ18rnWH5rOz1ZG7yiOA,1069
+proximl-0.5.5.dist-info/METADATA,sha256=ddskyt_dpYJW72fTThnwbbICvztVKb0eFgM2-crCFeg,7344
+proximl-0.5.5.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
+proximl-0.5.5.dist-info/entry_points.txt,sha256=HmI311IIabkZReMCXu-nGbvIEW-KfaduAOyfiSqt5SY,63
+proximl-0.5.5.dist-info/top_level.txt,sha256=-TWqc9tAaxmWmW4c7uYsmzPEYUIoh6z02xxqPbv7Kys,23
+proximl-0.5.5.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)

tests/integration/test_jobs_integration.py

@@ -269,7 +269,7 @@ class JobAPIDataValidationTests:
                 ),
             )
         assert (
-            "Invalid Request - Only regional datastore are allowed for Notebook and Endpoint output locations"
+            "Invalid Request - output_type invalid for Notebook and Endpoint jobs"
             in error.value.message
         )
 
@@ -298,7 +298,45 @@ class JobAPIDataValidationTests:
                 ),
             )
         assert (
-            "Invalid Request - Only regional datastore are allowed for Notebook and Endpoint output locations"
+            "Invalid Request - output_type invalid for Notebook and Endpoint jobs"
+            in error.value.message
+        )
+
+    async def test_invalid_volumes_for_training(self, proximl):
+        with raises(ApiError) as error:
+            await proximl.jobs.create(
+                name="Invalid Volumes for Training",
+                type="training",
+                gpu_types=["rtx3090"],
+                disk_size=10,
+                data=dict(
+                    output_uri="s3://proximl-examples/output/resnet_cifar10",
+                    output_type="aws",
+                    volumes=["volume-id"],
+                ),
+                workers=["python train.py"],
+            )
+        assert (
+            "Invalid Request - Only Notebook and Endpoint job types can use writable volumes"
+            in error.value.message
+        )
+
+    async def test_invalid_volumes_for_inference(self, proximl):
+        with raises(ApiError) as error:
+            await proximl.jobs.create(
+                name="Invalid Volumes for Inference",
+                type="inference",
+                gpu_types=["rtx3090"],
+                disk_size=10,
+                data=dict(
+                    output_uri="s3://proximl-examples/output/resnet_cifar10",
+                    output_type="aws",
+                    volumes=["volume-id"],
+                ),
+                workers=["python predict.py"],
+            )
+        assert (
+            "Invalid Request - Only Notebook and Endpoint job types can use writable volumes"
             in error.value.message
         )
 

tests/integration/test_models_integration.py

@@ -43,29 +43,27 @@ class GetModelTests:
 
     async def test_model_repr(self, model):
         string = repr(model)
-        regex = (
-            r"^Model\( proximl , \*\*{.*'model_uuid': '"
-            + model.id
-            + r"'.*}\)$"
-        )
+        regex = r"^Model\( proximl , \*\*{.*'model_uuid': '" + model.id + r"'.*}\)$"
         assert isinstance(string, str)
         assert re.match(regex, string)
 
 
 @mark.create
 @mark.asyncio
-async def test_model_aws(proximl, capsys):
+async def test_model_wasabi(proximl, capsys):
     model = await proximl.models.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"),
     )
     model = await model.wait_for("ready", 300)
     status = model.status
     size = model.size
     await model.remove()
     assert status == "ready"
-    assert size >= 1000000
+    assert size >= 500000
 
 
 @mark.create

tests/integration/test_projects_integration.py

@@ -11,9 +11,7 @@ pytestmark = [mark.sdk, mark.integration, mark.projects]
 class GetProjectsTests:
     @fixture(scope="class")
     async def project(self, proximl):
-        project = await proximl.projects.create(
-            name="New Project", copy_keys=False
-        )
+        project = await proximl.projects.create(name="New Project", copy_keys=False)
         yield project
         await project.remove()
 
@@ -41,8 +39,6 @@ class GetProjectsTests:
 
     async def test_project_repr(self, project):
         string = repr(project)
-        regex = (
-            r"^Project\( proximl , \*\*{.*'id': '" + project.id + r"'.*}\)$"
-        )
+        regex = r"^Project\( proximl , \*\*{.*'id': '" + project.id + r"'.*}\)$"
         assert isinstance(string, str)
         assert re.match(regex, string)

tests/integration/test_volumes_integration.py

@@ -0,0 +1,100 @@
+import re
+import sys
+import asyncio
+from pytest import mark, fixture
+
+pytestmark = [mark.sdk, mark.integration, mark.volumes]
+
+
+@mark.create
+@mark.asyncio
+class GetVolumeTests:
+    @fixture(scope="class")
+    async def volume(self, proximl):
+        volume = await proximl.volumes.create(
+            name="CLI Automated",
+            source_type="git",
+            source_uri="git@github.com:proxiML/environment-tests.git",
+            capacity="10G",
+        )
+        volume = await volume.wait_for("ready", 120)
+        yield volume
+        await volume.remove()
+        volume = await volume.wait_for("archived", 60)
+
+    async def test_get_volumes(self, proximl, volume):
+        volumes = await proximl.volumes.list()
+        assert len(volumes) > 0
+
+    async def test_get_volume(self, proximl, volume):
+        response = await proximl.volumes.get(volume.id)
+        assert response.id == volume.id
+
+    async def test_volume_properties(self, volume):
+        assert isinstance(volume.id, str)
+        assert isinstance(volume.status, str)
+        assert isinstance(volume.name, str)
+        assert isinstance(volume.capacity, str)
+        assert isinstance(volume.used_size, int)
+        assert isinstance(volume.billed_size, int)
+
+    async def test_volume_str(self, volume):
+        string = str(volume)
+        regex = r"^{.*\"id\": \"" + volume.id + r"\".*}$"
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    async def test_volume_repr(self, volume):
+        string = repr(volume)
+        regex = r"^Volume\( proximl , \*\*{.*'id': '" + volume.id + r"'.*}\)$"
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+
+@mark.create
+@mark.asyncio
+async def test_volume_wasabi(proximl, capsys):
+    volume = await proximl.volumes.create(
+        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"),
+    )
+    volume = await volume.wait_for("ready", 300)
+    status = volume.status
+    billed_size = volume.billed_size
+    used_size = volume.used_size
+    await volume.remove()
+    assert status == "ready"
+    assert billed_size >= 10000000
+    assert used_size >= 500000
+
+
+@mark.create
+@mark.asyncio
+async def test_volume_local(proximl, capsys):
+    volume = await proximl.volumes.create(
+        name="CLI Automated Local",
+        source_type="local",
+        source_uri="~/tensorflow-model",
+        capacity="10G",
+    )
+    attach_task = asyncio.create_task(volume.attach())
+    connect_task = asyncio.create_task(volume.connect())
+    await asyncio.gather(attach_task, connect_task)
+    await volume.disconnect()
+    await volume.refresh()
+    status = volume.status
+    billed_size = volume.billed_size
+    used_size = volume.used_size
+    await volume.remove()
+    assert status == "ready"
+    assert billed_size >= 10000000
+    assert used_size >= 1000000
+    captured = capsys.readouterr()
+    sys.stdout.write(captured.out)
+    sys.stderr.write(captured.err)
+    assert "Starting data upload from local" in captured.out
+    assert "official/LICENSE  11456 bytes" in captured.out
+    assert "Upload complete" in captured.out

tests/unit/cli/test_cli_volume_unit.py

@@ -0,0 +1,20 @@
+import re
+import json
+import click
+from unittest.mock import AsyncMock, patch
+from pytest import mark, fixture, raises
+
+pytestmark = [mark.cli, mark.unit, mark.volumes]
+
+from proximl.cli import volume as specimen
+from proximl.volumes import Volume
+
+
+def test_list(runner, mock_my_volumes):
+    with patch("proximl.cli.ProxiML", new=AsyncMock) as mock_proximl:
+        mock_proximl.volumes = AsyncMock()
+        mock_proximl.volumes.list = AsyncMock(return_value=mock_my_volumes)
+        result = runner.invoke(specimen, ["list"])
+        print(result)
+        assert result.exit_code == 0
+        mock_proximl.volumes.list.assert_called_once()