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

trainml/cli/model.py

@@ -36,15 +36,7 @@ def attach(config, model):
     if None is found:
         raise click.UsageError("Cannot find specified model.")
 
-    try:
-        config.trainml.run(found.attach())
-        return config.trainml.run(found.disconnect())
-    except:
-        try:
-            config.trainml.run(found.disconnect())
-        except:
-            pass
-        raise
+    config.trainml.run(found.attach())
 
 
 @model.command()
@@ -69,18 +61,10 @@ def connect(config, model, attach):
     if None is found:
         raise click.UsageError("Cannot find specified model.")
 
-    try:
-        if attach:
-            config.trainml.run(found.connect(), found.attach())
-            return config.trainml.run(found.disconnect())
-        else:
-            return config.trainml.run(found.connect())
-    except:
-        try:
-            config.trainml.run(found.disconnect())
-        except:
-            pass
-        raise
+    if attach:
+        config.trainml.run(found.connect(), found.attach())
+    else:
+        config.trainml.run(found.connect())
 
 
 @model.command()
@@ -125,41 +109,15 @@ def create(config, attach, connect, source, name, path):
             )
         )
 
-        try:
-            if connect and attach:
-                config.trainml.run(model.attach(), model.connect())
-                return config.trainml.run(model.disconnect())
-            elif connect:
-                return config.trainml.run(model.connect())
-            else:
-                raise click.UsageError(
-                    "Abort!\n"
-                    "No logs to show for local sourced model without connect."
-                )
-        except:
-            try:
-                config.trainml.run(model.disconnect())
-            except:
-                pass
-            raise
-
-
-@model.command()
-@click.argument("model", type=click.STRING)
-@pass_config
-def disconnect(config, model):
-    """
-    Disconnect and clean-up model upload.
-
-    MODEL may be specified by name or ID, but ID is preferred.
-    """
-    models = config.trainml.run(config.trainml.client.models.list())
-
-    found = search_by_id_name(model, models)
-    if None is found:
-        raise click.UsageError("Cannot find specified model.")
-
-    return config.trainml.run(found.disconnect())
+        if connect and attach:
+            config.trainml.run(model.attach(), model.connect())
+        elif connect:
+            config.trainml.run(model.connect())
+        else:
+            raise click.UsageError(
+                "Abort!\n"
+                "No logs to show for local sourced model without connect."
+            )
 
 
 @model.command()

trainml/cli/volume.py

@@ -35,15 +35,7 @@ def attach(config, volume):
     if None is found:
         raise click.UsageError("Cannot find specified volume.")
 
-    try:
-        config.trainml.run(found.attach())
-        return config.trainml.run(found.disconnect())
-    except:
-        try:
-            config.trainml.run(found.disconnect())
-        except:
-            pass
-        raise
+    config.trainml.run(found.attach())
 
 
 @volume.command()
@@ -67,18 +59,10 @@ def connect(config, volume, attach):
     if None is found:
         raise click.UsageError("Cannot find specified volume.")
 
-    try:
-        if attach:
-            config.trainml.run(found.connect(), found.attach())
-            return config.trainml.run(found.disconnect())
-        else:
-            return config.trainml.run(found.connect())
-    except:
-        try:
-            config.trainml.run(found.disconnect())
-        except:
-            pass
-        raise
+    if attach:
+        config.trainml.run(found.connect(), found.attach())
+    else:
+        config.trainml.run(found.connect())
 
 
 @volume.command()
@@ -120,45 +104,22 @@ def create(config, attach, connect, source, name, capacity, path):
     if source == "local":
         volume = config.trainml.run(
             config.trainml.client.volumes.create(
-                name=name, source_type="local", source_uri=path, capacity=capacity
+                name=name,
+                source_type="local",
+                source_uri=path,
+                capacity=capacity,
             )
         )
 
