What is the best way to ensure the shutdown of a kernel before initializing a new one?

[MarcosVn]

Description

I am experiencing an intermittent issue in my custom Notebook.tsx, which uses Cell.tsx from the library.

The problem is: intermittently, the shutdown of the current kernel doesn’t always seem to complete before the new kernel is created. This happens, for instance, when the user quickly exits the component and enters another notebook. Because of this, we got a disconnected kernel during the creation.

I had taken inspiration from the useKernel hook in NotebookInit.tsx from jupyter-react examples to handle the creation and shutdown process, and even with an additional lock, it still occasionally seems insufficient.

The current behavior below (using the isShuttingDown lock) seems to minimize the issue well, but is there a better-known way to handle this?

Additionally, I noticed that in version 0.17.x, some requests to the Jupyter server are being directed to oss.datalayer.run (even though my server configuration is set), and requests are subsequently made to it. I ran some tests, and this seems to be related to the default values set in Jupyter.tsx. When I removed them, the issue was resolved. Should I open a new issue for this?

Reproduce

1. Enter and exit repeatedly from a notebook using the code below.

Expected behavior

The kernel must be reliably shut down before initializing a new one.

Evidences

• Dead/disconnected kernel:

  • 

• Requests starting from oss.datalayer.. and after changing to my jupyter-server

  • 
  • 

iimport React, {useState, useEffect, useRef} from 'react';

import {useJupyter, Jupyter, Kernel} from '@datalayer/jupyter-react';
import {Box, CircularProgress, createTheme, Theme, ThemeProvider, Typography} from '@mui/material';
import {deepmerge} from '@mui/utils';

import {NotebookConfigProvider} from '../../config/NotebookConfig';
import {baseDarkTheme, NotebookStyledBox} from '../baseTheme';
import NotebookActionBar from './NotebookActionBar';
import {useNotebookStore, notebookStore as noteStore} from './NotebookState';
import NotebookController from './NotebookController';
import {ICellsConfig} from '../cell/factory/cell.type';
import {ICellState} from '../cell/CellState';
import CustomCell from '../cell/Cell';
import {NotebookEvent, notebookEventEmitter} from '../../events/NotebookEvents';
import {eventLogger} from '../../events/EventLogger';


interface INotebookConfig {
  cellsConfig?: ICellsConfig;
  muiTheme?: Theme;
  actionsDisabled?: boolean;
}

export interface ICustomNotebookProps {
  id: string;
  config?: INotebookConfig;
  autoStart?: boolean;
  initialCode?: string; // An optional and initial code to be executed before notebook start
}

export interface INotebookContainerProps {
  id: string;
  config?: INotebookConfig;
  autoStart?: boolean;
  kernel: Kernel;
  mergedTheme: Theme;
}

const JUPYTER_KERNEL_NAME = 'python';


const useKernel = () => {
  const {kernelManager, serviceManager} = useJupyter();
  const setIsShuttingDown = useNotebookStore(state => state.setIsShuttingDown);
  const [kernel, setKernel] = useState<Kernel>();
  useEffect(() => {
    if (!serviceManager) {
      return;
    }

    if (noteStore.getState().isShuttingDown) {
      return;
    }

    kernelManager?.ready.then(() => {
      const customKernel = new Kernel({
        kernelManager,
        kernelName: JUPYTER_KERNEL_NAME,
        kernelSpecName: JUPYTER_KERNEL_NAME,
        kernelspecsManager: serviceManager.kernelspecs,
        sessionManager: serviceManager.sessions,
      });
      customKernel.ready.then(() => {
        setKernel(customKernel);
      });
    });
    return () => {
      if (kernel) {
        setIsShuttingDown(true);
        kernelManager?.shutdown(kernel.id).then().finally(() => {
          setIsShuttingDown(false);
        });
      }
    };
  }, [kernelManager, serviceManager]);
  return kernel;
};

/**
 * Base component for Custom Notebooks container, it provide a flexible way to create cells grouped into a notebook
 * Is used inside CustomNotebook (explanation below)
 */
