[Bugfix] And PYTHONPATH in server launch. (dmlc#3352)

* put PYTHONPATH in server launch * remove prints Co-authored-by: xiangsx <[email protected]>
ssghost · Sep 14, 2021 · ad61a9a · ad61a9a
1 parent f4c79f7
commit ad61a9a
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/tools/launch.py b/tools/launch.py
@@ -271,6 +271,7 @@ def construct_dgl_server_env_vars(
     ip_config: str,
     num_servers: int,
     graph_format: str,
+    pythonpath: Optional[str] = "",
 ) -> str:
     """Constructs the DGL server-specific env vars string that are required for DGL code to behave in the correct
     server role.
@@ -286,6 +287,7 @@ def construct_dgl_server_env_vars(
             Relative path to workspace.
         num_servers:
         graph_format:
+        pythonpath: Optional. If given, this will pass this as PYTHONPATH.
 
     Returns:
         server_env_vars: The server-specific env-vars in a string format, friendly for CLI execution.
@@ -300,7 +302,11 @@ def construct_dgl_server_env_vars(
         "DGL_IP_CONFIG={DGL_IP_CONFIG} "
         "DGL_NUM_SERVER={DGL_NUM_SERVER} "
         "DGL_GRAPH_FORMAT={DGL_GRAPH_FORMAT} "
+        "{suffix_optional_envvars}"
     )
+    suffix_optional_envvars = ""
+    if pythonpath:
+        suffix_optional_envvars += f"PYTHONPATH={pythonpath} "
     return server_env_vars_template.format(
         DGL_ROLE="server",
         DGL_NUM_SAMPLER=num_samplers,
@@ -310,6 +316,7 @@ def construct_dgl_server_env_vars(
         DGL_IP_CONFIG=ip_config,
         DGL_NUM_SERVER=num_servers,
         DGL_GRAPH_FORMAT=graph_format,
+        suffix_optional_envvars=suffix_optional_envvars,
     )
 
 
@@ -440,6 +447,7 @@ def submit_jobs(args, udf_command):
         ip_config=args.ip_config,
         num_servers=args.num_servers,
         graph_format=args.graph_format,
+        pythonpath=os.environ.get("PYTHONPATH", ""),
     )
     for i in range(len(hosts) * server_count_per_machine):
         ip, _ = hosts[int(i / server_count_per_machine)]