feat: Allow passing workspace base directory via environment variable (#973)

I'm using sglang in kubernetes; sglang is importing flashinfer

I have rather strict setup that does not allow root user to run the
container, so I'm running the container with user `nobody` and I need to
specify various directories that are supposed to be writable.

With latest flashinfer, I noticed that it tries to create some
directories under `$HOME`. Common way to get around it is to path some
env variable to the libraries trying to write stuff so I can control the
target path and let it write to the writable directory (And in my case
this is not $HOME). So I'm proposing to add an option to explicitly
specify directory where flashinfer wants to write so that the user has
more control over it

Author: jsuchome

flashinfer/jit/env.py

@@ -14,6 +14,7 @@
 limitations under the License.
 """
 
+import os
 import pathlib
 import re
 import warnings
@@ -32,8 +33,11 @@ def _get_workspace_dir_name() -> pathlib.Path:
         arch = "_".join(sorted(set(re.findall(r"compute_(\d+)", "".join(flags)))))
     except Exception:
         arch = "noarch"
+    flashinfer_base = os.getenv(
+        "FLASHINFER_WORKSPACE_BASE", pathlib.Path.home().as_posix()
+    )
     # e.g.: $HOME/.cache/flashinfer/75_80_89_90/
-    return pathlib.Path.home() / ".cache" / "flashinfer" / arch
+    return pathlib.Path(flashinfer_base) / ".cache" / "flashinfer" / arch
 
 
 # use pathlib