trainml 0.5.12__tar.gz → 0.5.14__tar.gz

{trainml-0.5.12 → trainml-0.5.14}/PKG-INFO

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

trainml-0.5.14/tests/integration/projects/test_projects_members_integration.py

@@ -0,0 +1,49 @@
+import re
+import sys
+import asyncio
+from pytest import mark, fixture
+
+pytestmark = [mark.sdk, mark.integration, mark.projects]
+
+
+@mark.create
+@mark.asyncio
+class ProjectMembersTests:
+    @fixture(scope="class")
+    async def project_member(self, project):
+        member = await project.members.add("test.account@proximl.ai","read","read","read","read","read")
+        yield member
+        await project.members.remove("test.account@proximl.ai")
+
+    async def test_list_project_members(self, project):
+        members = await project.members.list()
+        assert len(members) > 0
+
+    async def test_project_member_properties(self, project, project_member):
+        assert isinstance(project_member.id, str)
+        assert isinstance(project_member.email, str)
+        assert isinstance(project_member.project_uuid, str)
+        assert isinstance(project_member.owner, bool)
+        assert isinstance(project_member.job, str)
+        assert isinstance(project_member.dataset, str)
+        assert isinstance(project_member.model, str)
+        assert isinstance(project_member.checkpoint, str)
+        assert isinstance(project_member.volume, str)
+        assert project.id == project_member.project_uuid
+        assert project_member.id == "test.account@proximl.ai"
+
+    async def test_project_member_str(self, project_member):
+        string = str(project_member)
+        regex = r"^{.*\"email\": \"" + project_member.email + r"\".*}$"
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    async def test_project_member_repr(self, project_member):
+        string = repr(project_member)
+        regex = (
+            r"^ProjectMember\( trainml , \*\*{.*'email': '"
+            + project_member.email
+            + r"'.*}\)$"
+        )
+        assert isinstance(string, str)
+        assert re.match(regex, string)

{trainml-0.5.12 → trainml-0.5.14}/tests/integration/test_jobs_integration.py

@@ -175,7 +175,7 @@ class JobAPIResourceValidationTests:
                 disk_size=10,
             )
         assert (
-            "Invalid Request - None (CPU Only) may be not be combined with other GPU Types"
+            "Invalid Request - CPU Only may be not be combined with other GPU Types"
             in error.value.message
         )
 
@@ -423,11 +423,11 @@ class JobIOTests:
             disk_size=10,
             workers=["python $ML_MODEL_PATH/tensorflow/main.py"],
             environment=dict(
-                type="DEEPLEARNING_PY310",
+                type="DEEPLEARNING_PY312",
                 env=[
                     dict(
                         key="CHECKPOINT_FILE",
-                        value="model.ckpt-0050",
+                        value="model.ckpt-0050.weights.h5",
                     )
                 ],
             ),
@@ -469,7 +469,7 @@ class JobIOTests:
         sys.stderr.write(captured.err)
         assert "Epoch 1/2" in captured.out
         assert "Epoch 2/2" in captured.out
-        assert "adding: model.ckpt-0001.data-00000-of-00001" in captured.out
+        assert "adding: model.ckpt-0001" in captured.out
         assert "Send complete" in captured.out
 
     async def test_job_model_input_and_output(self, trainml, capsys):
@@ -558,7 +558,7 @@ class JobTypeTests:
         assert job.url
         assert extract_domain_suffix(urlparse(job.url).hostname) == "proximl.cloud"
         tries = 0
-        await asyncio.sleep(30)
+        await asyncio.sleep(180) ## downloading weights can be slow
         async with aiohttp.ClientSession() as session:
             retry = True
             while retry:
