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

trainml/trainml.py

@@ -7,7 +7,7 @@ import traceback
 import random
 from importlib.metadata import version
 
-from trainml.auth import Auth
+from trainml.utils.auth import Auth
 from trainml.datasets import Datasets
 from trainml.models import Models
 from trainml.checkpoints import Checkpoints
@@ -16,7 +16,6 @@ from trainml.jobs import Jobs
 from trainml.gpu_types import GpuTypes
 from trainml.environments import Environments
 from trainml.exceptions import ApiError, TrainMLException
-from trainml.connections import Connections
 from trainml.projects import Projects
 from trainml.cloudbender import Cloudbender
 
@@ -34,13 +33,17 @@ class TrainML(object):
             os.environ.get("TRAINML_CONFIG_DIR") or "~/.trainml"
         )
         try:
-            with open(f"{CONFIG_DIR}/environment.json", "r") as file:
+            with open(
+                f"{CONFIG_DIR}/environment.json", "r", encoding="utf-8"
+            ) as file:
                 env_str = file.read().replace("\n", "")
             env = json.loads(env_str)
         except OSError:
             env = dict()
         try:
-            with open(f"{CONFIG_DIR}/config.json", "r") as file:
+            with open(
+                f"{CONFIG_DIR}/config.json", "r", encoding="utf-8"
+            ) as file:
                 config_str = file.read().replace("\n", "")
             config = json.loads(config_str)
         except OSError:
@@ -72,7 +75,6 @@ class TrainML(object):
         self.jobs = Jobs(self)
         self.gpu_types = GpuTypes(self)
         self.environments = Environments(self)
-        self.connections = Connections(self)
         self.projects = Projects(self)
         self.cloudbender = Cloudbender(self)
         self.api_url = (
@@ -92,7 +94,16 @@ class TrainML(object):
     def project(self) -> str:
         return self.active_project
 
-    async def _query(self, path, method, params=None, data=None, headers=None,max_retries=3, backoff_factor=0.5):
+    async def _query(
+        self,
+        path,
+        method,
+        params=None,
+        data=None,
+        headers=None,
+        max_retries=3,
+        backoff_factor=0.5,
+    ):
         try:
             tokens = self.auth.get_tokens()
         except TrainMLException as e:
@@ -120,7 +131,9 @@ class TrainML(object):
         )
         if params:
             if not isinstance(params, dict):
-                raise TrainMLException("Query parameters must be a valid dictionary")
+                raise TrainMLException(
+                    "Query parameters must be a valid dictionary"
+                )
             params = {
                 k: (str(v).lower() if isinstance(v, bool) else v)
                 for k, v in params.items()
@@ -154,27 +167,44 @@ class TrainML(object):
                         params=params,
                     ) as resp:
                         if (resp.status // 100) in [4, 5]:
-                            if resp.status == 502 and attempt < max_retries - 1:
-                                wait_time = (2 ** attempt) * backoff_factor * (random.random() + 0.5)
+                            if (
+                                resp.status == 502
+                                and attempt < max_retries - 1
+                            ):
+                                wait_time = (
+                                    (2**attempt)
+                                    * backoff_factor
+                                    * (random.random() + 0.5)
+                                )
                                 await asyncio.sleep(wait_time)
                                 continue
                             else:
                                 what = await resp.read()
-                                content_type = resp.headers.get("content-type", "")
+                                content_type = resp.headers.get(
+                                    "content-type", ""
+                                )
                                 resp.close()
                                 if content_type == "application/json":
-                                    raise ApiError(resp.status, json.loads(what.decode("utf8")))
+                                    raise ApiError(
+                                        resp.status,
+                                        json.loads(what.decode("utf8")),
+                                    )
                                 else:
-                                    raise ApiError(resp.status, {"message": what.decode("utf8")})
+                                    raise ApiError(
+                                        resp.status,
+                                        {"message": what.decode("utf8")},
+                                    )
                         results = await resp.json()
                         return results
             except aiohttp.ClientResponseError as e:
                 if e.status == 502 and attempt < max_retries - 1:
-                    wait_time = (2 ** attempt) * backoff_factor * (random.random() + 0.5)
+                    wait_time = (
+                        (2**attempt) * backoff_factor * (random.random() + 0.5)
+                    )
                     await asyncio.sleep(wait_time)
                     continue
                 else:
-                    raise  ApiError(e.status, f"Error {e.message}")
+                    raise ApiError(e.status, f"Error {e.message}")
 
         raise TrainMLException("Unexpected API failure")
 
@@ -184,114 +214,129 @@ class TrainML(object):
             "Content-Type": "application/json",
         }
         try:
