trainml 0.5.0__py3-none-any.whl → 0.5.2__py3-none-any.whl

trainml/datasets.py

@@ -37,8 +37,7 @@ class Datasets(object):
             source_type=source_type,
             source_uri=source_uri,
             source_options=kwargs.get("source_options"),
-            project_uuid=kwargs.get("project_uuid")
-            or self.trainml.active_project,
+            project_uuid=kwargs.get("project_uuid") or self.trainml.active_project,
         )
         payload = {k: v for k, v in data.items() if v is not None}
         logging.info(f"Creating Dataset {name}")
@@ -119,11 +118,13 @@ class Dataset:
                 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"),
-                output_path=None,
+                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()
@@ -133,7 +134,7 @@ class Dataset:
         if self.status in ["ready", "failed"]:
             raise SpecificationError(
                 "status",
-                f"You can only connect to new or downloading datasets.",
+                f"You can only connect to downloading or exporting datasets.",
             )
         if self.status == "new":
             await self.wait_for("downloading")
@@ -167,15 +168,27 @@ class Dataset:
         self.__init__(self.trainml, **resp)
         return self
 
+    async def export(self, output_type, output_uri, output_options=dict()):
+        resp = await self.trainml._query(
+            f"/dataset/{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):
             if data.get("type") == "subscription":
                 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()}"
                     )
@@ -212,9 +225,7 @@ class Dataset:
             return
         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/gpu_types.py

@@ -1,4 +1,5 @@
 import json
+from trainml.exceptions import TrainMLException
 
 
 class GpuTypes(object):
@@ -6,6 +7,10 @@ class GpuTypes(object):
         self.trainml = trainml
 
     async def list(self):
+        if not self.trainml.project:
+            raise TrainMLException(
+                "Active project not configured.  Run 'trainml configure' to select an active project."
+            )
         resp = await self.trainml._query(
             f"/project/{self.trainml.project}/gputypes", "GET"
         )

trainml/models.py

@@ -115,18 +115,23 @@ class Model:
                 project_uuid=self._model.get("project_uuid"),
                 cidr=self._model.get("vpn").get("cidr"),
                 ssh_port=self._model.get("vpn").get("client").get("ssh_port"),
-                input_path=self._model.get("source_uri"),
-                output_path=None,
+                input_path=self._model.get("source_uri")
+                if self.status in ["new", "downloading"]
+                else None,
+                output_path=self._model.get("output_uri")
+                if self.status == "exporting"
+                else None,
             )
         else:
             details = dict()
+        logging.debug(f"Connection Details: {details}")
         return details
 
     async def connect(self):
         if self.status in ["ready", "failed"]:
             raise SpecificationError(
                 "status",
-                f"You can only connect to new or downloading models.",
+                f"You can only connect to downloading or exporting models.",
             )
         if self.status == "new":
             await self.wait_for("downloading")
@@ -160,6 +165,20 @@ class Model:
         self.__init__(self.trainml, **resp)
         return self
 
