Skip to content

Commit bed3350

Browse files
authored
321 support getting spark config from be (#325)
* Refactor Notebook.js, SparkAppConfigModel.js, Config.js, spark_app.py, and spark_app.py to support saving Spark configuration to the database. * Refactor SparkAppConfigModel, Notebook.js, and spark_app.py to update Spark app configuration handling * Fix missing newline at end of file in test_spark_app_route.py * Refactor test_spark_app_route.py to use spark_app_blueprint in server/tests/routes/test_spark_app_route.py * Refactor test_spark_app_route.py to use POST method instead of GET in server/tests/routes/test_spark_app_route.py * Refactor test_spark_app_route.py to use POST method instead of GET in server/tests/routes/test_spark_app_route.py * Refactor test_spark_app_route.py to use POST method instead of GET in server/tests/routes/test_spark_app_route.py * Refactor test_spark_app_route.py to use POST method instead of GET in server/tests/routes/test_spark_app_route.py * Refactor test_spark_app_route.py to use POST method instead of GET in server/tests/routes/test_spark_app_route.py * Refactor spark_app.py and test_spark_app_route.py to add spark_app endpoint and update test case in server/tests/routes/test_spark_app_route.py * Refactor SparkAppConfigModel, Notebook.js, Config.js, and spark_app.py to update Spark app configuration handling
1 parent 2f15c64 commit bed3350

File tree

8 files changed

+314
-110
lines changed

8 files changed

+314
-110
lines changed

docker/postgres/init.sql

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ GRANT ALL PRIVILEGES ON SEQUENCE directories_id_seq TO server;
6464
GRANT ALL PRIVILEGES ON TABLE spark_apps TO server;
6565

6666
GRANT ALL PRIVILEGES ON TABLE spark_app_config TO server;
67+
GRANT ALL PRIVILEGES ON SEQUENCE spark_app_config_id_seq TO server;
6768

6869
-- Add some initial data
6970
-- user_0 -12345A

server/app/routes/spark_app.py

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
from flask import Blueprint, jsonify, request
22
from app.services.spark_app import SparkApp
3+
from flask_jwt_extended import jwt_required
4+
from app.auth.auth import identify_user
35
import logging
46

57
spark_app_blueprint = Blueprint('spark_app', __name__)
@@ -12,13 +14,15 @@ def create_spark_app(spark_app_id):
1214
notebook_path = data.get('notebookPath', None)
1315
return SparkApp.create_spark_app(spark_app_id=spark_app_id, notebook_path=notebook_path)
1416

15-
@spark_app_blueprint.route('/spark_app/<path:notbook_id>/config', methods=['GET'])
16-
def get_spark_app_config(notbook_id):
17-
logging.info(f"Getting spark app config for notebook id: {notbook_id}")
18-
return SparkApp.get_spark_app_config_by_notebook_id(notbook_id)
17+
# @jwt_required()
18+
# @identify_user
19+
@spark_app_blueprint.route('/spark_app/<path:notbook_path>/config', methods=['GET'])
20+
def get_spark_app_config(notbook_path):
21+
logging.info(f"Getting spark app config for notebook path: {notbook_path}")
22+
return SparkApp.get_spark_app_config_by_notebook_path(notbook_path)
1923

20-
@spark_app_blueprint.route('/spark_app/<path:notbook_id>/config', methods=['POST'])
21-
def update_spark_app_config(notbook_id):
22-
logging.info(f"Updating spark app config for notebook id: {notbook_id}")
24+
@spark_app_blueprint.route('/spark_app/<path:notbook_path>/config', methods=['POST'])
25+
def update_spark_app_config(notbook_path):
26+
logging.info(f"Updating spark app config for notebook path: {notbook_path}")
2327
data = request.get_json()
24-
return SparkApp.update_spark_app_config_by_notebook_id(notbook_id, data)
28+
return SparkApp.update_spark_app_config_by_notebook_path(notbook_path, data)

server/app/services/spark_app.py

Lines changed: 35 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -42,14 +42,35 @@ def get_spark_app_by_id(spark_app_id: str = None):
4242
)
4343