@@ -640,7 +640,7 @@ class JobTypeTests:
         assert "Epoch 2/2" in captured.out
         assert "Uploading s3://trainml-example/output/resnet_cifar10" in captured.out
         assert (
-            "upload: ./model.ckpt-0002.data-00000-of-00001 to s3://trainml-example/output/resnet_cifar10/model.ckpt-0002.data-00000-of-00001"
+            "upload: ./model.ckpt-0002.weights.h5 to s3://trainml-example/output/resnet_cifar10/model.ckpt-0002.weights.h5"
             in captured.out
         )
         assert "Upload complete" in captured.out
@@ -713,9 +713,9 @@ class JobFeatureTests:
         sys.stderr.write(captured.err)
         upload_contents = os.listdir(temp_dir.name)
         temp_dir.cleanup()
-        assert len(upload_contents) > 4
+        assert len(upload_contents) >= 3
         assert any(
-            "model.ckpt-0002.data-00000-of-00001" in content
+            "model.ckpt-0002" in content
             for content in upload_contents
         )
 
@@ -724,5 +724,5 @@ class JobFeatureTests:
         sys.stderr.write(captured.err)
         assert "Epoch 1/2" in captured.out
         assert "Epoch 2/2" in captured.out
-        assert "Number of regular files transferred: 7" in captured.out
+        assert "Number of regular files transferred: 4" in captured.out
         assert "Send complete" in captured.out

trainml-0.5.14/tests/unit/projects/test_project_members_unit.py

@@ -0,0 +1,107 @@
+import re
+import json
+import logging
+from unittest.mock import AsyncMock, patch
+from pytest import mark, fixture, raises
+from aiohttp import WSMessage, WSMsgType
+
+import trainml.projects.members as specimen
+from trainml.exceptions import (
+    ApiError,
+    SpecificationError,
+    TrainMLException,
+)
+
+pytestmark = [mark.sdk, mark.unit, mark.projects]
+
+
+@fixture
+def project_members(mock_trainml):
+    yield specimen.ProjectMembers(mock_trainml, project_id="1")
+
+
+@fixture
+def project_member(mock_trainml):
+    yield specimen.ProjectMember(
+        mock_trainml,
+        id="owner@gmail.com",
+        email="owner@gmail.com",
+        project_uuid="proj-id-1",
+       owner= True,
+                job= "all",
+                dataset= "all",
+                model= "all",
+                checkpoint="all",
+                volume= "all"
+    )
+
+
+class ProjectMembersTests:
+    @mark.asyncio
+    async def test_project_members_list(self, project_members, mock_trainml):
+        api_response = [
+            {
+                "project_uuid": "proj-id-1",
+                "email": "owner@gmail.com",
+                "createdAt": "2024-09-04T00:42:39.529Z",
+                "updatedAt": "2024-09-04T00:42:39.529Z",
+                "owner": True,
+                "job": "all",
+                "dataset": "all",
+                "model": "all",
+                "checkpoint": "all",
+                "volume": "all"
+            },
+            {
+                "project_uuid": "proj-id-1",
+                "email": "non-owner@gmail.com",
+                "createdAt": "2024-09-04T00:42:39.529Z",
+                "updatedAt": "2024-09-04T00:42:39.529Z",
+                "owner": False,
+                "job": "all",
+                "dataset": "all",
+                "model": "all",
+                "checkpoint": "read",
+                "volume": "read"
+            },
+        ]
+        mock_trainml._query = AsyncMock(return_value=api_response)
+        resp = await project_members.list()
+        mock_trainml._query.assert_called_once_with(
+            "/project/1/access", "GET", dict()
+        )
+        assert len(resp) == 2
+
+
+class ProjectMemberTests:
+    def test_project_member_properties(self, project_member):
+        assert isinstance(project_member.id, str)
+        assert isinstance(project_member.email, str)
+        assert isinstance(project_member.project_uuid, str)
+        assert isinstance(project_member.owner, bool)
+        assert isinstance(project_member.job, str)
+        assert isinstance(project_member.dataset, str)
+        assert isinstance(project_member.model, str)
+        assert isinstance(project_member.checkpoint, str)
+        assert isinstance(project_member.volume, str)
+
+    def test_project_member_str(self, project_member):
+        string = str(project_member)
+        regex = r"^{.*\"id\": \"" + project_member.id + r"\".*}$"
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    def test_project_member_repr(self, project_member):
+        string = repr(project_member)
+        regex = (
+            r"^ProjectMember\( trainml , \*\*{.*'id': '"
+            + project_member.id
+            + r"'.*}\)$"
+        )
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    def test_project_member_bool(self, project_member, mock_trainml):
+        empty_project_member = specimen.ProjectMember(mock_trainml)
+        assert bool(project_member)
+        assert not bool(empty_project_member)

