log jobs and job defs in mlflow

Author: andrii-i

jupyter_scheduler/executors.py

@@ -7,13 +7,15 @@
 from typing import Dict
 
 import fsspec
+import mlflow
 import nbconvert
 import nbformat
 from nbconvert.preprocessors import CellExecutionError, ExecutePreprocessor
 
 from jupyter_scheduler.models import DescribeJob, JobFeature, Status
 from jupyter_scheduler.orm import Job, create_session
 from jupyter_scheduler.parameterize import add_parameters
+from jupyter_scheduler.scheduler import MLFLOW_SERVER_URI
 from jupyter_scheduler.utils import get_utc_timestamp
 
 
@@ -172,6 +174,22 @@ def create_output_files(self, job: DescribeJob, notebook_node):
             output, _ = cls().from_notebook_node(notebook_node)
             with fsspec.open(self.staging_paths[output_format], "w", encoding="utf-8") as f:
                 f.write(output)
+        mlflow.set_tracking_uri(MLFLOW_SERVER_URI)
+        with mlflow.start_run(run_id=job.mlflow_run_id):
+            try:
+                ep.preprocess(nb)
+                if job.parameters:
+                    mlflow.log_params(job.parameters)
+            except CellExecutionError as e:
+                raise e
+            finally:
+                for output_format in job.output_formats:
+                    cls = nbconvert.get_exporter(output_format)
+                    output, resources = cls().from_notebook_node(nb)
+                    output_path = self.staging_paths[output_format]
+                    with fsspec.open(output_path, "w", encoding="utf-8") as f:
+                        f.write(output)
+                    mlflow.log_artifact(output_path)
 
     def supported_features(cls) -> Dict[JobFeature, bool]:
         return {

jupyter_scheduler/models.py

@@ -86,6 +86,9 @@ class CreateJob(BaseModel):
     output_filename_template: Optional[str] = OUTPUT_FILENAME_TEMPLATE
     compute_type: Optional[str] = None
     package_input_folder: Optional[bool] = None
+    mlflow_logging: Optional[bool] = None
+    mlflow_experiment_id: Optional[str] = None
+    mlflow_run_id: Optional[str] = None
 
     @root_validator
     def compute_input_filename(cls, values) -> Dict:
@@ -148,6 +151,9 @@ class DescribeJob(BaseModel):
     downloaded: bool = False
     package_input_folder: Optional[bool] = None
     packaged_files: Optional[List[str]] = []
+    mlflow_logging: Optional[bool] = None
+    mlflow_experiment_id: Optional[str] = None
+    mlflow_run_id: Optional[str] = None
 
     class Config:
         orm_mode = True
@@ -213,6 +219,8 @@ class CreateJobDefinition(BaseModel):
     schedule: Optional[str] = None
     timezone: Optional[str] = None
     package_input_folder: Optional[bool] = None
+    mlflow_logging: Optional[bool] = None
+    mlflow_experiment_id: Optional[str] = None
 
     @root_validator
     def compute_input_filename(cls, values) -> Dict:
@@ -240,6 +248,8 @@ class DescribeJobDefinition(BaseModel):
     active: bool
     package_input_folder: Optional[bool] = None
     packaged_files: Optional[List[str]] = []
+    mlflow_logging: Optional[bool] = None
+    mlflow_experiment_id: Optional[str] = None
 
     class Config:
         orm_mode = True

jupyter_scheduler/orm.py

@@ -87,6 +87,8 @@ class CommonColumns:
     create_time = Column(Integer, default=get_utc_timestamp)
     package_input_folder = Column(Boolean)
     packaged_files = Column(JsonType, default=[])
+    mlflow_logging = Column(Boolean)
+    mlflow_experiment_id = Column(String(256), nullable=True)
 
 
 class Job(CommonColumns, Base):
@@ -100,6 +102,7 @@ class Job(CommonColumns, Base):
     url = Column(String(256), default=generate_jobs_url)
     pid = Column(Integer)
     idempotency_token = Column(String(256))
+    mlflow_run_id = Column(String(256), nullable=True)
 
 
 class JobDefinition(CommonColumns, Base):

jupyter_scheduler/scheduler.py

@@ -4,8 +4,10 @@
 import shutil
 import subprocess
 from typing import Dict, List, Optional, Type, Union
+from uuid import uuid4
 
 import fsspec
+import mlflow
 import psutil
 from jupyter_core.paths import jupyter_data_dir
 from jupyter_server.transutils import _i18n
@@ -46,6 +48,10 @@
     create_output_filename,
 )
 
