trainml 0.5.17__py3-none-any.whl → 1.0.1__py3-none-any.whl

trainml/cli/job/create.py

@@ -275,7 +275,9 @@ def notebook(
         options["environment"]["type"] = environment
 
     try:
-        envs = [{"key": e.split("=")[0], "value": e.split("=")[1]} for e in env]
+        envs = [
+            {"key": e.split("=")[0], "value": e.split("=")[1]} for e in env
+        ]
         options["environment"]["env"] = envs
     except IndexError:
         raise click.UsageError(
@@ -289,21 +291,25 @@ def notebook(
         if pip_packages:
             options["environment"]["packages"]["pip"] = pip_packages.split(",")
         if conda_packages:
-            options["environment"]["packages"]["conda"] = conda_packages.split(",")
+            options["environment"]["packages"]["conda"] = conda_packages.split(
+                ","
+            )
 
     if data_dir:
         click.echo("Creating Dataset..", file=config.stdout)
         new_dataset = config.trainml.run(
-            config.trainml.client.datasets.create(f"Job - {name}", "local", data_dir)
+            config.trainml.client.datasets.create(
+                f"Job - {name}", "local", data_dir
+            )
         )
         if attach:
             config.trainml.run(new_dataset.attach(), new_dataset.connect())
-            config.trainml.run(new_dataset.disconnect())
         else:
             config.trainml.run(new_dataset.connect())
             config.trainml.run(new_dataset.wait_for("ready"))
-            config.trainml.run(new_dataset.disconnect())
-        options["data"]["datasets"].append(dict(id=new_dataset.id, type="existing"))
+        options["data"]["datasets"].append(
+            dict(id=new_dataset.id, type="existing")
+        )
 
     if git_uri:
         options["model"]["source_type"] = "git"
@@ -331,13 +337,11 @@ def notebook(
         if attach or connect:
             click.echo("Waiting for job to start...", file=config.stdout)
             config.trainml.run(job.connect(), job.attach())
-            config.trainml.run(job.disconnect())
             click.echo("Launching...", file=config.stdout)
             browse(job.notebook_url)
         else:
             config.trainml.run(job.connect())
             config.trainml.run(job.wait_for("running"))
-            config.trainml.run(job.disconnect())
     elif attach or connect:
         click.echo("Waiting for job to start...", file=config.stdout)
         config.trainml.run(job.wait_for("running", timeout))
@@ -626,15 +630,21 @@ def training(
     if output_type:
         options["data"]["output_type"] = output_type
         options["data"]["output_uri"] = output_uri
-        options["data"]["output_options"] = dict(archive=archive, save_model=save_model)
+        options["data"]["output_options"] = dict(
+            archive=archive, save_model=save_model
+        )
 
     if output_dir:
         options["data"]["output_type"] = "local"
         options["data"]["output_uri"] = output_dir
-        options["data"]["output_options"] = dict(archive=archive, save_model=save_model)
+        options["data"]["output_options"] = dict(
+            archive=archive, save_model=save_model
+        )
 
     try:
-        envs = [{"key": e.split("=")[0], "value": e.split("=")[1]} for e in env]
+        envs = [
+            {"key": e.split("=")[0], "value": e.split("=")[1]} for e in env
+        ]
         options["environment"]["env"] = envs
     except IndexError:
         raise click.UsageError(
@@ -648,21 +658,25 @@ def training(
         if pip_packages:
             options["environment"]["packages"]["pip"] = pip_packages.split(",")
         if conda_packages:
-            options["environment"]["packages"]["conda"] = conda_packages.split(",")
+            options["environment"]["packages"]["conda"] = conda_packages.split(
+                ","
+            )
 
     if data_dir:
         click.echo("Creating Dataset..", file=config.stdout)
         new_dataset = config.trainml.run(
-            config.trainml.client.datasets.create(f"Job - {name}", "local", data_dir)
+            config.trainml.client.datasets.create(
+                f"Job - {name}", "local", data_dir
+            )
         )
         if attach:
             config.trainml.run(new_dataset.attach(), new_dataset.connect())
-            config.trainml.run(new_dataset.disconnect())
         else:
             config.trainml.run(new_dataset.connect())
             config.trainml.run(new_dataset.wait_for("ready"))
-            config.trainml.run(new_dataset.disconnect())
-        options["data"]["datasets"].append(dict(id=new_dataset.id, type="existing"))
+        options["data"]["datasets"].append(
+            dict(id=new_dataset.id, type="existing")
+        )
 
     if git_uri:
         options["model"]["source_type"] = "git"
@@ -979,15 +993,21 @@ def inference(
     if output_type:
         options["data"]["output_type"] = output_type
         options["data"]["output_uri"] = output_uri
-        options["data"]["output_options"] = dict(archive=archive, save_model=save_model)
+        options["data"]["output_options"] = dict(
+            archive=archive, save_model=save_model
+        )
 
     if output_dir:
         options["data"]["output_type"] = "local"
         options["data"]["output_uri"] = output_dir
-        options["data"]["output_options"] = dict(archive=archive, save_model=save_model)
+        options["data"]["output_options"] = dict(
+            archive=archive, save_model=save_model
+        )
 
     try:
-        envs = [{"key": e.split("=")[0], "value": e.split("=")[1]} for e in env]
+        envs = [
+            {"key": e.split("=")[0], "value": e.split("=")[1]} for e in env
+        ]
         options["environment"]["env"] = envs
     except IndexError:
         raise click.UsageError(
@@ -1001,7 +1021,9 @@ def inference(
         if pip_packages:
             options["environment"]["packages"]["pip"] = pip_packages.split(",")
         if conda_packages:
-            options["environment"]["packages"]["conda"] = conda_packages.split(",")
+            options["environment"]["packages"]["conda"] = conda_packages.split(
+                ","
+            )
 
     if git_uri:
         options["model"]["source_type"] = "git"
@@ -1301,7 +1323,9 @@ def endpoint(
         options["environment"]["type"] = environment
 
     try:
-        envs = [{"key": e.split("=")[0], "value": e.split("=")[1]} for e in env]
+        envs = [
+            {"key": e.split("=")[0], "value": e.split("=")[1]} for e in env
+        ]
         options["environment"]["env"] = envs
     except IndexError:
         raise click.UsageError(
@@ -1315,7 +1339,9 @@ def endpoint(
         if pip_packages:
             options["environment"]["packages"]["pip"] = pip_packages.split(",")
         if conda_packages:
-            options["environment"]["packages"]["conda"] = conda_packages.split(",")
+            options["environment"]["packages"]["conda"] = conda_packages.split(
+                ","
+            )
 
     if git_uri:
         options["model"]["source_type"] = "git"
@@ -1349,7 +1375,6 @@ def endpoint(
             config.trainml.run(job.connect())
         click.echo("Waiting for job to start...", file=config.stdout)
         config.trainml.run(job.wait_for("running", timeout))
-        config.trainml.run(job.disconnect())
         config.trainml.run(job.refresh())
         click.echo(f"Endpoint is running at:  {job.url}", file=config.stdout)
     else:
@@ -1357,4 +1382,6 @@ def endpoint(
             click.echo("Waiting for job to start...", file=config.stdout)
             config.trainml.run(job.wait_for("running", timeout))
             config.trainml.run(job.refresh())
-            click.echo(f"Endpoint is running at:  {job.url}", file=config.stdout)
+            click.echo(
+                f"Endpoint is running at:  {job.url}", file=config.stdout
+            )

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: