From 9f02760bcf3c44b74f7e0a486522d1f4b1e7a487 Mon Sep 17 00:00:00 2001
From: Kiuk Chung <kiuk@meta.com>
Date: Tue, 30 Sep 2025 09:32:57 -0700
Subject: [PATCH] (torchx/specs) Add TORCHX_HOME function that returns the
 dot-torchx directory

Summary:
Adds a `TORCHX_HOME()` function to `torchx.specs.api` (re-exported from `torchx/specs/__init__.py` for ease of import) that returns the directory path to a `~/.torchx` (aka dot-directory) for torchx.

This directory is to be used by torchx to store temporary output such as workspace builds.

Usage:

```
from torchx.specs import TORCHX_HOME

outdir = TORCHX_HOME("out")
```

Reviewed By: hstonec

Differential Revision: D83575373
---
 torchx/specs/__init__.py      |  2 ++
 torchx/specs/api.py           | 28 ++++++++++++++++++++++++++++
 torchx/specs/test/api_test.py | 31 +++++++++++++++++++++++++++++++
 3 files changed, 61 insertions(+)

diff --git a/torchx/specs/__init__.py b/torchx/specs/__init__.py
index efb5df037..542cc8e6d 100644
--- a/torchx/specs/__init__.py
+++ b/torchx/specs/__init__.py
@@ -41,6 +41,7 @@
     RoleStatus,
     runopt,
     runopts,
+    TORCHX_HOME,
     UnknownAppException,
     UnknownSchedulerException,
     VolumeMount,
@@ -53,6 +54,7 @@
 
 GiB: int = 1024
 
+
 ResourceFactory = Callable[[], Resource]
 
 AWS_NAMED_RESOURCES: Mapping[str, ResourceFactory] = import_attr(
diff --git a/torchx/specs/api.py b/torchx/specs/api.py
index 02657aa4e..2d7907f55 100644
--- a/torchx/specs/api.py
+++ b/torchx/specs/api.py
@@ -11,6 +11,8 @@
 import inspect
 import json
 import logging as logger
+import os
+import pathlib
 import re
 import typing
 from dataclasses import asdict, dataclass, field
@@ -66,6 +68,32 @@
 RESET = "\033[0m"
 
 
+def TORCHX_HOME(*subdir_paths: str) -> pathlib.Path:
+    """
+    Path to the "dot-directory" for torchx.
+    Defaults to `~/.torchx` and is overridable via the `TORCHX_HOME` environment variable.
+
+    Usage:
+
+    .. doc-test::
+
+        from pathlib import Path
+        from torchx.specs import TORCHX_HOME
+
+        assert TORCHX_HOME() == Path.home() / ".torchx"
+        assert TORCHX_HOME("conda-pack-out") ==  Path.home() / ".torchx" / "conda-pack-out"
+    ```
+    """
+
+    default_dir = str(pathlib.Path.home() / ".torchx")
+    torchx_home = pathlib.Path(os.getenv("TORCHX_HOME", default_dir))
+
+    torchx_home = torchx_home / os.path.sep.join(subdir_paths)
+    torchx_home.mkdir(parents=True, exist_ok=True)
+
+    return torchx_home
+
+
 # ========================================
 # ==== Distributed AppDef API =======
 # ========================================
diff --git a/torchx/specs/test/api_test.py b/torchx/specs/test/api_test.py
index 2490e89b2..c99d6f700 100644
--- a/torchx/specs/test/api_test.py
+++ b/torchx/specs/test/api_test.py
@@ -10,10 +10,13 @@
 import asyncio
 import concurrent
 import os
+import tempfile
 import time
 import unittest
 from dataclasses import asdict
+from pathlib import Path
 from typing import Dict, List, Mapping, Tuple, Union
+from unittest import mock
 from unittest.mock import MagicMock
 
 import torchx.specs.named_resources_aws as named_resources_aws
@@ -40,9 +43,37 @@
     RoleStatus,
     runopt,
     runopts,
+    TORCHX_HOME,
 )
 
 
+class TorchXHomeTest(unittest.TestCase):
+    # guard against TORCHX_HOME set outside the test
+    @mock.patch.dict(os.environ, {}, clear=True)
+    def test_TORCHX_HOME_default(self) -> None:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            user_home = Path(tmpdir) / "sally"
+            with mock.patch("pathlib.Path.home", return_value=user_home):
+                torchx_home = TORCHX_HOME()
+                self.assertEqual(torchx_home, user_home / ".torchx")
+                self.assertTrue(torchx_home.exists())
+
+    def test_TORCHX_HOME_override(self) -> None:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            override_torchx_home = Path(tmpdir) / "test" / ".torchx"
+            with mock.patch.dict(
+                os.environ, {"TORCHX_HOME": str(override_torchx_home)}
+            ):
+                torchx_home = TORCHX_HOME()
+                conda_pack_out = TORCHX_HOME("conda-pack", "out")
+
+                self.assertEqual(override_torchx_home, torchx_home)
+                self.assertEqual(torchx_home / "conda-pack" / "out", conda_pack_out)
+
+                self.assertTrue(torchx_home.is_dir())
+                self.assertTrue(conda_pack_out.is_dir())
+
+
 class AppDryRunInfoTest(unittest.TestCase):
     def test_repr(self) -> None:
         request_mock = MagicMock()