fix: bringing vfs_cache env var into vfs launch environment (#262)

* fix: bringing vfs_cache env var into vfs launch environment Signed-off-by: Nathan MacLeod <[email protected]> * fix unit test Signed-off-by: Nathan MacLeod <[email protected]> --------- Signed-off-by: Nathan MacLeod <[email protected]>
aws-deadline · Mar 29, 2024 · 1a6b8c8 · 1a6b8c8
1 parent 435e07b
commit 1a6b8c8
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/src/deadline/job_attachments/vfs.py b/src/deadline/job_attachments/vfs.py
@@ -21,6 +21,7 @@
 log = logging.getLogger(__name__)
 
 DEADLINE_VFS_ENV_VAR = "DEADLINE_VFS_PATH"
+DEADLINE_VFS_CACHE_ENV_VAR = "DEADLINE_VFS_CACHE"
 DEADLINE_VFS_EXECUTABLE = "deadline_vfs"
 DEADLINE_VFS_INSTALL_PATH = "/opt/deadline_vfs"
 DEADLINE_VFS_EXECUTABLE_SCRIPT = "/scripts/production/al2/run_deadline_vfs_al2.sh"
@@ -428,6 +429,8 @@ def get_launch_environ(self) -> dict:
         my_env = {**self._os_env_vars}
         my_env["PATH"] = f"{VFSProcessManager.find_vfs_link_dir()}{os.pathsep}{os.environ['PATH']}"
         my_env["LD_LIBRARY_PATH"] = VFSProcessManager.get_library_path()  # type: ignore[assignment]
+        if os.environ.get(DEADLINE_VFS_CACHE_ENV_VAR) is not None:
+            my_env[DEADLINE_VFS_CACHE_ENV_VAR] = os.environ.get(DEADLINE_VFS_CACHE_ENV_VAR)  # type: ignore[assignment]
 
         return my_env
 

diff --git a/test/unit/deadline_job_attachments/test_vfs.py b/test/unit/deadline_job_attachments/test_vfs.py
@@ -24,6 +24,7 @@
 from deadline.job_attachments.vfs import (
     VFSProcessManager,
     DEADLINE_VFS_ENV_VAR,
+    DEADLINE_VFS_CACHE_ENV_VAR,
     DEADLINE_VFS_EXECUTABLE,
     DEADLINE_VFS_EXECUTABLE_SCRIPT,
     DEADLINE_VFS_INSTALL_PATH,
@@ -751,15 +752,21 @@ def test_manifest_group_set(
             mock_chmod.assert_called_with(manifest_path, DEADLINE_MANIFEST_GROUP_READ_PERMS)
 
 
+@pytest.mark.parametrize("vfs_cache_enabled", [True, False])
 def test_launch_environment_has_expected_settings(
     tmp_path: Path,
+    vfs_cache_enabled: bool,
 ):
     # Test to verify when retrieving the launch environment it does not contain os.environ variables (Unless passed in),
     # it DOES contain the VFSProcessManager's environment variables, and aws configuration variables aren't modified
     session_dir: str = str(tmp_path)
     test_mount: str = f"{session_dir}/test_mount"
     manifest_path1: str = f"{session_dir}/manifests/some_manifest.json"
     os.environ[DEADLINE_VFS_ENV_VAR] = str((Path(__file__) / "deadline_vfs").resolve())
+    if vfs_cache_enabled:
+        os.environ[DEADLINE_VFS_CACHE_ENV_VAR] = "V0"
+    else:
+        os.environ.pop(DEADLINE_VFS_CACHE_ENV_VAR, None)
 
     provided_vars = {
         "VFS_ENV_VAR": "test-vfs-env-var",
@@ -787,6 +794,11 @@ def test_launch_environment_has_expected_settings(
     for key, value in provided_vars.items():
         assert launch_env.get(key) == value
 
+    if vfs_cache_enabled:
+        assert launch_env.get(DEADLINE_VFS_CACHE_ENV_VAR) == "V0"
+    else:
+        assert launch_env.get(DEADLINE_VFS_CACHE_ENV_VAR) is None
+
     # Base environment variables are not passed through
     assert not launch_env.get(DEADLINE_VFS_ENV_VAR)