From 2e1af6fa27fbb9aa7e3f5e73c751db2b1a863523 Mon Sep 17 00:00:00 2001
From: Jake Schmidt <jacob.schmidt@noetik.ai>
Date: Thu, 2 Nov 2023 10:17:10 -0600
Subject: [PATCH] add mount args

---
 torchx/components/utils.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/torchx/components/utils.py b/torchx/components/utils.py
index f8502310b..a6fb13cbc 100644
--- a/torchx/components/utils.py
+++ b/torchx/components/utils.py
@@ -135,6 +135,7 @@ def python(
     memMB: int = 1024,
     h: Optional[str] = None,
     num_replicas: int = 1,
+    mounts: Optional[List[str]] = None,
 ) -> specs.AppDef:
     """
     Runs ``python`` with the specified module, command or script on the specified
@@ -157,6 +158,8 @@ def python(
         memMB: cpu memory in MB per replica
         h: a registered named resource (if specified takes precedence over cpu, gpu, memMB)
         num_replicas: number of copies to run (each on its own container)
+        mounts: mounts to mount into the worker environment/container (ex. type=<bind/volume>,src=/host,dst=/job[,readonly]).
+                See scheduler documentation for more info.
     :return:
     """
     if sum([m is not None, c is not None, script is not None]) != 1:
@@ -184,6 +187,7 @@ def python(
                 resource=specs.resource(cpu=cpu, gpu=gpu, memMB=memMB, h=h),
                 args=[*cmd, *args],
                 env={"HYDRA_MAIN_MODULE": m} if m else {},
+                mounts=specs.parse_mounts(mounts) if mounts else [],
             )
         ],
     )
@@ -227,7 +231,7 @@ def binary(
     )
 
 
-def copy(src: str, dst: str, image: str = torchx.IMAGE) -> specs.AppDef:
+def copy(src: str, dst: str, image: str = torchx.IMAGE, mounts: Optional[List[str]] = None,) -> specs.AppDef:
     """
     copy copies the file from src to dst. src and dst can be any valid fsspec
     url.
@@ -238,6 +242,8 @@ def copy(src: str, dst: str, image: str = torchx.IMAGE) -> specs.AppDef:
         src: the source fsspec file location
         dst: the destination fsspec file location
         image: the image that contains the copy app
+        mounts: mounts to mount into the worker environment/container (ex. type=<bind/volume>,src=/host,dst=/job[,readonly]).
+                See scheduler documentation for more info.
     """
 
     return specs.AppDef(
@@ -256,6 +262,7 @@ def copy(src: str, dst: str, image: str = torchx.IMAGE) -> specs.AppDef:
                     dst,
                 ],
                 resource=specs.Resource(cpu=1, gpu=0, memMB=1024),
+                mounts=specs.parse_mounts(mounts) if mounts else [],
             ),
         ],
     )