trainml 0.5.13__py3-none-any.whl → 0.5.15__py3-none-any.whl

tests/integration/cloudbender/conftest.py

@@ -0,0 +1,28 @@
+
+from pytest import fixture, mark
+
+
+pytestmark = [mark.integration,mark.cloudbender]
+
+@fixture(scope="session")
+@mark.create
+@mark.asyncio
+@mark.xdist_group("cloudbender_resources")
+async def provider( trainml):
+    provider = await trainml.cloudbender.providers.enable(type="test")
+    await provider.wait_for("ready")
+    yield provider
+    await provider.remove()
+
+@fixture(scope="session")
+@mark.create
+@mark.asyncio
+@mark.xdist_group("cloudbender_resources")
+async def region(trainml, provider):
+    region = await trainml.cloudbender.regions.create(provider_uuid=provider.id,name="test-region",
+        public=False,
+        storage=dict(mode="local"),)
+    await region.wait_for("healthy")
+    yield region
+    await region.remove()
+    await region.wait_for("archived")

tests/integration/cloudbender/test_providers_integration.py

@@ -8,14 +8,9 @@ pytestmark = [mark.sdk, mark.integration, mark.cloudbender, mark.providers]
 
 @mark.create
 @mark.asyncio
-class GetProvidersTests:
-    @fixture(scope="class")
-    async def provider(self, trainml):
-        provider = await trainml.cloudbender.providers.enable(type="test")
-        yield provider
-        await provider.remove()
-
-    async def test_get_providers(self, trainml):
+@mark.xdist_group("cloudbender_resources")
+class GetProviderTests:
+    async def test_get_providers(self, trainml, provider):
         providers = await trainml.cloudbender.providers.list()
         assert len(providers) > 0
 

tests/integration/cloudbender/test_regions_integration.py

@@ -0,0 +1,42 @@
+import re
+import sys
+import asyncio
+from pytest import mark, fixture
+
+pytestmark = [mark.sdk, mark.integration, mark.cloudbender, mark.regions]
+
+
+@mark.create
+@mark.asyncio
+@mark.xdist_group("cloudbender_resources")
+class GetRegionTests:
+    async def test_get_regions(self, trainml, provider, region):
+        regions = await trainml.cloudbender.regions.list(provider_uuid=provider.id)
+        assert len(regions) > 0
+
+    async def test_get_region(self, trainml, provider, region):
+        response = await trainml.cloudbender.regions.get(provider.id, region.id)
+        assert response.id == region.id
+
+    async def test_region_properties(self, provider, region):
+        assert isinstance(region.id, str)
+        assert isinstance(region.provider_uuid, str)
+        assert isinstance(region.type, str)
+        assert region.type == "test"
+        assert region.provider_uuid == provider.id
+
+    async def test_region_str(self, region):
+        string = str(region)
+        regex = r"^{.*\"region_uuid\": \"" + region.id + r"\".*}$"
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    async def test_region_repr(self, region):
+        string = repr(region)
+        regex = (
+            r"^Region\( trainml , \*\*{.*'region_uuid': '"
+            + region.id
+            + r"'.*}\)$"
+        )
+        assert isinstance(string, str)
+        assert re.match(regex, string)

tests/integration/cloudbender/test_services_integration.py

