include pythonpath and local dir in launch env (#93)

Co-authored-by: Shaden Smith <ShadenTSmith@gmail.com>

deepspeed/pt/deepspeed_run.py

@@ -15,6 +15,7 @@ import collections
 from copy import deepcopy
 
 DLTS_HOSTFILE = "/job/hostfile"
+EXPORT_ENVS = ["NCCL", "PYTHONPATH"]
 
 
 def parse_args(args=None):
@@ -305,13 +306,15 @@ def main(args=None):
         num_gpus_per_node = None
 
         curr_path = os.path.abspath('.')
+        env['PYTHONPATH'] = curr_path + ":" + env['PYTHONPATH']
 
-        nccl_export = ""
-        for nccl_var in filter(lambda x: "NCCL_" in x, env.keys()):
-            nccl_export += "export {}={}; ".format(nccl_var, env[nccl_var])
+        exports = ""
+        for var in env.keys():
+            if any(map(lambda name: name in var, EXPORT_ENVS)):
+                exports += "export {}={}; ".format(var, env[var])
 
         deepspeed_launch = [
-            nccl_export,
+            exports,
             "cd {};".format(curr_path),
             sys.executable,
             "-u",