class_docstr: This shows how to rewrite a monolith function so that it can handle preemption nicely,
  and not restart from scratch everytime it's preempted.
class_docstr_tok: ["This","shows","how","to","rewrite","a","monolith","function","so","that","it","can","handle","preemption","nicely",",","and","not","restart","from","scratch","everytime","it","'s","preempted","."]
code: def __call__(self, train_samples: int, model_path: Path = None):
      # `train_samples` and `model_path` will also be saved
      log = functools.partial(print, flush=True)
      log(f"*** Starting from stage '{self.stage}' ***")
  
      if self.train_test is None:
          self.stage = "Data Loading"
          t0 = time.time()
          log(f"*** Entering stage '{self.stage}' ***")
          # Load data from https://www.openml.org/d/554
          X, y = fetch_openml("mnist_784", version=1, return_X_y=True)
          X, y = X.numpy(), y.numpy()
  
          random_state = check_random_state(0)
          permutation = random_state.permutation(X.shape[0])
          X = X[permutation]
          y = y[permutation]
          X = X.reshape((X.shape[0], -1))
  
          # Checkpoint 1: save the train/test splits
          X_train, X_test, y_train, y_test = train_test_split(
              X, y, train_size=train_samples, test_size=10000
          )
          self.train_test = X_train, X_test, y_train, y_test
          log(
              f"Loaded data, shuffle and split in {time.time() - t0:.1f}s"
          )
  
      X_train, X_test, y_train, y_test = self.train_test
      if self.scaler is None:
          self.stage = "Data Cleaning"
          t0 = time.time()
          log(f"*** Entering stage '{self.stage}' ***")
          scaler = StandardScaler()
          X_train = scaler.fit_transform(X_train)
          X_test = scaler.transform(X_test)
          # Scaling is actual pretty fast, make it a bit slower to allow preemption to happen here
          time.sleep(10)
          # Checkpoint 2: save the scaler and the preprocessed data
          self.scaler = scaler
          self.train_test = X_train, X_test, y_train, y_test
          log(f"Scaled the data took {time.time() - t0:.0f}s")
  
      if not self.trained_clf:
          self.stage = "Model Training"
          t0 = time.time()
          log(f"*** Entering stage '{self.stage}' ***")
          self.clf.C = 50 / train_samples
          self.clf.fit(X_train, y_train)
          # Checkpoint 3: mark the classifier as trained
          self.trained_clf = True
          log(f"Training took {time.time() - t0:.0f}s")
  
      sparsity = np.mean(self.clf.coef_ == 0) * 100
      score = self.clf.score(X_test, y_test)
      log(f"Sparsity with L1 penalty: {sparsity / 100:.2%}")
      log(f"Test score with L1 penalty: {score:.4f}")
  
      if model_path:
          self.save(model_path)
      return score
  
file_path: submitit/docs/mnist.py
imports: import functools
  import pickle
  import time
  from pathlib import Path
  import numpy
  from sklearn.datasets import fetch_openml
  from sklearn.linear_model import LogisticRegression
  from sklearn.model_selection import train_test_split
  from sklearn.preprocessing import StandardScaler
  from sklearn.utils import check_random_state
  import submitit
label_desc: Use image node_id 2 for calling the MnistTrainer obj's underlying member method code with example usage: obj.__call__(train_samples, model_path) and returns: score