@@ -0,0 +1,87 @@
+import re
+import sys
+import asyncio
+from pytest import mark, fixture
+from cryptography import x509
+from cryptography.x509.oid import NameOID, ExtendedKeyUsageOID
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.hazmat.primitives import serialization, hashes
+
+pytestmark = [mark.sdk, mark.integration, mark.cloudbender, mark.regions]
+
+def get_csr(service_id):
+    private_key = rsa.generate_private_key(
+        public_exponent=65537,
+        key_size=4096
+    )
+    csr = x509.CertificateSigningRequestBuilder().subject_name(x509.Name([
+        x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
+        x509.NameAttribute(NameOID.ORGANIZATION_NAME, "proxiML"),
+        x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, service_id),
+        x509.NameAttribute(NameOID.COMMON_NAME, "test-client"),  # Client identity
+    ])).add_extension(
+        x509.ExtendedKeyUsage([
+            ExtendedKeyUsageOID.CLIENT_AUTH  # Client authentication usage
+        ]),
+        critical=True
+    ).sign(private_key, hashes.SHA256())
+    return csr.public_bytes(serialization.Encoding.PEM).decode("utf-8")
+
+@mark.create
+@mark.asyncio
+@mark.xdist_group("cloudbender_resources")
+class GetServiceTests:
+    @fixture(scope="class")
+    async def service(self, trainml, region):
+        service = await trainml.cloudbender.services.create(
+            provider_uuid=region.provider_uuid,
+            region_uuid=region.id,
+            name="CLI Automated Service",
+            type="tcp",
+            port="8989",
+            public=False,
+        )
+        await service.wait_for("active")
+        yield service
+        await service.remove()
+        await service.wait_for("archived")
+
+    async def test_get_services(self, trainml, region,service):
+        services = await trainml.cloudbender.services.list(provider_uuid=region.provider_uuid, region_uuid=region.id)
+        assert len(services) > 0
+
+    async def test_get_service(self, trainml, provider, region, service):
+        response = await trainml.cloudbender.services.get(provider.id, region.id, service.id)
+        assert response.id == service.id
+
+    async def test_service_properties(self, region, service):
+        assert isinstance(service.id, str)
+        assert isinstance(service.provider_uuid, str)
+        assert isinstance(service.region_uuid, str)
+        assert isinstance(service.public, bool)
+        assert service.port == "8989"
+        assert service.provider_uuid == region.provider_uuid
+        assert service.region_uuid == region.id
+
+    async def test_service_str(self, service):
+        string = str(service)
+        regex = r"^{.*\"service_id\": \"" + service.id + r"\".*}$"
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    async def test_service_repr(self, service):
+        string = repr(service)
+        regex = (
+            r"^Service\( trainml , \*\*{.*'service_id': '"
+            + service.id
+            + r"'.*}\)$"
+        )
+        assert isinstance(string, str)
+        assert re.match(regex, string)
+
+    async def test_service_certificate(self, service):
+        service = await service.generate_certificate()
+        assert isinstance(service._service.get("auth_cert"), str)
+        csr = get_csr(service.id)
+        certificate = await service.sign_client_certificate(csr)
+        assert isinstance(certificate, str)

tests/integration/conftest.py

@@ -40,7 +40,7 @@ def env(request):
     yield ENVS[env]
 
 
-@fixture(scope="module")
+@fixture(scope="session")
 def trainml(env):
     trainml = TrainML(**env)
     yield trainml

tests/integration/projects/conftest.py

@@ -1,7 +1,8 @@
-from pytest import fixture
+from pytest import fixture, mark
 
 
 @fixture(scope="module")
+@mark.xdist_group("project_resources")
 async def project(trainml):
     project = await trainml.projects.create(
         name="New Project", copy_credentials=False, copy_secrets=False

tests/integration/projects/test_projects_credentials_integration.py

@@ -8,6 +8,7 @@ pytestmark = [mark.sdk, mark.integration, mark.projects]
 
 @mark.create
 @mark.asyncio
+@mark.xdist_group("project_resources")
 class ProjectCredentialsTests:
     @fixture(scope="class")
     async def project_credential(self, project):

tests/integration/projects/test_projects_data_connectors_integration.py

@@ -8,6 +8,7 @@ pytestmark = [mark.sdk, mark.integration, mark.projects]
 
 @mark.create
 @mark.asyncio
+@mark.xdist_group("project_resources")
 class ProjectDataConnectorsTests:
     @fixture(scope="class")
     async def project_data_connector(self, project):

tests/integration/projects/test_projects_datastores_integration.py

@@ -8,6 +8,7 @@ pytestmark = [mark.sdk, mark.integration, mark.projects]
 
 @mark.create
 @mark.asyncio
+@mark.xdist_group("project_resources")
 class ProjectDatastoresTests:
     @fixture(scope="class")
     async def project_datastore(self, project):

tests/integration/projects/test_projects_integration.py

@@ -8,6 +8,7 @@ pytestmark = [mark.sdk, mark.integration, mark.projects]
 
 @mark.create
 @mark.asyncio
+@mark.xdist_group("project_resources")
 class GetProjectsTests:
     async def test_get_projects(self, trainml):
         projects = await trainml.projects.list()

tests/integration/projects/test_projects_members_integration.py

@@ -0,0 +1,50 @@
+import re
+import sys
+import asyncio
+from pytest import mark, fixture
+
+pytestmark = [mark.sdk, mark.integration, mark.projects]
+
+
+@mark.create
+@mark.asyncio
+@mark.xdist_group("project_resources")
+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)

