[App] Content for plugins (#17243)

Co-authored-by: Yurij Mikhalevich <[email protected]>
Co-authored-by: Luca Antiga <[email protected]>

(cherry picked from commit 2c3dfc0)

Author: ethanwharris

src/lightning/app/plugin/plugin.py

@@ -61,6 +61,8 @@ def run_job(self, name: str, app_entrypoint: str, env_vars: Dict[str, str] = {})
         """
         from lightning.app.runners.cloud import CloudRuntime
 
+        logger.info(f"Processing job run request. name: {name}, app_entrypoint: {app_entrypoint}, env_vars: {env_vars}")
+
         # Dispatch the job
         _set_flow_context()
 
@@ -123,6 +125,8 @@ def _run_plugin(run: _Run) -> Dict[str, Any]:
 
         # Download the tarball
         try:
+            logger.info(f"Downloading plugin source: {run.source_code_url}")
+
             # Sometimes the URL gets encoded, so we parse it here
             source_code_url = urlparse(run.source_code_url).geturl()
 
@@ -141,6 +145,8 @@ def _run_plugin(run: _Run) -> Dict[str, Any]:
 
         # Extract
         try:
+            logger.info("Extracting plugin source.")
+
             with tarfile.open(download_path, "r:gz") as tf:
                 tf.extractall(source_path)
         except Exception as ex:
@@ -151,6 +157,8 @@ def _run_plugin(run: _Run) -> Dict[str, Any]:
 
         # Import the plugin
         try:
+            logger.info(f"Importing plugin: {run.plugin_entrypoint}")
+
             plugin = _load_plugin_from_file(os.path.join(source_path, run.plugin_entrypoint))
         except Exception as ex:
             raise HTTPException(
@@ -163,6 +171,11 @@ def _run_plugin(run: _Run) -> Dict[str, Any]:
 
         # Setup and run the plugin
         try:
+            logger.info(
+                "Running plugin. "
+                f"project_id: {run.project_id}, cloudspace_id: {run.cloudspace_id}, cluster_id: {run.cluster_id}."
+            )
+
             plugin._setup(
                 project_id=run.project_id,
                 cloudspace_id=run.cloudspace_id,

src/lightning/app/runners/cloud.py

@@ -218,16 +218,22 @@ def cloudspace_dispatch(
         # Dispatch in four phases: resolution, validation, spec creation, API transactions
         # Resolution
         root = self._resolve_root()
-        repo = self._resolve_repo(root)
+        # If the root will already be there, we don't need to upload and preserve the absolute entrypoint
+        absolute_entrypoint = str(root).startswith("/project")
+        # If system customization files found, it will set their location path
+        sys_customizations_root = self._resolve_env_root()
+        repo = self._resolve_repo(
+            root,
+            default_ignore=False,
+            package_source=not absolute_entrypoint,
+            sys_customizations_root=sys_customizations_root,
+        )
         project = self._resolve_project(project_id=project_id)
         existing_instances = self._resolve_run_instances_by_name(project_id, name)
         name = self._resolve_run_name(name, existing_instances)
         cloudspace = self._resolve_cloudspace(project_id, cloudspace_id)
         queue_server_type = self._resolve_queue_server_type()
 
-        # If system customization files found, it will set their location path
-        sys_customizations_sync_root = self._resolve_env_root()
-
         self.app._update_index_file()
 
         # Validation
@@ -241,17 +247,26 @@ def cloudspace_dispatch(
         flow_servers = self._get_flow_servers()
         network_configs = self._get_network_configs(flow_servers)
         works = self._get_works(cloudspace=cloudspace)
-        run_body = self._get_run_body(cluster_id, flow_servers, network_configs, works, False, root, True)
+        run_body = self._get_run_body(
+            cluster_id,
+            flow_servers,
+            network_configs,
+            works,
+            False,
+            root,
+            True,
+            True,
+            absolute_entrypoint,
+        )
         env_vars = self._get_env_vars(self.env_vars, self.secrets, self.run_app_comment_commands)
 
-        # If the system customization root is set, prepare files for environment synchronization
-        if sys_customizations_sync_root is not None:
-            repo.prepare_sys_customizations_sync(sys_customizations_sync_root)
-
         # API transactions
+        logger.info(f"Creating cloudspace run. run_body: {run_body}")
         run = self._api_create_run(project_id, cloudspace_id, run_body)
+
         self._api_package_and_upload_repo(repo, run)
 
+        logger.info(f"Creating cloudspace run instance. name: {name}")
         run_instance = self._api_create_run_instance(
             cluster_id,
             project_id,
@@ -454,6 +469,9 @@ def _resolve_repo(
         self,
         root: Path,
         ignore_functions: Optional[List[_IGNORE_FUNCTION]] = None,
+        default_ignore: bool = True,
+        package_source: bool = True,
+        sys_customizations_root: Optional[Path] = None,
     ) -> LocalSourceCodeDir:
         """Gather and merge all lightningignores from the app children and create the ``LocalSourceCodeDir``
         object."""
@@ -470,7 +488,13 @@ def _resolve_repo(
                 patterns = _parse_lightningignore(merged)
                 ignore_functions = [*ignore_functions, partial(_filter_ignored, root, patterns)]
 
-        return LocalSourceCodeDir(path=root, ignore_functions=ignore_functions)
+        return LocalSourceCodeDir(
+            path=root,
+            ignore_functions=ignore_functions,
+            default_ignore=default_ignore,
+            package_source=package_source,
+            sys_customizations_root=sys_customizations_root,
+        )
 
     def _resolve_project(self, project_id: Optional[str] = None) -> V1Membership:
         """Determine the project to run on, choosing a default if multiple projects are found."""
@@ -785,7 +809,7 @@ def _get_works(self, cloudspace: Optional[V1CloudSpace] = None) -> List[V1Work]:
                 network_config=[V1NetworkConfig(name=random_name, port=work.port)],
                 data_connection_mounts=data_connection_mounts,
             )
-            works.append(V1Work(name=work.name, spec=work_spec))
+            works.append(V1Work(name=work.name, display_name=work.display_name, spec=work_spec))
 
         return works
 
@@ -798,12 +822,18 @@ def _get_run_body(
         no_cache: bool,
         root: Path,
         start_server: bool,
+        should_mount_cloudspace_content: bool = False,
+        absolute_entrypoint: bool = False,
     ) -> CloudspaceIdRunsBody:
         """Get the specification of the run creation request."""
-        # The entry point file needs to be relative to the root of the uploaded source file directory,
-        # because the backend will invoke the lightning commands relative said source directory
-        # TODO: we shouldn't set this if the entrypoint isn't a file but the backend gives an error if we don't
-        app_entrypoint_file = Path(self.entrypoint).absolute().relative_to(root)
+        if absolute_entrypoint:
+            # If the entrypoint will already exist in the cloud then we can choose to keep it as an absolute path.
+            app_entrypoint_file = Path(self.entrypoint).absolute()
+        else:
+            # The entry point file needs to be relative to the root of the uploaded source file directory,
+            # because the backend will invoke the lightning commands relative said source directory
+            # TODO: we shouldn't set this if the entrypoint isn't a file but the backend gives an error if we don't
+            app_entrypoint_file = Path(self.entrypoint).absolute().relative_to(root)
 
         run_body = CloudspaceIdRunsBody(
             cluster_id=cluster_id,
@@ -813,6 +843,7 @@ def _get_run_body(
             network_config=network_configs,
             works=works,
             local_source=True,
+            should_mount_cloudspace_content=should_mount_cloudspace_content,
         )
 
         if self.app is not None:
@@ -827,9 +858,10 @@ def _get_run_body(
         # if requirements file at the root of the repository is present,
         # we pass just the file name to the backend, so backend can find it in the relative path
         requirements_file = root / "requirements.txt"
-        if requirements_file.is_file():
+        if requirements_file.is_file() and requirements_file.exists():
+            requirements_path = requirements_file if absolute_entrypoint else "requirements.txt"
             run_body.image_spec = Gridv1ImageSpec(
-                dependency_file_info=V1DependencyFileInfo(package_manager=V1PackageManager.PIP, path="requirements.txt")
+                dependency_file_info=V1DependencyFileInfo(package_manager=V1PackageManager.PIP, path=requirements_path)
             )
             if not DISABLE_DEPENDENCY_CACHE and not no_cache:
                 # hash used for caching the dependencies
@@ -997,7 +1029,10 @@ def _api_create_run_instance(
         )
 
     @staticmethod
-    def _api_package_and_upload_repo(repo: LocalSourceCodeDir, run: V1LightningRun) -> None:
+    def _api_package_and_upload_repo(
+        repo: LocalSourceCodeDir,
+        run: V1LightningRun,
+    ) -> None:
         """Package and upload the provided local source code directory to the provided run."""
         if run.source_upload_url == "":
             raise RuntimeError("The source upload url is empty.")

src/lightning/app/source_code/local.py

@@ -13,14 +13,14 @@
 # limitations under the License.
 
 import os
+import uuid
 from contextlib import contextmanager
 from pathlib import Path
 from shutil import copytree, rmtree
 from typing import List, Optional
 
 from lightning.app.core.constants import DOT_IGNORE_FILENAME, SYS_CUSTOMIZATIONS_SYNC_PATH
 from lightning.app.source_code.copytree import _copytree, _IGNORE_FUNCTION
-from lightning.app.source_code.hashing import _get_hash
 from lightning.app.source_code.tar import _tar_path
 from lightning.app.source_code.uploader import FileUploader
 
@@ -30,20 +30,35 @@ class LocalSourceCodeDir:
 
     cache_location: Path = Path.home() / ".lightning" / "cache" / "repositories"
 
-    def __init__(self, path: Path, ignore_functions: Optional[List[_IGNORE_FUNCTION]] = None) -> None:
+    def __init__(
+        self,
+        path: Path,
+        ignore_functions: Optional[List[_IGNORE_FUNCTION]] = None,
+        default_ignore: bool = True,
+        package_source: bool = True,
+        sys_customizations_root: Optional[Path] = None,
+    ) -> None:
+        if "LIGHTNING_VSCODE_WORKSPACE" in os.environ:
+            # Don't use home to store the tar ball. This won't play nice with symlinks
+            self.cache_location: Path = Path("/tmp", ".lightning", "cache", "repositories")  # noqa: S108
+        else:
+            self.cache_location: Path = Path.home() / ".lightning" / "cache" / "repositories"
+
         self.path = path
         self.ignore_functions = ignore_functions
+        self.package_source = package_source
+        self.sys_customizations_root = sys_customizations_root
 
-        # cache checksum version
+        # cache version
         self._version: Optional[str] = None
         self._non_ignored_files: Optional[List[str]] = None
 
         # create global cache location if it doesn't exist
         if not self.cache_location.exists():
             self.cache_location.mkdir(parents=True, exist_ok=True)
 
-        # Create a default dotignore if it doesn't exist
-        if not (path / DOT_IGNORE_FILENAME).is_file():
+        # Create a default dotignore if requested and it doesn't exist
+        if default_ignore and not (path / DOT_IGNORE_FILENAME).is_file():
             with open(path / DOT_IGNORE_FILENAME, "w") as f:
                 f.write("venv/\n")
                 if (path / "bin" / "activate").is_file() or (path / "pyvenv.cfg").is_file():
@@ -57,7 +72,10 @@ def __init__(self, path: Path, ignore_functions: Optional[List[_IGNORE_FUNCTION]
     def files(self) -> List[str]:
         """Returns a set of files that are not ignored by .lightningignore."""
         if self._non_ignored_files is None:
-            self._non_ignored_files = _copytree(self.path, "", ignore_functions=self.ignore_functions, dry_run=True)
+            if self.package_source:
+                self._non_ignored_files = _copytree(self.path, "", ignore_functions=self.ignore_functions, dry_run=True)
+            else:
+                self._non_ignored_files = []
         return self._non_ignored_files
 
     @property
@@ -67,8 +85,8 @@ def version(self):
         if self._version is not None:
             return self._version
 
-        # stores both version and a set with the files used to generate the checksum
-        self._version = _get_hash(files=self.files, algorithm="blake2")
+        # create a random version ID and store it
+        self._version = uuid.uuid4().hex
         return self._version
 
     @property
@@ -83,7 +101,11 @@ def packaging_session(self) -> Path:
         session_path = self.cache_location / "packaging_sessions" / self.version
         try:
             rmtree(session_path, ignore_errors=True)
-            _copytree(self.path, session_path, ignore_functions=self.ignore_functions)
+            if self.package_source:
+                _copytree(self.path, session_path, ignore_functions=self.ignore_functions)
+            if self.sys_customizations_root is not None:
+                path_to_sync = Path(session_path, SYS_CUSTOMIZATIONS_SYNC_PATH)
+                copytree(self.sys_customizations_root, path_to_sync, dirs_exist_ok=True)
             yield session_path
         finally:
             rmtree(session_path, ignore_errors=True)
@@ -104,12 +126,6 @@ def package(self) -> Path:
             _tar_path(source_path=session_path, target_file=str(self.package_path), compression=True)
         return self.package_path
 
-    def prepare_sys_customizations_sync(self, sys_customizations_root: Path) -> None:
-        """Prepares files for system environment customization setup by copying conda and system environment files
-        to an app files directory."""
-        path_to_sync = Path(self.path, SYS_CUSTOMIZATIONS_SYNC_PATH)
-        copytree(sys_customizations_root, path_to_sync, dirs_exist_ok=True)
-
     def upload(self, url: str) -> None:
         """Uploads package to URL, usually pre-signed URL.