4444
@staticmethod
45-
def get_spark_app_config_by_notebook_id(notebook_id: str = None):
46-
logger.info(f"Getting spark app config for notebook id: {notebook_id}")
45+
def get_spark_app_config_by_notebook_path(notbook_path: str = None):
46+
# Get notebook id from path
47+
notebook = NotebookModel.query.filter_by(path=notbook_path).first()
48+
notebook_id = notebook.id
49+
4750
# Get the spark app config
4851
spark_app_config = SparkAppConfigModel.query.filter_by(notebook_id=notebook_id).first()
4952

5053
if (spark_app_config is not None):
54+
spark_app_config_dict = spark_app_config.to_dict()
55+
spark_app_config_dict_transformed = {
56+
'spark.driver.memory': spark_app_config_dict['driver_memory'],
57+
'spark.driver.memoryOverhead': spark_app_config_dict['driver_memory_overhead'],
58+
'spark.driver.cores': spark_app_config_dict['driver_cores'],
59+
'spark.executor.memory': spark_app_config_dict['executor_memory'],
60+
'spark.executor.memoryOverhead': spark_app_config_dict['executor_memory_overhead'],
61+
'spark.executor.memoryFraction': spark_app_config_dict['executor_memory_fraction'],
62+
'spark.executor.cores': spark_app_config_dict['executor_cores'],
63+
'spark.executor.instances': spark_app_config_dict['executor_instances'],
64+
'spark.dynamicAllocation.enabled': spark_app_config_dict['dynamic_allocation_enabled'],
65+
'spark.dynamicAllocation.minExecutors': spark_app_config_dict['executor_instances_min'],
66+
'spark.dynamicAllocation.maxExecutors': spark_app_config_dict['executor_instances_max'],
67+
'spark.shuffle.service.enabled': spark_app_config_dict['shuffle_service_enabled'],
68+
'spark.executor.idleTimeout': spark_app_config_dict['executor_idle_timeout'],
69+
'spark.queue': spark_app_config_dict['queue']
70+
}
71+
5172
return Response(
52-
response=json.dumps(spark_app_config.to_dict()),
73+
response=json.dumps(spark_app_config_dict_transformed),
5374
status=200
5475
)
5576
else:
@@ -69,23 +90,25 @@ def get_spark_app_config_by_notebook_id(notebook_id: str = None):
6990
)
7091

7192
@staticmethod
72-
def update_spark_app_config_by_notebook_id(notebook_id: str = None, data: dict = None):
73-
logger.info(f"Updating spark app config for notebook id: {notebook_id} with data: {data}")
93+
def update_spark_app_config_by_notebook_path(notebook_path: str = None, data: dict = None):
94+
logger.info(f"Updating spark app config for notebook path: {notebook_path} with data: {data}")
7495

75-
if notebook_id is None:
76-
logger.error("Notebook id is None")
96+
if notebook_path is None:
97+
logger.error("Notebook path is None")
7798
return Response(
78-
response=json.dumps({'message': 'Notebook id is None'}),
99+
response=json.dumps({'message': 'Notebook path is None'}),
79100
status=404)
80101

81102
# Get the notebook id
82-
notebook = NotebookModel.query.filter_by(id=notebook_id).first()
103+
notebook = NotebookModel.query.filter_by(path=notebook_path).first()
83104
if notebook is None:
84105
logger.error("Notebook not found")
85106
return Response(
86107
response=json.dumps({'message': 'Notebook not found'}),
87108
status=404)
88109

