Repository: deepmodeling/dpdispatcher

Issue:
Set useful environment variables
It is useful to set some environment variables like `DPDISPATCHER_NCPUS` and `DPDISPATCHER_NGPUS`. One can use them in the command, without needing to update the command when updating the resources.

Discussion:

Solution patch:
diff --git a/doc/conf.py b/doc/conf.py
index 3e395d4..d987ddb 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -21,7 +21,7 @@ from datetime import date
 
 project = "DPDispatcher"
 copyright = "2020-%d, Deep Modeling" % date.today().year
-author = "Deep Modeling"
+author = "DeepModeling"
 
 
 # -- General configuration ---------------------------------------------------
@@ -97,3 +97,7 @@ intersphinx_mapping = {
     "python": ("https://docs.python.org/", None),
     "dargs": ("https://docs.deepmodeling.com/projects/dargs/en/latest/", None),
 }
+
+myst_enable_extensions = [
+    "colon_fence",
+]
diff --git a/doc/context.md b/doc/context.md
index 884616c..c46a08b 100644
--- a/doc/context.md
+++ b/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
 
diff --git a/doc/env.md b/doc/env.md
new file mode 100644
index 0000000..6048217
--- /dev/null
+++ b/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`.
diff --git a/doc/index.rst b/doc/index.rst
index e498417..b614ff0 100644
--- a/doc/index.rst
+++ b/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
diff --git a/dpdispatcher/machine.py b/dpdispatcher/machine.py
index eb69130..352eade 100644
--- a/dpdispatcher/machine.py
+++ b/dpdispatcher/machine.py
@@ -265,6 +265,15 @@ class Machine(metaclass=ABCMeta):
 
         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:
diff --git a/dpdispatcher/submission.py b/dpdispatcher/submission.py
index 9c5c73c..4eb3836 100644
--- a/dpdispatcher/submission.py
+++ b/dpdispatcher/submission.py
@@ -1122,9 +1122,9 @@ class Resources:
 
     @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"


Test patch:
diff --git a/tests/test_lsf_script_generation.py b/tests/test_lsf_script_generation.py
index 48eb41d..aa1cb92 100755
--- a/tests/test_lsf_script_generation.py
+++ b/tests/test_lsf_script_generation.py
@@ -111,6 +111,11 @@ class TestLSFScriptGeneration(unittest.TestCase):
             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.'
diff --git a/tests/test_run_submission.py b/tests/test_run_submission.py
index 0831afe..51e9ec2 100644
--- a/tests/test_run_submission.py
+++ b/tests/test_run_submission.py
@@ -65,7 +65,14 @@ class RunSubmission:
         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"],