Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
lightning/tests/tests_app/runners/test_cloud.py

2114 lines
95 KiB
Python
Raw Blame History

import contextlib
import logging
import os
import pathlib
import re
import sys
from copy import copy
from pathlib import Path
from unittest import mock
from unittest.mock import MagicMock
import pytest
from lightning.app import BuildConfig, LightningApp, LightningFlow, LightningWork
from lightning.app.runners import CloudRuntime, backends, cloud
from lightning.app.source_code.copytree import _copytree, _parse_lightningignore
from lightning.app.source_code.local import LocalSourceCodeDir
from lightning.app.storage import Drive, Mount
from lightning.app.testing.helpers import EmptyWork
from lightning.app.utilities.cloud import _get_project
from lightning.app.utilities.dependency_caching import get_hash
from lightning.app.utilities.packaging.cloud_compute import CloudCompute
from lightning_cloud.openapi import (
CloudspaceIdRunsBody,
Externalv1Cluster,
Externalv1LightningappInstance,
Gridv1ImageSpec,
IdGetBody1,
ProjectIdProjectclustersbindingsBody,
V1BuildSpec,
V1CloudSpace,
V1CloudSpaceInstanceConfig,
V1ClusterSpec,
V1ClusterType,
V1DataConnectionMount,
V1DependencyFileInfo,
V1Drive,
V1DriveSpec,
V1DriveStatus,
V1DriveType,
V1EnvVar,
V1GetUserResponse,
V1LightningappInstanceSpec,
V1LightningappInstanceState,
V1LightningappInstanceStatus,
V1LightningAuth,
V1LightningBasicAuth,
V1LightningRun,
V1LightningworkDrives,
V1LightningworkSpec,
V1ListCloudSpacesResponse,
V1ListClustersResponse,
V1ListLightningappInstancesResponse,
V1ListMembershipsResponse,
V1ListProjectClusterBindingsResponse,
V1Membership,
V1Metadata,
V1NetworkConfig,
V1PackageManager,
V1ProjectClusterBinding,
V1PythonDependencyInfo,
V1QueueServerType,
V1SourceType,
V1UserFeatures,
V1UserRequestedComputeConfig,
V1UserRequestedFlowComputeConfig,
V1Work,
)
class MyWork(LightningWork):
def run(self):
print("my run")
class WorkWithSingleDrive(LightningWork):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.drive = None
def run(self):
pass
class WorkWithTwoDrives(LightningWork):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.lit_drive_1 = None
self.lit_drive_2 = None
def run(self):
pass
def get_cloud_runtime_request_body(**kwargs) -> "CloudspaceIdRunsBody":
default_request_body = {
"app_entrypoint_file": mock.ANY,
"enable_app_server": True,
"is_headless": True,
"should_mount_cloudspace_content": False,
"flow_servers": [],
"image_spec": None,
"works": [],
"local_source": True,
"dependency_cache_key": mock.ANY,
"user_requested_flow_compute_config": V1UserRequestedFlowComputeConfig(
name="flow-lite",
preemptible=False,
shm_size=0,
),
}
if kwargs.get("user_requested_flow_compute_config") is not None:
default_request_body["user_requested_flow_compute_config"] = kwargs["user_requested_flow_compute_config"]
return CloudspaceIdRunsBody(**default_request_body)
@pytest.fixture()
def cloud_backend(monkeypatch):
cloud_backend = mock.MagicMock()
monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock.MagicMock())
monkeypatch.setattr(cloud, "_prepare_lightning_wheels_and_requirements", mock.MagicMock())
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
return cloud_backend
@pytest.fixture()
def project_id():
return "test-project-id"
DEFAULT_CLUSTER = "litng-ai-03"
class TestAppCreationClient:
"""Testing the calls made using GridRestClient to create the app."""
def test_run_on_deleted_cluster(self, cloud_backend):
app_name = "test-app"
mock_client = mock.MagicMock()
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="Default Project", project_id=project_id)]
)
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([
Externalv1Cluster(id=DEFAULT_CLUSTER)
])
cloud_backend.client = mock_client
app = mock.MagicMock()
app.flows = []
app.frontend = {}
existing_instance = MagicMock()
existing_instance.status.phase = V1LightningappInstanceState.STOPPED
existing_instance.spec.cluster_id = DEFAULT_CLUSTER
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=[existing_instance])
)
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=Path("entrypoint.py"))
cloud_runtime._check_uploaded_folder = mock.MagicMock()
with pytest.raises(ValueError, match="that cluster doesn't exist"):
cloud_runtime.dispatch(name=app_name, cluster_id="unknown-cluster")
@pytest.mark.parametrize(
("old_cluster", "new_cluster"),
[
("test", "other"),
("test", "test"),
(None, None),
(None, "litng-ai-03"),
("litng-ai-03", None),
],
)
def test_new_instance_on_different_cluster(self, tmpdir, cloud_backend, project_id, old_cluster, new_cluster):
entrypoint = Path(tmpdir) / "entrypoint.py"
entrypoint.touch()
app_name = "test-app"
mock_client = mock.MagicMock()
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="Default Project", project_id=project_id)]
)
mock_client.lightningapp_v2_service_create_lightningapp_release.return_value = V1LightningRun(
cluster_id=new_cluster
)
# Note:
# backend converts "None" cluster to "litng-ai-03"
# dispatch should receive None, but API calls should return "litng-ai-03"
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([
Externalv1Cluster(id=old_cluster or DEFAULT_CLUSTER),
Externalv1Cluster(id=new_cluster or DEFAULT_CLUSTER),
])
mock_client.projects_service_list_project_cluster_bindings.return_value = V1ListProjectClusterBindingsResponse(
clusters=[
V1ProjectClusterBinding(cluster_id=old_cluster or DEFAULT_CLUSTER),
V1ProjectClusterBinding(cluster_id=new_cluster or DEFAULT_CLUSTER),
]
)
# Mock all clusters as global clusters
mock_client.cluster_service_get_cluster.side_effect = lambda cluster_id: Externalv1Cluster(
id=cluster_id, spec=V1ClusterSpec(cluster_type=V1ClusterType.GLOBAL)
)
cloud_backend.client = mock_client
app = mock.MagicMock()
app.flows = []
app.frontend = {}
existing_app = MagicMock()
existing_app.name = app_name
existing_app.id = "test-id"
mock_client.cloud_space_service_list_cloud_spaces.return_value = V1ListCloudSpacesResponse(
cloudspaces=[existing_app]
)
existing_instance = MagicMock()
existing_instance.name = app_name
existing_instance.status.phase = V1LightningappInstanceState.STOPPED
existing_instance.spec.cluster_id = old_cluster or DEFAULT_CLUSTER
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=[existing_instance])
)
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=entrypoint)
cloud_runtime._check_uploaded_folder = mock.MagicMock()
# This is the main assertion:
# we have an existing instance on `cluster-001`
# but we want to run this app on `cluster-002`
cloud_runtime.dispatch(name=app_name, cluster_id=new_cluster)
if new_cluster != old_cluster and None not in (old_cluster, new_cluster):
# If we switched cluster, check that a new name was used which starts with the old name
mock_client.cloud_space_service_create_lightning_run_instance.assert_called_once()
args = mock_client.cloud_space_service_create_lightning_run_instance.call_args
assert args[1]["body"].name != app_name
assert args[1]["body"].name.startswith(app_name)
assert args[1]["body"].cluster_id == new_cluster
def test_running_deleted_app(self, tmpdir, cloud_backend, project_id):
"""Deleted apps show up in list apps but not in list instances.
This tests that we don't try to reacreate a previously deleted app.
"""
entrypoint = Path(tmpdir) / "entrypoint.py"
entrypoint.touch()
app_name = "test-app"
mock_client = mock.MagicMock()
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="Default Project", project_id=project_id)]
)
mock_client.lightningapp_v2_service_create_lightningapp_release.return_value = V1LightningRun(
cluster_id=DEFAULT_CLUSTER
)
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([
Externalv1Cluster(id=DEFAULT_CLUSTER)
])
mock_client.projects_service_list_project_cluster_bindings.return_value = V1ListProjectClusterBindingsResponse(
clusters=[V1ProjectClusterBinding(cluster_id=DEFAULT_CLUSTER)]
)
# Mock all clusters as global clusters
mock_client.cluster_service_get_cluster.side_effect = lambda cluster_id: Externalv1Cluster(
id=cluster_id, spec=V1ClusterSpec(cluster_type=V1ClusterType.GLOBAL)
)
cloud_backend.client = mock_client
app = mock.MagicMock()
app.flows = []
app.frontend = {}
existing_app = MagicMock()
existing_app.name = app_name
existing_app.id = "test-id"
mock_client.cloud_space_service_list_cloud_spaces.return_value = V1ListCloudSpacesResponse(
cloudspaces=[existing_app]
)
# Simulate the app as deleted so no instance to return
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=[])
)
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=entrypoint)
cloud_runtime._check_uploaded_folder = mock.MagicMock()
cloud_runtime.dispatch(name=app_name)
# Check that a new name was used which starts with and does not equal the old name
mock_client.cloud_space_service_create_lightning_run_instance.assert_called_once()
args = mock_client.cloud_space_service_create_lightning_run_instance.call_args
assert args[1]["body"].name != app_name
assert args[1]["body"].name.startswith(app_name)
@pytest.mark.parametrize("flow_cloud_compute", [None, CloudCompute(name="t2.medium")])
@mock.patch("lightning.app.runners.backends.cloud.LightningClient", mock.MagicMock())
def test_run_with_default_flow_compute_config(self, tmpdir, monkeypatch, flow_cloud_compute):
entrypoint = Path(tmpdir) / "entrypoint.py"
entrypoint.touch()
mock_client = mock.MagicMock()
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="test-project", project_id="test-project-id")]
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=[])
)
mock_client.lightningapp_v2_service_create_lightningapp_release.return_value = V1LightningRun(cluster_id="test")
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([Externalv1Cluster(id="test")])
cloud_backend = mock.MagicMock()
cloud_backend.client = mock_client
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock.MagicMock())
dummy_flow = mock.MagicMock()
monkeypatch.setattr(dummy_flow, "run", lambda *args, **kwargs: None)
if flow_cloud_compute is None:
app = LightningApp(dummy_flow)
else:
app = LightningApp(dummy_flow, flow_cloud_compute=flow_cloud_compute)
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=entrypoint)
cloud_runtime._check_uploaded_folder = mock.MagicMock()
cloud_runtime.dispatch()
user_requested_flow_compute_config = None
if flow_cloud_compute is not None:
user_requested_flow_compute_config = V1UserRequestedFlowComputeConfig(
name=flow_cloud_compute.name, preemptible=False, shm_size=0
)
body = get_cloud_runtime_request_body(user_requested_flow_compute_config=user_requested_flow_compute_config)
cloud_runtime.backend.client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, body=body
)
@mock.patch("lightning.app.runners.backends.cloud.LightningClient", mock.MagicMock())
def test_run_on_byoc_cluster(self, tmpdir, monkeypatch):
entrypoint = Path(tmpdir) / "entrypoint.py"
entrypoint.touch()
mock_client = mock.MagicMock()
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="Default Project", project_id="default-project-id")]
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=[])
)
mock_client.cloud_space_service_create_lightning_run.return_value = V1LightningRun(cluster_id="test1234")
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([
Externalv1Cluster(id="test1234")
])
cloud_backend = mock.MagicMock()
cloud_backend.client = mock_client
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock.MagicMock())
monkeypatch.setattr(cloud, "_prepare_lightning_wheels_and_requirements", mock.MagicMock())
app = mock.MagicMock()
app.is_headless = False
app.flows = []
app.frontend = {}
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=entrypoint)
cloud_runtime._check_uploaded_folder = mock.MagicMock()
cloud_runtime.dispatch(cluster_id="test1234")
body = CloudspaceIdRunsBody(
cluster_id="test1234",
app_entrypoint_file=mock.ANY,
enable_app_server=True,
is_headless=False,
should_mount_cloudspace_content=False,
flow_servers=[],
image_spec=None,
works=[],
local_source=True,
dependency_cache_key=mock.ANY,
user_requested_flow_compute_config=mock.ANY,
)
cloud_runtime.backend.client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="default-project-id", cloudspace_id=mock.ANY, body=body
)
cloud_runtime.backend.client.projects_service_create_project_cluster_binding.assert_called_once_with(
project_id="default-project-id",
body=ProjectIdProjectclustersbindingsBody(cluster_id="test1234"),
)
@mock.patch("lightning.app.runners.backends.cloud.LightningClient", mock.MagicMock())
def test_requirements_file(self, tmpdir, monkeypatch):
entrypoint = Path(tmpdir) / "entrypoint.py"
entrypoint.touch()
mock_client = mock.MagicMock()
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="test-project", project_id="test-project-id")]
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=[])
)
mock_client.cloud_space_service_create_lightning_run.return_value = V1LightningRun()
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([Externalv1Cluster(id="test")])
cloud_backend = mock.MagicMock()
cloud_backend.client = mock_client
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock.MagicMock())
monkeypatch.setattr(cloud, "_prepare_lightning_wheels_and_requirements", mock.MagicMock())
app = mock.MagicMock()
app.is_headless = False
app.flows = []
app.frontend = {}
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=entrypoint)
cloud_runtime._check_uploaded_folder = mock.MagicMock()
# Without requirements file
cloud_runtime.dispatch()
body = CloudspaceIdRunsBody(
app_entrypoint_file=mock.ANY,
enable_app_server=True,
is_headless=False,
should_mount_cloudspace_content=False,
flow_servers=[],
image_spec=None,
works=[],
local_source=True,
dependency_cache_key=mock.ANY,
user_requested_flow_compute_config=mock.ANY,
)
cloud_runtime.backend.client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, body=body
)
# with requirements file
requirements = Path(tmpdir) / "requirements.txt"
requirements.touch()
cloud_runtime.dispatch(no_cache=True)
body.image_spec = Gridv1ImageSpec(
dependency_file_info=V1DependencyFileInfo(package_manager=V1PackageManager.PIP, path="requirements.txt")
)
cloud_runtime.backend.client.cloud_space_service_create_lightning_run.assert_called_with(
project_id="test-project-id", cloudspace_id=mock.ANY, body=body
)
@mock.patch("lightning.app.runners.backends.cloud.LightningClient", mock.MagicMock())
def test_basic_auth_enabled(self, tmpdir, monkeypatch):
entrypoint = Path(tmpdir) / "entrypoint.py"
entrypoint.touch()
mock_client = mock.MagicMock()
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="test-project", project_id="test-project-id")]
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=[])
)
mock_client.cloud_space_service_create_lightning_run.return_value = V1LightningRun()
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([Externalv1Cluster(id="test")])
cloud_backend = mock.MagicMock()
cloud_backend.client = mock_client
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock.MagicMock())
monkeypatch.setattr(cloud, "_prepare_lightning_wheels_and_requirements", mock.MagicMock())
app = mock.MagicMock()
app.is_headless = False
app.flows = []
app.frontend = {}
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=entrypoint)
cloud_runtime._check_uploaded_folder = mock.MagicMock()
# Set cloud_runtime.enable_basic_auth to be not empty:
cloud_runtime.enable_basic_auth = "username:password"
cloud_runtime.dispatch()
mock_client = cloud_runtime.backend.client
body = CloudspaceIdRunsBody(
app_entrypoint_file=mock.ANY,
enable_app_server=True,
is_headless=False,
should_mount_cloudspace_content=False,
flow_servers=[],
image_spec=None,
works=[],
local_source=True,
dependency_cache_key=mock.ANY,
user_requested_flow_compute_config=mock.ANY,
)
mock_client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, body=body
)
mock_client.cloud_space_service_create_lightning_run_instance.assert_called_once_with(
project_id="test-project-id",
cloudspace_id=mock.ANY,
id=mock.ANY,
body=IdGetBody1(
desired_state=mock.ANY,
name=mock.ANY,
env=mock.ANY,
queue_server_type=mock.ANY,
auth=V1LightningAuth(basic=V1LightningBasicAuth(username="username", password="password")),
),
)
@mock.patch("lightning.app.runners.backends.cloud.LightningClient", mock.MagicMock())
def test_no_cache(self, tmpdir, monkeypatch):
entrypoint = Path(tmpdir) / "entrypoint.py"
entrypoint.touch()
requirements = Path(tmpdir) / "requirements.txt"
requirements.touch()
mock_client = mock.MagicMock()
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="test-project", project_id="test-project-id")]
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=[])
)
mock_client.cloud_space_service_create_lightning_run.return_value = V1LightningRun(cluster_id="test")
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([Externalv1Cluster(id="test")])
cloud_backend = mock.MagicMock()
cloud_backend.client = mock_client
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock.MagicMock())
monkeypatch.setattr(cloud, "_prepare_lightning_wheels_and_requirements", mock.MagicMock())
monkeypatch.setattr(cloud, "get_hash", lambda *args, **kwargs: "dummy-hash")
app = mock.MagicMock()
app.flows = []
app.frontend = {}
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=entrypoint)
cloud_runtime._check_uploaded_folder = mock.MagicMock()
# testing with no-cache False
cloud_runtime.dispatch(no_cache=False)
_, _, kwargs = cloud_runtime.backend.client.cloud_space_service_create_lightning_run.mock_calls[0]
body = kwargs["body"]
assert body.dependency_cache_key == "dummy-hash"
# testing with no-cache True
mock_client.reset_mock()
cloud_runtime.dispatch(no_cache=True)
_, _, kwargs = cloud_runtime.backend.client.cloud_space_service_create_lightning_run.mock_calls[0]
body = kwargs["body"]
assert body.dependency_cache_key is None
@mock.patch("lightning.app.runners.backends.cloud.LightningClient", mock.MagicMock())
@pytest.mark.parametrize(
("lightningapps", "start_with_flow"),
[([], False), ([MagicMock()], False), ([MagicMock()], True)],
)
def test_call_with_work_app(self, lightningapps, start_with_flow, monkeypatch, tmpdir):
source_code_root_dir = Path(tmpdir / "src").absolute()
source_code_root_dir.mkdir()
Path(source_code_root_dir / ".lightning").write_text("name: myapp")
requirements_file = Path(source_code_root_dir / "requirements.txt")
Path(requirements_file).touch()
(source_code_root_dir / "entrypoint.py").touch()
mock_client = mock.MagicMock()
if lightningapps:
lightningapps[0].name = "myapp"
lightningapps[0].status.phase = V1LightningappInstanceState.STOPPED
lightningapps[0].spec.cluster_id = "test"
mock_client.cloud_space_service_list_cloud_spaces.return_value = V1ListCloudSpacesResponse(
cloudspaces=lightningapps
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=lightningapps)
)
mock_client.projects_service_list_project_cluster_bindings.return_value = V1ListProjectClusterBindingsResponse(
clusters=[V1ProjectClusterBinding(cluster_id="test")]
)
mock_client.cluster_service_get_cluster.side_effect = lambda cluster_id: Externalv1Cluster(
id=cluster_id, spec=V1ClusterSpec(cluster_type=V1ClusterType.GLOBAL)
)
mock_client.cloud_space_service_create_lightning_run_instance.return_value = V1LightningRun()
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([Externalv1Cluster(id="test")])
mock_client.cloud_space_service_create_lightning_run_instance.return_value = MagicMock()
existing_instance = MagicMock()
existing_instance.status.phase = V1LightningappInstanceState.STOPPED
mock_client.lightningapp_service_get_lightningapp = MagicMock(return_value=existing_instance)
cloud_backend = mock.MagicMock()
cloud_backend.client = mock_client
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock.MagicMock())
monkeypatch.setattr(cloud, "_prepare_lightning_wheels_and_requirements", mock.MagicMock())
app = mock.MagicMock()
app.is_headless = False
work = MyWork(start_with_flow=start_with_flow, cloud_compute=CloudCompute("custom"))
work._name = "test-work"
work._cloud_build_config.build_commands = lambda: ["echo 'start'"]
work._cloud_build_config.requirements = ["torch==1.0.0", "numpy==1.0.0"]
work._cloud_build_config.image = "random_base_public_image"
work._cloud_compute.disk_size = 0
work._port = 8080
app.works = [work]
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=(source_code_root_dir / "entrypoint.py"))
monkeypatch.setattr(
"lightning.app.runners.cloud._get_project",
lambda _, project_id: V1Membership(name="test-project", project_id="test-project-id"),
)
cloud_runtime.dispatch()
if lightningapps:
expected_body = CloudspaceIdRunsBody(
description=None,
local_source=True,
app_entrypoint_file="entrypoint.py",
enable_app_server=True,
is_headless=False,
should_mount_cloudspace_content=False,
flow_servers=[],
dependency_cache_key=get_hash(requirements_file),
user_requested_flow_compute_config=mock.ANY,
cluster_id="test",
image_spec=Gridv1ImageSpec(
dependency_file_info=V1DependencyFileInfo(
package_manager=V1PackageManager.PIP, path="requirements.txt"
)
),
)
if start_with_flow:
expected_body.works = [
V1Work(
name="test-work",
display_name="",
spec=V1LightningworkSpec(
build_spec=V1BuildSpec(
commands=["echo 'start'"],
python_dependencies=V1PythonDependencyInfo(
package_manager=V1PackageManager.PIP, packages="torch==1.0.0\nnumpy==1.0.0"
),
image="random_base_public_image",
),
drives=[],
user_requested_compute_config=V1UserRequestedComputeConfig(
name="custom",
count=1,
disk_size=0,
shm_size=0,
preemptible=False,
),
network_config=[V1NetworkConfig(name=mock.ANY, host=None, port=8080)],
data_connection_mounts=[],
),
)
]
else:
expected_body.works = []
mock_client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, body=expected_body
)
# running dispatch with disabled dependency cache
mock_client.reset_mock()
monkeypatch.setattr(cloud, "DISABLE_DEPENDENCY_CACHE", True)
expected_body.dependency_cache_key = None
cloud_runtime.dispatch()
mock_client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, body=expected_body
)
else:
mock_client.cloud_space_service_create_lightning_run_instance.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, id=mock.ANY, body=mock.ANY
)
@mock.patch("lightning.app.runners.backends.cloud.LightningClient", mock.MagicMock())
@pytest.mark.parametrize("lightningapps", [[], [MagicMock()]])
def test_call_with_queue_server_type_specified(self, tmpdir, lightningapps, monkeypatch):
entrypoint = Path(tmpdir) / "entrypoint.py"
entrypoint.touch()
mock_client = mock.MagicMock()
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="test-project", project_id="test-project-id")]
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=[])
)
mock_client.cloud_space_service_create_lightning_run.return_value = V1LightningRun()
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([Externalv1Cluster(id="test")])
cloud_backend = mock.MagicMock()
cloud_backend.client = mock_client
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock.MagicMock())
monkeypatch.setattr(cloud, "_prepare_lightning_wheels_and_requirements", mock.MagicMock())
app = mock.MagicMock()
app.flows = []
app.frontend = {}
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=entrypoint)
cloud_runtime._check_uploaded_folder = mock.MagicMock()
cloud_runtime.dispatch()
# calling with no env variable set
body = IdGetBody1(
desired_state=V1LightningappInstanceState.STOPPED,
env=[],
name=mock.ANY,
queue_server_type=V1QueueServerType.UNSPECIFIED,
)
client = cloud_runtime.backend.client
client.cloud_space_service_create_lightning_run_instance.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, id=mock.ANY, body=body
)
# calling with env variable set to http
monkeypatch.setitem(os.environ, "LIGHTNING_CLOUD_QUEUE_TYPE", "http")
cloud_runtime.backend.client.reset_mock()
cloud_runtime.dispatch()
body = IdGetBody1(
desired_state=V1LightningappInstanceState.STOPPED,
env=mock.ANY,
name=mock.ANY,
queue_server_type=V1QueueServerType.HTTP,
)
client = cloud_runtime.backend.client
client.cloud_space_service_create_lightning_run_instance.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, id=mock.ANY, body=body
)
@mock.patch("lightning.app.runners.backends.cloud.LightningClient", mock.MagicMock())
@pytest.mark.parametrize("lightningapps", [[], [MagicMock()]])
def test_call_with_work_app_and_attached_drives(self, lightningapps, monkeypatch, tmpdir):
source_code_root_dir = Path(tmpdir / "src").absolute()
source_code_root_dir.mkdir()
Path(source_code_root_dir / ".lightning").write_text("name: myapp")
requirements_file = Path(source_code_root_dir / "requirements.txt")
Path(requirements_file).touch()
(source_code_root_dir / "entrypoint.py").touch()
mock_client = mock.MagicMock()
if lightningapps:
lightningapps[0].name = "myapp"
lightningapps[0].status.phase = V1LightningappInstanceState.STOPPED
lightningapps[0].spec.cluster_id = "test"
mock_client.cloud_space_service_list_cloud_spaces.return_value = V1ListCloudSpacesResponse(
cloudspaces=lightningapps
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=lightningapps)
)
mock_client.projects_service_list_project_cluster_bindings.return_value = V1ListProjectClusterBindingsResponse(
clusters=[V1ProjectClusterBinding(cluster_id="test")]
)
mock_client.cluster_service_get_cluster.side_effect = lambda cluster_id: Externalv1Cluster(
id=cluster_id, spec=V1ClusterSpec(cluster_type=V1ClusterType.GLOBAL)
)
mock_client.cloud_space_service_create_lightning_run_instance.return_value = V1LightningRun()
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([Externalv1Cluster(id="test")])
lit_app_instance = MagicMock()
mock_client.cloud_space_service_create_lightning_run_instance = MagicMock(return_value=lit_app_instance)
existing_instance = MagicMock()
existing_instance.status.phase = V1LightningappInstanceState.STOPPED
mock_client.lightningapp_service_get_lightningapp = MagicMock(return_value=existing_instance)
cloud_backend = mock.MagicMock()
cloud_backend.client = mock_client
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock.MagicMock())
monkeypatch.setattr(cloud, "_prepare_lightning_wheels_and_requirements", mock.MagicMock())
app = mock.MagicMock()
app.is_headless = False
mocked_drive = MagicMock(spec=Drive)
setattr(mocked_drive, "id", "foobar")
setattr(mocked_drive, "protocol", "lit://")
setattr(mocked_drive, "component_name", "test-work")
setattr(mocked_drive, "allow_duplicates", False)
setattr(mocked_drive, "root_folder", tmpdir)
# deepcopy on a MagicMock instance will return an empty magicmock instance. To
# overcome this we set the __deepcopy__ method `return_value` to equal what
# should be the results of the deepcopy operation (an instance of the original class)
mocked_drive.__deepcopy__.return_value = copy(mocked_drive)
work = WorkWithSingleDrive(cloud_compute=CloudCompute("custom"))
monkeypatch.setattr(work, "drive", mocked_drive)
monkeypatch.setattr(work, "_state", {"_port", "drive"})
monkeypatch.setattr(work, "_name", "test-work")
monkeypatch.setattr(work._cloud_build_config, "build_commands", lambda: ["echo 'start'"])
monkeypatch.setattr(work._cloud_build_config, "requirements", ["torch==1.0.0", "numpy==1.0.0"])
monkeypatch.setattr(work._cloud_build_config, "image", "random_base_public_image")
monkeypatch.setattr(work._cloud_compute, "disk_size", 0)
monkeypatch.setattr(work, "_port", 8080)
app.works = [work]
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=(source_code_root_dir / "entrypoint.py"))
monkeypatch.setattr(
"lightning.app.runners.cloud._get_project",
lambda _, project_id: V1Membership(name="test-project", project_id="test-project-id"),
)
cloud_runtime.dispatch()
if lightningapps:
expected_body = CloudspaceIdRunsBody(
description=None,
local_source=True,
app_entrypoint_file="entrypoint.py",
enable_app_server=True,
is_headless=False,
should_mount_cloudspace_content=False,
flow_servers=[],
dependency_cache_key=get_hash(requirements_file),
user_requested_flow_compute_config=mock.ANY,
cluster_id="test",
image_spec=Gridv1ImageSpec(
dependency_file_info=V1DependencyFileInfo(
package_manager=V1PackageManager.PIP, path="requirements.txt"
)
),
works=[
V1Work(
name="test-work",
display_name="",
spec=V1LightningworkSpec(
build_spec=V1BuildSpec(
commands=["echo 'start'"],
python_dependencies=V1PythonDependencyInfo(
package_manager=V1PackageManager.PIP, packages="torch==1.0.0\nnumpy==1.0.0"
),
image="random_base_public_image",
),
drives=[
V1LightningworkDrives(
drive=V1Drive(
metadata=V1Metadata(
name="test-work.drive",
),
spec=V1DriveSpec(
drive_type=V1DriveType.NO_MOUNT_S3,
source_type=V1SourceType.S3,
source="lit://foobar",
),
status=V1DriveStatus(),
),
),
],
user_requested_compute_config=V1UserRequestedComputeConfig(
name="custom", count=1, disk_size=0, shm_size=0, preemptible=False
),
network_config=[V1NetworkConfig(name=mock.ANY, host=None, port=8080)],
data_connection_mounts=[],
),
)
],
)
mock_client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, body=expected_body
)
# running dispatch with disabled dependency cache
mock_client.reset_mock()
monkeypatch.setattr(cloud, "DISABLE_DEPENDENCY_CACHE", True)
expected_body.dependency_cache_key = None
cloud_runtime.dispatch()
mock_client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, body=expected_body
)
else:
mock_client.cloud_space_service_create_lightning_run_instance.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, id=mock.ANY, body=mock.ANY
)
@mock.patch("lightning.app.runners.backends.cloud.LightningClient", mock.MagicMock())
@mock.patch("lightning.app.core.constants.ENABLE_APP_COMMENT_COMMAND_EXECUTION", True)
@pytest.mark.parametrize("lightningapps", [[], [MagicMock()]])
def test_call_with_work_app_and_app_comment_command_execution_set(self, lightningapps, monkeypatch, tmpdir):
source_code_root_dir = Path(tmpdir / "src").absolute()
source_code_root_dir.mkdir()
Path(source_code_root_dir / ".lightning").write_text("name: myapp")
requirements_file = Path(source_code_root_dir / "requirements.txt")
Path(requirements_file).touch()
(source_code_root_dir / "entrypoint.py").touch()
mock_client = mock.MagicMock()
if lightningapps:
lightningapps[0].name = "myapp"
lightningapps[0].status.phase = V1LightningappInstanceState.STOPPED
lightningapps[0].spec.cluster_id = "test"
mock_client.projects_service_list_project_cluster_bindings.return_value = (
V1ListProjectClusterBindingsResponse(clusters=[V1ProjectClusterBinding(cluster_id="test")])
)
mock_client.cluster_service_get_cluster.side_effect = lambda cluster_id: Externalv1Cluster(
id=cluster_id, spec=V1ClusterSpec(cluster_type=V1ClusterType.GLOBAL)
)
mock_client.cloud_space_service_list_cloud_spaces.return_value = V1ListCloudSpacesResponse(
cloudspaces=lightningapps
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=lightningapps)
)
mock_client.cloud_space_service_create_lightning_run_instance.return_value = V1LightningRun()
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([Externalv1Cluster(id="test")])
lit_app_instance = MagicMock()
mock_client.cloud_space_service_create_lightning_run_instance = MagicMock(return_value=lit_app_instance)
existing_instance = MagicMock()
existing_instance.status.phase = V1LightningappInstanceState.STOPPED
mock_client.lightningapp_service_get_lightningapp = MagicMock(return_value=existing_instance)
cloud_backend = mock.MagicMock()
cloud_backend.client = mock_client
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock.MagicMock())
monkeypatch.setattr(cloud, "_prepare_lightning_wheels_and_requirements", mock.MagicMock())
app = mock.MagicMock()
app.is_headless = False
work = MyWork(cloud_compute=CloudCompute("custom"))
work._state = {"_port"}
work._name = "test-work"
work._cloud_build_config.build_commands = lambda: ["echo 'start'"]
work._cloud_build_config.requirements = ["torch==1.0.0", "numpy==1.0.0"]
work._cloud_build_config.image = "random_base_public_image"
work._cloud_compute.disk_size = 0
work._port = 8080
app.works = [work]
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=(source_code_root_dir / "entrypoint.py"))
monkeypatch.setattr(
"lightning.app.runners.cloud._get_project",
lambda _, project_id: V1Membership(name="test-project", project_id="test-project-id"),
)
cloud_runtime.run_app_comment_commands = True
cloud_runtime.dispatch()
if lightningapps:
expected_body = CloudspaceIdRunsBody(
description=None,
local_source=True,
app_entrypoint_file="entrypoint.py",
enable_app_server=True,
is_headless=False,
should_mount_cloudspace_content=False,
flow_servers=[],
dependency_cache_key=get_hash(requirements_file),
user_requested_flow_compute_config=mock.ANY,
cluster_id="test",
image_spec=Gridv1ImageSpec(
dependency_file_info=V1DependencyFileInfo(
package_manager=V1PackageManager.PIP, path="requirements.txt"
)
),
works=[
V1Work(
name="test-work",
display_name="",
spec=V1LightningworkSpec(
build_spec=V1BuildSpec(
commands=["echo 'start'"],
python_dependencies=V1PythonDependencyInfo(
package_manager=V1PackageManager.PIP, packages="torch==1.0.0\nnumpy==1.0.0"
),
image="random_base_public_image",
),
drives=[],
user_requested_compute_config=V1UserRequestedComputeConfig(
name="custom", count=1, disk_size=0, shm_size=0, preemptible=mock.ANY
),
network_config=[V1NetworkConfig(name=mock.ANY, host=None, port=8080)],
cluster_id=mock.ANY,
data_connection_mounts=[],
),
)
],
)
mock_client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, body=expected_body
)
# running dispatch with disabled dependency cache
mock_client.reset_mock()
monkeypatch.setattr(cloud, "DISABLE_DEPENDENCY_CACHE", True)
expected_body.dependency_cache_key = None
cloud_runtime.dispatch()
mock_client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, body=expected_body
)
else:
mock_client.cloud_space_service_create_lightning_run_instance.assert_called_once_with(
project_id="test-project-id",
cloudspace_id=mock.ANY,
id=mock.ANY,
body=IdGetBody1(
desired_state=V1LightningappInstanceState.STOPPED,
name=mock.ANY,
env=[V1EnvVar(name="ENABLE_APP_COMMENT_COMMAND_EXECUTION", value="1")],
queue_server_type=mock.ANY,
),
)
@mock.patch("lightning.app.runners.backends.cloud.LightningClient", mock.MagicMock())
@pytest.mark.parametrize("lightningapps", [[], [MagicMock()]])
def test_call_with_work_app_and_multiple_attached_drives(self, lightningapps, monkeypatch, tmpdir):
source_code_root_dir = Path(tmpdir / "src").absolute()
source_code_root_dir.mkdir()
Path(source_code_root_dir / ".lightning").write_text("name: myapp")
requirements_file = Path(source_code_root_dir / "requirements.txt")
Path(requirements_file).touch()
(source_code_root_dir / "entrypoint.py").touch()
mock_client = mock.MagicMock()
if lightningapps:
lightningapps[0].name = "myapp"
lightningapps[0].status.phase = V1LightningappInstanceState.STOPPED
lightningapps[0].spec.cluster_id = "test"
mock_client.projects_service_list_project_cluster_bindings.return_value = (
V1ListProjectClusterBindingsResponse(
clusters=[
V1ProjectClusterBinding(cluster_id="test"),
]
)
)
mock_client.cluster_service_get_cluster.side_effect = lambda cluster_id: Externalv1Cluster(
id=cluster_id, spec=V1ClusterSpec(cluster_type=V1ClusterType.GLOBAL)
)
mock_client.cloud_space_service_list_cloud_spaces.return_value = V1ListCloudSpacesResponse(
cloudspaces=lightningapps
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=lightningapps)
)
mock_client.cloud_space_service_create_lightning_run_instance.return_value = V1LightningRun(cluster_id="test")
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([Externalv1Cluster(id="test")])
lit_app_instance = MagicMock()
mock_client.cloud_space_service_create_lightning_run_instance = MagicMock(return_value=lit_app_instance)
existing_instance = MagicMock()
existing_instance.status.phase = V1LightningappInstanceState.STOPPED
mock_client.lightningapp_service_get_lightningapp = MagicMock(return_value=existing_instance)
cloud_backend = mock.MagicMock()
cloud_backend.client = mock_client
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock.MagicMock())
monkeypatch.setattr(cloud, "_prepare_lightning_wheels_and_requirements", mock.MagicMock())
app = mock.MagicMock()
app.is_headless = False
mocked_lit_drive = MagicMock(spec=Drive)
setattr(mocked_lit_drive, "id", "foobar")
setattr(mocked_lit_drive, "protocol", "lit://")
setattr(mocked_lit_drive, "component_name", "test-work")
setattr(mocked_lit_drive, "allow_duplicates", False)
setattr(mocked_lit_drive, "root_folder", tmpdir)
# deepcopy on a MagicMock instance will return an empty magicmock instance. To
# overcome this we set the __deepcopy__ method `return_value` to equal what
# should be the results of the deepcopy operation (an instance of the original class)
mocked_lit_drive.__deepcopy__.return_value = copy(mocked_lit_drive)
work = WorkWithTwoDrives(cloud_compute=CloudCompute("custom"))
work.lit_drive_1 = mocked_lit_drive
work.lit_drive_2 = mocked_lit_drive
work._state = {"_port", "_name", "lit_drive_1", "lit_drive_2"}
work._name = "test-work"
work._cloud_build_config.build_commands = lambda: ["echo 'start'"]
work._cloud_build_config.requirements = ["torch==1.0.0", "numpy==1.0.0"]
work._cloud_build_config.image = "random_base_public_image"
work._cloud_compute.disk_size = 0
work._port = 8080
app.works = [work]
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=(source_code_root_dir / "entrypoint.py"))
monkeypatch.setattr(
"lightning.app.runners.cloud._get_project",
lambda _, project_id: V1Membership(name="test-project", project_id="test-project-id"),
)
cloud_runtime.dispatch()
if lightningapps:
lit_drive_1_spec = V1LightningworkDrives(
drive=V1Drive(
metadata=V1Metadata(
name="test-work.lit_drive_1",
),
spec=V1DriveSpec(
drive_type=V1DriveType.NO_MOUNT_S3,
source_type=V1SourceType.S3,
source="lit://foobar",
),
status=V1DriveStatus(),
),
)
lit_drive_2_spec = V1LightningworkDrives(
drive=V1Drive(
metadata=V1Metadata(
name="test-work.lit_drive_2",
),
spec=V1DriveSpec(
drive_type=V1DriveType.NO_MOUNT_S3,
source_type=V1SourceType.S3,
source="lit://foobar",
),
status=V1DriveStatus(),
),
)
# order of drives in the spec is non-deterministic, so there are two options
# depending for the expected body value on which drive is ordered in the list first.
expected_body_option_1 = CloudspaceIdRunsBody(
description=None,
local_source=True,
app_entrypoint_file="entrypoint.py",
enable_app_server=True,
is_headless=False,
should_mount_cloudspace_content=False,
flow_servers=[],
dependency_cache_key=get_hash(requirements_file),
user_requested_flow_compute_config=mock.ANY,
cluster_id="test",
image_spec=Gridv1ImageSpec(
dependency_file_info=V1DependencyFileInfo(
package_manager=V1PackageManager.PIP, path="requirements.txt"
)
),
works=[
V1Work(
name="test-work",
display_name="",
spec=V1LightningworkSpec(
build_spec=V1BuildSpec(
commands=["echo 'start'"],
python_dependencies=V1PythonDependencyInfo(
package_manager=V1PackageManager.PIP, packages="torch==1.0.0\nnumpy==1.0.0"
),
image="random_base_public_image",
),
drives=[lit_drive_2_spec, lit_drive_1_spec],
user_requested_compute_config=V1UserRequestedComputeConfig(
name="custom",
count=1,
disk_size=0,
shm_size=0,
preemptible=False,
),
network_config=[V1NetworkConfig(name=mock.ANY, host=None, port=8080)],
data_connection_mounts=[],
),
)
],
)
expected_body_option_2 = CloudspaceIdRunsBody(
description=None,
local_source=True,
app_entrypoint_file="entrypoint.py",
enable_app_server=True,
is_headless=False,
should_mount_cloudspace_content=False,
flow_servers=[],
dependency_cache_key=get_hash(requirements_file),
user_requested_flow_compute_config=mock.ANY,
cluster_id="test",
image_spec=Gridv1ImageSpec(
dependency_file_info=V1DependencyFileInfo(
package_manager=V1PackageManager.PIP, path="requirements.txt"
)
),
works=[
V1Work(
name="test-work",
display_name="",
spec=V1LightningworkSpec(
build_spec=V1BuildSpec(
commands=["echo 'start'"],
python_dependencies=V1PythonDependencyInfo(
package_manager=V1PackageManager.PIP, packages="torch==1.0.0\nnumpy==1.0.0"
),
image="random_base_public_image",
),
drives=[lit_drive_1_spec, lit_drive_2_spec],
user_requested_compute_config=V1UserRequestedComputeConfig(
name="custom",
count=1,
disk_size=0,
shm_size=0,
preemptible=False,
),
network_config=[V1NetworkConfig(name=mock.ANY, host=None, port=8080)],
data_connection_mounts=[],
),
)
],
)
# try both options for the expected body to avoid false
# positive test failures depending on system randomness
expected_body = expected_body_option_1
try:
mock_client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, body=expected_body
)
except Exception:
expected_body = expected_body_option_2
mock_client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, body=expected_body
)
# running dispatch with disabled dependency cache
mock_client.reset_mock()
monkeypatch.setattr(cloud, "DISABLE_DEPENDENCY_CACHE", True)
expected_body.dependency_cache_key = None
cloud_runtime.dispatch()
mock_client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, body=expected_body
)
else:
mock_client.cloud_space_service_create_lightning_run_instance.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, id=mock.ANY, body=mock.ANY
)
@mock.patch("lightning.app.runners.backends.cloud.LightningClient", mock.MagicMock())
@pytest.mark.parametrize("lightningapps", [[], [MagicMock()]])
def test_call_with_work_app_and_attached_mount_and_drive(self, lightningapps, monkeypatch, tmpdir):
source_code_root_dir = Path(tmpdir / "src").absolute()
source_code_root_dir.mkdir()
Path(source_code_root_dir / ".lightning").write_text("name: myapp")
requirements_file = Path(source_code_root_dir / "requirements.txt")
Path(requirements_file).touch()
(source_code_root_dir / "entrypoint.py").touch()
mock_client = mock.MagicMock()
if lightningapps:
lightningapps[0].name = "myapp"
lightningapps[0].status.phase = V1LightningappInstanceState.STOPPED
lightningapps[0].spec.cluster_id = "test"
mock_client.projects_service_list_project_cluster_bindings.return_value = (
V1ListProjectClusterBindingsResponse(clusters=[V1ProjectClusterBinding(cluster_id="test")])
)
mock_client.cluster_service_get_cluster.side_effect = lambda cluster_id: Externalv1Cluster(
id=cluster_id, spec=V1ClusterSpec(cluster_type=V1ClusterType.GLOBAL)
)
mock_client.cloud_space_service_list_cloud_spaces.return_value = V1ListCloudSpacesResponse(
cloudspaces=lightningapps
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=lightningapps)
)
mock_client.cloud_space_service_create_lightning_run_instance.return_value = V1LightningRun(cluster_id="test")
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([Externalv1Cluster(id="test")])
lit_app_instance = MagicMock()
mock_client.cloud_space_service_create_lightning_run_instance = MagicMock(return_value=lit_app_instance)
existing_instance = MagicMock()
existing_instance.status.phase = V1LightningappInstanceState.STOPPED
existing_instance.spec.cluster_id = None
mock_client.lightningapp_service_get_lightningapp = MagicMock(return_value=existing_instance)
cloud_backend = mock.MagicMock()
cloud_backend.client = mock_client
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock.MagicMock())
monkeypatch.setattr(cloud, "_prepare_lightning_wheels_and_requirements", mock.MagicMock())
app = mock.MagicMock()
app.is_headless = False
mocked_drive = MagicMock(spec=Drive)
setattr(mocked_drive, "id", "foobar")
setattr(mocked_drive, "protocol", "lit://")
setattr(mocked_drive, "component_name", "test-work")
setattr(mocked_drive, "allow_duplicates", False)
setattr(mocked_drive, "root_folder", tmpdir)
# deepcopy on a MagicMock instance will return an empty magicmock instance. To
# overcome this we set the __deepcopy__ method `return_value` to equal what
# should be the results of the deepcopy operation (an instance of the original class)
mocked_drive.__deepcopy__.return_value = copy(mocked_drive)
mocked_mount = MagicMock(spec=Mount)
setattr(mocked_mount, "source", "s3://foo/")
setattr(mocked_mount, "mount_path", "/content/foo")
setattr(mocked_mount, "protocol", "s3://")
work = WorkWithSingleDrive(cloud_compute=CloudCompute("custom"))
monkeypatch.setattr(work, "drive", mocked_drive)
monkeypatch.setattr(work, "_state", {"_port", "drive"})
monkeypatch.setattr(work, "_name", "test-work")
monkeypatch.setattr(work._cloud_build_config, "build_commands", lambda: ["echo 'start'"])
monkeypatch.setattr(work._cloud_build_config, "requirements", ["torch==1.0.0", "numpy==1.0.0"])
monkeypatch.setattr(work._cloud_build_config, "image", "random_base_public_image")
monkeypatch.setattr(work._cloud_compute, "disk_size", 0)
monkeypatch.setattr(work._cloud_compute, "mounts", mocked_mount)
monkeypatch.setattr(work, "_port", 8080)
app.works = [work]
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=(source_code_root_dir / "entrypoint.py"))
monkeypatch.setattr(
"lightning.app.runners.cloud._get_project",
lambda _, project_id: V1Membership(name="test-project", project_id="test-project-id"),
)
cloud_runtime.dispatch()
if lightningapps:
expected_body = CloudspaceIdRunsBody(
description=None,
local_source=True,
app_entrypoint_file="entrypoint.py",
enable_app_server=True,
is_headless=False,
should_mount_cloudspace_content=False,
flow_servers=[],
dependency_cache_key=get_hash(requirements_file),
image_spec=Gridv1ImageSpec(
dependency_file_info=V1DependencyFileInfo(
package_manager=V1PackageManager.PIP, path="requirements.txt"
)
),
user_requested_flow_compute_config=mock.ANY,
cluster_id="test",
works=[
V1Work(
name="test-work",
display_name="",
spec=V1LightningworkSpec(
build_spec=V1BuildSpec(
commands=["echo 'start'"],
python_dependencies=V1PythonDependencyInfo(
package_manager=V1PackageManager.PIP, packages="torch==1.0.0\nnumpy==1.0.0"
),
image="random_base_public_image",
),
drives=[
V1LightningworkDrives(
drive=V1Drive(
metadata=V1Metadata(
name="test-work.drive",
),
spec=V1DriveSpec(
drive_type=V1DriveType.NO_MOUNT_S3,
source_type=V1SourceType.S3,
source="lit://foobar",
),
status=V1DriveStatus(),
),
),
V1LightningworkDrives(
drive=V1Drive(
metadata=V1Metadata(
name="test-work",
),
spec=V1DriveSpec(
drive_type=V1DriveType.INDEXED_S3,
source_type=V1SourceType.S3,
source="s3://foo/",
),
status=V1DriveStatus(),
),
mount_location="/content/foo",
),
],
user_requested_compute_config=V1UserRequestedComputeConfig(
name="custom",
count=1,
disk_size=0,
shm_size=0,
preemptible=False,
),
network_config=[V1NetworkConfig(name=mock.ANY, host=None, port=8080)],
data_connection_mounts=[],
),
)
],
)
mock_client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, body=expected_body
)
# running dispatch with disabled dependency cache
mock_client.reset_mock()
monkeypatch.setattr(cloud, "DISABLE_DEPENDENCY_CACHE", True)
expected_body.dependency_cache_key = None
cloud_runtime.dispatch()
mock_client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, body=expected_body
)
else:
mock_client.cloud_space_service_create_lightning_run_instance.assert_called_once_with(
project_id="test-project-id", cloudspace_id=mock.ANY, id=mock.ANY, body=mock.ANY
)
class TestOpen:
def test_open(self, monkeypatch):
"""Tests that the open method calls the expected API endpoints."""
mock_client = mock.MagicMock()
mock_client.auth_service_get_user.return_value = V1GetUserResponse(
username="tester", features=V1UserFeatures(code_tab=True)
)
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="test-project", project_id="test-project-id")]
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=[])
)
mock_client.cloud_space_service_create_cloud_space.return_value = V1CloudSpace(id="cloudspace_id")
mock_client.cloud_space_service_create_lightning_run.return_value = V1LightningRun(id="run_id")
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([Externalv1Cluster(id="test")])
cloud_backend = mock.MagicMock()
cloud_backend.client = mock_client
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
mock_local_source = mock.MagicMock()
monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock_local_source)
cloud_runtime = cloud.CloudRuntime(entrypoint=Path("."))
cloud_runtime.open("test_space")
mock_client.cloud_space_service_create_cloud_space.assert_called_once_with(
project_id="test-project-id", body=mock.ANY
)
mock_client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="test-project-id", cloudspace_id="cloudspace_id", body=mock.ANY
)
assert mock_client.cloud_space_service_create_cloud_space.call_args.kwargs["body"].name == "test_space"
@pytest.mark.parametrize(
("path", "expected_root", "entries", "expected_filtered_entries"),
[(".", ".", ["a.py", "b.ipynb"], ["a.py", "b.ipynb"]), ("a.py", ".", ["a.py", "b.ipynb"], ["a.py"])],
)
def test_open_repo(self, tmpdir, monkeypatch, path, expected_root, entries, expected_filtered_entries):
"""Tests that the local source code repo is set up with the correct path and ignore functions."""
tmpdir = Path(tmpdir)
for entry in entries:
(tmpdir / entry).touch()
mock_client = mock.MagicMock()
mock_client.auth_service_get_user.return_value = V1GetUserResponse(
username="tester", features=V1UserFeatures(code_tab=True)
)
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="test-project", project_id="test-project-id")]
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=[])
)
mock_client.lightningapp_v2_service_create_lightningapp_release.return_value = V1LightningRun(cluster_id="test")
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([Externalv1Cluster(id="test")])
cloud_backend = mock.MagicMock()
cloud_backend.client = mock_client
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
mock_local_source = mock.MagicMock()
monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock_local_source)
cloud_runtime = cloud.CloudRuntime(entrypoint=tmpdir / path)
cloud_runtime.open("test_space")
mock_local_source.assert_called_once()
repo_call = mock_local_source.call_args
assert repo_call.kwargs["path"] == (tmpdir / expected_root).absolute()
ignore_functions = repo_call.kwargs["ignore_functions"]
if len(ignore_functions) > 0:
filtered = ignore_functions[0]("", [tmpdir / entry for entry in entries])
else:
filtered = [tmpdir / entry for entry in entries]
filtered = [entry.absolute() for entry in filtered]
expected_filtered_entries = [(tmpdir / entry).absolute() for entry in expected_filtered_entries]
assert filtered == expected_filtered_entries
def test_reopen(self, monkeypatch, capsys):
"""Tests that the open method calls the expected API endpoints when the CloudSpace already exists."""
mock_client = mock.MagicMock()
mock_client.auth_service_get_user.return_value = V1GetUserResponse(
username="tester", features=V1UserFeatures(code_tab=True)
)
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="test-project", project_id="test-project-id")]
)
mock_client.cloud_space_service_list_cloud_spaces.return_value = V1ListCloudSpacesResponse(
cloudspaces=[V1CloudSpace(id="cloudspace_id", name="test_space")]
)
running_instance = Externalv1LightningappInstance(
id="instance_id",
name="test_space",
spec=V1LightningappInstanceSpec(cluster_id="test"),
status=V1LightningappInstanceStatus(phase=V1LightningappInstanceState.RUNNING),
)
stopped_instance = Externalv1LightningappInstance(
id="instance_id",
name="test_space",
spec=V1LightningappInstanceSpec(cluster_id="test"),
status=V1LightningappInstanceStatus(phase=V1LightningappInstanceState.STOPPED),
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=[running_instance])
)
mock_client.lightningapp_instance_service_update_lightningapp_instance.return_value = running_instance
mock_client.lightningapp_instance_service_get_lightningapp_instance.return_value = stopped_instance
mock_client.cloud_space_service_create_cloud_space.return_value = V1CloudSpace(id="cloudspace_id")
mock_client.cloud_space_service_create_lightning_run.return_value = V1LightningRun(id="run_id")
cluster = Externalv1Cluster(id="test", spec=V1ClusterSpec(cluster_type=V1ClusterType.GLOBAL))
mock_client.projects_service_list_project_cluster_bindings.return_value = V1ListProjectClusterBindingsResponse(
clusters=[V1ProjectClusterBinding(cluster_id="test")],
)
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([cluster])
mock_client.cluster_service_get_cluster.return_value = cluster
cloud_backend = mock.MagicMock()
cloud_backend.client = mock_client
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
mock_local_source = mock.MagicMock()
monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock_local_source)
cloud_runtime = cloud.CloudRuntime(entrypoint=Path("."))
cloud_runtime.open("test_space")
mock_client.cloud_space_service_create_lightning_run_instance.assert_not_called()
mock_client.cloud_space_service_create_cloud_space.assert_not_called()
mock_client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="test-project-id", cloudspace_id="cloudspace_id", body=mock.ANY
)
def test_not_enabled(self, monkeypatch, capsys):
"""Tests that an error is printed and the call exits if the feature isn't enabled for the user."""
mock_client = mock.MagicMock()
mock_client.auth_service_get_user.return_value = V1GetUserResponse(
username="tester",
features=V1UserFeatures(code_tab=False),
)
cloud_backend = mock.MagicMock()
cloud_backend.client = mock_client
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
cloud_runtime = cloud.CloudRuntime(entrypoint=Path("."))
monkeypatch.setattr(cloud, "Path", Path)
exited = False
try:
cloud_runtime.open("test_space")
except SystemExit:
# Expected behaviour
exited = True
out, _ = capsys.readouterr()
assert exited
assert "`lightning_app open` command has not been enabled" in out
class TestCloudspaceDispatch:
@mock.patch.object(pathlib.Path, "exists")
@pytest.mark.parametrize(
("custom_env_sync_path_value", "cloudspace"),
[
(None, V1CloudSpace(id="test_id", code_config=V1CloudSpaceInstanceConfig())),
(
Path("/tmp/sys-customizations-sync"),
V1CloudSpace(id="test_id", code_config=V1CloudSpaceInstanceConfig()),
),
(
Path("/tmp/sys-customizations-sync"),
V1CloudSpace(
id="test_id",
code_config=V1CloudSpaceInstanceConfig(data_connection_mounts=[V1DataConnectionMount(id="test")]),
),
),
],
)
def test_cloudspace_dispatch(self, custom_env_sync_root, custom_env_sync_path_value, cloudspace, monkeypatch):
"""Tests that the cloudspace_dispatch method calls the expected API endpoints."""
mock_client = mock.MagicMock()
mock_client.auth_service_get_user.return_value = V1GetUserResponse(
username="tester",
features=V1UserFeatures(),
)
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="project", project_id="project_id")]
)
mock_client.cloud_space_service_create_lightning_run.return_value = V1LightningRun(id="run_id")
mock_client.cloud_space_service_create_lightning_run_instance.return_value = Externalv1LightningappInstance(
id="instance_id"
)
cluster = Externalv1Cluster(id="test", spec=V1ClusterSpec(cluster_type=V1ClusterType.GLOBAL))
mock_client.projects_service_list_project_cluster_bindings.return_value = V1ListProjectClusterBindingsResponse(
clusters=[V1ProjectClusterBinding(cluster_id="cluster_id")],
)
mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([cluster])
mock_client.cluster_service_get_cluster.return_value = cluster
mock_client.cloud_space_service_get_cloud_space.return_value = cloudspace
cloud_backend = mock.MagicMock()
cloud_backend.client = mock_client
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
mock_repo = mock.MagicMock()
mock_local_source = mock.MagicMock(return_value=mock_repo)
monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock_local_source)
custom_env_sync_root.return_value = custom_env_sync_path_value
mock_app = mock.MagicMock()
mock_app.works = [mock.MagicMock()]
cloud_runtime = cloud.CloudRuntime(app=mock_app, entrypoint=Path("."))
app = cloud_runtime.cloudspace_dispatch("project_id", "cloudspace_id", "run_name", "cluster_id")
assert app.id == "instance_id"
mock_client.cloud_space_service_get_cloud_space.assert_called_once_with(
project_id="project_id", id="cloudspace_id"
)
mock_client.cloud_space_service_create_lightning_run.assert_called_once_with(
project_id="project_id", cloudspace_id="cloudspace_id", body=mock.ANY
)
assert (
mock_client.cloud_space_service_create_lightning_run.call_args.kwargs["body"]
.works[0]
.spec.data_connection_mounts
== cloudspace.code_config.data_connection_mounts
)
mock_client.cloud_space_service_create_lightning_run_instance.assert_called_once_with(
project_id="project_id", cloudspace_id="cloudspace_id", id="run_id", body=mock.ANY
)
assert mock_client.cloud_space_service_create_lightning_run_instance.call_args.kwargs["body"].name == "run_name"
@mock.patch("lightning.app.core.queues.QueuingSystem", MagicMock())
@mock.patch("lightning.app.runners.backends.cloud.LightningClient", MagicMock())
def test_get_project(monkeypatch):
mock_client = mock.MagicMock()
monkeypatch.setattr(cloud, "CloudBackend", mock.MagicMock(return_value=mock_client))
app = mock.MagicMock(spec=LightningApp)
cloud.CloudRuntime(app=app, entrypoint=Path("entrypoint.py"))
# No valid projects
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(memberships=[])
with pytest.raises(ValueError, match="No valid projects found"):
_get_project(mock_client)
# One valid project
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="test-project", project_id="test-project-id")]
)
ret = _get_project(mock_client)
assert ret.project_id == "test-project-id"
# Multiple valid projects
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[
V1Membership(name="test-project1", project_id="test-project-id1"),
V1Membership(name="test-project2", project_id="test-project-id2"),
]
)
ret = _get_project(mock_client)
assert ret.project_id == "test-project-id1"
def write_file_of_size(path, size):
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "wb") as f:
f.seek(size)
f.write(b"\0")
@mock.patch("lightning.app.core.queues.QueuingSystem", MagicMock())
@mock.patch("lightning.app.runners.backends.cloud.LightningClient", MagicMock())
def test_check_uploaded_folder(monkeypatch, tmpdir, caplog):
app = MagicMock()
root = Path(tmpdir)
repo = LocalSourceCodeDir(root)
backend = cloud.CloudRuntime(app)
with caplog.at_level(logging.WARN):
backend._validate_repo(root, repo)
assert caplog.messages == []
# write some files to assert the message below.
write_file_of_size(root / "a.png", 4 * 1000 * 1000)
write_file_of_size(root / "b.txt", 5 * 1000 * 1000)
write_file_of_size(root / "c.jpg", 6 * 1000 * 1000)
repo._non_ignored_files = None # force reset
with caplog.at_level(logging.WARN):
backend._validate_repo(root, repo)
assert f"Your application folder '{root.absolute()}' is more than 2 MB" in caplog.text
assert "The total size is 15.0 MB" in caplog.text
assert "4 files were uploaded" in caplog.text
assert "files:\n6.0 MB: c.jpg\n5.0 MB: b.txt\n4.0 MB: a.png\nPerhaps" in caplog.text # tests the order
assert "adding them to `.lightningignore`." in caplog.text
assert "lightningingore` attribute in a Flow or Work" in caplog.text
@mock.patch("lightning.app.core.queues.QueuingSystem", MagicMock())
@mock.patch("lightning.app.runners.backends.cloud.LightningClient", MagicMock())
def test_project_has_sufficient_credits():
app = mock.MagicMock(spec=LightningApp)
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=Path("entrypoint.py"))
credits_and_test_value = [
[0.3, True],
[1, False],
[1.1, False],
]
for balance, result in credits_and_test_value:
project = V1Membership(name="test-project1", project_id="test-project-id1", balance=balance)
assert cloud_runtime._resolve_needs_credits(project) is result
@pytest.mark.parametrize(
"lines",
[
[
"import this_package_is_not_real",
"from lightning.app import LightningApp",
"from lightning.app.testing.helpers import EmptyWork",
"app = LightningApp(EmptyWork())",
],
[
"from this_package_is_not_real import this_module_is_not_real",
"from lightning.app import LightningApp",
"from lightning.app.testing.helpers import EmptyWork",
"app = LightningApp(EmptyWork())",
],
[
"import this_package_is_not_real",
"from this_package_is_not_real import this_module_is_not_real",
"from lightning.app import LightningApp",
"from lightning.app.testing.helpers import EmptyWork",
"app = LightningApp(EmptyWork())",
],
[
"import this_package_is_not_real",
"from lightning.app import LightningApp",
"from lightning.app.core.flow import _RootFlow",
"from lightning.app.testing.helpers import EmptyWork",
"class MyFlow(_RootFlow):",
" def configure_layout(self):",
" return [{'name': 'test', 'content': this_package_is_not_real()}]",
"app = LightningApp(MyFlow(EmptyWork()))",
],
],
)
@pytest.mark.skipif(sys.platform != "linux", reason="Causing conflicts on non-linux")
def test_load_app_from_file_mock_imports(tmpdir, lines):
path = copy(sys.path)
app_file = os.path.join(tmpdir, "app.py")
with open(app_file, "w") as f:
f.write("\n".join(lines))
app = CloudRuntime.load_app_from_file(app_file)
assert isinstance(app, LightningApp)
assert isinstance(app.root.work, EmptyWork)
# Cleanup PATH to prevent conflict with other tests
sys.path = path
os.remove(app_file)
def test_load_app_from_file():
test_script_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "core", "scripts")
app = CloudRuntime.load_app_from_file(
os.path.join(test_script_dir, "app_with_env.py"),
)
assert app.works[0].cloud_compute.name == "cpu-small"
app = CloudRuntime.load_app_from_file(
os.path.join(test_script_dir, "app_with_env.py"),
env_vars={"COMPUTE_NAME": "foo"},
)
assert app.works[0].cloud_compute.name == "foo"
@pytest.mark.parametrize(
("print_format", "expected"),
[
(
"web",
[
{
"displayName": "",
"name": "root.work",
"spec": {
"buildSpec": {
"commands": [],
"pythonDependencies": {"packageManager": "PACKAGE_MANAGER_PIP", "packages": ""},
},
"dataConnectionMounts": [],
"drives": [],
"networkConfig": [{"name": "*", "port": "*"}],
"userRequestedComputeConfig": {
"count": 1,
"diskSize": 0,
"name": "cpu-small",
"preemptible": "*",
"shmSize": 0,
},
},
}
],
),
(
"gallery",
[
{
"display_name": "",
"name": "root.work",
"spec": {
"build_spec": {
"commands": [],
"python_dependencies": {"package_manager": "PACKAGE_MANAGER_PIP", "packages": ""},
},
"data_connection_mounts": [],
"drives": [],
"network_config": [{"name": "*", "port": "*"}],
"user_requested_compute_config": {
"count": 1,
"disk_size": 0,
"name": "cpu-small",
"preemptible": "*",
"shm_size": 0,
},
},
}
],
),
],
)
def test_print_specs(tmpdir, caplog, monkeypatch, print_format, expected):
entrypoint = Path(tmpdir) / "entrypoint.py"
entrypoint.touch()
mock_client = mock.MagicMock()
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="test-project", project_id="test-project-id")]
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=[])
)
cloud_backend = mock.MagicMock(client=mock_client)
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
cloud_runtime = cloud.CloudRuntime(app=LightningApp(EmptyWork()), entrypoint=entrypoint)
cloud.LIGHTNING_CLOUD_PRINT_SPECS = print_format
try:
with caplog.at_level(logging.INFO), contextlib.suppress(SystemExit):
cloud_runtime.dispatch()
lines = caplog.text.split("\n")
expected = re.escape(str(expected).replace("'", '"').replace(" ", "")).replace('"\\*"', "(.*)")
expected = "INFO(.*)works: " + expected
assert any(re.fullmatch(expected, line) for line in lines)
finally:
cloud.LIGHTNING_CLOUD_PRINT_SPECS = None
def test_incompatible_cloud_compute_and_build_config(monkeypatch):
"""Test that an exception is raised when a build config has a custom image defined, but the cloud compute is the
default.
This combination is not supported by the platform.
"""
mock_client = mock.MagicMock()
cloud_backend = mock.MagicMock(client=mock_client)
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
class Work(LightningWork):
def __init__(self):
super().__init__()
self.cloud_compute = CloudCompute(name="default")
# TODO: Remove me
self.cloud_compute.name = "default"
self.cloud_build_config = BuildConfig(image="custom")
def run(self):
pass
app = MagicMock()
app.works = [Work()]
with pytest.raises(ValueError, match="You requested a custom base image for the Work with name"):
CloudRuntime(app=app)._validate_work_build_specs_and_compute()
def test_programmatic_lightningignore(monkeypatch, caplog, tmpdir):
path = Path(tmpdir)
entrypoint = path / "entrypoint.py"
entrypoint.touch()
mock_client = mock.MagicMock()
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="test-project", project_id="test-project-id")]
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=[])
)
mock_client.cloud_space_service_create_lightning_run.return_value = V1LightningRun(cluster_id="test")
cloud_backend = mock.MagicMock(client=mock_client)
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
class MyWork(LightningWork):
def __init__(self):
super().__init__()
self.lightningignore += ("foo", "lightning_logs")
def run(self):
with pytest.raises(RuntimeError, match="w.lightningignore` does not"):
self.lightningignore += ("foobar",)
class MyFlow(LightningFlow):
def __init__(self):
super().__init__()
self.lightningignore = ("foo",)
self.w = MyWork()
def run(self):
with pytest.raises(RuntimeError, match="root.lightningignore` does not"):
self.lightningignore = ("baz",)
self.w.run()
flow = MyFlow()
app = LightningApp(flow)
monkeypatch.setattr(app, "_update_index_file", mock.MagicMock())
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=entrypoint)
monkeypatch.setattr(LocalSourceCodeDir, "upload", mock.MagicMock())
# write some files
write_file_of_size(path / "a.txt", 5 * 1000 * 1000)
write_file_of_size(path / "foo.png", 4 * 1000 * 1000)
write_file_of_size(path / "lightning_logs" / "foo.ckpt", 6 * 1000 * 1000)
# also an actual .lightningignore file
(path / ".lightningignore").write_text("foo.png")
with mock.patch(
"lightning.app.runners.cloud._parse_lightningignore", wraps=_parse_lightningignore
) as parse_mock, mock.patch(
"lightning.app.source_code.local._copytree", wraps=_copytree
) as copy_mock, caplog.at_level(logging.WARN):
cloud_runtime.dispatch()
parse_mock.assert_called_once_with(("foo", "foo", "lightning_logs"))
assert copy_mock.mock_calls[0].kwargs["ignore_functions"][0].args[1] == {"lightning_logs", "foo"}
assert f"Your application folder '{path.absolute()}' is more than 2 MB" in caplog.text
assert "The total size is 5.0 MB" in caplog.text
assert "2 files were uploaded" # a.txt and .lightningignore
assert "files:\n5.0 MB: a.txt\nPerhaps" in caplog.text # only this file appears
flow.run()
def test_default_lightningignore(monkeypatch, caplog, tmpdir):
path = Path(tmpdir)
entrypoint = path / "entrypoint.py"
entrypoint.touch()
mock_client = mock.MagicMock()
mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
memberships=[V1Membership(name="test-project", project_id="test-project-id")]
)
mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
V1ListLightningappInstancesResponse(lightningapps=[])
)
mock_client.cloud_space_service_create_lightning_run.return_value = V1LightningRun(cluster_id="test")
cloud_backend = mock.MagicMock(client=mock_client)
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
class MyWork(LightningWork):
def run(self):
pass
app = LightningApp(MyWork())
cloud_runtime = cloud.CloudRuntime(app=app, entrypoint=entrypoint)
monkeypatch.setattr(LocalSourceCodeDir, "upload", mock.MagicMock())
# write some files
write_file_of_size(path / "a.txt", 5 * 1000 * 1000)
write_file_of_size(path / "venv" / "foo.txt", 4 * 1000 * 1000)
assert not (path / ".lightningignore").exists()
with mock.patch(
"lightning.app.runners.cloud._parse_lightningignore", wraps=_parse_lightningignore
) as parse_mock, mock.patch(
"lightning.app.source_code.local._copytree", wraps=_copytree
) as copy_mock, caplog.at_level(logging.WARN):
cloud_runtime.dispatch()
parse_mock.assert_called_once_with(())
assert copy_mock.mock_calls[0].kwargs["ignore_functions"][0].args[1] == set()
assert (path / ".lightningignore").exists()
assert f"Your application folder '{path.absolute()}' is more than 2 MB" in caplog.text
assert "The total size is 5.0 MB" in caplog.text
assert "2 files were uploaded" # a.txt and .lightningignore
assert "files:\n5.0 MB: a.txt\nPerhaps" in caplog.text # only this file appears
@pytest.mark.parametrize(
("project", "run_instance", "user", "tab", "lightning_cloud_url", "expected_url"),
[
# Old style
(
V1Membership(),
Externalv1LightningappInstance(id="test-app-id"),
V1GetUserResponse(username="tester", features=V1UserFeatures()),
"logs",
"https://lightning.ai",
"https://lightning.ai/tester/apps/test-app-id/logs",
),
(
V1Membership(),
Externalv1LightningappInstance(id="test-app-id"),
V1GetUserResponse(username="tester", features=V1UserFeatures()),
"logs",
"http://localhost:9800",
"http://localhost:9800/tester/apps/test-app-id/logs",
),
# New style
(
V1Membership(name="tester's project"),
Externalv1LightningappInstance(name="test/job"),
V1GetUserResponse(username="tester", features=V1UserFeatures(project_selector=True)),
"logs",
"https://lightning.ai",
"https://lightning.ai/tester/tester%27s%20project/jobs/test%2Fjob/logs",
),
(
V1Membership(name="tester's project"),
Externalv1LightningappInstance(name="test/job"),
V1GetUserResponse(username="tester", features=V1UserFeatures(project_selector=True)),
"logs",
"https://localhost:9800",
"https://localhost:9800/tester/tester%27s%20project/jobs/test%2Fjob/logs",
),
],
)
def test_get_app_url(monkeypatch, project, run_instance, user, tab, lightning_cloud_url, expected_url):
mock_client = mock.MagicMock()
mock_client.auth_service_get_user.return_value = user
cloud_backend = mock.MagicMock(client=mock_client)
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
runtime = CloudRuntime()
with mock.patch(
"lightning.app.runners.cloud.get_lightning_cloud_url", mock.MagicMock(return_value=lightning_cloud_url)
):
assert runtime._get_app_url(project, run_instance, tab) == expected_url
@pytest.mark.parametrize(
("user", "project", "cloudspace_name", "tab", "lightning_cloud_url", "expected_url"),
[
(
V1GetUserResponse(username="tester", features=V1UserFeatures()),
V1Membership(name="default-project"),
"test/cloudspace",
"code",
"https://lightning.ai",
"https://lightning.ai/tester/default-project/apps/test%2Fcloudspace/code",
),
(
V1GetUserResponse(username="tester", features=V1UserFeatures()),
V1Membership(name="Awesome Project"),
"The Best CloudSpace ever",
"web-ui",
"http://localhost:9800",
"http://localhost:9800/tester/Awesome%20Project/apps/The%20Best%20CloudSpace%20ever/web-ui",
),
],
)
def test_get_cloudspace_url(monkeypatch, user, project, cloudspace_name, tab, lightning_cloud_url, expected_url):
mock_client = mock.MagicMock()
mock_client.auth_service_get_user.return_value = user
cloud_backend = mock.MagicMock(client=mock_client)
monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
runtime = CloudRuntime()
with mock.patch(
"lightning.app.runners.cloud.get_lightning_cloud_url", mock.MagicMock(return_value=lightning_cloud_url)
):
assert runtime._get_cloudspace_url(project, cloudspace_name, tab) == expected_url
Reference in New Issue View Git Blame Copy Permalink