feat: export resources information to the environment variables (#478)

Author: njzjz

doc/conf.py

@@ -21,7 +21,7 @@
 
 project = "DPDispatcher"
 copyright = "2020-%d, Deep Modeling" % date.today().year
-author = "Deep Modeling"
+author = "DeepModeling"
 
 
 # -- General configuration ---------------------------------------------------
@@ -97,3 +97,7 @@ def setup(app):
     "python": ("https://docs.python.org/", None),
     "dargs": ("https://docs.deepmodeling.com/projects/dargs/en/latest/", None),
 }
+
+myst_enable_extensions = [
+    "colon_fence",
+]

doc/context.md

@@ -9,7 +9,7 @@ One needs to set {dargs:argument}`context_type <machine/context_type>` to one of
 
 `LazyLocal` directly runs jobs in the local server and local directory.
 
-Since [`bash -l`](https://www.gnu.org/software/bash/manual/bash.html#Invoking-Bash) is used in the shebang line of the submission scripts, the [login shell startup files](https://www.gnu.org/software/bash/manual/bash.html#Invoking-Bash) will be executed, potentially overriding the current environmental variables. Therefore, it's advisable to explicitly set the environmental variables using {dargs:argument}`envs <resources/envs>` or {dargs:argument}`source_list <resources/source_list>`.
+Since [`bash -l`](https://www.gnu.org/software/bash/manual/bash.html#Invoking-Bash) is used in the shebang line of the submission scripts, the [login shell startup files](https://www.gnu.org/software/bash/manual/bash.html#Invoking-Bash) will be executed, potentially overriding the current environment variables. Therefore, it's advisable to explicitly set the environment variables using {dargs:argument}`envs <resources/envs>` or {dargs:argument}`source_list <resources/source_list>`.
 
 ## Local
 
@@ -18,7 +18,7 @@ Since [`bash -l`](https://www.gnu.org/software/bash/manual/bash.html#Invoking-Ba
 `Local` runs jobs in the local server, but in a different directory.
 Files will be copied to the remote directory before jobs start and copied back after jobs finish.
 
-Since [`bash -l`](https://www.gnu.org/software/bash/manual/bash.html#Invoking-Bash) is used in the shebang line of the submission scripts, the [login shell startup files](https://www.gnu.org/software/bash/manual/bash.html#Invoking-Bash) will be executed, potentially overriding the current environmental variables. Therefore, it's advisable to explicitly set the environmental variables using {dargs:argument}`envs <resources/envs>` or {dargs:argument}`source_list <resources/source_list>`.
+Since [`bash -l`](https://www.gnu.org/software/bash/manual/bash.html#Invoking-Bash) is used in the shebang line of the submission scripts, the [login shell startup files](https://www.gnu.org/software/bash/manual/bash.html#Invoking-Bash) will be executed, potentially overriding the current environment variables. Therefore, it's advisable to explicitly set the environment variables using {dargs:argument}`envs <resources/envs>` or {dargs:argument}`source_list <resources/source_list>`.
 
 ## SSH
 

doc/env.md

@@ -0,0 +1,35 @@
+# Environment variables
+
+When launching a job, dpdispatcher sets the following environment variables according to the resources, in addition to user-defined environment variables:
+
+:::{envvar} DPDISPATCHER_NUMBER_NODE
+
+The number of nodes required for each job.
+
+:::
+
+:::{envvar} DPDISPATCHER_CPU_PER_NODE
+
+CPU numbers of each node assigned to each job.
+
+:::
+
+:::{envvar} DPDISPATCHER_GPU_PER_NODE
+
+GPU numbers of each node assigned to each job.
+
+:::
+
+:::{envvar} DPDISPATCHER_QUEUE_NAME
+
+The queue name of batch job scheduler system.
+
+:::
+
+:::{envvar} DPDISPATCHER_GROUP_SIZE
+
+The number of tasks in a job. 0 means infinity.
+
+:::
+
+These environment variables can be used in the {dargs:argument}`command <task/command>`, for example, `mpirun -n ${DPDISPATCHER_CPU_PER_NODE} xx.run`.

doc/index.rst

@@ -22,6 +22,7 @@ DPDispatcher will monitor (poke) until these jobs finish and download the result
    machine
    resources
    task
+   env
    run
    cli
    api/api

dpdispatcher/machine.py

@@ -265,6 +265,15 @@ def gen_script_env(self, job):
 
         export_envs_part = ""
         envs = job.resources.envs
+        envs = {
+            # export resources information to the environment variables
+            "DPDISPATCHER_NUMBER_NODE": job.resources.number_node,
+            "DPDISPATCHER_CPU_PER_NODE": job.resources.cpu_per_node,
+            "DPDISPATCHER_GPU_PER_NODE": job.resources.gpu_per_node,
+            "DPDISPATCHER_QUEUE_NAME": job.resources.queue_name,
+            "DPDISPATCHER_GROUP_SIZE": job.resources.group_size,
+            **envs,
+        }
         for k, v in envs.items():
             if isinstance(v, list):
                 for each_value in v:

dpdispatcher/submission.py

@@ -1122,9 +1122,9 @@ def load_from_dict(cls, resources_dict):
 
     @staticmethod
     def arginfo(detail_kwargs=True):
-        doc_number_node = "The number of node need for each `job`"
-        doc_cpu_per_node = "cpu numbers of each node assigned to each job."
-        doc_gpu_per_node = "gpu numbers of each node assigned to each job."
+        doc_number_node = "The number of nodes required for each `job`."
+        doc_cpu_per_node = "CPU numbers of each node assigned to each job."
+        doc_gpu_per_node = "GPU numbers of each node assigned to each job."
         doc_queue_name = "The queue name of batch job scheduler system."
         doc_group_size = "The number of `tasks` in a `job`. 0 means infinity."
         doc_custom_flags = "The extra lines pass to job submitting script header"

tests/test_lsf_script_generation.py

@@ -111,6 +111,11 @@ def test_shell_trival(self):
             source /data/home/ypliu/scripts/avail_gpu.sh
             source /data/home/ypliu/dprun/tf_envs.sh
 
+            export DPDISPATCHER_NUMBER_NODE=1
+            export DPDISPATCHER_CPU_PER_NODE=4
+            export DPDISPATCHER_GPU_PER_NODE=1
+            export DPDISPATCHER_QUEUE_NAME=gpu
+            export DPDISPATCHER_GROUP_SIZE=4
             export DP_DISPATCHER_EXPORT=test_foo_bar_baz
 
             echo 'The summer you were there.'

tests/test_run_submission.py

@@ -65,7 +65,14 @@ def test_run_submission(self):
         task_list = []
         for ii in range(4):
             task = Task(
-                command=f"echo dpdispatcher_unittest_{ii}",
+                command=(
+                    f"[[ $DPDISPATCHER_NUMBER_NODE -eq {self.resources_dict['number_node']} ]] &&"
+                    f"[[ $DPDISPATCHER_CPU_PER_NODE -eq {self.resources_dict['cpu_per_node']} ]] &&"
+                    f"[[ $DPDISPATCHER_GPU_PER_NODE -eq {self.resources_dict['gpu_per_node']} ]] &&"
+                    f"[[ $DPDISPATCHER_QUEUE_NAME == {self.resources_dict['queue_name']} ]] &&"
+                    f"[[ $DPDISPATCHER_GROUP_SIZE -eq {self.resources_dict['group_size']} ]] &&"
+                    f"echo dpdispatcher_unittest_{ii}"
+                ),
                 task_work_path="./",
                 forward_files=[],
                 backward_files=[f"out{ii}.txt"],