{trainml-0.5.12 → trainml-0.5.14}/trainml/__init__.py

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

{trainml-0.5.12 → trainml-0.5.14}/trainml/auth.py

@@ -508,7 +508,7 @@ class AWSSRP(object):
 
 
 class Auth(object):
-    def __init__(self, config_dir, domain_suffix="trainml.ai", **kwargs):
+    def __init__(self, config_dir, domain_suffix="proximl.ai", **kwargs):
         try:
             with open(f"{config_dir}/environment.json", "r") as file:
                 env_str = file.read().replace("\n", "")

{trainml-0.5.12 → trainml-0.5.14}/trainml/cli/job/create.py

@@ -180,6 +180,12 @@ def create(config):
     help="Third Party Credentials to add to the job environment",
     multiple=True,
 )
+@click.option(
+    "--secret",
+    type=click.STRING,
+    help="Project secrets to add to the job environment",
+    multiple=True,
+)
 @click.option(
     "--git-uri",
     type=click.STRING,
@@ -222,6 +228,7 @@ def notebook(
     custom_image,
     env,
     credential,
+    secret,
     apt_packages,
     pip_packages,
     conda_packages,
@@ -254,7 +261,7 @@ def notebook(
         data=dict(datasets=datasets),
         model=dict(checkpoints=checkpoints),
         environment=dict(
-            credentials=[k for k in credential],
+            credentials=[k for k in credential], secrets=[s for s in secret]
         ),
     )
 
@@ -508,6 +515,12 @@ def notebook(
     help="Third Party Credentials to add to the job environment",
     multiple=True,
 )
+@click.option(
+    "--secret",
+    type=click.STRING,
+    help="Project secrets to add to the job environment",
+    multiple=True,
+)
 @click.option(
     "--apt-packages",
     type=click.STRING,
@@ -564,6 +577,7 @@ def training(
     custom_image,
     env,
     credential,
+    secret,
     apt_packages,
     pip_packages,
     conda_packages,
@@ -596,7 +610,7 @@ def training(
         data=dict(datasets=datasets),
         model=dict(checkpoints=checkpoints),
         environment=dict(
-            credentials=[k for k in credential],
+            credentials=[k for k in credential], secrets=[s for s in secret]
         ),
     )
 
@@ -852,6 +866,12 @@ def training(
     help="Third Party Credentials to add to the job environment",
     multiple=True,
 )
+@click.option(
+    "--secret",
+    type=click.STRING,
+    help="Project secrets to add to the job environment",
+    multiple=True,
+)
 @click.option(
     "--apt-packages",
     type=click.STRING,
@@ -908,6 +928,7 @@ def inference(
     custom_image,
     env,
     credential,
+    secret,
     apt_packages,
     pip_packages,
     conda_packages,
@@ -934,6 +955,7 @@ def inference(
         model=dict(checkpoints=checkpoints),
         environment=dict(
             credentials=[k for k in credential],
+            secrets=[s for s in secret],
         ),
     )
 
@@ -1171,6 +1193,12 @@ def from_json(config, attach, connect, file):
     help="Third Party Credentials to add to the job environment",
     multiple=True,
 )
+@click.option(
+    "--secret",
+    type=click.STRING,
+    help="Project secrets to add to the job environment",
+    multiple=True,
+)
 @click.option(
     "--apt-packages",
     type=click.STRING,
@@ -1231,6 +1259,7 @@ def endpoint(
     custom_image,
     env,
     credential,
+    secret,
     apt_packages,
     pip_packages,
     conda_packages,
@@ -1258,6 +1287,7 @@ def endpoint(
         model=dict(checkpoints=checkpoints),
         environment=dict(
             credentials=[k for k in credential],
+            secrets=[s for s in secret],
         ),
     )
 

{trainml-0.5.12 → trainml-0.5.14}/trainml/jobs.py

@@ -530,6 +530,15 @@ class Job:
                         "waiting for data/model download",
                     ]
                 )
+                or (
+                    status
+                    == "running"  ## this status could be too short for polling could miss it
+                    and self.status
+                    in [
+                        "uploading",
+                        "finished"
+                    ]
+                )
             ):
                 return self
             elif self.status == "failed":

{trainml-0.5.12 → trainml-0.5.14}/trainml/projects/data_connectors.py

@@ -7,6 +7,12 @@ class ProjectDataConnectors(object):
         self.trainml = trainml
         self.project_id = project_id
 
+    async def get(self, id, **kwargs):
+        resp = await self.trainml._query(
+            f"/project/{self.project_id}/data_connectors/{id}", "GET", kwargs
+        )
+        return ProjectDataConnector(self.trainml, **resp)
+
     async def list(self, **kwargs):
         resp = await self.trainml._query(
             f"/project/{self.project_id}/data_connectors", "GET", kwargs
@@ -61,3 +67,13 @@ class ProjectDataConnector:
 
     def __bool__(self):
         return bool(self._id)
+
+    async def enable(self):
+        await self.trainml._query(
+            f"/project/{self._project_uuid}/data_connectors/{self._id}/enable", "PATCH"
+        )
+
+    async def disable(self):
+        await self.trainml._query(
+            f"/project/{self._project_uuid}/data_connectors/{self._id}/disable", "PATCH"
+        )

{trainml-0.5.12 → trainml-0.5.14}/trainml/projects/datastores.py

@@ -7,6 +7,12 @@ class ProjectDatastores(object):
         self.trainml = trainml
         self.project_id = project_id
 
+    async def get(self, id, **kwargs):
+        resp = await self.trainml._query(
+            f"/project/{self.project_id}/datastores/{id}", "GET", kwargs
+        )
+        return ProjectDatastore(self.trainml, **resp)
+
     async def list(self, **kwargs):
         resp = await self.trainml._query(
             f"/project/{self.project_id}/datastores", "GET", kwargs
@@ -56,3 +62,13 @@ class ProjectDatastore:
 
     def __bool__(self):
         return bool(self._id)
+
+    async def enable(self):
+        await self.trainml._query(
+            f"/project/{self._project_uuid}/datastores/{self._id}/enable", "PATCH"
+        )
+
+    async def disable(self):
+        await self.trainml._query(
+            f"/project/{self._project_uuid}/datastores/{self._id}/disable", "PATCH"
+        )

trainml-0.5.14/trainml/projects/members.py

@@ -0,0 +1,98 @@
+import json
+import logging
+from typing import Literal
+
+class ProjectMembers(object):
+    def __init__(self, trainml, project_id):
+        self.trainml = trainml
+        self.project_id = project_id
+
+    async def list(self, **kwargs):
+        resp = await self.trainml._query(
+            f"/project/{self.project_id}/access", "GET", kwargs
+        )
+        members = [ProjectMember(self.trainml, **member) for member in resp]
+        return members
+    
+    async def add(self, email: str, job: Literal["all", "read"], dataset: Literal["all", "read"], model: Literal["all", "read"], checkpoint: Literal["all", "read"], volume: Literal["all", "read"],  **kwargs):
+        data = dict(
+            email=email,
+            job=job,
+            dataset=dataset,
+            model=model,
+            checkpoint=checkpoint,
+            volume=volume,
+        )
+        payload = {k: v for k, v in data.items() if v is not None}
+        resp = await self.trainml._query(
+            f"/project/{self.project_id}/access", "POST",kwargs, payload)
+        member = ProjectMember(self.trainml, **resp)
+        logging.info(f"Added Project Member {email} to project {self.project_id}")
+        return member
+
+    
+    async def remove(self, email, **kwargs):
+        await self.trainml._query(
+            f"/project/{self.project_id}/access", "DELETE", dict(**kwargs, email=email)
+        )
+
+
+class ProjectMember:
+    def __init__(self, trainml, **kwargs):
+        self.trainml = trainml
+        self._entity = kwargs
+        self._id = self._entity.get("email")
+        self._project_uuid = self._entity.get("project_uuid")
+        self._owner = self._entity.get("owner")
+        self._job = self._entity.get("job")
+        self._dataset = self._entity.get("dataset")
+        self._model = self._entity.get("model")
+        self._checkpoint = self._entity.get("checkpoint")
+        self._volume = self._entity.get("volume")
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    @property
+    def project_uuid(self) -> str:
+        return self._project_uuid
+    
+    @property
+    def email(self) -> str:
+        return self._id
+
+    @property
+    def owner(self) -> bool:
+        return self._owner
+
+    @property
+    def job(self) -> str:
+        return self._job
+
+    @property
+    def dataset(self) -> str:
+        return self._dataset
+    
+    @property
+    def model(self) -> str:
+        return self._model
+    
+    @property
+    def checkpoint(self) -> str:
+        return self._checkpoint
+    
+    @property
+    def volume(self) -> str:
+        return self._volume
+
+    def __str__(self):
+        return json.dumps({k: v for k, v in self._entity.items()})
+
+    def __repr__(self):
+        return f"ProjectMember( trainml , **{self._entity.__repr__()})"
+
+    def __bool__(self):
+        return bool(self._id)
+
+

{trainml-0.5.12 → trainml-0.5.14}/trainml/projects/projects.py

@@ -5,6 +5,7 @@ from .data_connectors import ProjectDataConnectors
 from .services import ProjectServices
 from .credentials import ProjectCredentials
 from .secrets import ProjectSecrets
+from .members import ProjectMembers
 
 
 class Projects(object):
@@ -56,6 +57,7 @@ class Project:
         self.services = ProjectServices(self.trainml, self._id)
         self.credentials = ProjectCredentials(self.trainml, self._id)
         self.secrets = ProjectSecrets(self.trainml, self._id)
+        self.members = ProjectMembers(self.trainml, self._id)
 
     @property
     def id(self) -> str:

{trainml-0.5.12 → trainml-0.5.14}/trainml/projects/services.py

@@ -7,6 +7,12 @@ class ProjectServices(object):
         self.trainml = trainml
         self.project_id = project_id
 
+    async def get(self, id, **kwargs):
+        resp = await self.trainml._query(
+            f"/project/{self.project_id}/services/{id}", "GET", kwargs
+        )
+        return ProjectService(self.trainml, **resp)
+
     async def list(self, **kwargs):
         resp = await self.trainml._query(
             f"/project/{self.project_id}/services", "GET", kwargs
@@ -61,3 +67,13 @@ class ProjectService:
 
     def __bool__(self):
         return bool(self._id)
+
+    async def enable(self):
+        await self.trainml._query(
+            f"/project/{self._project_uuid}/services/{self._id}/enable", "PATCH"
+        )
+
+    async def disable(self):
+        await self.trainml._query(
+            f"/project/{self._project_uuid}/services/{self._id}/disable", "PATCH"
+        )

{trainml-0.5.12 → trainml-0.5.14}/trainml/trainml.py

@@ -4,6 +4,7 @@ import asyncio
 import aiohttp
 import logging
 import traceback
+import random
 from importlib.metadata import version
 
 from trainml.auth import Auth
@@ -48,7 +49,7 @@ class TrainML(object):
             kwargs.get("domain_suffix")
             or os.environ.get("TRAINML_DOMAIN_SUFFIX")
             or env.get("domain_suffix")
-            or "trainml.ai"
+            or "proximl.ai"
         )
         self.auth = Auth(
             config_dir=CONFIG_DIR,
@@ -91,7 +92,7 @@ class TrainML(object):
     def project(self) -> str:
         return self.active_project
 
-    async def _query(self, path, method, params=None, data=None, headers=None):
+    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:
@@ -142,24 +143,40 @@ class TrainML(object):
         logging.debug(
             f"Request - Url: {url}, Method: {method}, Params: {params}, Body: {data}, Headers: {headers}"
         )
-        async with aiohttp.ClientSession() as session:
-            async with session.request(
-                method,
-                url,
-                data=json.dumps(data),
-                headers=headers,
-                params=params,
-            ) as resp:
-                if (resp.status // 100) in [4, 5]:
-                    what = await resp.read()
-                    content_type = resp.headers.get("content-type", "")
-                    resp.close()
-                    if content_type == "application/json":
-                        raise ApiError(resp.status, json.loads(what.decode("utf8")))
-                    else:
-                        raise ApiError(resp.status, {"message": what.decode("utf8")})
-                results = await resp.json()
-                return results
+        for attempt in range(max_retries):
+            try:
+                async with aiohttp.ClientSession() as session:
+                    async with session.request(
+                        method,
+                        url,
+                        data=json.dumps(data),
+                        headers=headers,
+                        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)
+                                await asyncio.sleep(wait_time)
+                                continue
+                            else:
+                                what = await resp.read()
+                                content_type = resp.headers.get("content-type", "")
+                                resp.close()
+                                if content_type == "application/json":
+                                    raise ApiError(resp.status, json.loads(what.decode("utf8")))
+                                else:
+                                    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)
+                    await asyncio.sleep(wait_time)
+                    continue
+                else:
+                    raise  ApiError(e.status, f"Error {e.message}")
+
+        raise TrainMLException("Unexpected API failure")
 
     async def _ws_subscribe(self, entity, project_uuid, id, msg_handler):
         headers = {

{trainml-0.5.12 → trainml-0.5.14}/trainml.egg-info/PKG-INFO

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

{trainml-0.5.12 → trainml-0.5.14}/trainml.egg-info/SOURCES.txt

@@ -23,6 +23,7 @@ tests/integration/projects/test_projects_credentials_integration.py
 tests/integration/projects/test_projects_data_connectors_integration.py
 tests/integration/projects/test_projects_datastores_integration.py
 tests/integration/projects/test_projects_integration.py
+tests/integration/projects/test_projects_members_integration.py
 tests/integration/projects/test_projects_secrets_integration.py
 tests/integration/projects/test_projects_services_integration.py
 tests/unit/__init__.py
@@ -74,6 +75,7 @@ tests/unit/projects/__init__.py
 tests/unit/projects/test_project_credentials_unit.py
 tests/unit/projects/test_project_data_connectors_unit.py
 tests/unit/projects/test_project_datastores_unit.py
+tests/unit/projects/test_project_members_unit.py
 tests/unit/projects/test_project_secrets_unit.py
 tests/unit/projects/test_project_services_unit.py
 tests/unit/projects/test_projects_unit.py
@@ -134,6 +136,7 @@ trainml/projects/__init__.py
 trainml/projects/credentials.py
 trainml/projects/data_connectors.py
 trainml/projects/datastores.py
+trainml/projects/members.py
 trainml/projects/projects.py
 trainml/projects/secrets.py
 trainml/projects/services.py