allow arbitrary expressions as job parameters (#247)

* allow arbitrary expressions as job parameters * restore EnvironmentParameterValues * fix UI type Co-authored-by: Piyush Jain <[email protected]> Co-authored-by: Piyush Jain <[email protected]>
jupyter-server · Nov 1, 2022 · 0b6c21b · 0b6c21b
1 parent b15f9fe
commit 0b6c21b
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 18 deletions.
diff --git a/jupyter_scheduler/models.py b/jupyter_scheduler/models.py
@@ -5,8 +5,7 @@
 from pydantic import BaseModel, root_validator
 
 Tags = List[str]
-ParameterValues = Union[int, str, float, bool]
-EnvironmentParameterValues = Union[int, str, float, bool]
+EnvironmentParameterValues = Union[int, float, bool, str]
 
 EMAIL_RE = ""
 SCHEDULE_RE = ""
@@ -80,7 +79,7 @@ class CreateJob(BaseModel):
  output_formats: Optional[List[str]] = None
  idempotency_token: Optional[str] = None
  job_definition_id: Optional[str] = None
- parameters: Optional[Dict[str, ParameterValues]] = None
+ parameters: Optional[Dict[str, str]] = None
  tags: Optional[Tags] = None
  name: Optional[str] = None
  output_filename_template: Optional[str] = OUTPUT_FILENAME_TEMPLATE
@@ -130,7 +129,7 @@ class DescribeJob(BaseModel):
  output_formats: Optional[List[str]] = None
  idempotency_token: Optional[str] = None
  job_definition_id: Optional[str] = None
- parameters: Optional[Dict[str, ParameterValues]] = None
+ parameters: Optional[Dict[str, str]] = None
  tags: Optional[Tags] = None
  name: Optional[str] = None
  output_filename_template: Optional[str] = OUTPUT_FILENAME_TEMPLATE
@@ -202,7 +201,7 @@ class CreateJobDefinition(BaseModel):
  runtime_environment_name: str
  runtime_environment_parameters: Optional[Dict[str, EnvironmentParameterValues]]
  output_formats: Optional[List[str]] = None
- parameters: Optional[Dict[str, ParameterValues]] = None
+ parameters: Optional[Dict[str, str]] = None
  tags: Optional[Tags] = None
  name: Optional[str] = None
  output_filename_template: Optional[str] = OUTPUT_FILENAME_TEMPLATE
@@ -223,7 +222,7 @@ class DescribeJobDefinition(BaseModel):
  runtime_environment_name: str
  runtime_environment_parameters: Optional[Dict[str, EnvironmentParameterValues]]
  output_formats: Optional[List[str]] = None
- parameters: Optional[Dict[str, ParameterValues]] = None
+ parameters: Optional[Dict[str, str]] = None
  tags: Optional[Tags] = None
  name: Optional[str] = None
  output_filename_template: Optional[str] = OUTPUT_FILENAME_TEMPLATE
@@ -243,7 +242,7 @@ class UpdateJobDefinition(BaseModel):
  runtime_environment_name: Optional[str]
  runtime_environment_parameters: Optional[Dict[str, EnvironmentParameterValues]]
  output_formats: Optional[List[str]] = None
- parameters: Optional[Dict[str, ParameterValues]] = None
+ parameters: Optional[Dict[str, str]] = None
  tags: Optional[Tags] = None
  name: Optional[str] = None
  url: Optional[str] = None
@@ -270,7 +269,7 @@ class ListJobDefinitionsResponse(BaseModel):
 
 
 class CreateJobFromDefinition(BaseModel):
- parameters: Optional[Dict[str, ParameterValues]] = None
+ parameters: Optional[Dict[str, str]] = None
 
 
 class JobFeature(str, Enum):

diff --git a/jupyter_scheduler/parameterize.py b/jupyter_scheduler/parameterize.py
@@ -2,18 +2,14 @@
 
 from nbformat import NotebookNode, v4
 
-from jupyter_scheduler.models import ParameterValues
 from jupyter_scheduler.utils import find_cell_index_with_tag
 
 
-def add_parameters(nb: NotebookNode, parameters: Dict[str, ParameterValues]) -> NotebookNode:
+def add_parameters(nb: NotebookNode, parameters: Dict[str, str]) -> NotebookNode:
  content = []
 
  for key, value in parameters.items():
- if type(value) == str:
- content.append(f"{key} = '{value}'")
- else:
- content.append(f"{key} = {value}")
+ content.append(f"{key} = {value}")
 
  new_cell = v4.new_code_cell(source="\n".join(content))
  new_cell.metadata["tags"] = ["injected-parameters"]

diff --git a/src/handler.ts b/src/handler.ts
@@ -345,10 +345,8 @@ export namespace SchedulerService {
 }
 
 export namespace Scheduler {
- export type RuntimeEnvironmentParameters = {
- [key: string]: number | string | boolean;
- };
- export type Parameters = { [key: string]: number | string | boolean };
+ export type RuntimeEnvironmentParameters = Record<string, number | string | boolean>;
+ export type Parameters = Record<string, string>;
 
  export interface ICreateJobDefinition {
  input_uri: string;