def test_deploy_inference_endpoint(mocker):
service = create_job_service(mocker)
mocker.patch.object(service, 'build_inference_with_model_id')
data_bytes, temp_dir, zip_filename = create_inference_zip()
service.submit_training_data = mocker.Mock()
with flask.Flask("Test").app_context():
controller = InferenceController(service)
controller.deploy_inference_endpoint("test_workspace",
"test_user",
"test_model_id",
data_bytes=data_bytes)
service.build_inference_with_model_id.assert_called_with(
"test_workspace",
"test_user",
"test_model_id",
t_any(str),
gpu=0,
cpu=None,
memory=None)
temp_dir.cleanup()
os.remove(zip_filename)
def test_invalid_put_notebook_data(mocker):
service = create_job_service(mocker)
controller = DataController(service)
with pytest.raises(InvalidBundleError, match="Bundle is malformed"):
with flask.Flask("Test").app_context():
controller.put_notebook_data("test_workspace", "test_user", data_bytes=b'0x1')
def test_invalid_put_features(mocker):
service = create_job_service(mocker)
mocker.patch.object(service, 'submit_training_data')
controller = DataController(service)
with pytest.raises(InvalidBundleError, match="Bundle is malformed"):
with flask.Flask("Test").app_context():
controller.put_features("test_workspace", "test_user", data_bytes=b'0x1')
def test_invalid_deploy_inference_endpoint(mocker):
service = create_job_service(mocker)
with pytest.raises(InvalidBundleError, match="Bundle is malformed"):
with flask.Flask("Test").app_context():
controller = InferenceController(service)
controller.deploy_inference_endpoint("test_workspace",
"test_user",
"test_model_id",
data_bytes=b'0x1')
def test_put_notebook_data(mocker):
service = create_job_service(mocker)
mocker.patch.object(service, 'submit_notebook_data')
controller = DataController(service)
data_bytes, _, temp_dir, zip_filename = create_zip()
with flask.Flask("Test").app_context():
controller.put_notebook_data("test_workspace", "test_user", data_bytes=data_bytes)
service.submit_notebook_data.assert_called_with("test_workspace", "test_user", t_any(str))
temp_dir.cleanup()
os.remove(zip_filename)
def test_submit_notebook_without_data(mocker):
service = create_job_service(mocker)
mocker.patch.object(service, 'submit_notebook_code')
service.create_notebook_data = mocker.Mock()
with flask.Flask("Test").app_context():
controller = NotebookController(service)
controller.submit_notebook("test_workspace", "test_user")
service.submit_notebook_code.assert_called_with("test_workspace",
"test_user",
t_any(str),
gpu=0,
cpu=None,
memory=None)
def test_submit_notebook_missing_root(mocker):
service = create_job_service(mocker)
data_bytes, _, temp_dir, zip_filename = create_zip()
with pytest.raises(InvalidBundleError,
match="Missing root directory in source-code bundle"):
with flask.Flask("Test").app_context():
controller = NotebookController(service)
controller.submit_notebook("test_workspace",
"test_user",
data_bytes=data_bytes)
temp_dir.cleanup()
os.remove(zip_filename)
def test_deploy_inference_endpoint_missing_dockerfile(mocker):
service = create_job_service(mocker)
data_bytes, _, temp_dir, zip_filename = create_zip()
with pytest.raises(InvalidBundleError,
match="Missing root directory in source-code bundle"):
with flask.Flask("Test").app_context():
controller = InferenceController(service)
controller.deploy_inference_endpoint("test_workspace",
"test_user",
"test_model_id",
data_bytes=data_bytes)
temp_dir.cleanup()
os.remove(zip_filename)
def test_put_manifest_features(mocker):
service = create_job_service(mocker)
mocker.patch.object(service, 'submit_manifest')
mocker.patch.object(service, 'define_train_pipeline')
mocker.patch.object(service, 'define_manifest_ingestion_pipeline')
controller = DataController(service)
data_bytes = json.dumps({'path': 'a', 'url': 'b'}).encode('utf-8')
hash_hex = hashlib.sha256(data_bytes).hexdigest()[:5]
data_name = f"manifest:{hash_hex}"
with flask.Flask("Test").app_context():
controller.put_manifest_features("test_workspace", "test_user",
data_bytes=data_bytes, cpu=None, gpu=0, memory=None)
service.submit_manifest.assert_called_with("test_workspace", "test_user", t_any(bytes), data_name)
service.define_train_pipeline.assert_called_with("test_workspace", "test_user",
data_name=data_name, cpu=None, gpu=0, memory=None)
def test_put_features(mocker):
service = create_job_service(mocker)
mocker.patch.object(service, 'submit_training_data')
mocker.patch.object(service, 'define_train_pipeline')
mocker.patch.object(service, 'define_bundle_ingestion_pipeline')
controller = DataController(service)
data_bytes, temp_file, temp_dir, zip_filename = create_zip()
hash_hex = build_hash(data_bytes)
data_name = f"{os.path.split(temp_file.name)[-1]}:{hash_hex}"
with flask.Flask("Test").app_context():
controller.put_features("test_workspace", "test_user", data_bytes=data_bytes)
service.submit_training_data.assert_called_with("test_workspace", "test_user", t_any(str))
service.define_train_pipeline.assert_called_with("test_workspace", "test_user",
data_name=data_name, cpu=None, memory=None, gpu=0)
temp_dir.cleanup()
os.remove(zip_filename)
def test_submit_notebook(mocker):
service = create_job_service(mocker)
mocker.patch.object(service, 'submit_notebook_code')
data_bytes, temp_dir, zip_filename = create_train_zip()
service.create_notebook_data = mocker.Mock()
with flask.Flask("Test").app_context():
controller = NotebookController(service)
controller.submit_notebook("test_workspace",
"test_user",
data_bytes=data_bytes)
service.submit_notebook_code.assert_called_with("test_workspace",
"test_user",
t_any(str),
gpu=0,
cpu=None,
memory=None)
temp_dir.cleanup()
os.remove(zip_filename)