Enable and disable Rules

.gitignore

@@ -15,3 +15,4 @@ MANIFEST
 manual_dist/**
 sagemaker_run_notebook/cloudformation.yml
 docs/build/**
+.DS_Store

container/build_and_push.sh

@@ -50,8 +50,14 @@ region=$(aws configure get region)
 region=${region:-us-west-2}
 echo "Region ${region}"
 
+if [[ $region == cn-* ]]
+then
+    amazon_extension="amazonaws.com.cn"
+else
+    amazon_extension="amazonaws.com"
+fi
 
-fullname="${account}.dkr.ecr.${region}.amazonaws.com/${image}:latest"
+fullname="${account}.dkr.ecr.${region}.${amazon_extension}/${image}:latest"
 
 # If the repository doesn't exist in ECR, create it.
 
@@ -73,7 +79,7 @@ then
   # Get the login command from ECR and execute it directly
   $(aws ecr get-login --region ${region} --no-include-email)
 else
-  aws ecr get-login-password --region ${region} | docker login --username AWS --password-stdin ${account}.dkr.ecr.${region}.amazonaws.com
+  aws ecr get-login-password --region ${region} | docker login --username AWS --password-stdin ${account}.dkr.ecr.${region}.${amazon_extension}
 fi
 
 # Build the docker image locally with the image name and then push it to ECR

labextension/src/components/RuleList.tsx

@@ -60,7 +60,10 @@ export class RulesList extends React.Component<RuleListProps, RuleListState> {
       rule.image,
       rule.instance,
       rule.role,
-      rule.state,
+      // rule.state,
+      <a key="state" onClick={this.toggleState(rule.name, rule.state)} className={tableLinkClass}>
+        {rule.state}
+      </a>,
       <a key={rule.name} onClick={this.deleteRule(rule.name)} className={tableLinkClass}>
         Delete
       </a>,
@@ -96,6 +99,40 @@ export class RulesList extends React.Component<RuleListProps, RuleListState> {
     return <SimpleTablePage title="Schedule and Event Rules">{content}</SimpleTablePage>;
   }
 
+  private toggleState(rule: string, state: string) {
+    return async (): Promise<boolean> => {
+      console.log(`toggle rule ${rule} state from ${state}`);
+      if (!['DISABLED', 'ENABLED'].includes(state)) {
+        console.log(`unknown state: ${state}`);
+        return;
+      }
+      const settings = ServerConnection.makeSettings();
+      const response = await ServerConnection.makeRequest(
+        URLExt.join(settings.baseUrl, 'sagemaker-scheduler', 'schedule', rule),
+        { method: 'PATCH', body: JSON.stringify({ action: state === 'DISABLED' ? 'enable' : 'disable' }) },
+        settings,
+      );
+
+      if (!response.ok) {
+        const error = (await response.json()) as ErrorResponse;
+        let errorMessage: string;
+        if (error.error) {
+          errorMessage = error.error.message;
+        } else {
+          errorMessage = JSON.stringify(error);
+        }
+        showDialog({
+          title: 'Error toggling schedule state',
+          body: <p>{errorMessage}</p>,
+          buttons: [Dialog.okButton({ label: 'Close' })],
+        });
+        return;
+      }
+      this.props.model.refresh();
+      return true;
+    };
+  }
+
   private deleteRule(rule: string) {
     return async (): Promise<boolean> => {
       const deleteBtn = Dialog.warnButton({ label: 'Delete' });

sagemaker_run_notebook/__init__.py

@@ -35,6 +35,8 @@
     create_lambda_role,
     schedule,
     unschedule,
+    enable_schedule,
+    disable_schedule,
     describe_schedule,
     describe_schedules,
     list_schedules,

sagemaker_run_notebook/run_notebook.py

@@ -982,7 +982,7 @@ def unschedule(rule_name, session=None):
     prefixed_rule_name = RULE_PREFIX + rule_name
 
     session = ensure_session(session)
-    events = boto3.client("events")
+    events = session.client("events")
     lambda_ = session.client("lambda")
 
     try:
@@ -1000,6 +1000,32 @@ def unschedule(rule_name, session=None):
 
     events.delete_rule(Name=prefixed_rule_name)
 
+def enable_schedule(rule_name, session=None):
+    """Enable an existing notebook schedule rule.
+
+    Args:
+        rule_name (str): The name of the rule for CloudWatch Events (required).
+        session (boto3.Session): The boto3 session to use. Will create a default session if not supplied (default: None).
+    """
+    prefixed_rule_name = RULE_PREFIX + rule_name
+
+    session = ensure_session(session)
+    events = session.client("events")
+    events.enable_rule(Name=prefixed_rule_name)
+
+def disable_schedule(rule_name, session=None):
+    """Disable an existing notebook schedule rule.
+
+    Args:
+        rule_name (str): The name of the rule for CloudWatch Events (required).
+        session (boto3.Session): The boto3 session to use. Will create a default session if not supplied (default: None).
+    """
+    prefixed_rule_name = RULE_PREFIX + rule_name
+
+    session = ensure_session(session)
+    events = session.client("events")
+    events.disable_rule(Name=prefixed_rule_name)
+
 
 def describe_schedules(n=0, rule_prefix=None, session=None):
     """A generator that returns descriptions of all the notebook schedule rules

sagemaker_run_notebook/server_extension/handlers.py

@@ -122,7 +122,7 @@ def load_params(self, required):
         try:
             data = json.loads(self.request.body.decode("utf-8"))
 
-            if not self.required_params(data, ["image", "input_path", "notebook"]):
+            if not self.required_params(data, required):
                 return None
             else:
                 return data
@@ -296,6 +296,22 @@ def delete(self, rule_name):
         except botocore.exceptions.BotoCoreError as e:
             self.botocore_error_response(e)
 
+    def patch(self, rule_name):
+        """Enable or disable a rule"""
+        try:
+            data = self.load_params(["action"])
+            if data is None:
+                return
+
+            if data['action'] == 'enable':
+                run.enable_schedule(rule_name=rule_name)
+            elif data['action'] == 'disable':
+                run.disable_schedule(rule_name=rule_name)
+        except botocore.exceptions.ClientError as e:
+            self.client_error_response(e)
+        except botocore.exceptions.BotoCoreError as e:
+            self.botocore_error_response(e)
+
 
 class InvokeHandler(BaseHandler):
     def post(self):