Document and test Metaflow tempdir defaults

Author: codex

metaflow/plugins/aws/batch/batch_decorator.py

@@ -299,8 +299,9 @@ def task_pre_step(
         self.metadata = metadata
         self.task_datastore = task_datastore
 
-        # current.tempdir reflects the value of METAFLOW_TEMPDIR (the current working
-        # directory by default), or the value of tmpfs_path if tmpfs_tempdir=False.
+        # current.tempdir reflects the value of METAFLOW_TEMPDIR (the system
+        # temporary directory from tempfile.gettempdir() by default), or the
+        # value of tmpfs_path if tmpfs_tempdir=False.
         if not self.attributes["tmpfs_tempdir"]:
             current._update_env({"tempdir": self.attributes["tmpfs_path"]})
 

metaflow/plugins/datatools/s3/s3.py

@@ -479,8 +479,9 @@ class S3(object):
     data = [obj.blob for obj in s3.get_many(urls)]
     s3.close()
     ```
-    You can customize the location of the temporary directory with `tmproot`. It
-    defaults to the current working directory.
+    You can customize the location of the temporary directory with `tmproot`, or
+    with the `METAFLOW_TEMPDIR` configuration variable. If neither is specified,
+    the system temporary directory returned by `tempfile.gettempdir()` is used.
 
     To make it easier to deal with object locations, the client can be initialized
     with an S3 path prefix. There are three ways to handle locations:
@@ -499,7 +500,7 @@ class S3(object):
 
     Parameters
     ----------
-    tmproot : str, default '.'
+    tmproot : str, default METAFLOW_TEMPDIR or tempfile.gettempdir()
         Where to store the temporary directory.
     bucket : str, optional, default None
         Override the bucket from `DATATOOLS_S3ROOT` when `run` is specified.

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -529,8 +529,9 @@ def task_pre_step(
         self.metadata = metadata
         self.task_datastore = task_datastore
 
-        # current.tempdir reflects the value of METAFLOW_TEMPDIR (the current working
-        # directory by default), or the value of tmpfs_path if tmpfs_tempdir=False.
+        # current.tempdir reflects the value of METAFLOW_TEMPDIR (the system
+        # temporary directory from tempfile.gettempdir() by default), or the
+        # value of tmpfs_path if tmpfs_tempdir=False.
         if not self.attributes["tmpfs_tempdir"]:
             current._update_env({"tempdir": self.attributes["tmpfs_path"]})
 

test/unit/test_s3_readonly_cwd.py

@@ -1,13 +1,13 @@
 """Test that S3 client works when current directory is not writeable."""
 
+import importlib
 import os
 import tempfile
+
 import pytest
 
 
-def test_s3_client_readonly_cwd(mocker):
-    """Test that S3 client doesn't require writable CWD (issue #854)."""
-    # Mock boto3 and dependencies before importing S3
+def _mock_s3_dependencies(mocker):
     mock_boto3 = mocker.MagicMock()
     mock_transfer_config = mocker.MagicMock()
     mocker.patch.dict("sys.modules", {"boto3": mock_boto3})
@@ -16,33 +16,82 @@ def test_s3_client_readonly_cwd(mocker):
         {"boto3.s3.transfer": mocker.MagicMock(TransferConfig=mock_transfer_config)},
     )
 
-    from metaflow.plugins.datatools.s3 import S3
 
-    # Create a temporary read-only directory
-    with tempfile.TemporaryDirectory() as tmpdir:
-        readonly_dir = os.path.join(tmpdir, "readonly")
-        os.makedirs(readonly_dir, mode=0o555)
+def _load_s3_class():
+    """Reload configuration and S3 so environment changes affect the default."""
+    metaflow_config = importlib.import_module("metaflow.metaflow_config")
+    s3_module = importlib.import_module("metaflow.plugins.datatools.s3.s3")
+    importlib.reload(metaflow_config)
+    importlib.reload(s3_module)
+    # Keep the package-level export in sync for tests that import S3 later.
+    s3_package = importlib.import_module("metaflow.plugins.datatools.s3")
+    s3_package.S3 = s3_module.S3
+    return s3_module.S3
 
-        # Change to read-only directory
-        original_cwd = os.getcwd()
-        try:
-            os.chdir(readonly_dir)
 
-            # This should not raise an exception about permissions
-            # The S3 client should use system temp dir, not CWD
-            s3_client = S3()
+def test_s3_client_readonly_cwd(mocker, monkeypatch):
+    """Test that S3 client doesn't require writable CWD (issue #854)."""
+    _mock_s3_dependencies(mocker)
+    original_tempdir = os.environ.get("METAFLOW_TEMPDIR")
+
+    try:
+        with monkeypatch.context() as env:
+            env.delenv("METAFLOW_TEMPDIR", raising=False)
+            S3 = _load_s3_class()
+
+            # Create a temporary read-only directory
+            with tempfile.TemporaryDirectory() as tmpdir:
+                readonly_dir = os.path.join(tmpdir, "readonly")
+                os.makedirs(readonly_dir, mode=0o555)
+
+                # Change to read-only directory
+                original_cwd = os.getcwd()
+                try:
+                    os.chdir(readonly_dir)
 
-            # Verify that tmpdir was created in system temp, not CWD
-            assert s3_client._tmpdir is not None
-            # The temp dir should NOT be in the current (read-only) directory
-            assert not s3_client._tmpdir.startswith(readonly_dir)
-            # The temp dir should be in the system temp directory
-            system_temp = tempfile.gettempdir()
-            assert s3_client._tmpdir.startswith(system_temp)
-
-            # Clean up
-            s3_client.close()
-        finally:
-            os.chdir(original_cwd)
-            # Make directory writable again so it can be deleted
-            os.chmod(readonly_dir, 0o755)
+                    # This should not raise an exception about permissions
+                    # The S3 client should use system temp dir, not CWD
+                    s3_client = S3()
+
+                    # Verify that tmpdir was created in system temp, not CWD
+                    assert s3_client._tmpdir is not None
+                    # The temp dir should NOT be in the current (read-only) directory
+                    assert not s3_client._tmpdir.startswith(readonly_dir)
+                    # The temp dir should be in the system temp directory
+                    system_temp = tempfile.gettempdir()
+                    assert s3_client._tmpdir.startswith(system_temp)
+
+                    # Clean up
+                    s3_client.close()
+                finally:
+                    os.chdir(original_cwd)
+                    # Make directory writable again so it can be deleted
+                    os.chmod(readonly_dir, 0o755)
+        assert os.environ.get("METAFLOW_TEMPDIR") == original_tempdir
+    finally:
+        # Restore the S3 module's import-time default after the isolated test.
+        _load_s3_class()
+
+
+def test_s3_client_respects_metaflow_tempdir(mocker, monkeypatch, tmp_path):
+    """Test that METAFLOW_TEMPDIR overrides the system temporary directory."""
+    _mock_s3_dependencies(mocker)
+    configured_dir = tmp_path / "configured"
+    configured_dir.mkdir()
+    original_tempdir = os.environ.get("METAFLOW_TEMPDIR")
+
+    try:
+        with monkeypatch.context() as env:
+            env.setenv("METAFLOW_TEMPDIR", str(configured_dir))
+            S3 = _load_s3_class()
+
+            s3_client = S3()
+            try:
+                assert s3_client._tmproot == str(configured_dir)
+                assert os.path.dirname(s3_client._tmpdir) == str(configured_dir)
+            finally:
+                s3_client.close()
+        assert os.environ.get("METAFLOW_TEMPDIR") == original_tempdir
+    finally:
+        # Restore the S3 module's import-time default after the isolated test.
+        _load_s3_class()