trainml 0.4.13__py3-none-any.whl → 0.4.15__py3-none-any.whl

tests/unit/test_jobs_unit.py

@@ -314,7 +314,13 @@ class JobTests:
 
     @mark.asyncio
     async def test_job_start(self, job, mock_trainml):
-        api_response = None
+        api_response = {
+            "customer_uuid": "cus-id-1",
+            "job_uuid": "job-id-1",
+            "name": "test notebook",
+            "type": "notebook",
+            "status": "starting",
+        }
         mock_trainml._query = AsyncMock(return_value=api_response)
         await job.start()
         mock_trainml._query.assert_called_once_with(
@@ -326,7 +332,13 @@ class JobTests:
 
     @mark.asyncio
     async def test_job_stop(self, job, mock_trainml):
-        api_response = None
+        api_response = {
+            "customer_uuid": "cus-id-1",
+            "job_uuid": "job-id-1",
+            "name": "test notebook",
+            "type": "notebook",
+            "status": "stopping",
+        }
         mock_trainml._query = AsyncMock(return_value=api_response)
         await job.stop()
         mock_trainml._query.assert_called_once_with(

trainml/__init__.py

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

trainml/checkpoints.py

@@ -126,8 +126,12 @@ class Checkpoint:
                 ssh_port=self._checkpoint.get("vpn")
                 .get("client")
                 .get("ssh_port"),
-                input_path=self._checkpoint.get("source_uri"),
-                output_path=None,
+                input_path=self._checkpoint.get("source_uri")
+                if self.status in ["new", "downloading"]
+                else None,
+                output_path=self._checkpoint.get("output_uri")
+                if self.status == "exporting"
+                else None,
             )
         else:
             details = dict()
@@ -137,7 +141,7 @@ class Checkpoint:
         if self.status in ["ready", "failed"]:
             raise SpecificationError(
                 "status",
-                f"You can only connect to new or downloading checkpoints.",
+                f"You can only connect to downloading or exporting checkpoints.",
             )
         if self.status == "new":
             await self.wait_for("downloading")
@@ -162,12 +166,28 @@ class Checkpoint:
         )
 
     async def rename(self, name):
-        await self.trainml._query(
+        resp = await self.trainml._query(
             f"/checkpoint/{self._id}",
             "PATCH",
-            None,
+            dict(project_uuid=self._project_uuid),
             dict(name=name),
         )
+        self.__init__(self.trainml, **resp)
+        return self
+
+    async def export(self, output_type, output_uri, output_options=dict()):
+        resp = await self.trainml._query(
+            f"/checkpoint/{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.trainml, **resp)
+        return self
 
     def _get_msg_handler(self, msg_handler):
         def handler(data):

trainml/datasets.py

@@ -158,12 +158,14 @@ class Dataset:
         )
 
     async def rename(self, name):
-        await self.trainml._query(
+        resp = await self.trainml._query(
             f"/dataset/{self._id}",
             "PATCH",
             None,
             dict(name=name),
         )
+        self.__init__(self.trainml, **resp)
+        return self
 
     def _get_msg_handler(self, msg_handler):
         def handler(data):

trainml/jobs.py

@@ -253,20 +253,24 @@ class Job:
         return create_json
 
     async def start(self):
-        await self.trainml._query(
+        resp = await self.trainml._query(
             f"/job/{self._id}",
             "PATCH",
             dict(project_uuid=self._project_uuid),
             dict(command="start"),
         )
+        self.__init__(self.trainml, **resp)
+        return self
 
     async def stop(self):
-        await self.trainml._query(
+        resp = await self.trainml._query(
             f"/job/{self._id}",
             "PATCH",
             dict(project_uuid=self._project_uuid),
             dict(command="stop"),
         )
+        self.__init__(self.trainml, **resp)
+        return self
 
     async def update(self, data):
         if self.type != "notebook":
@@ -274,12 +278,14 @@ class Job:
                 "type",
                 "Only notebook jobs can be modified.",
             )
-        await self.trainml._query(
+        resp = await self.trainml._query(
             f"/job/{self._id}",
             "PATCH",
             dict(project_uuid=self._project_uuid),
             data,
         )
+        self.__init__(self.trainml, **resp)
+        return self
 
     async def get_worker_log_url(self, job_worker_uuid):
         resp = await self.trainml._query(

trainml/models.py

@@ -151,12 +151,14 @@ class Model:
         )
 
     async def rename(self, name):
-        await self.trainml._query(
+        resp = await self.trainml._query(
             f"/model/{self._id}",
             "PATCH",
             None,
             dict(name=name),
         )
+        self.__init__(self.trainml, **resp)
+        return self
 
     def _get_msg_handler(self, msg_handler):
         def handler(data):

{trainml-0.4.13.dist-info → trainml-0.4.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: trainml
-Version: 0.4.13
+Version: 0.4.15
 Summary: trainML client SDK and command line utilities
 Home-page: https://github.com/trainML/trainml-cli
 Author: trainML

{trainml-0.4.13.dist-info → trainml-0.4.15.dist-info}/RECORD

@@ -21,7 +21,7 @@ tests/unit/test_datasets_unit.py,sha256=lVNoBZu4RIiJK26gbUPOUAra_k0YS2GcnjJDnT7U
 tests/unit/test_environments_unit.py,sha256=1QFGf1xwM0yKCyVHT_Xi0DX8g0Aelr0mcqAImEXJfQU,1882
 tests/unit/test_exceptions.py,sha256=3tAok6kAU1QRjN7qTNVYuSGWDg7IEoK__OXFLyzLr7k,906
 tests/unit/test_gpu_types_unit.py,sha256=6v_n_AytYjQZxv2OtcUYBxQz7iRjigSS7xmBo6wSJk0,1703
-tests/unit/test_jobs_unit.py,sha256=zG--HP1NOWDNx-5ZYWfv4R10RW9s-DR77OQC2EN6kaM,26459
+tests/unit/test_jobs_unit.py,sha256=N-51OkTU_nHdPxCcLzGyz4A0LD7dmupE5TGr7JXjWVM,26833
 tests/unit/test_models_unit.py,sha256=uezWF7FUHGmCSQBtpyyKhBttTnCTRjxU22NsHdJLYYg,15064
 tests/unit/test_projects_unit.py,sha256=mV0CejcTSNJEYfFl-vYcPbQ2HnDw31xc04Asn6-jZrM,3871
 tests/unit/test_providers_unit.py,sha256=nEizghnC8pfDubkCw-kMmS_QQOUUWBk3i8D44pnyljo,3700
@@ -29,17 +29,17 @@ tests/unit/test_trainml.py,sha256=8vAKvFD1xYsx_VY4HFVa0b1MUlMoNApY6TO8r7vI-UQ,17
 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_environment_unit.py,sha256=7FFKFPVa6yqJqujTXD_tsSUDruW5KhlLgShPZTn1Eio,683
-trainml/__init__.py,sha256=YCDyPnS-RTLikmOhBm90HYDsjdcnIut8_uloxGBQr08,433
+trainml/__init__.py,sha256=SZqtnWUI0eSxdRWMWS8SVB1kXKfEJSLzw5GhbC3r4SA,433
 trainml/__main__.py,sha256=JgErYkiskih8Y6oRwowALtR-rwQhAAdqOYWjQraRIPI,59
 trainml/auth.py,sha256=VI8wXgmXbwiBpgw0Sm-DzC4VmNqKg_zniua-Z61PafA,26646
-trainml/checkpoints.py,sha256=2Mq7VWs1_Jn9hhTwylrbh5XMocDup4cb1Urweod_X_E,7497
+trainml/checkpoints.py,sha256=-P8FOkw1ihA5KCSQll-FDY1VtuKZUOboBjjEipUltF4,8262
 trainml/connections.py,sha256=xkcFoyB5AXztRh4DPkpg3GPIZPxBoDvb5FZfON46aVs,20059
-trainml/datasets.py,sha256=jhzA0t72oGiYYXOiMckJECap57TlfFUCYCvZZhaQaFU,7311
+trainml/datasets.py,sha256=eImi3BO74M7mV9kmUsH03O1ybP_mcMA_FPdLMC--ioM,7382
 trainml/environments.py,sha256=hctwxZX7zNO-StgMnD9Zh4JS8fc3mUc98Ws3xM4c9QQ,1537
 trainml/exceptions.py,sha256=hhR78fI8rbU3fWQ-kUsgxOLyYP8D2bQcjLjrskqPG0Q,3724
 trainml/gpu_types.py,sha256=VXB4XqE0vEwtRyj5nOhWHlR2hRh_lU1VuURU3vyO_gE,1831
-trainml/jobs.py,sha256=YxRta4vVIob3OTRkX0vz45F9uXRtEJ34rUDYuScZR1Q,17594
-trainml/models.py,sha256=nZ1GzzvdXUyt30riev6mC4K3DG7gLu4BBJCFavGzbz0,6960
+trainml/jobs.py,sha256=cUuiMvtJSFSupv2jy7fqAvUOkO8FRLfBDe3RuSZUjfE,17807
+trainml/models.py,sha256=jQ84rqrTbOuCstcZYlh8agWeZOQkrApZEpy_m9bfKWU,7031
 trainml/projects.py,sha256=-lFWwii5cnaqm6vIVZbiVsolZ9kNRsawdd8HjF4BLJE,2063
 trainml/providers.py,sha256=97VegYVSeK0BuYv04hfBY_awNBbGz_GR_mdDrknfO-A,1844
 trainml/trainml.py,sha256=1xGKaI7mAUkB4I7bDyPVDNNrDVmgC6Sg8tJx2UXkKH4,10974
@@ -55,9 +55,9 @@ trainml/cli/project.py,sha256=sBId3S4K7kRwIqURkabSXU3iL2TXbg8XMcjMxHwtt7g,1654
 trainml/cli/provider.py,sha256=eaklYo0IUOGeyZ0ziZZz6UhOZ50Py4ewaPtWA9UDCqU,1594
 trainml/cli/job/__init__.py,sha256=KP3_j6aaokVYNJyi8BUReUdJ6WbZ3ObEBmRJTFWtS10,6544
 trainml/cli/job/create.py,sha256=JyyqVH0Oe9wktRRxxqMC5GEVWd6nzEdKkE42jn1jjtc,35272
-trainml-0.4.13.dist-info/LICENSE,sha256=s0lpBxhSSUEpMavwde-Vb6K_K7xDCTTvSpNznVqVGR0,1069
-trainml-0.4.13.dist-info/METADATA,sha256=qTw565C9HMTZ3zDTOx6ktUiki5EdBk7i02v_Q-JINiM,7346
-trainml-0.4.13.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
-trainml-0.4.13.dist-info/entry_points.txt,sha256=OzBDm2wXby1bSGF02jTVxzRFZLejnbFiLHXhKdW3Bds,63
-trainml-0.4.13.dist-info/top_level.txt,sha256=Y1kLFRWKUW7RG8BX7cvejHF_yW8wBOaRYF1JQHENY4w,23
-trainml-0.4.13.dist-info/RECORD,,
+trainml-0.4.15.dist-info/LICENSE,sha256=s0lpBxhSSUEpMavwde-Vb6K_K7xDCTTvSpNznVqVGR0,1069
+trainml-0.4.15.dist-info/METADATA,sha256=XY6ZStH-E7tb-tBeJ4zPPjHYbUhYRh2GkFrwqe1Pk4s,7346
+trainml-0.4.15.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
+trainml-0.4.15.dist-info/entry_points.txt,sha256=OzBDm2wXby1bSGF02jTVxzRFZLejnbFiLHXhKdW3Bds,63
+trainml-0.4.15.dist-info/top_level.txt,sha256=Y1kLFRWKUW7RG8BX7cvejHF_yW8wBOaRYF1JQHENY4w,23
+trainml-0.4.15.dist-info/RECORD,,