Getting rid of profiling and adding route for index (it sees the root of the...

Getting rid of profiling and adding route for index (it sees the root of the project as the megatron subdirectory for some reason)

megatron/api_server.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import torch
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, current_app
 from flask_restful import Resource, Api
 
 from megatron import get_args
@@ -103,9 +103,13 @@ class MegatronGenerate(Resource):
         return jsonify({"sentences": resp_sentences})
     
 
+def index():
+    return current_app.send_static_file('index.html')
+
 class MegatronServer(object):
     def __init__(self, model):
         self.app = Flask(__name__)
+        self.app.add_url_rule('/', 'index', index)
         api = Api(self.app)
         api.add_resource(MegatronGenerate, '/generate', resource_class_args=[model])
 

run_api_server_530B.sh

@@ -4,7 +4,7 @@ DATA_PATH="/home/universal-lm-data.cosmos549/scratch/mshoeybi/data/gpt2"
 VOCAB_FILE="${DATA_PATH}/bpe/gpt2-vocab.json"
 MERGE_FILE="${DATA_PATH}/bpe/gpt2-merges.txt"
 RUN_CMD=(
-python -m cProfile -s cumtime tools/run_api_server.py 
+python tools/run_api_server.py 
        --tensor-model-parallel-size 16 
        --pipeline-model-parallel-size 3 
        --num-layers 105