-            tokens = self.auth.get_tokens()
-        except TrainMLException as e:
-            raise e
-        except Exception:
-            raise TrainMLException(
-                f"Error getting authorization tokens.  Verify configured credentials. Error: {traceback.format_exc()}"
-            )
-        async with aiohttp.ClientSession() as session:
-            done = False
-            async with session.ws_connect(
-                f"wss://{self.ws_url}?Authorization={tokens.get('id_token')}",
-                headers=headers,
-                heartbeat=30,
-            ) as ws:
-                asyncio.create_task(
-                    ws.send_json(
-                        dict(
-                            action="getlogs",
-                            data=dict(
-                                type="init",
-                                entity=entity,
-                                id=id,
-                                project_uuid=project_uuid,
-                            ),
-                        )
-                    )
+            try:
+                tokens = self.auth.get_tokens()
+            except TrainMLException as e:
+                raise e
+            except Exception:
+                raise TrainMLException(
+                    f"Error getting authorization tokens.  Verify configured credentials. Error: {traceback.format_exc()}"
                 )
-                asyncio.create_task(
-                    ws.send_json(
-                        dict(
-                            action="subscribe",
-                            data=dict(
-                                type="logs",
-                                entity=entity,
-                                id=id,
-                                project_uuid=project_uuid,
-                            ),
+            async with aiohttp.ClientSession() as session:
+                done = False
+                async with session.ws_connect(
+                    f"wss://{self.ws_url}?Authorization={tokens.get('id_token')}",
+                    headers=headers,
+                    heartbeat=30,
+                ) as ws:
+                    asyncio.create_task(
+                        ws.send_json(
+                            dict(
+                                action="getlogs",
+                                data=dict(
+                                    type="init",
+                                    entity=entity,
+                                    id=id,
+                                    project_uuid=project_uuid,
+                                ),
+                            )
                         )
                     )
-                )
-                async for msg in ws:
-                    if msg.type in (
-                        aiohttp.WSMsgType.CLOSED,
-                        aiohttp.WSMsgType.ERROR,
-                        aiohttp.WSMsgType.CLOSE,
-                    ):
-                        logging.debug(
-                            f"Websocket Received Closed Message.  Done? {done}"
+                    asyncio.create_task(
+                        ws.send_json(
+                            dict(
+                                action="subscribe",
+                                data=dict(
+                                    type="logs",
+                                    entity=entity,
+                                    id=id,
+                                    project_uuid=project_uuid,
+                                ),
+                            )
                         )
-                        await ws.close()
-                        break
-                    data = json.loads(msg.data)
-                    if data.get("type") == "end":
-                        done = True
-                        asyncio.create_task(delayed_close(ws))
-                    else:
-                        msg_handler(data)
-            logging.debug(f"Websocket Disconnected.  Done? {done}")
+                    )
+                    async for msg in ws:
+                        if msg.type in (
+                            aiohttp.WSMsgType.CLOSED,
+                            aiohttp.WSMsgType.ERROR,
+                            aiohttp.WSMsgType.CLOSE,
+                        ):
+                            logging.debug(
+                                f"Websocket Received Closed Message.  Done? {done}"
+                            )
+                            await ws.close()
+                            break
+                        data = json.loads(msg.data)
+                        if data.get("type") == "end":
+                            done = True
+                            asyncio.create_task(delayed_close(ws))
+                        else:
+                            msg_handler(data)
+                logging.debug(f"Websocket Disconnected.  Done? {done}")
 
-            connection_tries = 0
-            while not done:
-                tokens = self.auth.get_tokens()
-                try:
-                    async with session.ws_connect(
-                        f"wss://{self.ws_url}?Authorization={tokens.get('id_token')}",
-                        headers=headers,
-                        heartbeat=30,
-                    ) as ws:
-                        asyncio.create_task(
-                            ws.send_json(
-                                dict(
-                                    action="subscribe",
-                                    data=dict(
-                                        type="logs",
-                                        entity=entity,
-                                        id=id,
-                                        project_uuid=project_uuid,
-                                    ),
+                connection_tries = 0
+                while not done:
+                    tokens = self.auth.get_tokens()
+                    try:
+                        async with session.ws_connect(
+                            f"wss://{self.ws_url}?Authorization={tokens.get('id_token')}",
+                            headers=headers,
+                            heartbeat=30,
+                        ) as ws:
+                            asyncio.create_task(
+                                ws.send_json(
+                                    dict(
+                                        action="subscribe",
+                                        data=dict(
+                                            type="logs",
+                                            entity=entity,
+                                            id=id,
+                                            project_uuid=project_uuid,
+                                        ),
+                                    )
                                 )
                             )
+                            async for msg in ws:
+                                if msg.type in (
+                                    aiohttp.WSMsgType.CLOSED,
+                                    aiohttp.WSMsgType.ERROR,
+                                    aiohttp.WSMsgType.CLOSE,
+                                ):
+                                    logging.debug(
+                                        f"Websocket Received Closed Message.  Done? {done}"
+                                    )
+                                    await ws.close()
+                                    break
+                                data = json.loads(msg.data)
+                                if data.get("type") == "end":
+                                    done = True
+                                    asyncio.create_task(delayed_close(ws))
+                                else:
+                                    msg_handler(data)
+                        connection_tries = 0
+                        logging.debug(f"Websocket Disconnected.  Done? {done}")
+                    except Exception as e:
+                        connection_tries += 1
+                        logging.debug(
+                            f"Connection error: {traceback.format_exc()}"
                         )
-                        async for msg in ws:
-                            if msg.type in (
-                                aiohttp.WSMsgType.CLOSED,
-                                aiohttp.WSMsgType.ERROR,
-                                aiohttp.WSMsgType.CLOSE,
-                            ):
-                                logging.debug(
-                                    f"Websocket Received Closed Message.  Done? {done}"
-                                )
-                                await ws.close()
-                                break
-                            data = json.loads(msg.data)
-                            if data.get("type") == "end":
-                                done = True
-                                asyncio.create_task(delayed_close(ws))
-                            else:
-                                msg_handler(data)
-                    connection_tries = 0
-                    logging.debug(f"Websocket Disconnected.  Done? {done}")
-                except Exception as e:
-                    connection_tries += 1
-                    logging.debug(f"Connection error: {traceback.format_exc()}")
-                    if connection_tries == 5:
-                        raise ApiError(
-                            500,
-                            {"message": f"Connection error: {traceback.format_exc()}"},
-                        )
+                        if connection_tries == 5:
+                            raise ApiError(
+                                500,
+                                {
+                                    "message": f"Connection error: {traceback.format_exc()}"
+                                },
+                            )
+        except GeneratorExit:
+            # Handle graceful shutdown - GeneratorExit is raised during
+            # event loop cleanup. Don't re-raise to avoid "coroutine ignored"
+            # warnings.
+            logging.debug("Websocket subscription cancelled during shutdown")
+            return
+        except asyncio.CancelledError:
+            # Re-raise CancelledError to properly propagate task cancellation
+            logging.debug("Websocket subscription task cancelled")
+            raise
 
     def set_active_project(self, project_uuid):
         CONFIG_DIR = os.path.expanduser(

trainml/utils/__init__.py

@@ -0,0 +1 @@
+"""Utility modules for trainml SDK."""