const NotebookContainer: React.FC<INotebookContainerProps> = ({id, config, autoStart = false, kernel, mergedTheme}) => {
  const notebookStore = useNotebookStore();
  const {kernelManager, serviceManager} = useJupyter();
  const [controller, setController] = useState<NotebookController>();
  const [isLoaded, setIsLoaded] = useState(false);
  const hasLoadedRef = useRef(false);

  useEffect(() => {
    // defaultKernel triggers the base jupyter-cell rerendering, so we wait for it here also
    if (!serviceManager) {
      console.log('Service manager is not ready yet');
      return;
    }

    if (controller) {
      console.log('Controller is already set');
      return;
    }

    setController(new NotebookController(id, serviceManager, autoStart));
  }, [serviceManager, kernelManager]);

  /**
   * Calls the controller promise to load the notebooks content. After the loads, set isLoaded state
   */
  const loadNotebook = async () => {
    // Skip if already loaded or controller is not set
    if (hasLoadedRef.current || !controller) {
      return;
    }

    await controller.load();
    setIsLoaded(true);
  };

  useEffect(() => {
    loadNotebook();

    // Component unmount - resets the notebook store state
    return () => {
      console.log(`Unmounting notebook ${id}...`);
      notebookStore.setCells([]);
      notebookStore.setActiveCell(undefined);
    };
  }, [kernel, controller]);

  return (
    <NotebookConfigProvider config={config}>
      <ThemeProvider theme={mergedTheme}>
        <Box
          display="flex"
          flexDirection="column"
          minHeight="100vh"
          height="100%"
          sx={{
            borderTop: `1px solid ${mergedTheme.custom.darkerBorder}`,
            backgroundColor: mergedTheme.custom.darkerBackground,
          }}
        >
          {/* Notebook action bar */}
          {controller && <NotebookActionBar controller={controller} isNotebookLoaded={isLoaded && !config?.actionsDisabled} />}

          {/* Cells list */}
          <NotebookStyledBox flexGrow={1}>
            {notebookStore.cells.map((cell: ICellState) => (
              <CustomCell key={cell.id} cell={cell} config={config?.cellsConfig || {}} kernel={kernel} />
            ))}
          </NotebookStyledBox>
        </Box>
      </ThemeProvider>
    </NotebookConfigProvider>
  );
};


/**
 * Wraper for Notebook with upper Jupyter connection context
 * It is currently necessary to ensure that the Notebook will have serviceManager attached, due to it's async auth
 * TODO: Review it's there's no better approach to do this in the future
 */
const CustomNotebook: React.FC<ICustomNotebookProps> = ({id, config, autoStart = false, initialCode}) => {
  const kernel = useKernel();

  /**
   * Merge user provided theme and base theme, there is,
   *  - user prefs will be applied and any other not configured will be inherited from base
   */
  const mergedTheme = createTheme(deepmerge(baseDarkTheme, config?.muiTheme || {}));

  useEffect(() => {
    // Execute the inicial user code (async to better UX)
    if (kernel && initialCode) {
      const executionResult = kernel.execute(initialCode);
      executionResult?.done.catch(error => {
        console.error('Failed to execute the initial code', error);
      });
    }
  }, [kernel]);

  // Render the kernel if it exists, otherwise, render the loading feedback
  return kernel ? (
    <Jupyter
      colorMode={config?.muiTheme?.palette?.mode || 'dark'}
      jupyterServerUrl={process.env.REACT_APP_JUPYTER_SERVER_HTTP}
      jupyterServerToken={process.env.REACT_APP_JUPYTER_TOKEN}
      startDefaultKernel={false}
      useRunningKernelId={kernel.id}>
      <NotebookContainer
        id={id}
        config={config}
        autoStart={autoStart}
        kernel={kernel}
        mergedTheme={mergedTheme}
      />
    </Jupyter>
  ) : (
    // Loading section until kernel is not ready
    <ThemeProvider theme={mergedTheme}>
      <Box
        display="flex"
        flexDirection="column"
        minHeight="100vh"
        height="100%"
        sx={{
          borderTop: `1px solid ${mergedTheme.custom.darkerBorder}`,
          backgroundColor: mergedTheme.custom.darkerBackground,
        }}
      >
        <Box
          display="flex"
          flexDirection="row"
          alignItems="center"
          justifyContent="center"
          sx={{gap: mergedTheme.spacing(1)}}>
          <Typography sx={{
            textAlign: 'center',
            fontSize: 16,
            color: mergedTheme.custom.darkerText,
            padding: mergedTheme.spacing(2)
          }}>
            O seu notebook está sendo carregado. Por favor, aguarde...
          </Typography>
          <CircularProgress
            size={12}
            sx={{color: '#fff', marginRight: mergedTheme.spacing(2)}} />
        </Box>
      </Box>
    </ThemeProvider>
  );
};

export default CustomNotebook;

• Datalayer version: 0.17.1
• Operating System and version: Ubuntu 22.04
• Browser and version: Chrome 130.0.6723.116

Browser Output

Paste the output from your browser Javascript console here.

[MarcosVn]

Hey @echarles, after some additional testing on my side, I believe we can close this issue.

• Regarding the problem with the kernel being marked as dead, this was caused by an incorrect unmounting on our side (it was being called multiple times, even during initialization). I've adjusted it here, and so far, I haven’t been able to reproduce the issue anymore.

• The other matter, related to requests to oss.datalayer.run, has been addressed with the integration of #341.

Finally, I have one last question/confirmation regarding the Cell requests to /jupyter-server/api/sessions?{id} and /jupyter-server/api/kernels?{id}, but maybe we can close this one and I could open a new issue for that. What do you think?

[echarles]

Sure @MarcosVn I am closing this issue and let you open a new one regarding the request to sessoins/kernels.

[MarcosVn]

Thx, @echarles