-        try:
-            if connect and attach:
-                config.trainml.run(volume.attach(), volume.connect())
-                return config.trainml.run(volume.disconnect())
-            elif connect:
-                return config.trainml.run(volume.connect())
-            else:
-                raise click.UsageError(
-                    "Abort!\n"
-                    "No logs to show for local sourced volume without connect."
-                )
-        except:
-            try:
-                config.trainml.run(volume.disconnect())
-            except:
-                pass
-            raise
-
-
-@volume.command()
-@click.argument("volume", type=click.STRING)
-@pass_config
-def disconnect(config, volume):
-    """
-    Disconnect and clean-up volume upload.
-
-    VOLUME may be specified by name or ID, but ID is preferred.
-    """
-    volumes = config.trainml.run(config.trainml.client.volumes.list())
-
-    found = search_by_id_name(volume, volumes)
-    if None is found:
-        raise click.UsageError("Cannot find specified volume.")
-
-    return config.trainml.run(found.disconnect())
+        if connect and attach:
+            config.trainml.run(volume.attach(), volume.connect())
+        elif connect:
+            config.trainml.run(volume.connect())
+        else:
+            raise click.UsageError(
+                "Abort!\n"
+                "No logs to show for local sourced volume without connect."
+            )
 
 
 @volume.command()

trainml/datasets.py

@@ -10,7 +10,7 @@ from .exceptions import (
     SpecificationError,
     TrainMLException,
 )
-from .connections import Connection
+from trainml.utils.transfer import upload, download
 
 
 class Datasets(object):
@@ -71,10 +71,12 @@ class Dataset:
         self._id = self._dataset.get("id", self._dataset.get("dataset_uuid"))
         self._status = self._dataset.get("status")
         self._name = self._dataset.get("name")
-        self._size = self._dataset.get("size") or self._dataset.get("used_size")
-        self._billed_size = self._dataset.get("billed_size") or self._dataset.get(
-            "size"
+        self._size = self._dataset.get("size") or self._dataset.get(
+            "used_size"
         )
+        self._billed_size = self._dataset.get(
+            "billed_size"
+        ) or self._dataset.get("size")
         self._project_uuid = self._dataset.get("project_uuid")
 
     @property
@@ -122,56 +124,45 @@ class Dataset:
         )
         return resp
 
-    async def get_connection_utility_url(self):
-        resp = await self.trainml._query(
-            f"/dataset/{self._id}/download",
-            "GET",
-            dict(project_uuid=self._project_uuid),
-        )
-        return resp
-
-    def get_connection_details(self):
-        if self._dataset.get("vpn"):
-            details = dict(
-                entity_type="dataset",
-                project_uuid=self._dataset.get("project_uuid"),
-                cidr=self._dataset.get("vpn").get("cidr"),
-                ssh_port=self._dataset.get("vpn").get("client").get("ssh_port"),
-                input_path=(
-                    self._dataset.get("source_uri")
-                    if self.status in ["new", "downloading"]
-                    else None
-                ),
-                output_path=(
-                    self._dataset.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 datasets.",
-            )
-        if self.status == "new":
-            await self.wait_for("downloading")
-        connection = Connection(
-            self.trainml, entity_type="dataset", 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 datasets.",
+                )
 
-    async def disconnect(self):
-        connection = Connection(
-            self.trainml, entity_type="dataset", 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 dataset
+            auth_token = self._dataset.get("auth_token")
+            hostname = self._dataset.get("hostname")
+            source_uri = self._dataset.get("source_uri")
+
+            if not auth_token or not hostname or not source_uri:
+                raise SpecificationError(
+                    "status",
+                    f"Dataset 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 dataset
+            auth_token = self._dataset.get("auth_token")
+            hostname = self._dataset.get("hostname")
+            output_uri = self._dataset.get("output_uri")
+
+            if not auth_token or not hostname or not output_uri:
+                raise SpecificationError(
+                    "status",
+                    f"Dataset 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.trainml._query(
@@ -210,7 +201,9 @@ class Dataset:
                 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()}"
                     )
@@ -239,7 +232,7 @@ class Dataset:
     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",
@@ -254,7 +247,9 @@ class Dataset:
 
         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: