trainml 0.4.13__tar.gz → 0.4.15__tar.gz

{trainml-0.4.13 → trainml-0.4.15}/PKG-INFO

@@ -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 → trainml-0.4.15}/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-0.4.13 → trainml-0.4.15}/trainml/__init__.py

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

{trainml-0.4.13 → trainml-0.4.15}/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-0.4.13 → trainml-0.4.15}/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-0.4.13 → trainml-0.4.15}/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-0.4.13 → trainml-0.4.15}/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 → trainml-0.4.15}/trainml.egg-info/PKG-INFO

@@ -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