diff --git a/aixplain/factories/model_factory.py b/aixplain/factories/model_factory.py index 0fb845f1..c11d837a 100644 --- a/aixplain/factories/model_factory.py +++ b/aixplain/factories/model_factory.py @@ -361,7 +361,7 @@ def create_asset_repo( if function_dict["name"] == function: function_id = function_dict["id"] if function_id is None: - raise Exception("Invalid function name") + raise Exception(f"Invalid function name {function}") create_url = urljoin(config.BACKEND_URL, f"sdk/models/onboard") logging.debug(f"URL: {create_url}") if api_key: diff --git a/tests/functional/model/hf_onboarding_test.py b/tests/functional/model/hf_onboarding_test.py index 47a38361..fa68d2e8 100644 --- a/tests/functional/model/hf_onboarding_test.py +++ b/tests/functional/model/hf_onboarding_test.py @@ -13,7 +13,7 @@ def test_deploy_model(): # Start the deployment model_name = "Test Model" repo_id = "tiiuae/falcon-7b" - response = ModelFactory.deploy_huggingface_model(model_name, repo_id, config.HF_TOKEN) + response = ModelFactory.deploy_huggingface_model(model_name, repo_id, hf_token=config.HF_TOKEN) assert "id" in response.keys() # Check for status @@ -30,31 +30,31 @@ def test_deploy_model(): delete_asset(model_id, config.TEAM_API_KEY) -@pytest.mark.skip(reason="Model Deployment is deactivated for improvements.") +# @pytest.mark.skip(reason="Model Deployment is deactivated for improvements.") def test_nonexistent_model(): # Start the deployment model_name = "Test Model" repo_id = "nonexistent-supplier/nonexistent-model" - response = ModelFactory.deploy_huggingface_model(model_name, repo_id, config.HF_TOKEN) + response = ModelFactory.deploy_huggingface_model(model_name, repo_id, hf_token=config.HF_TOKEN) assert response["statusCode"] == 400 assert response["message"] == "err.unable_to_onboard_model" -@pytest.mark.skip(reason="Model Deployment is deactivated for improvements.") +# @pytest.mark.skip(reason="Model Deployment is deactivated for improvements.") def test_size_limit(): # Start the deployment model_name = "Test Model" repo_id = "tiiuae/falcon-40b" - response = ModelFactory.deploy_huggingface_model(model_name, repo_id, config.HF_TOKEN) + response = ModelFactory.deploy_huggingface_model(model_name, repo_id, hf_token=config.HF_TOKEN) assert response["statusCode"] == 400 assert response["message"] == "err.unable_to_onboard_model" -@pytest.mark.skip(reason="Model Deployment is deactivated for improvements.") +# @pytest.mark.skip(reason="Model Deployment is deactivated for improvements.") def test_gated_model(): # Start the deployment model_name = "Test Model" repo_id = "meta-llama/Llama-2-7b-hf" - response = ModelFactory.deploy_huggingface_model(model_name, repo_id, "mock_key") + response = ModelFactory.deploy_huggingface_model(model_name, repo_id, hf_token="mock_key") assert response["statusCode"] == 400 assert response["message"] == "err.unable_to_onboard_model" diff --git a/tests/image_upload_e2e_test.py b/tests/functional/model/image_upload_e2e_test.py similarity index 72% rename from tests/image_upload_e2e_test.py rename to tests/functional/model/image_upload_e2e_test.py index 0e2ccbc5..7c7efbcc 100644 --- a/tests/image_upload_e2e_test.py +++ b/tests/functional/model/image_upload_e2e_test.py @@ -9,7 +9,6 @@ import pytest -@pytest.mark.skip(reason="Model Upload is deactivated for improvements.") def test_create_and_upload_model(): # List the host machines host_response = ModelFactory.list_host_machines() @@ -31,14 +30,15 @@ def test_create_and_upload_model(): # Register the model, and create an image repository for it. with open(Path("tests/test_requests/create_asset_request.json")) as f: - register_payload = json.load(f) - name = register_payload["name"] - host_machine = register_payload["hostingMachine"] - version = register_payload["version"] - description = register_payload["description"] - function = register_payload["function"] - source_language = register_payload["sourceLanguage"] - register_response = ModelFactory.create_asset_repo(name, host_machine, version, description, function, source_language) + mock_register_payload = json.load(f) + name = mock_register_payload["name"] + description = mock_register_payload["description"] + function = mock_register_payload["function"] + source_language = mock_register_payload["sourceLanguage"] + input_modality = mock_register_payload["input_modality"] + output_modality = mock_register_payload["output_modality"] + documentation_url = mock_register_payload["documentation_url"] + register_response = ModelFactory.create_asset_repo(name, description, function, source_language, input_modality, output_modality, documentation_url, config.TEAM_API_KEY) assert "id" in register_response.keys() assert "repositoryName" in register_response.keys() model_id = register_response["id"] @@ -56,10 +56,7 @@ def test_create_and_upload_model(): registry = login_response["registry"] # Push an image to ECR - # os.system("aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 535945872701.dkr.ecr.us-east-1.amazonaws.com") low_level_client = docker.APIClient(base_url="unix://var/run/docker.sock") - # low_level_client.pull("535945872701.dkr.ecr.us-east-1.amazonaws.com/bash") - # low_level_client.tag("535945872701.dkr.ecr.us-east-1.amazonaws.com/bash", f"{registry}/{repo_name}") low_level_client.pull("bash") low_level_client.tag("bash", f"{registry}/{repo_name}") low_level_client.push(f"{registry}/{repo_name}", auth_config={"username": username, "password": password}) diff --git a/tests/image_upload_functional_test.py b/tests/functional/model/image_upload_functional_test.py similarity index 81% rename from tests/image_upload_functional_test.py rename to tests/functional/model/image_upload_functional_test.py index b9dd3ebf..60d1d3f0 100644 --- a/tests/image_upload_functional_test.py +++ b/tests/functional/model/image_upload_functional_test.py @@ -1,13 +1,12 @@ __author__ = "michaellam" from pathlib import Path import json -from aixplain.utils import config -from tests.test_utils import delete_asset, delete_service_account from aixplain.factories.model_factory import ModelFactory +from tests.test_utils import delete_asset, delete_service_account +from aixplain.utils import config +import docker import pytest - -@pytest.mark.skip(reason="Model Upload is deactivated for improvements.") def test_login(): response = ModelFactory.asset_repo_login() assert response["username"] == "AWS" @@ -18,18 +17,17 @@ def test_login(): delete_service_account(config.TEAM_API_KEY) -@pytest.mark.skip(reason="Model Upload is deactivated for improvements.") def test_create_asset_repo(): with open(Path("tests/test_requests/create_asset_request.json")) as f: mock_register_payload = json.load(f) name = mock_register_payload["name"] - host_machine = mock_register_payload["hostingMachine"] - version = mock_register_payload["version"] description = mock_register_payload["description"] function = mock_register_payload["function"] source_language = mock_register_payload["sourceLanguage"] - response = ModelFactory.create_asset_repo(name, host_machine, version, description, function, source_language) - print(response) + input_modality = mock_register_payload["input_modality"] + output_modality = mock_register_payload["output_modality"] + documentation_url = mock_register_payload["documentation_url"] + response = ModelFactory.create_asset_repo(name, description, function, source_language, input_modality, output_modality, documentation_url, config.TEAM_API_KEY) response_dict = dict(response) assert "id" in response_dict.keys() assert "repositoryName" in response_dict.keys() @@ -38,7 +36,6 @@ def test_create_asset_repo(): delete_asset(response["id"], config.TEAM_API_KEY) -@pytest.mark.skip(reason="Model Upload is deactivated for improvements.") def test_list_host_machines(): response = ModelFactory.list_host_machines() for hosting_machine_dict in response: @@ -49,7 +46,6 @@ def test_list_host_machines(): assert "hourlyCost" in hosting_machine_dict.keys() -@pytest.mark.skip(reason="Model Upload is deactivated for improvements.") def test_get_functions(): # Verbose response = ModelFactory.list_functions(True) diff --git a/tests/test_requests/create_asset_request.json b/tests/test_requests/create_asset_request.json index 4683e526..688dd33a 100644 --- a/tests/test_requests/create_asset_request.json +++ b/tests/test_requests/create_asset_request.json @@ -1,8 +1,9 @@ { "name": "mock_name", - "hostingMachine": "aix-2c-8g-od", - "version": "mock_version", "description": "mock_description", - "function": "Speech Recognition", - "sourceLanguage": "en" + "function": "Text Generation", + "sourceLanguage": "en", + "input_modality": "text", + "output_modality": "text", + "documentation_url": "" } \ No newline at end of file diff --git a/tests/image_upload_test.py b/tests/unit/image_upload_test.py similarity index 85% rename from tests/image_upload_test.py rename to tests/unit/image_upload_test.py index fb919171..4b192292 100644 --- a/tests/image_upload_test.py +++ b/tests/unit/image_upload_test.py @@ -13,7 +13,6 @@ API_FIXED_HEADER = {"x-api-key": f"{config.TEAM_API_KEY}", "Content-Type": "application/json"} -@pytest.mark.skip(reason="Model Upload is deactivated for improvements.") def test_login(): url = urljoin(config.BACKEND_URL, f"sdk/ecr/login") with requests_mock.Mocker() as mock: @@ -24,24 +23,26 @@ def test_login(): assert creds == mock_json -@pytest.mark.skip(reason="Model Upload is deactivated for improvements.") def test_create_asset_repo(): - url_register = urljoin(config.BACKEND_URL, f"sdk/models/register") + url_register = urljoin(config.BACKEND_URL, f"sdk/models/onboard") url_function = urljoin(config.BACKEND_URL, f"sdk/functions") + print(f"URL_Register {url_register}") with requests_mock.Mocker() as mock: with open(Path("tests/mock_responses/create_asset_repo_response.json")) as f: mock_json_register = json.load(f) - mock.post(url_register, headers=API_FIXED_HEADER, json=mock_json_register) + mock.post(url_register, headers=API_FIXED_HEADER, json=mock_json_register, status_code=201) + with open(Path("tests/mock_responses/list_functions_response.json")) as f: mock_json_functions = json.load(f) mock.get(url_function, headers=AUTH_FIXED_HEADER, json=mock_json_functions) + model_id = ModelFactory.create_asset_repo( - "mock_name", "mock_machines", "mock_version", "mock_description", "Speech Recognition", "en", config.TEAM_API_KEY + "mock_name", "mock_description", "Text Generation", "en", "text", "text", api_key=config.TEAM_API_KEY ) + # print(f"Model ID {model_id}") assert model_id == mock_json_register -@pytest.mark.skip(reason="Model Upload is deactivated for improvements.") def test_list_host_machines(): url = urljoin(config.BACKEND_URL, f"sdk/hosting-machines") with requests_mock.Mocker() as mock: @@ -55,8 +56,6 @@ def test_list_host_machines(): for key in machine_dict.keys(): assert machine_dict[key] == mock_json_dict[key] - -@pytest.mark.skip(reason="Model Upload is deactivated for improvements.") def test_get_functions(): url = urljoin(config.BACKEND_URL, f"sdk/functions") with requests_mock.Mocker() as mock: @@ -66,7 +65,6 @@ def test_get_functions(): functions = ModelFactory.list_functions(config.TEAM_API_KEY) assert functions == mock_json - @pytest.mark.skip(reason="Not currently supported.") def test_list_image_repo_tags(): model_id = "mock_id"