tests/integration/projects/test_projects_secrets_integration.py

@@ -8,6 +8,7 @@ pytestmark = [mark.sdk, mark.integration, mark.projects]
 
 @mark.create
 @mark.asyncio
+@mark.xdist_group("project_resources")
 class ProjectSecretsTests:
     @fixture(scope="class")
     async def project_secret(self, project):

tests/integration/projects/test_projects_services_integration.py

@@ -8,6 +8,7 @@ pytestmark = [mark.sdk, mark.integration, mark.projects]
 
 @mark.create
 @mark.asyncio
+@mark.xdist_group("project_resources")
 class ProjectServicesTests:
     @fixture(scope="class")
     async def project_service(self, project):

tests/integration/test_checkpoints_integration.py

@@ -8,6 +8,7 @@ pytestmark = [mark.sdk, mark.integration, mark.checkpoints]
 
 @mark.create
 @mark.asyncio
+@mark.xdist_group("checkpoints")
 class GetCheckpointTests:
     @fixture(scope="class")
     async def checkpoint(self, trainml):

tests/integration/test_datasets_integration.py

@@ -8,6 +8,7 @@ pytestmark = [mark.sdk, mark.integration, mark.datasets]
 
 @mark.create
 @mark.asyncio
+@mark.xdist_group("datasets")
 class GetDatasetTests:
     @fixture(scope="class")
     async def dataset(self, trainml):

tests/integration/test_jobs_integration.py

@@ -20,6 +20,7 @@ def extract_domain_suffix(hostname):
 
 
 @fixture(scope="class")
+@mark.xdist_group("job_lifecycle")
 async def job(trainml):
     job = await trainml.jobs.create(
         name="CLI Automated Tests - Job Lifecycle",
@@ -38,6 +39,7 @@ async def job(trainml):
 
 @mark.create
 @mark.asyncio
+@mark.xdist_group("job_lifecycle")
 class JobLifeCycleTests:
     async def test_wait_for_running(self, job):
         assert job.status != "running"
@@ -413,6 +415,7 @@ class JobAPIWorkerValidationTests:
 
 @mark.create
 @mark.asyncio
+@mark.xdist_group("job_io")
 class JobIOTests:
     async def test_job_local_output(self, trainml, capsys):
         temp_dir = tempfile.TemporaryDirectory()
@@ -423,11 +426,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 +472,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):
@@ -527,6 +530,7 @@ class JobIOTests:
 
 @mark.create
 @mark.asyncio
+@mark.xdist_group("job_types")
 class JobTypeTests:
     async def test_endpoint(self, trainml):
 
@@ -558,7 +562,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 +644,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
@@ -648,6 +652,7 @@ class JobTypeTests:
 
 @mark.create
 @mark.asyncio
+@mark.xdist_group("job_features")
 class JobFeatureTests:
     async def test_cpu_instance(self, trainml, capsys):
         job = await trainml.jobs.create(
@@ -713,9 +718,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 +729,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

tests/integration/test_models_integration.py

@@ -8,6 +8,7 @@ pytestmark = [mark.sdk, mark.integration, mark.models]
 
 @mark.create
 @mark.asyncio
+@mark.xdist_group("models")
 class GetModelTests:
     @fixture(scope="class")
     async def model(self, trainml):

tests/integration/test_volumes_integration.py

@@ -8,6 +8,7 @@ pytestmark = [mark.sdk, mark.integration, mark.volumes]
 
 @mark.create
 @mark.asyncio
+@mark.xdist_group("volumes")
 class GetVolumeTests:
     @fixture(scope="class")
     async def volume(self, trainml):

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/__init__.py

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

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/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],
         ),
     )