+    async def export(self, output_type, output_uri, output_options=dict()):
+        resp = await self.trainml._query(
+            f"/model/{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):
             if data.get("type") == "subscription":

trainml/trainml.py

@@ -28,7 +28,7 @@ async def delayed_close(ws):
 class TrainML(object):
     def __init__(self, **kwargs):
         self._version = version("trainml")
-        CONFIG_DIR = os.path.expanduser(
+        CONFIG_DIR = kwargs.get("config_dir") or os.path.expanduser(
             os.environ.get("TRAINML_CONFIG_DIR") or "~/.trainml"
         )
         try:
@@ -50,6 +50,7 @@ class TrainML(object):
             or "trainml.ai"
         )
         self.auth = Auth(
+            config_dir=CONFIG_DIR,
             domain_suffix=self.domain_suffix,
             user=kwargs.get("user"),
             key=kwargs.get("key"),

{trainml-0.5.0.dist-info → trainml-0.5.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: trainml
-Version: 0.5.0
+Version: 0.5.2
 Summary: trainML client SDK and command line utilities
 Home-page: https://github.com/trainML/trainml-cli
 Author: trainML

{trainml-0.5.0.dist-info → trainml-0.5.2.dist-info}/RECORD

@@ -8,15 +8,15 @@ tests/integration/test_checkpoints_integration.py,sha256=Ev-GxXiBupOi3KduYoEMhb2
 tests/integration/test_datasets_integration.py,sha256=jskX8y9moLvAkZLcQL4iSBrUXSGGuWJ95Vw1JuxxwD8,3424
 tests/integration/test_environments_integration.py,sha256=0IckhJvQhd8j4Ouiu0hMq2b7iA1dbZpZYmknyfWjsFM,1403
 tests/integration/test_gpu_types_integration.py,sha256=V2OncokZWWVq_l5FSmKEDM4EsWrmpB-zKiVPt-we0aY,1256
-tests/integration/test_jobs_integration.py,sha256=YP_cpKMXqP3eMlcVR0dtVsHkRY2neYsdwcPw5BOXkzI,23496
+tests/integration/test_jobs_integration.py,sha256=0fMg7kOjQMotEo40XnerAamqhhRbb4pz0ZhmvVrwlb0,23342
 tests/integration/test_models_integration.py,sha256=xJPq_3m0Cf1liMH8e49ON_L3MO5XcPtJIz_MD9plkyU,2848
 tests/integration/test_projects_integration.py,sha256=_tmMRFFBe29WaWWuEy3_0j7SKyJc_60JS99hgJUHTG4,1492
 tests/integration/test_providers_integration.py,sha256=bBVvlSDLCnofSaD36OB21KlinEKeK04gPB--bbHJiP0,1419
 tests/integration/cloudbender/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/integration/cloudbender/test_providers_integration.py,sha256=oV8ydFsosDZ_Z1Dkg2IN-ZhWuIl5e_HkHAORMsOsAJc,1473
 tests/unit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/unit/conftest.py,sha256=viwLZR1Aja7soy5whypAWoZ-2avEWpel_43fooLA6OY,27336
-tests/unit/test_auth.py,sha256=m5dGjVo5uGrMVbfiwvtePHeG0mouDE8vwci_DsUlVFo,815
+tests/unit/conftest.py,sha256=AZfqMjB6qXyJdH6cRpxrZK4ciioG6HJTQIIDDkJB2H8,29059
+tests/unit/test_auth.py,sha256=nfhlOCR7rUsn_MaD8QQtBc2v0k8pIxqbzGgRAZK1WGc,858
 tests/unit/test_checkpoints_unit.py,sha256=4Add2DXZCuriSZ0atvOXc8fsEGMaEfPhYmT8Q3UgP5E,16008
 tests/unit/test_connections_unit.py,sha256=FzN2ddQxNpjxzNGUsXhjTk0HnD24wSPelPTL4o_r-Ho,5507
 tests/unit/test_datasets_unit.py,sha256=lVNoBZu4RIiJK26gbUPOUAra_k0YS2GcnjJDnT7UV6Y,15879
@@ -39,60 +39,64 @@ tests/unit/cli/test_cli_model_unit.py,sha256=fE-CRVg8gbtDlwrKBkf-hc9x7EhFlYeE3jl
 tests/unit/cli/test_cli_project_unit.py,sha256=1tEfwXJqnac41bDtvYcN_gpScnb4UY6s5dHfTl-0YoQ,1756
 tests/unit/cli/cloudbender/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/unit/cli/cloudbender/test_cli_datastore_unit.py,sha256=DQWDjqg4viBZRONi00nVzqF9rJ5qKOKRub9pKbTmMWU,1381
+tests/unit/cli/cloudbender/test_cli_device_unit.py,sha256=2BSMyXQ8fOzNKh_-pa_tx7fy_GCdNlGLNvkiA7vuV7s,1342
 tests/unit/cli/cloudbender/test_cli_node_unit.py,sha256=KbK7axJ1L4y4sN7KQRpOVIqphnNpi0aFW4-HlYLtwnI,1316
 tests/unit/cli/cloudbender/test_cli_provider_unit.py,sha256=Rm-tRNPbTTB7ZzkkIpLfDp_pEYfqihjB0ZYk_EPQUfs,781
 tests/unit/cli/cloudbender/test_cli_region_unit.py,sha256=iH5AbrzZ-R2EJ-Bd2HFN7FN2lTpkr3-pCLR59ZVvdQU,1262
 tests/unit/cli/cloudbender/test_cli_reservation_unit.py,sha256=9y4FhcUwti7OzKYaJ23YAgOpx-dBIYj_0i--xaIk8YU,1407
 tests/unit/cloudbender/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/unit/cloudbender/test_datastores_unit.py,sha256=VVbINi31I3lf85JsBodqI-GZhfnGY5OIIlTC-8QFmIc,4644
-tests/unit/cloudbender/test_device_configs_unit.py,sha256=TYCmykeInX9t6PjXccdmZKuR26VomRsesCTQPdTq6gg,4858
-tests/unit/cloudbender/test_nodes_unit.py,sha256=dl5ufBbzvI559Y0AgdIjxF3QfQkw20R7tOx6N0NqFxI,4686
-tests/unit/cloudbender/test_providers_unit.py,sha256=dv0pzFIdswr1euB7k_nwMVcCUhuAAX8afWrRzxwKz98,3730
-tests/unit/cloudbender/test_regions_unit.py,sha256=BYX1NF6vI2sPj3aUbm-Pf1QvsEwwQZ3oDoZbfygYEy0,5635
-tests/unit/cloudbender/test_reservations_unit.py,sha256=JthNLYCQPV4RSECziuQ3Mqf-DOsOCEEbxtMfkwGeeH0,4827
-trainml/__init__.py,sha256=bDSwVwfrPDcEGAaZlSl9EWcFbz7F1UK7BvQWLfWNEHE,432
+tests/unit/cloudbender/test_datastores_unit.py,sha256=54mPokxhrRjlkBfqpmeA_q-PLml-HUNNit91aQVTpCg,5398
+tests/unit/cloudbender/test_device_configs_unit.py,sha256=lzyCuF7MRoQrtJVTQFL27lqPnRwQFv25htPgJqDuQI8,5714
+tests/unit/cloudbender/test_devices_unit.py,sha256=QBWnlOe1tw_XNA_i-yDHkmpGvtK36f2u1HhoXquoVaE,9103
+tests/unit/cloudbender/test_nodes_unit.py,sha256=BDpfJXCBNNpLt5rhJMk2BVXDQ_4QSmxoVrO_YPs6xBU,6231
+tests/unit/cloudbender/test_providers_unit.py,sha256=OgxifgC1IqLH8DNMKXy1Ne9_7a75ea6kHEOfRSRoQuQ,4373
+tests/unit/cloudbender/test_regions_unit.py,sha256=BbJICLIQmlotpA1UmLD0KTW_H9g2UW0J8ZYzQk1_Xjc,6299
+tests/unit/cloudbender/test_reservations_unit.py,sha256=nWEZ_p9EF2C49nbgL7Dt4NG2Irmyt94ZqJJQDyNfGFI,5624
+trainml/__init__.py,sha256=PACqklwvG2CiPJqOAaEZLAgJyM6nllgNp0p12pqDYUc,432
 trainml/__main__.py,sha256=JgErYkiskih8Y6oRwowALtR-rwQhAAdqOYWjQraRIPI,59
-trainml/auth.py,sha256=VI8wXgmXbwiBpgw0Sm-DzC4VmNqKg_zniua-Z61PafA,26646
+trainml/auth.py,sha256=gruZv27nhttrCbhcVQTH9kZkF2uMm1E06SwA_2pQAHQ,26565
 trainml/checkpoints.py,sha256=726yaeTHzIPvkQ-BhqmvJ6u0s1Oh732N6UAFbuWEv0Q,8274
 trainml/connections.py,sha256=h-S1NZbOkaXpIlpRStA6q-3eXc_OMlFWOLzF8R9SVG8,20029
-trainml/datasets.py,sha256=9Hb9AqGHu-mWfRWI2n3PqeUjKHc6GvowMyBPkFiRRq8,7394
+trainml/datasets.py,sha256=5zhxpmK12aA7iXp4n-8bD00Tie7bi3LU0tPNJKdnZlY,7935
 trainml/environments.py,sha256=OH4o08zXZ7IJ2CiA1rPnys2Fl45r8qvQHfM2mCBRAIc,1507
 trainml/exceptions.py,sha256=hhR78fI8rbU3fWQ-kUsgxOLyYP8D2bQcjLjrskqPG0Q,3724
-trainml/gpu_types.py,sha256=Nfaojyw4e-AeMPaIiExeOiewUSFWa1JVd65DZ8b_Hvk,1663
+trainml/gpu_types.py,sha256=mm-dwfYc02192bmYPIJmzesndyBcoOdkKYBaYZXOUwU,1901
 trainml/jobs.py,sha256=cUuiMvtJSFSupv2jy7fqAvUOkO8FRLfBDe3RuSZUjfE,17807
-trainml/models.py,sha256=laxVTyLqPNRSdst3rC0ql6HygCfl0ffcgXHF0t2Q1Bg,7043
+trainml/models.py,sha256=Lqs3OJMuOZXx8cfGFNC3JZ8nVK6l9_jU1xM1Uw1c5UQ,7750
 trainml/projects.py,sha256=Jk-0xhEgBFPIVMAK_Szxp3B-h-tJJSIEtTF8JWAYzo8,5209
 trainml/providers.py,sha256=97VegYVSeK0BuYv04hfBY_awNBbGz_GR_mdDrknfO-A,1844
-trainml/trainml.py,sha256=W7R0wn4XXHxm8ZdHGJeFJ89xnQ2rcTRWbrf_PTv6fg4,10982
+trainml/trainml.py,sha256=OrDXkJiStA_fTkwl0HX529n-cPCQUTVfdMOX-w0u1Oo,11045
 trainml/cli/__init__.py,sha256=-lA19Djkvqr1I-5PWoONgzHia-Lkg0J45dzt7qbetD8,4338
 trainml/cli/checkpoint.py,sha256=8Rh4bmFwJ4DKlIjHK-FLTeRynABqKCgIUGRtbQhAsX4,7170
 trainml/cli/connection.py,sha256=ELV6bPL30dzttFNxDU7Fb74R8oPL_E70k7TcJEzbwtQ,1700
-trainml/cli/dataset.py,sha256=joVsNxnCw23s1T6rTJMJeYRFHM-EIWmKFMsB7YyY1zs,6893
+trainml/cli/dataset.py,sha256=Pc00M6t7hGoRzCxznmmkijsWhG4PhIfG7UkrwtwykTY,6871
 trainml/cli/environment.py,sha256=dfm_T8OlCM5M8dLyOQBapJl3eFuVIku2P4JO6V0cYVQ,1019
 trainml/cli/gpu.py,sha256=CMcQyl2qbUgc2bc-gvUVT6X7bq2-sgiCHl3hyZ6kFWM,883
 trainml/cli/job.py,sha256=Xh_542a2rjAjhXyq8WG9ZEySboq_MmFQ7sCqYJYRXcM,4302
-trainml/cli/model.py,sha256=W3ah6Zy3whDrfUgvdaE3B7kVhPsLXBl_4Pr7dElb9oo,6071
+trainml/cli/model.py,sha256=hR23E6ttRXcLk-RofkPK6wUXMO7OU6sT6jTEHTmUg9Q,6111
 trainml/cli/project.py,sha256=iar4-Aic6k0O6HWfxB5vJ5EtjS4aAaw2g8wW7H-3Bjk,3522
 trainml/cli/provider.py,sha256=eaklYo0IUOGeyZ0ziZZz6UhOZ50Py4ewaPtWA9UDCqU,1594
-trainml/cli/cloudbender/__init__.py,sha256=xPo_OEqCBaU4xPnEpM4D8CBfANoO2qA2FtyV-J3XtdY,484
+trainml/cli/cloudbender/__init__.py,sha256=ewat7I7PvHQEW0a8zuky28kuMwtjL-gFQOAQmG0mmJA,534
 trainml/cli/cloudbender/datastore.py,sha256=gJ-comfAq65uiPoONQ35UIDLNVN7QKMf3l_2EcTN6zY,3478
+trainml/cli/cloudbender/device.py,sha256=KGZCFwwvS4tWsWuudrhlvquu_IFtV7LCUAOmCajicic,3453
 trainml/cli/cloudbender/node.py,sha256=iN_WaPCxOhtgDtnSsIFAEMGADG4MKiLjWoez6YSYwZI,3843
 trainml/cli/cloudbender/provider.py,sha256=oFjZWKfFQjNY7OtDu7nUdfv-RTmQc_Huuug963D3BdA,1726
 trainml/cli/cloudbender/region.py,sha256=X6-FYOb-pGpOEazn-NbsYSwa9ergB7FGATFkTe4a8Pk,2892
 trainml/cli/cloudbender/reservation.py,sha256=z2oMYwp-w_Keo1DepKUtuRnwiGz2VscVHDYWEFap1gs,3569
 trainml/cli/job/__init__.py,sha256=ljY-ELeXhXQ7txASbJEKGBom7OXfNyy7sWILz3nxRAE,6545
-trainml/cli/job/create.py,sha256=eaHOLRT0_ujp_xuMjWlR5GcSPpJ9x5AoS_Vdn6b55hU,34652
+trainml/cli/job/create.py,sha256=QDxHTSpiAjdnmGxUhBASwSNbL1pYqGoCV12UnB-c4N4,34655
 trainml/cloudbender/__init__.py,sha256=iE29obtC0_9f0IhRvHQcG5aY58fVhVYipTakpjAhdss,64
-trainml/cloudbender/cloudbender.py,sha256=ViVvOQTefXgWlyuLFQbLhVkTgRo0q2o4qxJQHBNtNv8,565
-trainml/cloudbender/datastores.py,sha256=1QoWT-cSKBA2-BA9703rJDsQ-mu7CI7RtKknSSp-6oU,3245
-trainml/cloudbender/device_configs.py,sha256=JBBb-33OD304Js7GsnBOwVCOs97TqcJ6qaw1gr-9oFk,3134
-trainml/cloudbender/nodes.py,sha256=dQva4u-ejTpKJAT51V4LxwZ2ovioyzVA5ysmQYUU6rc,3608
-trainml/cloudbender/providers.py,sha256=FyIF-FD9K9hkAD1NYJzRxW2ioNFw1bCWcadPVardv6w,1834
-trainml/cloudbender/regions.py,sha256=lRnBIYPn3DppVG2FuidK0o4apsCSg6ziixld3j5IYgg,3339
-trainml/cloudbender/reservations.py,sha256=D4ZlLudyNnbR2EhMuQWl14QoMqGWc3h-SgV3q75wL2w,3323
-trainml-0.5.0.dist-info/LICENSE,sha256=s0lpBxhSSUEpMavwde-Vb6K_K7xDCTTvSpNznVqVGR0,1069
-trainml-0.5.0.dist-info/METADATA,sha256=IBdBZCdm9Nio49kvHvlN161zyW6tv_CHu0s0H9aIdBA,7345
-trainml-0.5.0.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
-trainml-0.5.0.dist-info/entry_points.txt,sha256=OzBDm2wXby1bSGF02jTVxzRFZLejnbFiLHXhKdW3Bds,63
-trainml-0.5.0.dist-info/top_level.txt,sha256=Y1kLFRWKUW7RG8BX7cvejHF_yW8wBOaRYF1JQHENY4w,23
-trainml-0.5.0.dist-info/RECORD,,
+trainml/cloudbender/cloudbender.py,sha256=rAc93mtMaa3tWLi-NT4tiqO664MCPkhlGG9oNYSmimQ,634
+trainml/cloudbender/datastores.py,sha256=biVGifedc3r1DcuxsfCQh-f1Tw4HcJMMJfdgHxPfkKM,3506
+trainml/cloudbender/device_configs.py,sha256=DJWiGFaOE4C4xLE1BLDAiEjeL4T00R3FA_pb1xnSOr4,3399
+trainml/cloudbender/devices.py,sha256=QORNmKdLJoqGZmeWXRnivC1JmNBIw-ebvf4bsoem3r8,5660
+trainml/cloudbender/nodes.py,sha256=7HV2VLmxiUcJ-Kc6AAXS3M8C_XO-HKmaVgJpPdVnBQk,4332
+trainml/cloudbender/providers.py,sha256=-gkdiTu6Ah2znUuyyc3ZuRALagW8s1-OgqVjtlvc1AU,2036
+trainml/cloudbender/regions.py,sha256=Aqc_MeLVAeEv21e-lR5u8x1eintqUhZT2DBiQG3AcEE,3570
+trainml/cloudbender/reservations.py,sha256=rOrGXWIUHON4ad2aufEcvK4Yv_Mv3dDoScUtLJE8LWw,3586
+trainml-0.5.2.dist-info/LICENSE,sha256=s0lpBxhSSUEpMavwde-Vb6K_K7xDCTTvSpNznVqVGR0,1069
+trainml-0.5.2.dist-info/METADATA,sha256=RGNR1MBae8_fsnjuBRWPTi0d2fGR7ctzdSP5YOwC2Og,7345
+trainml-0.5.2.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
+trainml-0.5.2.dist-info/entry_points.txt,sha256=OzBDm2wXby1bSGF02jTVxzRFZLejnbFiLHXhKdW3Bds,63
+trainml-0.5.2.dist-info/top_level.txt,sha256=Y1kLFRWKUW7RG8BX7cvejHF_yW8wBOaRYF1JQHENY4w,23
+trainml-0.5.2.dist-info/RECORD,,