slurmray 3.5.5__tar.gz → 3.5.7__tar.gz

{slurmray-3.5.5 → slurmray-3.5.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: slurmray
-Version: 3.5.5
+Version: 3.5.7
 Summary: SlurmRay is a module for effortlessly distributing tasks on a Slurm cluster using the Ray library. 
 Home-page: https://henri-jamet.vercel.app/
 License: Apache License

{slurmray-3.5.5 → slurmray-3.5.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "slurmray"
-version = "3.5.5"
+version = "3.5.7"
 description = "SlurmRay is a module for effortlessly distributing tasks on a Slurm cluster using the Ray library. "
 authors = ["Henri Jamet <henri.jamet@unil.ch>"]
 license = "Apache License"

{slurmray-3.5.5 → slurmray-3.5.7}/slurmray/RayLauncher.py

@@ -25,6 +25,7 @@ class RayLauncher:
         use_gpu: bool = False,
         memory: int = 64,
         max_running_time: int = 60,
+        runtime_env: dict = {"env_vars": {}},
         server_run: bool = True,
         server_ssh: str = "curnagl.dcsr.unil.ch",
         server_username: str = "hjamet",
@@ -41,6 +42,7 @@ class RayLauncher:
             use_gpu (bool, optional): Use GPU or not. Defaults to False.
             memory (int, optional): Amount of RAM to use per node in GigaBytes. Defaults to 64.
             max_running_time (int, optional): Maximum running time of the job in minutes. Defaults to 60.
+            runtime_env (dict, optional): Environment variables to share between all the workers. Can be useful for issues like https://github.com/ray-project/ray/issues/418. Default to empty.
             server_run (bool, optional): If you run the launcher from your local machine, you can use this parameter to execute your function using online cluster ressources. Defaults to True.
             server_ssh (str, optional): If `server_run` is set to true, the addess of the **SLURM** server to use.
             server_username (str, optional): If `server_run` is set to true, the username with which you wish to connect.
@@ -54,6 +56,7 @@ class RayLauncher:
         self.use_gpu = use_gpu
         self.memory = memory
         self.max_running_time = max_running_time
+        self.runtime_env = runtime_env
         self.server_run = server_run
         self.server_ssh = server_ssh
         self.server_username = server_username
@@ -196,7 +199,7 @@ class RayLauncher:
         text = text.replace("{{PROJECT_PATH}}", f'"{self.project_path}"')
         local_mode = ""
         if self.cluster or self.server_run:
-            "\n\taddress='auto',\n\tinclude_dashboard=True,\n\tdashboard_host='0.0.0.0',\n\tdashboard_port=8888,\n"
+            f"\n\taddress='auto',\n\tinclude_dashboard=True,\n\tdashboard_host='0.0.0.0',\n\tdashboard_port=8888,\nruntime_env = {self.runtime_env},\n"
         text = text.replace(
             "{{LOCAL_MODE}}",
             local_mode,
@@ -434,7 +437,8 @@ class RayLauncher:
             lines.append("slurmray --pre\n")
             # Solve torch buf (https://github.com/pytorch/pytorch/issues/111469)
             if "torchaudio\n" or "torchvision\n" in lines:
-                lines.append("torch==2.0.1\n")
+                lines.append("torch==2.1.1\n")
+                lines.append("--index-url https://download.pytorch.org/whl/cu121\n")
 
         with open(f"{self.project_path}/requirements.txt", "w") as file:
             file.writelines(lines)
@@ -540,6 +544,7 @@ if __name__ == "__main__":
         use_gpu=True,
         memory=8,
         max_running_time=5,
+        runtime_env={"env_vars": {"NCCL_SOCKET_IFNAME": "eno1"}},
         server_run=True,
         server_ssh="curnagl.dcsr.unil.ch",
         server_username="hjamet",