110+
notebook_id = notebook.id
111+
89112
# Transform data
90113
transformed_data = {
91114
'driver_memory': data.get('spark.driver.memory', None),
@@ -97,8 +120,8 @@ def update_spark_app_config_by_notebook_id(notebook_id: str = None, data: dict =
97120
'executor_cores': data.get('spark.executor.cores', None),
98121
'executor_instances': data.get('spark.executor.instances', None),
99122
'dynamic_allocation_enabled': data.get('spark.dynamicAllocation.enabled', None),
100-
'executor_instances_min': data.get('spark.executor.instancesMin', None),
101-
'executor_instances_max': data.get('spark.executor.instancesMax', None),
123+
'executor_instances_min': data.get('spark.dynamicAllocation.minExecutors', None),
124+
'executor_instances_max': data.get('spark.dynamicAllocation.maxExecutors', None),
102125
'shuffle_service_enabled': data.get('spark.shuffle.service.enabled', None),
103126
'executor_idle_timeout': data.get('spark.executor.idleTimeout', None),
104127
'queue': data.get('spark.queue', None)
@@ -113,6 +136,7 @@ def update_spark_app_config_by_notebook_id(notebook_id: str = None, data: dict =
113136
setattr(config, key, value)
114137

115138
db.session.commit()
139+
db.session.refresh(config)
116140

117141
return Response(
118142
response=json.dumps({'message': 'Updated spark app config'}),
Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
# import unittest
2+
# import json
3+
# from flask_cors import CORS
4+
# from flask import g
5+
# from database import db
6+
# from run import create_app
7+
# from app.routes.spark_app import spark_app_blueprint
8+
# from app.routes.login import login_blueprint
9+
# from app.services.directory import Directory
10+
# from app.models.user import UserModel
11+
# from app.services.user import User
12+
# from app.models.spark_app import SparkAppModel
13+
# from app.models.notebook import NotebookModel
14+
15+
# class SparkAppRouteTestCase(unittest.TestCase):
16+
17+
# def setUp(self):
18+
# self.app = create_app()
19+
# self.app.register_blueprint(spark_app_blueprint)
20+
# self.app.register_blueprint(login_blueprint)
21+
# self.client = self.app.test_client()
22+
# with self.app.app_context():
23+
# db.create_all()
24+
# user = UserModel(name='test_user', email='test_email')
25+
# user.set_password('test_password')
26+
# db.session.add(user)
27+
# db.session.commit()
28+
29+
# def tearDown(self):
30+
# with self.app.app_context():
31+
# db.session.remove()
32+
# db.drop_all()
33+
34+
# def login_and_get_token(self):
35+
# with self.app.app_context():
36+
# response = self.client.post('/login', auth=('test_user', 'test_password'))
37+
# return json.loads(response.data)['access_token']
38+
39+
# # def test_create_spark_app(self):
40+
# # with self.app.app_context():
41+
# # # Create Notebook
42+
# # notebook = NotebookModel(name='Test Notebook', path='/path/to/notebook', user_id=1)
43+
# # db.session.add(notebook)
44+
# # db.session.commit()
45+
46+
# # # Create Spark App
47+
# # spark_app_id = 'app_0001'
48+
# # path = f'/spark-app/app_0001'
49+
50+
# # # data = {
51+
# # # 'notebookPath': notebook.path
52+
# # # }
53+
54+
# # # token = self.login_and_get_token()
55+
# # # headers = {
56+
# # # 'Authorization': f'Bearer {token}',
57+
# # # }
58+
59+
# # response = self.client.post(
60+
# # path,
61+
# # # headers=headers,
62+
# # # json=json.dumps(data),
63+
# # )
64+
65+
# # print(response.data)
66+
# # # self.assertEqual(response.status_code, 200)
67+
# # # self.assertEqual(json.loads(response.data)['spark_app_id'], spark_app_id)
68+
# # # self.assertEqual(json.loads(response.data)['notebook_id'], notebook.id)
69+
# # # self.assertEqual(json.loads(response.data)['user_id'], notebook.user_id)
70+
71+
# def test_get_spark_app_config_by_notebook_path(self):
72+
# with self.app.app_context():
73+
# token = self.login_and_get_token()
74+
# headers = {
75+
# 'Authorization': f'Bearer {token}',
76+
# }
77+
78+
# # response = self.client.get('/spark-app/path_to_notebook/config', headers=headers)
79+
# # print(response.data)
80+
81+
# response = self.client.get('/spark-app')
82+
# print(response.data)

server/tests/services/test_spark_app_service.py

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ def test_get_spark_by_id(self):
4949
spark_app_dict = json.loads(response.data)
5050
self.assertEqual(spark_app_dict['spark_app_id'], '1234')
5151

52-
def test_get_spark_app_config_by_notebook_id(self):
52+
def test_get_spark_app_config_by_notebook_path(self):
5353
with self.app.app_context():
5454
# Create User
5555
user_0 = UserModel(name='testuser0', email='[email protected]')
@@ -65,7 +65,7 @@ def test_get_spark_app_config_by_notebook_id(self):
6565
db.session.commit()
6666

6767
# Get spark app config by notebook path
68-
response = SparkApp.get_spark_app_config_by_notebook_id(notebook_0.id)
68+
response = SparkApp.get_spark_app_config_by_notebook_path('/path/to/notebook')
6969
spark_app_config_dict = json.loads(response.data)
7070

7171
self.assertEqual(spark_app_config_dict['spark.driver.memory'], '1g')
@@ -75,7 +75,7 @@ def test_get_spark_app_config_by_notebook_id(self):
7575
self.assertEqual(spark_app_config_dict['spark.executor.instances'], 1)
7676
self.assertEqual(spark_app_config_dict['spark.dynamicAllocation.enabled'], False)
7777

78-
def test_update_spark_app_config_by_notebook_id(self):
78+
def test_update_spark_app_config(self):
7979
with self.app.app_context():
8080
# Create User
8181
user_0 = UserModel(name='testuser0', email='[email protected]')
@@ -100,15 +100,15 @@ def test_update_spark_app_config_by_notebook_id(self):
100100
'spark.dynamicAllocation.enabled': True,
101101
}
102102

103-
response_0 = SparkApp.update_spark_app_config_by_notebook_id(None, data=data)
103+
response_0 = SparkApp.update_spark_app_config_by_notebook_path(None, data=data)
104104
self.assertEqual(response_0.status_code, 404)
105-
self.assertEqual(json.loads(response_0.data)['message'], 'Notebook id is None')
105+
self.assertEqual(json.loads(response_0.data)['message'], 'Notebook path is None')
106106

107-
response_1 = SparkApp.update_spark_app_config_by_notebook_id(999, data=data)
107+
response_1 = SparkApp.update_spark_app_config_by_notebook_path('path_not_found', data=data)
108108
self.assertEqual(response_1.status_code, 404)
109109
self.assertEqual(json.loads(response_1.data)['message'], 'Notebook not found')
110110

111-
response_2 = SparkApp.update_spark_app_config_by_notebook_id(notebook_0.id, data=data)
111+
response_2 = SparkApp.update_spark_app_config_by_notebook_path('/path/to/notebook', data=data)
112112
self.assertEqual(response_2.status_code, 200)
113113
self.assertEqual(json.loads(response_2.data)['message'], 'Updated spark app config')
114114

@@ -121,6 +121,8 @@ def test_update_spark_app_config_by_notebook_id(self):
121121
self.assertEqual(spark_app_config.executor_cores, 2)
122122
self.assertEqual(spark_app_config.executor_instances, 2)
123123
self.assertEqual(spark_app_config.dynamic_allocation_enabled, True)
124+
125+
124126

125127
def test_create_spark_app(self):
126128
with self.app.app_context():

webapp/src/components/notebook/Notebook.js

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -323,7 +323,10 @@ function Notebook({
323323
saveNotebook={handleUpdateNotebook}
324324
deleteNotebook={handleDeleteNotebook}
325325
/> : contentType === ContentType.Config ?
326-
<Config /> :
326+
<Config
327+
notebook={notebook}
328+
notebookPath={notebook.path}
329+
/> :
327330
<Runs
328331
notebook={notebook}
329332
contentType={contentType}

0 commit comments

Comments
 (0)