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

trainml/__init__.py

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

trainml/checkpoints.py

@@ -10,7 +10,7 @@ from .exceptions import (
     SpecificationError,
     TrainMLException,
 )
-from .connections import Connection
+from trainml.utils.transfer import upload, download
 
 
 class Checkpoints(object):
@@ -23,7 +23,9 @@ class Checkpoints(object):
 
     async def list(self, **kwargs):
         resp = await self.trainml._query(f"/checkpoint", "GET", kwargs)
-        checkpoints = [Checkpoint(self.trainml, **checkpoint) for checkpoint in resp]
+        checkpoints = [
+            Checkpoint(self.trainml, **checkpoint) for checkpoint in resp
+        ]
         return checkpoints
 
     async def list_public(self, **kwargs):
@@ -68,13 +70,17 @@ class Checkpoint:
     def __init__(self, trainml, **kwargs):
         self.trainml = trainml
         self._checkpoint = kwargs
-        self._id = self._checkpoint.get("id", self._checkpoint.get("checkpoint_uuid"))
+        self._id = self._checkpoint.get(
+            "id", self._checkpoint.get("checkpoint_uuid")
+        )
         self._status = self._checkpoint.get("status")
         self._name = self._checkpoint.get("name")
-        self._size = self._checkpoint.get("size") or self._checkpoint.get("used_size")
-        self._billed_size = self._checkpoint.get("billed_size") or self._checkpoint.get(
-            "size"
+        self._size = self._checkpoint.get("size") or self._checkpoint.get(
+            "used_size"
         )
+        self._billed_size = self._checkpoint.get(
+            "billed_size"
+        ) or self._checkpoint.get("size")
         self._project_uuid = self._checkpoint.get("project_uuid")
 
     @property
@@ -122,56 +128,45 @@ class Checkpoint:
         )
         return resp
 
-    async def get_connection_utility_url(self):
-        resp = await self.trainml._query(
-            f"/checkpoint/{self._id}/download",
-            "GET",
-            dict(project_uuid=self._project_uuid),
-        )
-        return resp
-
-    def get_connection_details(self):
-        if self._checkpoint.get("vpn"):
-            details = dict(
-                entity_type="checkpoint",
-                project_uuid=self._checkpoint.get("project_uuid"),
-                cidr=self._checkpoint.get("vpn").get("cidr"),
-                ssh_port=self._checkpoint.get("vpn").get("client").get("ssh_port"),
-                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()
-        return details
-
     async def connect(self):
-        if self.status in ["ready", "failed"]:
-            raise SpecificationError(
-                "status",
-                f"You can only connect to downloading or exporting checkpoints.",
-            )
-        if self.status == "new":
-            await self.wait_for("downloading")
-        connection = Connection(
-            self.trainml, entity_type="checkpoint", 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 checkpoints.",
+                )
 
-    async def disconnect(self):
-        connection = Connection(
-            self.trainml, entity_type="checkpoint", 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 checkpoint
+            auth_token = self._checkpoint.get("auth_token")
+            hostname = self._checkpoint.get("hostname")
+            source_uri = self._checkpoint.get("source_uri")
+
+            if not auth_token or not hostname or not source_uri:
+                raise SpecificationError(
+                    "status",
+                    f"Checkpoint 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 checkpoint
+            auth_token = self._checkpoint.get("auth_token")
+            hostname = self._checkpoint.get("hostname")
+            output_uri = self._checkpoint.get("output_uri")
+
+            if not auth_token or not hostname or not output_uri:
+                raise SpecificationError(
+                    "status",
+                    f"Checkpoint 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 +205,9 @@ class Checkpoint:
                 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 +236,7 @@ class Checkpoint:
     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 +251,9 @@ class Checkpoint:
             )
         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:

trainml/cli/__init__.py

@@ -142,7 +142,9 @@ def configure(config):
         project for project in projects if project.id == active_project_id
     ]
 
-    active_project_name = active_project[0].name if len(active_project) else "UNSET"
+    active_project_name = (
+        active_project[0].name if len(active_project) else "UNSET"
+    )
 
     click.echo(f"Current Active Project: {active_project_name}")
 
@@ -152,11 +154,12 @@ def configure(config):
         show_choices=True,
         default=active_project_name,
     )
-    selected_project = [project for project in projects if project.name == name]
+    selected_project = [
+        project for project in projects if project.name == name
+    ]
     config.trainml.client.set_active_project(selected_project[0].id)
 
 
-from trainml.cli.connection import connection
 from trainml.cli.dataset import dataset
 from trainml.cli.model import model
 from trainml.cli.checkpoint import checkpoint

trainml/cli/checkpoint.py

@@ -35,15 +35,7 @@ def attach(config, checkpoint):
     if None is found:
         raise click.UsageError("Cannot find specified checkpoint.")
 
-    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())
 
 
 @checkpoint.command()
@@ -67,18 +59,10 @@ def connect(config, checkpoint, attach):
     if None is found:
         raise click.UsageError("Cannot find specified checkpoint.")
 
-    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())
 
 
 @checkpoint.command()
@@ -123,41 +107,15 @@ def create(config, attach, connect, source, name, path):
             )
         )
 
