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

examples/local_storage.py

@@ -19,7 +19,6 @@ async def create_dataset():
     attach_task = asyncio.create_task(dataset.attach())
     connect_task = asyncio.create_task(dataset.connect())
     await asyncio.gather(attach_task, connect_task)
-    await dataset.disconnect()
     return dataset
 
 
@@ -55,7 +54,6 @@ async def run_job(dataset):
     await asyncio.gather(attach_task, connect_task)
 
     # Cleanup job
-    await job.disconnect()
     await job.remove()
 
 

proximl/__init__.py

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

proximl/checkpoints.py

@@ -10,7 +10,7 @@ from .exceptions import (
     SpecificationError,
     ProxiMLException,
 )
-from .connections import Connection
+from proximl.utils.transfer import upload, download
 
 
 class Checkpoints(object):
@@ -23,7 +23,9 @@ class Checkpoints(object):
 
     async def list(self, **kwargs):
         resp = await self.proximl._query(f"/checkpoint", "GET", kwargs)
-        checkpoints = [Checkpoint(self.proximl, **checkpoint) for checkpoint in resp]
+        checkpoints = [
+            Checkpoint(self.proximl, **checkpoint) for checkpoint in resp
+        ]
         return checkpoints
 
     async def list_public(self, **kwargs):
@@ -68,13 +70,17 @@ class Checkpoint:
     def __init__(self, proximl, **kwargs):
         self.proximl = proximl
         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.proximl._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.proximl, 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.proximl, 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.proximl._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:

proximl/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.proximl.client.set_active_project(selected_project[0].id)
 
 
-from proximl.cli.connection import connection
 from proximl.cli.dataset import dataset
 from proximl.cli.model import model
 from proximl.cli.checkpoint import checkpoint

proximl/cli/checkpoint.py

@@ -35,15 +35,7 @@ def attach(config, checkpoint):
     if None is found:
         raise click.UsageError("Cannot find specified checkpoint.")
 
-    try:
-        config.proximl.run(found.attach())
-        return config.proximl.run(found.disconnect())
-    except:
-        try:
-            config.proximl.run(found.disconnect())
-        except:
-            pass
-        raise
+    config.proximl.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.proximl.run(found.connect(), found.attach())
-            return config.proximl.run(found.disconnect())
-        else:
-            return config.proximl.run(found.connect())
-    except:
-        try:
-            config.proximl.run(found.disconnect())
-        except:
-            pass
-        raise
+    if attach:
+        config.proximl.run(found.connect(), found.attach())
+    else:
+        config.proximl.run(found.connect())
 
 
 @checkpoint.command()
@@ -123,41 +107,15 @@ def create(config, attach, connect, source, name, path):
             )
         )
 
-        try:
-            if connect and attach:
-                config.proximl.run(checkpoint.attach(), checkpoint.connect())
-                return config.proximl.run(checkpoint.disconnect())
-            elif connect:
-                return config.proximl.run(checkpoint.connect())
-            else:
-                raise click.UsageError(
-                    "Abort!\n"
-                    "No logs to show for local sourced checkpoint without connect."
-                )
-        except:
-            try:
-                config.proximl.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.proximl.run(config.proximl.client.checkpoints.list())
-
-    found = search_by_id_name(checkpoint, checkpoints)
-    if None is found:
-        raise click.UsageError("Cannot find specified checkpoint.")
-
-    return config.proximl.run(found.disconnect())
+        if connect and attach:
+            config.proximl.run(checkpoint.attach(), checkpoint.connect())
+        elif connect:
+            config.proximl.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.proximl.run(found.client.checkpoints.remove(checkpoint))
+            config.proximl.run(
+                config.proximl.client.checkpoints.remove(checkpoint)
+            )
+            return
         else:
             raise click.UsageError("Cannot find specified checkpoint.")
 

proximl/cli/dataset.py

@@ -35,15 +35,7 @@ def attach(config, dataset):
     if None is found:
         raise click.UsageError("Cannot find specified dataset.")
 
-    try:
-        config.proximl.run(found.attach())
-        return config.proximl.run(found.disconnect())
-    except:
-        try:
-            config.proximl.run(found.disconnect())
-        except:
-            pass
-        raise
+    config.proximl.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.proximl.run(found.connect(), found.attach())
-            return config.proximl.run(found.disconnect())
-        else:
-            return config.proximl.run(found.connect())
-    except:
-        try:
-            config.proximl.run(found.disconnect())
-        except:
-            pass
-        raise
+    if attach:
+        config.proximl.run(found.connect(), found.attach())
+    else:
+        config.proximl.run(found.connect())
 
 
 @dataset.command()
@@ -123,41 +107,15 @@ def create(config, attach, connect, source, name, path):
             )
         )
 
-        try:
-            if connect and attach:
-                config.proximl.run(dataset.attach(), dataset.connect())
-                return config.proximl.run(dataset.disconnect())
-            elif connect:
-                return config.proximl.run(dataset.connect())
-            else:
-                raise click.UsageError(
-                    "Abort!\n"
-                    "No logs to show for local sourced dataset without connect."
-                )
-        except:
-            try:
-                config.proximl.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.proximl.run(config.proximl.client.datasets.list())
-
-    found = search_by_id_name(dataset, datasets)
-    if None is found:
-        raise click.UsageError("Cannot find specified dataset.")
-
-    return config.proximl.run(found.disconnect())
+        if connect and attach:
+            config.proximl.run(dataset.attach(), dataset.connect())
+        elif connect:
+            config.proximl.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.proximl.run(config.proximl.client.datasets.get(dataset))
+        dataset = config.proximl.run(
+            config.proximl.client.datasets.get(dataset)
+        )
         if dataset is None:
             raise click.UsageError("Cannot find specified dataset.")
     except:

proximl/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.proximl.run(found.attach())
-        return config.proximl.run(found.disconnect())
-    except:
-        try:
-            config.proximl.run(found.disconnect())
-        except:
-            pass
-        raise
+    config.proximl.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.proximl.run(found.connect(), found.attach())
-                return config.proximl.run(found.disconnect())
-            else:
-                return config.proximl.run(found.connect())
-        except:
-            try:
-                config.proximl.run(found.disconnect())
-            except:
-                pass
-            raise
+        if attach:
+            config.proximl.run(found.connect(), found.attach())
+        else:
+            config.proximl.run(found.connect())
     else:
         if found.status in [
             "new",
             "waiting for data/model download",
             "waiting for GPUs",
         ]:
-            try:
-                if attach:
-                    config.proximl.run(found.connect(), found.attach())
-                    config.proximl.run(found.disconnect())
-                    click.echo("Launching...", file=config.stdout)
-                    browse(found.notebook_url)
-                else:
-                    return config.proximl.run(found.connect())
-            except:
-                try:
-                    config.proximl.run(found.disconnect())
-                except:
-                    pass
-                raise
+            if attach:
+                config.proximl.run(found.connect(), found.attach())
+                click.echo("Launching...", file=config.stdout)
+                browse(found.notebook_url)
+            else:
+                config.proximl.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.proximl.run(config.proximl.client.jobs.list())
-
-    found = search_by_id_name(job, jobs)
-    if None is found:
-        raise click.UsageError("Cannot find specified job.")
-
-    return config.proximl.run(found.disconnect())
-
-
 @job.command()
 @click.option(
     "--wait/--no-wait",