proximl 0.5.16__py3-none-any.whl → 1.0.0__py3-none-any.whl

proximl/volumes.py

@@ -10,7 +10,7 @@ from .exceptions import (
     SpecificationError,
     ProxiMLException,
 )
-from .connections import Connection
+from proximl.utils.transfer import upload, download
 
 
 class Volumes(object):
@@ -56,7 +56,9 @@ class Volumes(object):
         return volume
 
     async def remove(self, id, **kwargs):
-        await self.proximl._query(f"/volume/{id}", "DELETE", dict(**kwargs, force=True))
+        await self.proximl._query(
+            f"/volume/{id}", "DELETE", dict(**kwargs, force=True)
+        )
 
 
 class Volume:
@@ -120,56 +122,45 @@ class Volume:
         )
         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
+        if self.status not in ["downloading", "exporting"]:
+            if self.status == "new":
+                await self.wait_for("downloading")
+            else:
+                raise SpecificationError(
+                    "status",
+                    f"You can only connect to downloading or exporting volumes.",
+                )
 
-    async def disconnect(self):
-        connection = Connection(
-            self.proximl, entity_type="volume", id=self.id, entity=self
-        )
-        await connection.stop()
-        return connection.status
+        # Refresh to get latest entity data
+        await self.refresh()
+
+        if self.status == "downloading":
+            # Upload task - get auth_token, hostname, and source_uri from volume
+            auth_token = self._volume.get("auth_token")
+            hostname = self._volume.get("hostname")
+            source_uri = self._volume.get("source_uri")
+
+            if not auth_token or not hostname or not source_uri:
+                raise SpecificationError(
+                    "status",
+                    f"Volume in downloading status missing required connection properties (auth_token, hostname, source_uri).",
+                )
+
+            await upload(hostname, auth_token, source_uri)
+        elif self.status == "exporting":
+            # Download task - get auth_token, hostname, and output_uri from volume
+            auth_token = self._volume.get("auth_token")
+            hostname = self._volume.get("hostname")
+            output_uri = self._volume.get("output_uri")
+
+            if not auth_token or not hostname or not output_uri:
+                raise SpecificationError(
+                    "status",
+                    f"Volume in exporting status missing required connection properties (auth_token, hostname, output_uri).",
+                )
+
+            await download(hostname, auth_token, output_uri)
 
     async def remove(self, force=False):
         await self.proximl._query(
@@ -208,7 +199,9 @@ class Volume:
                 if msg_handler:
                     msg_handler(data)
                 else:
-                    timestamp = datetime.fromtimestamp(int(data.get("time")) / 1000)
+                    timestamp = datetime.fromtimestamp(
+                        int(data.get("time")) / 1000
+                    )
                     print(
                         f"{timestamp.strftime('%m/%d/%Y, %H:%M:%S')}: {data.get('msg').rstrip()}"
                     )
@@ -237,7 +230,7 @@ class Volume:
     async def wait_for(self, status, timeout=300):
         if self.status == status:
             return
-        valid_statuses = ["downloading", "ready", "archived"]
+        valid_statuses = ["downloading", "ready", "exporting", "archived"]
         if not status in valid_statuses:
             raise SpecificationError(
                 "status",
@@ -252,7 +245,9 @@ class Volume:
             )
         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:

{proximl-0.5.16.dist-info → proximl-1.0.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: proximl
-Version: 0.5.16
+Version: 1.0.0
 Summary: proxiML client SDK and command line utilities
 Home-page: https://github.com/proxiML/python-sdk
 Author: proxiML
@@ -12,8 +12,8 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.8
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
-Requires-Dist: aiodocker
 Requires-Dist: aiohttp
+Requires-Dist: aiofiles
 Requires-Dist: boto3
 Requires-Dist: Click (>8)
 Requires-Dist: python-jose[cryptography]
@@ -190,7 +190,7 @@ To list all datasets:
 proximl dataset list
 ```
 
-To connect to a job that requires the [connection capability](https://docs.proximl.ai/reference/connection-capability):
+To connect to a job that uses the "local" file transfer method:
 
 ```
 proximl job connect <job ID or name>

{proximl-0.5.16.dist-info → proximl-1.0.0.dist-info}/RECORD

@@ -1,28 +1,25 @@
 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/local_storage.py,sha256=1_bvGoE04ttvKT5MH7Vm8VFiJrw6nF_S7FojlcdC3kY,1687
 examples/training_inference_pipeline.py,sha256=pxux0QUUtRXxKj2rX-6fPEKBIi43mhRd1zJ-Lf1ZJGI,2334
-proximl/__init__.py,sha256=9y__q8yiD7SWW8PzY4Cx6ZFTodmOM-e-L572DUjwATg,433
+proximl/__init__.py,sha256=6bDDwARPoJxFWPJex__P6-Qdw49mnh_w-YXWnlVLSYg,432
 proximl/__main__.py,sha256=JgErYkiskih8Y6oRwowALtR-rwQhAAdqOYWjQraRIPI,59
-proximl/auth.py,sha256=LacGBDAVel5HcJx7JXp1wVn3s0gZc7nf--vDfSFOXYU,26565
-proximl/checkpoints.py,sha256=Ezpiab9Wfcmd2h5Slm9U5lekZGXiPzvhDKxXXgla1yo,8790
-proximl/connections.py,sha256=0C8VJSkDturQVlzg_3yAkQz8W9NAmgzDCBRK_SrzeUI,20035
-proximl/datasets.py,sha256=dBDzf7DZ1rS5LmJPy5oV-JC0PXGTf-NZuBDG1KPyoA4,8600
+proximl/checkpoints.py,sha256=SClBmtCUM4Ql8DWEDw4mpDERJAXE7z9ZnYqEcTwsr-g,8924
+proximl/datasets.py,sha256=m2rza1wdRx0DPWc7cHWcf7J1FNdMrWV59Yxlq7ouVYU,8690
 proximl/environments.py,sha256=L_cRmau1wJxpGvnJAqgms-GiNdDhiuOntrlBqsdoE3A,1507
 proximl/exceptions.py,sha256=q1YfsqbLw3y1DJHP1FcMM8xG5CEhQsTIvRfq4v5PyVw,5468
 proximl/gpu_types.py,sha256=V-EZzE-hDLi5eVQ2_9yGLTm8-Qk1AnnzctfSVC44yLY,1901
-proximl/jobs.py,sha256=HiGs2gRKPFZReHLg9DLPsDWEecxZpafWv6XFBSGXYj0,17947
-proximl/models.py,sha256=gJQ1C3HGEHScW41PAQS3Q8IKSQSauIUIJxOZvkDzizQ,8284
-proximl/proximl.py,sha256=7UyqYbGO9edFJec54xqk_wwahKldXQlY5112Izg3saA,11902
-proximl/volumes.py,sha256=qRIgzvJ4NTHuLi7QN_6v61Urns8l4Y9KjWladCIQ2gk,8443
-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
+proximl/jobs.py,sha256=lBj1UGPKDQAPbQEAUQ5lPL4-35wHNVDQVUGdtoE8G5Q,24928
+proximl/models.py,sha256=ZigU31SdpsxsnkW5wMt4narvlqVO5MbxTQuV__Mndqo,8339
+proximl/proximl.py,sha256=x0P7D6iMKl1nXWCWMzaHhFo63oxb7RWhf1KPa9zUOkA,12724
+proximl/volumes.py,sha256=uQDVLLJa_-GNdxwo8xh67KcGhlJ8EMHtW_V3B9IfgWA,8533
+proximl/cli/__init__.py,sha256=vtNYq0HJ491CWMjIevCVSnLilCIDQv9LBhZmVWv9N3c,4330
+proximl/cli/checkpoint.py,sha256=mgEQ13XO9FS1SOXW4hTfbN2cqAYY7izSWt2KyrEITJQ,6020
+proximl/cli/dataset.py,sha256=MpAC-AvXXEZXr1PE0ch5TPq1U8v2fecbrTQcvJ8Lhcc,5721
 proximl/cli/environment.py,sha256=nh7oYbG5oOrZEpZkMkKgvzFXmQJWnFTMw1-YbuvkdFU,1019
 proximl/cli/gpu.py,sha256=xL8eqM5ca_Ueaj8cWit1iKn34KhaR0StrubVeRU2YQY,883
-proximl/cli/model.py,sha256=xdjveIaRPK7MdfrnFygPEuwYRJRW9VqheZ-11XnXDcE,6111
-proximl/cli/volume.py,sha256=uyIrKov4zwCjyLyZrEJYoEbIkS0zdU3xSyWZk2BM1kA,6246
+proximl/cli/model.py,sha256=n6hghRf3Krj1K39UD6xg8UvTGIax0XxXrx7hGe8TiSY,4963
+proximl/cli/volume.py,sha256=lACn2vG0R1NMfGLRHl6t9oSUZI2twouv3Y-N3kwri_Q,5135
 proximl/cli/cloudbender/__init__.py,sha256=ATPj85Be6o4xw8M6vvHfkuLyw1drc_0IlCMh1wSebYQ,592
 proximl/cli/cloudbender/data_connector.py,sha256=D5xbUzpUGpn_NN-NZQ9vSxOZY610c1LsoBC40DkPwkw,3606
 proximl/cli/cloudbender/datastore.py,sha256=YBFbobEKjw3s5uYC1krD_Oa-VEz3kjQ5UvPOqkq-NJE,3323
@@ -31,8 +28,8 @@ 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/service.py,sha256=6NuwlFULJLtOmMy9OFA8GkW4MLvNemAcd_KitQyDXV8,3147
-proximl/cli/job/__init__.py,sha256=s8mU2PvCWDcv4gGT3EmjHn8MIZlXBAoayoZKmnKpXnY,6545
-proximl/cli/job/create.py,sha256=plRiX9lRqt6lrjWvTIVbLC-EoDgSL0-oYqwIfuWb_ek,35098
+proximl/cli/job/__init__.py,sha256=IaoNWZUn4VHHahTaAk9z7Vr4CXzHk4v3RWHzPcEZP-8,5379
+proximl/cli/job/create.py,sha256=EmIxMO9WyN06fwsAGl7zlUJURa-9z4n3jWOA38ONfcw,35157
 proximl/cli/project/__init__.py,sha256=ZSiXhsq7O1JwGecsBKBOD1W0VuiwErs-FFsfrXMOiFs,1918
 proximl/cli/project/credential.py,sha256=5vPI6VBIo_YyvAzfUa9ARmz-LUnkdIVW0zbko9t0Q38,3443
 proximl/cli/project/data_connector.py,sha256=CKF4snfrzDmeDXb4au8nV2W9jTweRMDZ89U2GAlBedQ,1411
@@ -54,18 +51,21 @@ proximl/projects/credentials.py,sha256=hWz6EUEAgltkAQMDjQVsFlVvfWUbZfY3KoWpldnCr
 proximl/projects/data_connectors.py,sha256=o2-K40BdF85TBXJgj1UMsznBmNHmnYh1qbZtHmmOa2Y,2216
 proximl/projects/datastores.py,sha256=8nHpMEHn3UzsIhmPMoXympbRRJ10XkHdyRhBXUVXyeE,2095
 proximl/projects/members.py,sha256=siREAa-Oac1pLGRHMXmj3ShoB0DVl8uY34eFdTlSezA,2826
-proximl/projects/projects.py,sha256=PJYUaXbYMINcvlDQePKUSAz6uO11jg9tj5gJW-cNjGQ,2890
+proximl/projects/projects.py,sha256=2lYStCIQ6_KB411NS_u1Ia2uatKxre1jFS6Lbr2LqzY,2840
 proximl/projects/secrets.py,sha256=xlfzRONjT2ySJDcq8FYRQmSU2ZJMdPLS8NbePWYNCeM,2160
 proximl/projects/services.py,sha256=aCESZjqfI9K-iXl8uJQsjEyoX29LuPaNv_TFjK6B2Ts,2771
+proximl/utils/__init__.py,sha256=omDnxAdZ_MWqiK0cfcqPEp1EfRTHUCAC_kJYQa77K6M,39
+proximl/utils/auth.py,sha256=K-JqDmifY83hVqGJvk1qlEEK8fb5KmVqpqQrEX9X-wA,26573
+proximl/utils/transfer.py,sha256=EEQOUVbdu12P4gWcECuZYjozoi8lMcak9dVIz4guGqY,22395
 tests/integration/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/integration/conftest.py,sha256=rqOKgDZ37SZH6CxXZIQSHtx8SF5GC1MlC49Zb5xrzsk,1135
-tests/integration/test_checkpoints_integration.py,sha256=2gKhbax96ePIEZqcM4IeqSB0fTOYPh9nexNJtC96LuU,3239
-tests/integration/test_datasets_integration.py,sha256=y8NL7fKQvZhJs1r9UAcLjXI0832tHr9XfIaRhHX9aOk,3529
+tests/integration/test_checkpoints_integration.py,sha256=qHD-tzQc6QclqKf5o3pH2phcUps7K0PS3ctqPJls8M0,3219
+tests/integration/test_datasets_integration.py,sha256=o1abXw5vvZ3RvKA-sQNBO3xJSDVgPaX-P4lNTeRnhWk,3518
 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=mSgPERDycfphLuLd7Hp3tSo-SpnbACp963x99K3XbtU,25263
-tests/integration/test_models_integration.py,sha256=du9i3GsWjwEWei_-5BzfoBneBNcZ-9NQsDdOfn4dzd8,2906
-tests/integration/test_volumes_integration.py,sha256=fTsP5_wyspzR_DX_ArXz5Fflq2dEmXtJbhwqCjSQDOA,3299
+tests/integration/test_jobs_integration.py,sha256=sX_DYopEW8mwKepMteyiaHoZMnd5Ct_ZKOjtU3uHZZw,25386
+tests/integration/test_models_integration.py,sha256=2GY_RF_fvARIO67YS4qfAMfndk3NO8abcXMiEhvjH7g,2924
+tests/integration/test_volumes_integration.py,sha256=ecgwQRc83glqHoRwohrDxvNHQ0XLFDwCFbVLWSN6C3M,3268
 tests/integration/cloudbender/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/integration/cloudbender/conftest.py,sha256=K1EQC7U7aBnywJq-QvBZINqHtQm7kq5feDWVKAtp22E,777
 tests/integration/cloudbender/test_providers_integration.py,sha256=8_1TBxcIpQcET7VStHIsi_-_3j5hs8H-icsKN3qnNa8,1327
@@ -81,21 +81,21 @@ tests/integration/projects/test_projects_members_integration.py,sha256=4i4SeDCDZ
 tests/integration/projects/test_projects_secrets_integration.py,sha256=TyqItJs-SBzOUwDy3vs_-yrDXeDF8r1_bAovr8C1a_s,1520
 tests/integration/projects/test_projects_services_integration.py,sha256=cHOGSUpl3rqYJnCrbYp-2Zmn1ngFJTDV9Io7K8fDOp0,1569
 tests/unit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/unit/conftest.py,sha256=11Z49O4mKXtm6xl18kWPomZvbZJQpqdaDYeJHFQbO1A,35588
-tests/unit/test_auth_unit.py,sha256=IJZHT5CZLNfu3MybTdEWIsKlvxNfFYpZ6oWYzbS856g,858
-tests/unit/test_checkpoints_unit.py,sha256=FbBM8Ffqy8JGTpWq2zWAWle2CI8A0MB_zABL_9Mezbk,15965
-tests/unit/test_connections_unit.py,sha256=LFAZzlrvL9oM8rZJTiC1oA9quw1KA2vMUCc3LV6SjXs,5507
-tests/unit/test_datasets_unit.py,sha256=c3suBgQMtzJckHlY8MFwU38WundXobbjb-KXmUboulI,15910
+tests/unit/conftest.py,sha256=LRrAR1KxlVDVuiSgS7krd12CxCaEXU4nG1fV2lUjjY0,35592
+tests/unit/test_auth_unit.py,sha256=3oqdvz2jM-G37QM9FK6oCvS_w3EcVadz7rKawFqt-OU,1413
+tests/unit/test_checkpoints_unit.py,sha256=OcPL61j9AqZWWPQsDxME5GQCud0xowkJuAB-__3OkB0,21914
+tests/unit/test_datasets_unit.py,sha256=rv4qQoIjQaPCMPOGkH4T5sZyi5-cJlIDwtWr3OC0tK0,21135
 tests/unit/test_environments_unit.py,sha256=HS4NN_cRWGQNDfovCxNGrVRoZI4NcUgsXyzvapNi6Ms,1882
-tests/unit/test_exceptions.py,sha256=ASiwzV9z-6jk0Y2qSHEBmGOwKx-hXz5vaHcyApbjSKA,906
-tests/unit/test_gpu_types_unit.py,sha256=d2JCQcFgBpkHEBUhKeiyngioEs9cz2qMGM2MZdwg-1o,2054
-tests/unit/test_jobs_unit.py,sha256=kCrm2FdogqYgji0Coe9KGaODYEsXmvuV46FN2mUE1Ts,26835
-tests/unit/test_models_unit.py,sha256=e9tcTcojLLTPKgx1wdAwP3j1k6C2aBnCbdxBhwKiRoQ,15091
-tests/unit/test_proximl_unit.py,sha256=E-er5V1O-4ZVfTO2R-2wGsKN80q0_BVG5ho9Ab9lnE4,1701
-tests/unit/test_volumes_unit.py,sha256=6FwEfANDZ2Z_ec3OL3Q12H_ml3TK9q2WgJW8hN-qwrc,15382
+tests/unit/test_exceptions.py,sha256=qD6YV1wwAx-taCXSaBisc1w88-pKujJJPQykjebOIk8,5749
+tests/unit/test_gpu_types_unit.py,sha256=eCPzi3EhKupDhlfIMPCIHCP3HXDM-KiheSsigxKJGIw,2526
+tests/unit/test_jobs_unit.py,sha256=V_9xCupqQwcg_Cz3dtbZDsB9sknlgKJ-Hm6mElR_6zk,63964
+tests/unit/test_main_unit.py,sha256=DSa7gE9D1ccxrPZCC0mUqWqTm7otC-OOGBfHaI_998I,627
+tests/unit/test_models_unit.py,sha256=6YKMHHrAnhxU7RvVpqsf4elzcK8TOaQsuw0dscIUiu8,20127
+tests/unit/test_proximl_unit.py,sha256=q0Sifg6uu8F_l5lUw6Sxifw-q5Y45eI7Bfk3A8o5WmA,23756
+tests/unit/test_volumes_unit.py,sha256=2w3Nv7IqMMK3nXgmEN2LNNmt67hp2omeKNe1CB4_amo,20185
 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
+tests/unit/cli/test_cli_checkpoint_unit.py,sha256=9-mtZxHb06b5QyyAXurcbOdRzdSfqBfUGkmZRiiKTgI,11830
 tests/unit/cli/test_cli_datasets_unit.py,sha256=Hq9kWHMXIG2rI5GWjN-QuCquZ1D4Qaa62Bow0iPDOLk,653
 tests/unit/cli/test_cli_environment_unit.py,sha256=iiNj_fxiCinq9FZgkwzcO_OLu9O279Hk8nJpFJwDpPg,683
 tests/unit/cli/test_cli_gpu_unit.py,sha256=xdwIZEZJcJlWSuLBEcLhZXXH9EogZoKRiJlMRxaDa3o,650
@@ -121,21 +121,23 @@ tests/unit/cloudbender/test_data_connectors_unit.py,sha256=ZCnI9Gr2RAcpHiPZ5t2ka
 tests/unit/cloudbender/test_datastores_unit.py,sha256=8PKa6XRdZ1MoVEwbRxA67qKeI6DrycpLU1RxWeokMnw,5281
 tests/unit/cloudbender/test_device_configs_unit.py,sha256=wBDnTfrP6yl0K_jWcNU2ADR6iHoe1sRzMpvNfs_dwXs,5714
 tests/unit/cloudbender/test_devices_unit.py,sha256=C2YTnfIpHlxdidgfbTnlzl72r5O7kqKStQUhWcTTXDg,9103
-tests/unit/cloudbender/test_nodes_unit.py,sha256=ehOHkNroiLKNTR09SbnBPpwELE72GcGT6le_YN5aVwQ,6231
-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_services_unit.py,sha256=V1JHfqNcLjz7d3qtVxPBCKzd5_lcL5iwhvWqQ0kw0J0,5220
+tests/unit/cloudbender/test_nodes_unit.py,sha256=qsiqAVDU_TaLStSBVwKJeK4Y1GzmNCTtGqeSYJJMtP0,10928
+tests/unit/cloudbender/test_providers_unit.py,sha256=acCN_8sjmzg9Wc87Erd6M48-ngWhcsWziPMsIbRDeQM,8588
+tests/unit/cloudbender/test_regions_unit.py,sha256=J_Qn_Vi-FQfMLqIy5T4MwEPUxjZW8Z2cDMpkX7esJSE,10961
+tests/unit/cloudbender/test_services_unit.py,sha256=ZmFaleqzIeIXXDGAqYlochdoVMwgiHQDBg3vxRutc4Q,10956
 tests/unit/projects/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/unit/projects/test_project_credentials_unit.py,sha256=y528DyvIm7uc5mFuwiUPjOMO1dkVcMEeYQqzMmKLOuk,3480
-tests/unit/projects/test_project_data_connectors_unit.py,sha256=Yx2yCUgcuN_KfTKynZxse0Nx8jRrSlMygI7R8a-NjE4,3385
-tests/unit/projects/test_project_datastores_unit.py,sha256=WZMKkhpEg2dcevJqT7k6QRcYJF6FJfq177BDk2GWOBk,3129
-tests/unit/projects/test_project_members_unit.py,sha256=dtbRfP454mviQYn6Dc0Uzyb3BBIuR3Ive6g3mnUAR6M,3466
+tests/unit/projects/test_project_data_connectors_unit.py,sha256=C7IPXiZ6ZhAcgLFbjkkI8JDLt13DLI6eyZZFB1i3tGE,5004
+tests/unit/projects/test_project_datastores_unit.py,sha256=9s_bzl2WrEyysyCaFI77kyY9h-wDGwRnvzHvKO4OiSk,4646
+tests/unit/projects/test_project_members_unit.py,sha256=J9l1sS1U4Uj4zy89xP_3ZOjz5bEklLk4sWJ51H11H3E,5097
 tests/unit/projects/test_project_secrets_unit.py,sha256=VE9L91FJodcwVGizfF65WYMiHZaF0s2AdW1aiJ3z7xA,3276
-tests/unit/projects/test_project_services_unit.py,sha256=PzeNuJRuAG7RkrPWX0FfgFTt6-63FviecrDY06rLQ6A,3331
-tests/unit/projects/test_projects_unit.py,sha256=4vMo7TF-SjduoT2ejpuyM3dULslpdqE8-c25M9X-iYc,3810
-proximl-0.5.16.dist-info/LICENSE,sha256=ADFxLEZDxKY0j4MdyUd5GNuhQ18rnWH5rOz1ZG7yiOA,1069
-proximl-0.5.16.dist-info/METADATA,sha256=Pm7udhIj-sHr1U3trGMxzbBclXIjHue5W9U-uRKxMxc,7345
-proximl-0.5.16.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
-proximl-0.5.16.dist-info/entry_points.txt,sha256=HmI311IIabkZReMCXu-nGbvIEW-KfaduAOyfiSqt5SY,63
-proximl-0.5.16.dist-info/top_level.txt,sha256=-TWqc9tAaxmWmW4c7uYsmzPEYUIoh6z02xxqPbv7Kys,23
-proximl-0.5.16.dist-info/RECORD,,
+tests/unit/projects/test_project_services_unit.py,sha256=8CHd6w2koDo7cwIba6_nGuHe5wCnCRv9Q8xsv_F9I7U,6140
+tests/unit/projects/test_projects_unit.py,sha256=J6nXrigjMJONSfKpNzgWOUl0CrB3GYaLmcjnale5zvQ,4409
+tests/unit/utils/__init__.py,sha256=VDy3m4Lfazpey8NMzHWdgcs3T9aaXWuH7pglIjCnZW0,38
+tests/unit/utils/test_transfer_unit.py,sha256=hfPv-6sTnUlwrZdxTzciaF_huXOVkyX662UIGzmO994,188309
+proximl-1.0.0.dist-info/LICENSE,sha256=ADFxLEZDxKY0j4MdyUd5GNuhQ18rnWH5rOz1ZG7yiOA,1069
+proximl-1.0.0.dist-info/METADATA,sha256=9jMrrMKtP_j3VHdT6wukLnztZDDRHzou0NifyfaOz3Y,7287
+proximl-1.0.0.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
+proximl-1.0.0.dist-info/entry_points.txt,sha256=HmI311IIabkZReMCXu-nGbvIEW-KfaduAOyfiSqt5SY,63
+proximl-1.0.0.dist-info/top_level.txt,sha256=-TWqc9tAaxmWmW4c7uYsmzPEYUIoh6z02xxqPbv7Kys,23
+proximl-1.0.0.dist-info/RECORD,,

tests/integration/test_checkpoints_integration.py

@@ -23,6 +23,7 @@ class GetCheckpointTests:
         checkpoint = await checkpoint.wait_for("archived", 60)
 
     async def test_get_checkpoints(self, proximl, checkpoint):
+        _ = checkpoint
         checkpoints = await proximl.checkpoints.list()
         assert len(checkpoints) > 0
 
@@ -55,7 +56,7 @@ class GetCheckpointTests:
 
 @mark.create
 @mark.asyncio
-async def test_checkpoint_wasabi(proximl, capsys):
+async def test_checkpoint_wasabi(proximl):
     checkpoint = await proximl.checkpoints.create(
         name="CLI Automated Wasabi",
         source_type="wasabi",
@@ -72,6 +73,7 @@ async def test_checkpoint_wasabi(proximl, capsys):
 
 @mark.create
 @mark.asyncio
+@mark.local
 async def test_checkpoint_local(proximl, capsys):
     checkpoint = await proximl.checkpoints.create(
         name="CLI Automated Local",
@@ -81,7 +83,6 @@ async def test_checkpoint_local(proximl, capsys):
     attach_task = asyncio.create_task(checkpoint.attach())
     connect_task = asyncio.create_task(checkpoint.connect())
     await asyncio.gather(attach_task, connect_task)
-    await checkpoint.disconnect()
     await checkpoint.refresh()
     status = checkpoint.status
     size = checkpoint.size
@@ -92,5 +93,5 @@ async def test_checkpoint_local(proximl, capsys):
     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 "official/LICENSE" in captured.out
     assert "Upload complete" in captured.out

tests/integration/test_datasets_integration.py

@@ -50,7 +50,9 @@ 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)
@@ -79,6 +81,7 @@ class GetDatasetTests:
 
 @mark.create
 @mark.asyncio
+@mark.local
 async def test_dataset_local(proximl, capsys):
     dataset = await proximl.datasets.create(
         name="CLI Automated Local",
@@ -88,7 +91,6 @@ async def test_dataset_local(proximl, capsys):
     attach_task = asyncio.create_task(dataset.attach())
     connect_task = asyncio.create_task(dataset.connect())
     await asyncio.gather(attach_task, connect_task)
-    await dataset.disconnect()
     await dataset.refresh()
     status = dataset.status
     size = dataset.size
@@ -99,5 +101,5 @@ async def test_dataset_local(proximl, capsys):
     sys.stdout.write(captured.out)
     sys.stderr.write(captured.err)
     assert "Starting data upload from local" in captured.out
-    assert "data_batch_1.bin  30733788 bytes" in captured.out
+    assert "data_batch_1.bin" in captured.out
     assert "Upload complete" in captured.out

tests/integration/test_jobs_integration.py

@@ -46,7 +46,10 @@ class JobLifeCycleTests:
         job = await job.wait_for("running")
         assert job.status == "running"
         assert job.url
-        assert extract_domain_suffix(urlparse(job.url).hostname) == "proximl.cloud"
+        assert (
+            extract_domain_suffix(urlparse(job.url).hostname)
+            == "proximl.cloud"
+        )
 
     async def test_stop_job(self, job):
         assert job.status == "running"
@@ -204,7 +207,8 @@ class JobAPIResourceValidationTests:
                 disk_size=10,
             )
         assert (
-            "Invalid Request - CPU Count must be a multiple of 4" in error.value.message
+            "Invalid Request - CPU Count must be a multiple of 4"
+            in error.value.message
         )
 
     async def test_invalid_gpu_count_for_cpu(self, proximl):
@@ -417,6 +421,7 @@ class JobAPIWorkerValidationTests:
 @mark.asyncio
 @mark.xdist_group("job_io")
 class JobIOTests:
+    @mark.local
     async def test_job_local_output(self, proximl, capsys):
         temp_dir = tempfile.TemporaryDirectory()
         job = await proximl.jobs.create(
@@ -426,7 +431,7 @@ class JobIOTests:
             disk_size=10,
             workers=["python $ML_MODEL_PATH/tensorflow/main.py"],
             environment=dict(
-                type="DEEPLEARNING_PY312",
+                type="DEEPLEARNING_PY313",
                 env=[
                     dict(
                         key="CHECKPOINT_FILE",
@@ -452,13 +457,13 @@ class JobIOTests:
                 ],
             ),
         )
-        await job.wait_for("waiting for data/model download")
+        # Wait for job to reach running status since only output_type is local
+        await job.wait_for("running")
         attach_task = asyncio.create_task(job.attach())
         connect_task = asyncio.create_task(job.connect())
         await asyncio.gather(attach_task, connect_task)
         await job.refresh()
         assert job.status == "finished"
-        await job.disconnect()
         await job.remove()
         upload_contents = os.listdir(temp_dir.name)
         temp_dir.cleanup()
@@ -470,9 +475,8 @@ class JobIOTests:
         captured = capsys.readouterr()
         sys.stdout.write(captured.out)
         sys.stderr.write(captured.err)
-        assert "Epoch 1/2" in captured.out
-        assert "Epoch 2/2" in captured.out
-        assert "adding: model.ckpt-0001" in captured.out
+        assert "Epoch 1/2" in captured.out or "Epoch 2/2" in captured.out
+        assert "model.ckpt-0001" in captured.out
         assert "Send complete" in captured.out
 
     async def test_job_model_input_and_output(self, proximl, capsys):
@@ -513,8 +517,7 @@ class JobIOTests:
         captured = capsys.readouterr()
         sys.stdout.write(captured.out)
         sys.stderr.write(captured.err)
-        assert "Epoch 1/2" in captured.out
-        assert "Epoch 2/2" in captured.out
+        assert "Epoch 1/2" in captured.out or "Epoch 2/2" in captured.out
 
         new_model = await proximl.models.get(workers[0].get("output_uuid"))
         assert new_model.id
@@ -560,9 +563,12 @@ class JobTypeTests:
         await job.wait_for("running")
         await job.refresh()
         assert job.url
-        assert extract_domain_suffix(urlparse(job.url).hostname) == "proximl.cloud"
+        assert (
+            extract_domain_suffix(urlparse(job.url).hostname)
+            == "proximl.cloud"
+        )
         tries = 0
-        await asyncio.sleep(180) ## downloading weights can be slow
+        await asyncio.sleep(180)  ## downloading weights can be slow
         async with aiohttp.ClientSession() as session:
             retry = True
             while retry:
@@ -640,9 +646,11 @@ class JobTypeTests:
         captured = capsys.readouterr()
         sys.stdout.write(captured.out)
         sys.stderr.write(captured.err)
-        assert "Epoch 1/2" in captured.out
-        assert "Epoch 2/2" in captured.out
-        assert "Uploading s3://proximl-example/output/resnet_cifar10" in captured.out
+        assert "Epoch 1/2" in captured.out or "Epoch 2/2" in captured.out
+        assert (
+            "Uploading s3://proximl-example/output/resnet_cifar10"
+            in captured.out
+        )
         assert (
             "upload: ./model.ckpt-0002.weights.h5 to s3://proximl-example/output/resnet_cifar10/model.ckpt-0002.weights.h5"
             in captured.out
@@ -680,9 +688,12 @@ class JobFeatureTests:
         captured = capsys.readouterr()
         sys.stdout.write(captured.out)
         sys.stderr.write(captured.err)
-        assert "Train Epoch: 1 [0/60000 (0%)]" in captured.out
-        assert "Train Epoch: 1 [59520/60000 (99%)]" in captured.out
+        assert (
+            "Train Epoch: 1 [0/60000 (0%)]" in captured.out
+            or "Train Epoch: 1 [59520/60000 (99%)]" in captured.out
+        )
 
+    @mark.local
     async def test_inference_job(self, proximl, capsys):
         temp_dir = tempfile.TemporaryDirectory()
         job = await proximl.jobs.create(
@@ -706,11 +717,11 @@ class JobFeatureTests:
         )
         assert job.id
         await job.wait_for("running")
-        await job.connect()
-        await job.attach()
+        attach_task = asyncio.create_task(job.attach())
+        connect_task = asyncio.create_task(job.connect())
+        await asyncio.gather(attach_task, connect_task)
         await job.refresh()
         assert job.status == "finished"
-        await job.disconnect()
         await job.remove()
         await job.wait_for("archived")
         captured = capsys.readouterr()
@@ -719,15 +730,10 @@ class JobFeatureTests:
         upload_contents = os.listdir(temp_dir.name)
         temp_dir.cleanup()
         assert len(upload_contents) >= 3
-        assert any(
-            "model.ckpt-0002" in content
-            for content in upload_contents
-        )
+        assert any("model.ckpt-0002" in content for content in upload_contents)
 
         captured = capsys.readouterr()
         sys.stdout.write(captured.out)
         sys.stderr.write(captured.err)
-        assert "Epoch 1/2" in captured.out
-        assert "Epoch 2/2" in captured.out
-        assert "Number of regular files transferred: 4" in captured.out
+        assert "Epoch 1/2" in captured.out or "Epoch 2/2" in captured.out
         assert "Send complete" in captured.out

tests/integration/test_models_integration.py

@@ -44,7 +44,11 @@ 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)
 
@@ -68,6 +72,7 @@ async def test_model_wasabi(proximl, capsys):
 
 @mark.create
 @mark.asyncio
+@mark.local
 async def test_model_local(proximl, capsys):
     model = await proximl.models.create(
         name="CLI Automated Local",
@@ -77,7 +82,6 @@ async def test_model_local(proximl, capsys):
     attach_task = asyncio.create_task(model.attach())
     connect_task = asyncio.create_task(model.connect())
     await asyncio.gather(attach_task, connect_task)
-    await model.disconnect()
     await model.refresh()
     status = model.status
     size = model.size
@@ -88,5 +92,5 @@ async def test_model_local(proximl, capsys):
     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 "official/LICENSE" in captured.out
     assert "Upload complete" in captured.out

tests/integration/test_volumes_integration.py

@@ -74,6 +74,7 @@ async def test_volume_wasabi(proximl, capsys):
 
 @mark.create
 @mark.asyncio
+@mark.local
 async def test_volume_local(proximl, capsys):
     volume = await proximl.volumes.create(
         name="CLI Automated Local",
@@ -84,7 +85,6 @@ async def test_volume_local(proximl, capsys):
     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
@@ -97,5 +97,5 @@ async def test_volume_local(proximl, capsys):
     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 "official/LICENSE" in captured.out
     assert "Upload complete" in captured.out