-        try:
-            if connect and attach:
-                config.trainml.run(checkpoint.attach(), checkpoint.connect())
-                return config.trainml.run(checkpoint.disconnect())
-            elif connect:
-                return config.trainml.run(checkpoint.connect())
-            else:
-                raise click.UsageError(
-                    "Abort!\n"
-                    "No logs to show for local sourced checkpoint without connect."
-                )
-        except:
-            try:
-                config.trainml.run(checkpoint.disconnect())
-            except:
-                pass
-            raise
-
-
-@checkpoint.command()
-@click.argument("checkpoint", type=click.STRING)
-@pass_config
-def disconnect(config, checkpoint):
-    """
-    Disconnect and clean-up checkpoint upload.
-
-    CHECKPOINT may be specified by name or ID, but ID is preferred.
-    """
-    checkpoints = config.trainml.run(config.trainml.client.checkpoints.list())
-
-    found = search_by_id_name(checkpoint, checkpoints)
-    if None is found:
-        raise click.UsageError("Cannot find specified checkpoint.")
-
-    return config.trainml.run(found.disconnect())
+        if connect and attach:
+            config.trainml.run(checkpoint.attach(), checkpoint.connect())
+        elif connect:
+            config.trainml.run(checkpoint.connect())
+        else:
+            raise click.UsageError(
+                "Abort!\n"
+                "No logs to show for local sourced checkpoint without connect."
+            )
 
 
 @checkpoint.command()
@@ -236,7 +194,10 @@ def remove(config, checkpoint, force):
     found = search_by_id_name(checkpoint, checkpoints)
     if None is found:
         if force:
-            config.trainml.run(found.client.checkpoints.remove(checkpoint))
+            config.trainml.run(
+                config.trainml.client.checkpoints.remove(checkpoint)
+            )
+            return
         else:
             raise click.UsageError("Cannot find specified checkpoint.")
 

trainml/cli/dataset.py

@@ -35,15 +35,7 @@ def attach(config, dataset):
     if None is found:
         raise click.UsageError("Cannot find specified dataset.")
 
-    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())
 
 
 @dataset.command()
@@ -67,18 +59,10 @@ def connect(config, dataset, attach):
     if None is found:
         raise click.UsageError("Cannot find specified dataset.")
 
-    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())
 
 
 @dataset.command()
@@ -123,41 +107,15 @@ def create(config, attach, connect, source, name, path):
             )
         )
 
-        try:
-            if connect and attach:
-                config.trainml.run(dataset.attach(), dataset.connect())
-                return config.trainml.run(dataset.disconnect())
-            elif connect:
-                return config.trainml.run(dataset.connect())
-            else:
-                raise click.UsageError(
-                    "Abort!\n"
-                    "No logs to show for local sourced dataset without connect."
-                )
-        except:
-            try:
-                config.trainml.run(dataset.disconnect())
-            except:
-                pass
-            raise
-
-
-@dataset.command()
-@click.argument("dataset", type=click.STRING)
-@pass_config
-def disconnect(config, dataset):
-    """
-    Disconnect and clean-up dataset upload.
-
-    DATASET may be specified by name or ID, but ID is preferred.
-    """
-    datasets = config.trainml.run(config.trainml.client.datasets.list())
-
-    found = search_by_id_name(dataset, datasets)
-    if None is found:
-        raise click.UsageError("Cannot find specified dataset.")
-
-    return config.trainml.run(found.disconnect())
+        if connect and attach:
+            config.trainml.run(dataset.attach(), dataset.connect())
+        elif connect:
+            config.trainml.run(dataset.connect())
+        else:
+            raise click.UsageError(
+                "Abort!\n"
+                "No logs to show for local sourced dataset without connect."
+            )
 
 
 @dataset.command()
@@ -252,7 +210,9 @@ def rename(config, dataset, name):
     DATASET may be specified by name or ID, but ID is preferred.
     """
     try:
-        dataset = config.trainml.run(config.trainml.client.datasets.get(dataset))
+        dataset = config.trainml.run(
+            config.trainml.client.datasets.get(dataset)
+        )
         if dataset is None:
             raise click.UsageError("Cannot find specified dataset.")
     except:

trainml/cli/job/__init__.py

@@ -25,15 +25,7 @@ def attach(config, job):
     if None is found:
         raise click.UsageError("Cannot find specified job.")
 
-    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())
 
 
 @job.command()
@@ -58,38 +50,22 @@ def connect(config, job, attach):
         raise click.UsageError("Cannot find specified job.")
 
     if found.type != "notebook":
-        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())
     else:
         if found.status in [
             "new",
             "waiting for data/model download",
             "waiting for GPUs",
         ]:
-            try:
-                if attach:
-                    config.trainml.run(found.connect(), found.attach())
-                    config.trainml.run(found.disconnect())
-                    click.echo("Launching...", file=config.stdout)
-                    browse(found.notebook_url)
-                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())
+                click.echo("Launching...", file=config.stdout)
+                browse(found.notebook_url)
+            else:
+                config.trainml.run(found.connect())
         elif found.status not in [
             "starting",
             "running",
@@ -103,24 +79,6 @@ def connect(config, job, attach):
             browse(found.notebook_url)
 
 
-@job.command()
-@click.argument("job", type=click.STRING)
-@pass_config
-def disconnect(config, job):
-    """
-    Disconnect and clean-up job.
-
-    JOB may be specified by name or ID, but ID is preferred.
-    """
-    jobs = config.trainml.run(config.trainml.client.jobs.list())
-
-    found = search_by_id_name(job, jobs)
-    if None is found:
-        raise click.UsageError("Cannot find specified job.")
-
-    return config.trainml.run(found.disconnect())
-
-
 @job.command()
 @click.option(
     "--wait/--no-wait",

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
+            )