+MLFLOW_SERVER_HOST = "127.0.0.1"
+MLFLOW_SERVER_PORT = "5000"
+MLFLOW_SERVER_URI = f"http://{MLFLOW_SERVER_HOST}:{MLFLOW_SERVER_PORT}"
+
 
 class BaseScheduler(LoggingConfigurable):
     """Base class for schedulers. A default implementation
@@ -405,16 +411,13 @@ def start_mlflow_server(self):
             [
                 "mlflow",
                 "server",
-                "--backend-store-uri",
-                "./mlruns",
-                "--default-artifact-root",
-                "./mlartifacts",
                 "--host",
-                "0.0.0.0",
+                MLFLOW_SERVER_HOST,
                 "--port",
-                "5000",
+                MLFLOW_SERVER_PORT,
             ]
         )
+        mlflow.set_tracking_uri(MLFLOW_SERVER_URI)
 
     def __init__(
         self,
@@ -481,6 +484,19 @@ def create_job(self, model: CreateJob) -> str:
             if not model.output_formats:
                 model.output_formats = []
 
+            mlflow_client = mlflow.MlflowClient()
+
+            if model.job_definition_id and model.mlflow_experiment_id:
+                experiment_id = model.mlflow_experiment_id
+            else:
+                experiment_id = mlflow_client.create_experiment(f"{model.name}-{uuid4()}")
+                model.mlflow_experiment_id = experiment_id
+                input_file_path = os.path.join(self.root_dir, model.input_uri)
+                mlflow.log_artifact(input_file_path, "input")
+
+            mlflow_run = mlflow_client.create_run(experiment_id=experiment_id, run_name=model.name)
+            model.mlflow_run_id = mlflow_run.info.run_id
+
             job = Job(**model.dict(exclude_none=True, exclude={"input_uri"}))
 
             session.add(job)
@@ -628,6 +644,12 @@ def create_job_definition(self, model: CreateJobDefinition) -> str:
             if not self.file_exists(model.input_uri):
                 raise InputUriError(model.input_uri)
 
+            mlflow_client = mlflow.MlflowClient()
+            experiment_id = mlflow_client.create_experiment(f"{model.name}-{uuid4()}")
+            model.mlflow_experiment_id = experiment_id
+            input_file_path = os.path.join(self.root_dir, model.input_uri)
+            mlflow.log_artifact(input_file_path, "input")
+
             job_definition = JobDefinition(**model.dict(exclude_none=True, exclude={"input_uri"}))
             session.add(job_definition)
             session.commit()

src/components/mlflow-checkbox.tsx

@@ -2,15 +2,13 @@ import React, { ChangeEvent } from 'react';
 
 import { Checkbox, FormControlLabel, FormGroup } from '@mui/material';
 
-export type MLFlowCheckboxProps = {
+export function MLFlowLoggingControl(props: {
   onChange: (event: ChangeEvent<HTMLInputElement>) => void;
-};
-
-export function MLFlowCheckbox(props: MLFlowCheckboxProps): JSX.Element {
+}): JSX.Element {
   return (
     <FormGroup>
       <FormControlLabel
-        control={<Checkbox onChange={props.onChange} value={'mlflowLogging'} />}
+        control={<Checkbox onChange={props.onChange} name={'mlflowLogging'} />}
         label="Log with MLFlow"
       />
     </FormGroup>

src/handler.ts

@@ -364,6 +364,8 @@ export namespace Scheduler {
     schedule?: string;
     timezone?: string;
     package_input_folder?: boolean;
+    mlflow_logging?: boolean;
+    mlflow_experiment_id?: string;
   }
 
   export interface IUpdateJobDefinition {
@@ -391,6 +393,8 @@ export namespace Scheduler {
     update_time: number;
     active: boolean;
     package_input_folder?: boolean;
+    mlflow_logging: boolean;
+    mlflow_experiment_id?: string;
   }
 
   export interface IEmailNotifications {
@@ -418,6 +422,9 @@ export namespace Scheduler {
     output_formats?: string[];
     compute_type?: string;
     package_input_folder?: boolean;
+    mlflow_logging?: boolean;
+    mlflow_experiment_id?: string;
+    mlflow_run_id?: string;
   }
 
   export interface ICreateJobFromDefinition {
@@ -467,6 +474,9 @@ export namespace Scheduler {
     end_time?: number;
     downloaded: boolean;
     package_input_folder?: boolean;
+    mlflow_logging?: boolean;
+    mlflow_experiment_id?: string;
+    mlflow_run_id?: string;
   }
 
   export interface ICreateJobResponse {

src/mainviews/create-job.tsx

@@ -43,7 +43,7 @@ import {
 import { Box, Stack } from '@mui/system';
 import { getErrorMessage } from '../util/errors';
 import { PackageInputFolderControl } from '../components/input-folder-checkbox';
-import { MLFlowCheckbox } from '../components/mlflow-checkbox';
+import { MLFlowLoggingControl } from '../components/mlflow-checkbox';
 
 export interface ICreateJobProps {
   model: ICreateJobModel;
@@ -512,7 +512,7 @@ export function CreateJob(props: ICreateJobProps): JSX.Element {
             onChange={handleInputChange}
             inputFile={props.model.inputFile}
           />
-          <MLFlowCheckbox onChange={handleInputChange} />
+          <MLFlowLoggingControl onChange={handleInputChange} />
           <OutputFormatPicker
             label={trans.__('Output formats')}
             name="outputFormat"

src/mainviews/detail-view/job-definition.tsx

@@ -32,6 +32,7 @@ import { Scheduler as SchedulerTokens } from '../../tokens';
 
 import { timestampLocalize } from './job-detail';
 import { getErrorMessage } from '../../util/errors';
+import { OpenInNew } from '@mui/icons-material';
 
 export interface IJobDefinitionProps {
   app: JupyterFrontEnd;
@@ -175,6 +176,18 @@ export function JobDefinition(props: IJobDefinitionProps): JSX.Element {
       >
         {trans.__('Edit Job Definition')}
       </Button>
+      {model.mlflowLogging === true && (
+        <Button
+          variant="outlined"
+          onClick={() => {
+            const mlFlowUrl = `http://127.0.0.1:5000/#/experiments/${props.model?.mlflowExperimentId}`;
+            window.open(mlFlowUrl);
+          }}
+          endIcon={<OpenInNew />}
+        >
+          {trans.__('Open in MLFlow')}
+        </Button>
+      )}
       <ConfirmDialogDeleteButton
         handleDelete={async () => {
           log('job-definition-detail.delete');
@@ -231,6 +244,16 @@ export function JobDefinition(props: IJobDefinitionProps): JSX.Element {
         label: trans.__('Time zone')
       }
     ],
+    [
+      {
+        value: model.mlflowLogging ? trans.__('Yes') : trans.__('No'),
+        label: trans.__('MLFlow Logging')
+      },
+      {
+        value: props.model.mlflowExperimentId,
+        label: trans.__('MLFLow Experiment Id')
+      }
+    ],
     [
       {
         value: model.packageInputFolder ? trans.__('Yes') : trans.__('No'),

src/mainviews/detail-view/job-detail.tsx

@@ -39,6 +39,11 @@ import {
   LabeledValue
 } from '../../components/labeled-value';
 import { getErrorMessage } from '../../util/errors';
+import { OpenInNew } from '@mui/icons-material';
+
+const MLFLOW_SERVER_HOST = '127.0.0.1';
+const MLFLOW_SERVER_PORT = '5000';
+const MLFLOW_SERVER_URI = `http://${MLFLOW_SERVER_HOST}:${MLFLOW_SERVER_PORT}`;
 
 export interface IJobDetailProps {
   app: JupyterFrontEnd;
@@ -167,6 +172,18 @@ export function JobDetail(props: IJobDetailProps): JSX.Element {
             {trans.__('Download Job Files')}
           </Button>
         )}
+      {props.model?.mlflowLogging === true && (
+        <Button
+          variant="outlined"
+          onClick={() => {
+            const mlFlowUrl = `${MLFLOW_SERVER_URI}/#/experiments/${props.model?.mlflowExperimentId}/runs/${props.model?.mlflowRunId}`;
+            window.open(mlFlowUrl);
+          }}
+          endIcon={<OpenInNew />}
+        >
+          {trans.__('Open in MLFlow')}
+        </Button>
+      )}
       {props.model !== null && props.model.status === 'IN_PROGRESS' && (
         <ConfirmDialogStopButton
           handleStop={handleStopJob}
@@ -250,6 +267,24 @@ export function JobDetail(props: IJobDetailProps): JSX.Element {
         value: timestampLocalize(props.model.endTime ?? ''),
         label: trans.__('End time')
       },
+      {
+        value: props.model.mlflowLogging ? trans.__('Yes') : trans.__('No'),
+
+        label: trans.__('MLFlow Logging')
+      }
+    ],
+    [
+      {
+        value: props.model.mlflowExperimentId,
+        label: trans.__('MLFLow Experiment Id')
+      },
+      {
+        value: props.model.mlflowRunId,
+
+        label: trans.__('MLFlow Run Id')
+      }
+    ],
+    [
       {
         value: props.model.packageInputFolder
           ? trans.__('Yes')

src/model.ts

@@ -100,6 +100,9 @@ export interface ICreateJobModel
   // Is the create button disabled due to a submission in progress?
   createInProgress?: boolean;
   packageInputFolder?: boolean;
+  mlflowLogging?: boolean;
+  mlflowExperimentId?: string;
+  mlflowRunId?: string;
 }
 
 export const defaultScheduleFields: ModelWithScheduleFields = {
@@ -312,6 +315,9 @@ export interface IJobDetailModel {
   job_files: Scheduler.IJobFile[];
   downloaded: boolean;
   packageInputFolder?: boolean;
+  mlflowLogging?: boolean;
+  mlflowExperimentId?: string;
+  mlflowRunId?: string;
 }
 
 export interface IJobDefinitionModel {
@@ -339,6 +345,8 @@ export interface IJobDefinitionModel {
   endTime?: number;
   outputPrefix?: string;
   packageInputFolder?: boolean;
+  mlflowLogging?: boolean;
+  mlflowExperimentId?: string;
 }
 
 const convertParameters = (parameters: {
@@ -388,7 +396,10 @@ export function convertDescribeJobtoJobDetail(
     startTime: describeJob.start_time,
     endTime: describeJob.end_time,
     downloaded: describeJob.downloaded,
-    packageInputFolder: describeJob.package_input_folder
+    packageInputFolder: describeJob.package_input_folder,
+    mlflowLogging: describeJob.mlflow_logging,
+    mlflowExperimentId: describeJob.mlflow_experiment_id,
+    mlflowRunId: describeJob.mlflow_run_id
   };
 }
 
@@ -417,7 +428,9 @@ export function convertDescribeDefinitiontoDefinition(
     updateTime: describeDefinition.update_time,
     schedule: describeDefinition.schedule,
     timezone: describeDefinition.timezone,
-    packageInputFolder: describeDefinition.package_input_folder
+    packageInputFolder: describeDefinition.package_input_folder,
+    mlflowLogging: describeDefinition.mlflow_logging,
+    mlflowExperimentId: describeDefinition.mlflow_experiment_id
   };
 }