From 6c95a762031de3e5c7b434d2112de2329262cae0 Mon Sep 17 00:00:00 2001
From: illia-prokopchuk <78905712+illia-prokopchuk@users.noreply.github.com>
Date: Tue, 7 Nov 2023 18:18:45 +0200
Subject: [PATCH] Impl [Batch Run] Model Training (#2066)

---
 src/actions/functions.js                      |  4 +-
 src/common/NameFilter/nameFilter.scss         |  2 +
 src/common/TargetPath/targetPath.util.js      |  2 +-
 .../ArtifactsActionBar/ArtifactsActionBar.js  | 39 ++++++++++-------
 src/components/Datasets/Datasets.js           |  8 ++--
 src/components/Datasets/DatasetsView.js       | 14 ++++---
 src/components/Datasets/datasets.util.js      | 29 ++++++++++++-
 src/components/Files/FilesView.js             | 14 ++++---
 src/components/JobWizard/JobWizard.js         | 24 ++++++++++-
 src/components/JobWizard/JobWizard.util.js    | 28 +++++++++++--
 .../JobWizardFunctionSelection.js             | 42 ++++++++++++-------
 .../jobWizardFunctionSelection.util.js        | 18 ++++++++
 .../JobWizardRunDetails.js                    |  4 +-
 src/components/ModelsPage/Models/Models.js    | 10 +++++
 .../ModelsPage/Models/ModelsView.js           | 28 +++++++++----
 .../ProjectOverview/ProjectOverview.util.js   |  5 +--
 src/scss/main.scss                            |  6 ++-
 src/utils/generateTemplatesCategories.js      | 16 ++++---
 src/utils/panelPathScheme.js                  |  2 +-
 19 files changed, 217 insertions(+), 78 deletions(-)

diff --git a/src/actions/functions.js b/src/actions/functions.js
index 42cbe89cd..01ccb07ee 100644
--- a/src/actions/functions.js
+++ b/src/actions/functions.js
@@ -287,13 +287,13 @@ const functionsActions = {
     type: FETCH_HUB_FUNCTION_TEMPLATE_FAILURE,
     payload: err
   }),
-  fetchHubFunctions: () => dispatch => {
+  fetchHubFunctions: allowedHubFunctions => dispatch => {
     dispatch(functionsActions.fetchHubFunctionsBegin())
 
     return functionsApi
       .getHubFunctions()
       .then(({ data: functionTemplates }) => {
-        const templatesData = generateHubCategories(functionTemplates.catalog)
+        const templatesData = generateHubCategories(functionTemplates.catalog, allowedHubFunctions)
 
         dispatch(functionsActions.setHubFunctions(templatesData))
 
diff --git a/src/common/NameFilter/nameFilter.scss b/src/common/NameFilter/nameFilter.scss
index 54382f0c9..fe294be76 100644
--- a/src/common/NameFilter/nameFilter.scss
+++ b/src/common/NameFilter/nameFilter.scss
@@ -16,6 +16,8 @@ such restriction.
 */
 
 .name-filter {
+  min-width: 180px;
+
   &__icon {
     cursor: pointer;
   }
diff --git a/src/common/TargetPath/targetPath.util.js b/src/common/TargetPath/targetPath.util.js
index b21f46d96..16def23a2 100644
--- a/src/common/TargetPath/targetPath.util.js
+++ b/src/common/TargetPath/targetPath.util.js
@@ -275,7 +275,7 @@ export const generateComboboxMatchesList = (
 export const generateArtifactsList = artifacts => {
   const generatedArtifacts = artifacts
     .map(artifact => {
-      const key = artifact.link_iteration ? artifact.link_iteration.db_key : artifact.key ?? ''
+      const key = artifact.link_iteration?.db_key || artifact.db_key || artifact.key || ''
       return {
         label: key,
         id: key
diff --git a/src/components/ArtifactsActionBar/ArtifactsActionBar.js b/src/components/ArtifactsActionBar/ArtifactsActionBar.js
index da408556c..894298281 100644
--- a/src/components/ArtifactsActionBar/ArtifactsActionBar.js
+++ b/src/components/ArtifactsActionBar/ArtifactsActionBar.js
@@ -40,7 +40,7 @@ import detailsActions from '../../actions/details'
 import { ReactComponent as RefreshIcon } from 'igz-controls/images/refresh.svg'
 
 function ArtifactsActionBar({
-  actionButton,
+  actionButtons,
   cancelRequest,
   filterMenuName,
   handleRefresh,
@@ -154,13 +154,18 @@ function ArtifactsActionBar({
             </FilterMenuModal>
           </div>
           <div className="action-bar__actions">
-            {actionButton && !actionButton.hidden && (
-              <Button
-                variant={actionButton.variant}
-                label={actionButton.label}
-                className={actionButton.className}
-                onClick={actionButton.onClick}
-              />
+            {actionButtons.map(
+              (actionButton, index) =>
+                actionButton &&
+                !actionButton.hidden && (
+                  <Button
+                    key={index}
+                    variant={actionButton.variant}
+                    label={actionButton.label}
+                    className={actionButton.className}
+                    onClick={actionButton.onClick}
+                  />
+                )
             )}
 
             <RoundedIcon tooltipText="Refresh" onClick={() => refresh(formState)} id="refresh">
@@ -174,19 +179,21 @@ function ArtifactsActionBar({
 }
 
 ArtifactsActionBar.defaultProps = {
-  actionButton: null,
+  actionButtons: [],
   cancelRequest: null,
   tab: ''
 }
 
 ArtifactsActionBar.propTypes = {
-  actionButton: PropTypes.shape({
-    className: PropTypes.string,
-    hidden: PropTypes.bool,
-    label: PropTypes.string.isRequired,
-    onClick: PropTypes.func.isRequired,
-    variant: PropTypes.string.isRequired
-  }),
+  actionButtons: PropTypes.arrayOf(
+    PropTypes.shape({
+      className: PropTypes.string,
+      hidden: PropTypes.bool,
+      label: PropTypes.string.isRequired,
+      onClick: PropTypes.func.isRequired,
+      variant: PropTypes.string.isRequired
+    })
+  ),
   cancelRequest: PropTypes.func,
   filterMenuName: PropTypes.string.isRequired,
   handleRefresh: PropTypes.func.isRequired,
diff --git a/src/components/Datasets/Datasets.js b/src/components/Datasets/Datasets.js
index 9371724f5..08c02dff3 100644
--- a/src/components/Datasets/Datasets.js
+++ b/src/components/Datasets/Datasets.js
@@ -83,10 +83,6 @@ const Datasets = () => {
   const filtersStore = useSelector(store => store.filtersStore)
   const datasetsRef = useRef(null)
   const viewMode = getViewMode(window.location.search)
-  const pageData = useMemo(
-    () => generatePageData(selectedDataset, viewMode),
-    [selectedDataset, viewMode]
-  )
   const params = useParams()
   const navigate = useNavigate()
   const location = useLocation()
@@ -96,6 +92,10 @@ const Datasets = () => {
     () => filtersStore[FILTER_MENU_MODAL][DATASETS_FILTERS].values,
     [filtersStore]
   )
+  const pageData = useMemo(
+    () => generatePageData(selectedDataset, viewMode, params),
+    [selectedDataset, viewMode, params]
+  )
 
   const detailsFormInitialValues = useMemo(
     () => ({
diff --git a/src/components/Datasets/DatasetsView.js b/src/components/Datasets/DatasetsView.js
index 364dfde61..5b13a37ba 100644
--- a/src/components/Datasets/DatasetsView.js
+++ b/src/components/Datasets/DatasetsView.js
@@ -75,12 +75,14 @@ const DatasetsView = React.forwardRef(
             <div className="table-container">
               <div className="content__action-bar-wrapper">
                 <ArtifactsActionBar
-                  actionButton={{
-                    variant: SECONDARY_BUTTON,
-                    label: registerDatasetTitle,
-                    className: 'register-button',
-                    onClick: handleRegisterDataset
-                  }}
+                  actionButtons={[
+                    {
+                      variant: SECONDARY_BUTTON,
+                      label: registerDatasetTitle,
+                      className: 'action-button',
+                      onClick: handleRegisterDataset
+                    }
+                  ]}
                   filterMenuName={DATASETS_FILTERS}
                   handleRefresh={handleRefresh}
                   page={DATASETS_PAGE}
diff --git a/src/components/Datasets/datasets.util.js b/src/components/Datasets/datasets.util.js
index 3d00719a6..8c1f2b514 100644
--- a/src/components/Datasets/datasets.util.js
+++ b/src/components/Datasets/datasets.util.js
@@ -18,6 +18,8 @@ under the Apache 2.0 license is conditioned upon your compliance with
 such restriction.
 */
 
+import JobWizard from '../JobWizard/JobWizard'
+import { openPopUp } from 'igz-controls/utils/common.util'
 import { applyTagChanges } from '../../utils/artifacts.util'
 import { getArtifactIdentifier } from '../../utils/getUniqueIdentifier'
 import {
@@ -34,6 +36,8 @@ import { searchArtifactItem } from '../../utils/searchArtifactItem'
 import { sortListByDate } from '../../utils'
 import { fetchDataSet } from '../../reducers/artifactsReducer'
 
+import { PRIMARY_BUTTON } from 'igz-controls/constants'
+
 export const infoHeaders = [
   {
     label: 'Hash',
@@ -85,17 +89,38 @@ export const generateDataSetsDetailsMenu = selectedItem => [
   }
 ]
 
-export const generatePageData = (selectedItem, viewMode) => ({
+export const generatePageData = (selectedItem, viewMode, params) => ({
   page: DATASETS_PAGE,
   details: {
     menu: generateDataSetsDetailsMenu(selectedItem),
     infoHeaders,
     type: DATASETS,
     hideBackBtn: viewMode === FULL_VIEW_MODE,
-    withToggleViewBtn: true
+    withToggleViewBtn: true,
+    actionButton: {
+      label: 'Train',
+      variant: PRIMARY_BUTTON,
+      onClick: () => handleTrainDataset(selectedItem, params)
+    }
   }
 })
 
+const handleTrainDataset = (selectedItem, params) => {
+  const prePopulatedDataInputs = [{
+    name: selectedItem.db_key || selectedItem.key || 'dataset',
+    path: selectedItem.URI
+  }]
+
+  openPopUp(JobWizard, {
+    params,
+    isTrain: true,
+    wizardTitle: 'Train model',
+    prePopulatedData: {
+      dataInputs: prePopulatedDataInputs
+    }
+  })
+}
+
 export const fetchDataSetRowData = async (
   dispatch,
   dataSet,
diff --git a/src/components/Files/FilesView.js b/src/components/Files/FilesView.js
index 34ab14a34..b73b94320 100644
--- a/src/components/Files/FilesView.js
+++ b/src/components/Files/FilesView.js
@@ -75,12 +75,14 @@ const FilesView = React.forwardRef(
             <div className="table-container">
               <div className="content__action-bar-wrapper">
                 <ArtifactsActionBar
-                  actionButton={{
-                    variant: SECONDARY_BUTTON,
-                    label: registerArtifactTitle,
-                    className: 'register-button',
-                    onClick: handleRegisterArtifact
-                  }}
+                  actionButtons={[
+                    {
+                      variant: SECONDARY_BUTTON,
+                      label: registerArtifactTitle,
+                      className: 'action-button',
+                      onClick: handleRegisterArtifact
+                    }
+                  ]}
                   filterMenuName={FILES_FILTERS}
                   handleRefresh={handleRefresh}
                   page={FILES_PAGE}
diff --git a/src/components/JobWizard/JobWizard.js b/src/components/JobWizard/JobWizard.js
index ade3ce45f..a72ff50f9 100644
--- a/src/components/JobWizard/JobWizard.js
+++ b/src/components/JobWizard/JobWizard.js
@@ -84,12 +84,14 @@ const JobWizard = ({
   functionsStore,
   isBatchInference,
   isOpen,
+  isTrain,
   jobsStore,
   mode,
   onResolve,
   onSuccessRequest,
   onWizardClose,
   params,
+  prePopulatedData,
   removeJobFunction,
   removeHubFunctions,
   runNewJob,
@@ -286,7 +288,11 @@ const JobWizard = ({
 
     return [
       {
-        label: isBatchInference ? 'Schedule Infer' : 'Schedule for later',
+        label: isBatchInference
+          ? 'Schedule Infer'
+          : isTrain
+          ? 'Schedule training job'
+          : 'Schedule for later',
         onClick: () => {
           formState.handleSubmit()
 
@@ -301,7 +307,14 @@ const JobWizard = ({
         ref: scheduleButtonRef
       },
       {
-        label: mode === PANEL_EDIT_MODE ? 'Save' : isBatchInference ? 'Infer now' : 'Run',
+        label:
+          mode === PANEL_EDIT_MODE
+            ? 'Save'
+            : isBatchInference
+            ? 'Infer now'
+            : isTrain
+            ? 'Run training now'
+            : 'Run',
         onClick: () => {
           formState.handleSubmit()
 
@@ -435,7 +448,9 @@ const JobWizard = ({
                 frontendSpec={frontendSpec}
                 functions={functions}
                 isEditMode={isEditMode}
+                isTrain={isTrain}
                 params={params}
+                prePopulatedData={prePopulatedData}
                 selectedFunctionData={selectedFunctionData}
                 selectedFunctionTab={selectedFunctionTab}
                 setActiveTab={setActiveTab}
@@ -456,6 +471,7 @@ const JobWizard = ({
                 isBatchInference={isBatchInference}
                 isEditMode={isEditMode}
                 jobAdditionalData={jobAdditionalData}
+                prePopulatedData={prePopulatedData}
                 selectedFunctionData={selectedFunctionData}
               />
               <JobWizardDataInputs formState={formState} />
@@ -505,9 +521,11 @@ const JobWizard = ({
 JobWizard.defaultProps = {
   defaultData: {},
   isBatchInference: false,
+  isTrain: false,
   mode: PANEL_CREATE_MODE,
   onSuccessRequest: () => {},
   onWizardClose: () => {},
+  prePopulatedData: {},
   wizardTitle: 'Batch run'
 }
 
@@ -515,11 +533,13 @@ JobWizard.propTypes = {
   defaultData: PropTypes.shape({}),
   isBatchInference: PropTypes.bool,
   isOpen: PropTypes.bool.isRequired,
+  isTrain: PropTypes.bool,
   mode: JOB_WIZARD_MODE,
   onResolve: PropTypes.func.isRequired,
   onSuccessRequest: PropTypes.func,
   onWizardClose: PropTypes.func,
   params: PropTypes.shape({}).isRequired,
+  prePopulatedData: PropTypes.shape({}),
   wizardTitle: PropTypes.string
 }
 
diff --git a/src/components/JobWizard/JobWizard.util.js b/src/components/JobWizard/JobWizard.util.js
index 1a309711f..d9f671e7f 100644
--- a/src/components/JobWizard/JobWizard.util.js
+++ b/src/components/JobWizard/JobWizard.util.js
@@ -102,7 +102,8 @@ export const generateJobWizardData = (
   selectedFunctionData,
   defaultData,
   currentProjectName,
-  isEditMode
+  isEditMode,
+  prePopulatedData
 ) => {
   const functions = selectedFunctionData.functions
   const functionInfo = getFunctionInfo(selectedFunctionData)
@@ -181,12 +182,18 @@ export const generateJobWizardData = (
     jobFormData[RESOURCES_STEP].jobPriorityClassName = jobPriorityClassName
   }
 
+  if (!isEmpty(functionParameters) || !isEmpty(prePopulatedData.dataInputs)) {
+    jobFormData[DATA_INPUTS_STEP].dataInputsTable = parseDataInputs(
+      functionParameters,
+      prePopulatedData.dataInputs
+    )
+  }
+
   if (!isEmpty(functionParameters)) {
     jobFormData[PARAMETERS_STEP].parametersTable = {
       predefined: parsePredefinedParameters(functionParameters),
       custom: []
     }
-    jobFormData[DATA_INPUTS_STEP].dataInputsTable = parseDataInputs(functionParameters)
   }
 
   return [jobFormData, jobAdditionalData]
@@ -611,8 +618,8 @@ const getDataInputData = (dataInputName, dataInputValue, dataInputIsChecked) =>
 
 const sortParameters = (parameter, nextParameter) => nextParameter.isRequired - parameter.isRequired
 
-export const parseDataInputs = functionParameters => {
-  return functionParameters
+export const parseDataInputs = (functionParameters = [], prePopulatedDataInputs) => {
+  const parsedDataInputs = functionParameters
     .filter(dataInputs => dataInputs.type?.includes('DataItem'))
     .map(dataInput => {
       return {
@@ -624,6 +631,19 @@ export const parseDataInputs = functionParameters => {
       }
     })
     .sort(sortParameters)
+
+  if (!isEmpty(prePopulatedDataInputs)) {
+    prePopulatedDataInputs.forEach(dataInput => {
+      parsedDataInputs.unshift({
+        data: getDataInputData(dataInput.name, dataInput.path, true),
+        isRequired: true,
+        isDefault: true,
+        isPredefined: true
+      })
+    })
+  }
+
+  return parsedDataInputs
 }
 
 export const parseDefaultDataInputs = (funcParams, runDataInputs = {}) => {
diff --git a/src/components/JobWizard/JobWizardSteps/JobWizardFunctionSelection/JobWizardFunctionSelection.js b/src/components/JobWizard/JobWizardSteps/JobWizardFunctionSelection/JobWizardFunctionSelection.js
index 6c043f952..e670d9e99 100644
--- a/src/components/JobWizard/JobWizardSteps/JobWizardFunctionSelection/JobWizardFunctionSelection.js
+++ b/src/components/JobWizard/JobWizardSteps/JobWizardFunctionSelection/JobWizardFunctionSelection.js
@@ -45,11 +45,13 @@ import { generateProjectsList } from '../../../../utils/projects'
 import { functionRunKinds } from '../../../Jobs/jobs.util'
 import { openConfirmPopUp } from 'igz-controls/utils/common.util'
 import {
+  filterTrainFunctionMethods,
   FUNCTIONS_SELECTION_FUNCTIONS_TAB,
   FUNCTIONS_SELECTION_HUB_TAB,
   functionsSelectionTabs,
   generateFunctionCardData,
-  generateFunctionTemplateCardData
+  generateFunctionTemplateCardData,
+  trainModelAllowedHubFunctions
 } from './jobWizardFunctionSelection.util'
 
 import './jobWizardFunctionSelection.scss'
@@ -63,7 +65,9 @@ const JobWizardFunctionSelection = ({
   frontendSpec,
   functions,
   isEditMode,
+  isTrain,
   params,
+  prePopulatedData,
   selectedFunctionData,
   selectedFunctionTab,
   setActiveTab,
@@ -215,7 +219,8 @@ const JobWizardFunctionSelection = ({
         functionData,
         defaultData,
         params.projectName,
-        isEditMode
+        isEditMode,
+        prePopulatedData
       )
 
       const newInitial = {
@@ -272,7 +277,9 @@ const JobWizardFunctionSelection = ({
       activeTab === FUNCTIONS_SELECTION_HUB_TAB &&
       (isEmpty(hubFunctions) || isEmpty(hubFunctionsCatalog))
     ) {
-      dispatch(functionsActions.fetchHubFunctions()).then(templatesObject => {
+      dispatch(
+        functionsActions.fetchHubFunctions(isTrain ? trainModelAllowedHubFunctions : {})
+      ).then(templatesObject => {
         if (templatesObject) {
           setTemplatesCategories(templatesObject.hubFunctionsCategories)
           setTemplates(templatesObject.hubFunctions)
@@ -298,6 +305,7 @@ const JobWizardFunctionSelection = ({
     formState.initialValues,
     hubFunctions,
     hubFunctionsCatalog,
+    isTrain,
     setTemplates,
     setTemplatesCategories
   ])
@@ -327,8 +335,10 @@ const JobWizardFunctionSelection = ({
       const functionTemplatePath = `${functionData.spec.item_uri}${functionData.spec.assets.function}`
 
       dispatch(functionsActions.fetchFunctionTemplate(functionTemplatePath)).then(result => {
-        setSelectedFunctionData(result)
-        generateData(result)
+        const resultData = isTrain ? filterTrainFunctionMethods(result) : result
+
+        setSelectedFunctionData(resultData)
+        generateData(resultData)
         setSelectedFunctionTab(FUNCTIONS_SELECTION_HUB_TAB)
         setShowSchedule(false)
       })
@@ -416,15 +426,17 @@ const JobWizardFunctionSelection = ({
               placeholder="Search functions..."
               setMatches={setFilterMatches}
             />
-            <FilterMenuModal
-              header="Filter by category"
-              wizardClassName="hub-filter"
-              filterMenuName={JOB_WIZARD_FILTERS}
-              initialValues={hubFiltersInitialValues}
-              values={hubFiltersInitialValues}
-            >
-              <HubCategoriesFilter templates={filterTemplates} />
-            </FilterMenuModal>
+            {!isTrain && (
+              <FilterMenuModal
+                header="Filter by category"
+                wizardClassName="hub-filter"
+                filterMenuName={JOB_WIZARD_FILTERS}
+                initialValues={hubFiltersInitialValues}
+                values={hubFiltersInitialValues}
+              >
+                <HubCategoriesFilter templates={filterTemplates} />
+              </FilterMenuModal>
+            )}
           </div>
           {(filterByName.length > 0 &&
             (filterMatches.length === 0 || isEmpty(filteredTemplates))) ||
@@ -476,7 +488,9 @@ JobWizardFunctionSelection.propTypes = {
   frontendSpec: PropTypes.shape({}).isRequired,
   functions: PropTypes.arrayOf(PropTypes.shape({})).isRequired,
   isEditMode: PropTypes.bool.isRequired,
+  isTrain: PropTypes.bool.isRequired,
   params: PropTypes.shape({}).isRequired,
+  prePopulatedData: PropTypes.shape({}).isRequired,
   selectedFunctionData: PropTypes.shape({}).isRequired,
   selectedFunctionTab: PropTypes.string.isRequired,
   setActiveTab: PropTypes.func.isRequired,
diff --git a/src/components/JobWizard/JobWizardSteps/JobWizardFunctionSelection/jobWizardFunctionSelection.util.js b/src/components/JobWizard/JobWizardSteps/JobWizardFunctionSelection/jobWizardFunctionSelection.util.js
index 5eb5e6926..d9e2a8586 100644
--- a/src/components/JobWizard/JobWizardSteps/JobWizardFunctionSelection/jobWizardFunctionSelection.util.js
+++ b/src/components/JobWizard/JobWizardSteps/JobWizardFunctionSelection/jobWizardFunctionSelection.util.js
@@ -24,6 +24,11 @@ import { FUNCTION_SELECTION_STEP } from '../../../../constants'
 export const FUNCTIONS_SELECTION_FUNCTIONS_TAB = 'functions'
 export const FUNCTIONS_SELECTION_HUB_TAB = 'hub'
 
+export const trainModelAllowedHubFunctions = {
+  'auto-trainer': ['train'],
+  'azureml-utils': ['submit_training_job', 'train']
+}
+
 export const functionsSelectionTabs = [
   {
     id: FUNCTIONS_SELECTION_FUNCTIONS_TAB,
@@ -63,3 +68,16 @@ export const generateFunctionTemplateCardData = templateData => {
 
   return functionTemplateCardData
 }
+
+export const filterTrainFunctionMethods = result => {
+  const allowedMethods = trainModelAllowedHubFunctions[result.name]
+  const { entry_points } = result.functions[0].spec
+
+  if (entry_points) {
+    result.functions[0].spec.entry_points = Object.fromEntries(
+      Object.entries(entry_points).filter(([key]) => allowedMethods.includes(key))
+    )
+  }
+
+  return result
+}
diff --git a/src/components/JobWizard/JobWizardSteps/JobWizardRunDetails/JobWizardRunDetails.js b/src/components/JobWizard/JobWizardSteps/JobWizardRunDetails/JobWizardRunDetails.js
index c88ab9a28..747278ef7 100644
--- a/src/components/JobWizard/JobWizardSteps/JobWizardRunDetails/JobWizardRunDetails.js
+++ b/src/components/JobWizard/JobWizardSteps/JobWizardRunDetails/JobWizardRunDetails.js
@@ -59,6 +59,7 @@ const JobWizardRunDetails = ({
   isBatchInference,
   isEditMode,
   jobAdditionalData,
+  prePopulatedData,
   selectedFunctionData,
   stepIsActive
 }) => {
@@ -82,7 +83,7 @@ const JobWizardRunDetails = ({
     setSpyOnMethodChange(true)
 
     const functionParameters = getFunctionParameters(selectedFunctionData.functions, method)
-    const dataInputs = parseDataInputs(functionParameters)
+    const dataInputs = parseDataInputs(functionParameters, prePopulatedData.dataInputs)
     const predefinedParameters = parsePredefinedParameters(functionParameters)
     const methodData = getMethodData(selectedFunctionData, method)
 
@@ -317,6 +318,7 @@ JobWizardRunDetails.propTypes = {
   isBatchInference: PropTypes.bool.isRequired,
   isEditMode: PropTypes.bool.isRequired,
   jobAdditionalData: PropTypes.shape({}).isRequired,
+  prePopulatedData: PropTypes.shape({}).isRequired,
   selectedFunctionData: PropTypes.shape({}).isRequired
 }
 
diff --git a/src/components/ModelsPage/Models/Models.js b/src/components/ModelsPage/Models/Models.js
index e1c1174cb..46abf1a63 100644
--- a/src/components/ModelsPage/Models/Models.js
+++ b/src/components/ModelsPage/Models/Models.js
@@ -26,6 +26,7 @@ import AddArtifactTagPopUp from '../../../elements/AddArtifactTagPopUp/AddArtifa
 import DeployModelPopUp from '../../../elements/DeployModelPopUp/DeployModelPopUp'
 import ModelsView from './ModelsView'
 import RegisterModelModal from '../../../elements/RegisterModelModal/RegisterModelModal'
+import JobWizard from '../../JobWizard/JobWizard'
 
 import {
   fetchModel,
@@ -380,6 +381,14 @@ const Models = ({ fetchModelFeatureVector }) => {
     openPopUp(RegisterModelModal, { projectName: params.projectName, refresh: handleRefresh })
   }, [handleRefresh, params.projectName])
 
+  const handleTrainModel = () => {
+    openPopUp(JobWizard, {
+      params,
+      isTrain: true,
+      wizardTitle: 'Train model'
+    })
+  }
+
   return (
     <ModelsView
       actionsMenu={actionsMenu}
@@ -391,6 +400,7 @@ const Models = ({ fetchModelFeatureVector }) => {
       handleExpandRow={handleExpandRow}
       handleRefresh={handleRefresh}
       handleRegisterModel={handleRegisterModel}
+      handleTrainModel={handleTrainModel}
       isDemoMode={isDemoMode}
       models={models}
       pageData={pageData}
diff --git a/src/components/ModelsPage/Models/ModelsView.js b/src/components/ModelsPage/Models/ModelsView.js
index 24f2aa2c5..90314557a 100644
--- a/src/components/ModelsPage/Models/ModelsView.js
+++ b/src/components/ModelsPage/Models/ModelsView.js
@@ -34,7 +34,7 @@ import { getNoDataMessage } from '../../../utils/getNoDataMessage'
 import { removeModel } from '../../../reducers/artifactsReducer'
 import { filters } from './models.util'
 import { ACTIONS_MENU } from '../../../types'
-import { SECONDARY_BUTTON } from 'igz-controls/constants'
+import { SECONDARY_BUTTON, PRIMARY_BUTTON } from 'igz-controls/constants'
 
 const ModelsView = React.forwardRef(
   (
@@ -48,6 +48,7 @@ const ModelsView = React.forwardRef(
       handleExpandRow,
       handleRefresh,
       handleRegisterModel,
+      handleTrainModel,
       isDemoMode,
       models,
       pageData,
@@ -69,15 +70,23 @@ const ModelsView = React.forwardRef(
           <div className="table-container">
             <div className="content__action-bar-wrapper">
               <ModelsPageTabs />
-              {/* TODO: remove from demo in 1.4 */}
+              {/* TODO: remove from demo in 1.7 */}
               <ArtifactsActionBar
-                actionButton={{
-                  variant: SECONDARY_BUTTON,
-                  label: 'Register model',
-                  className: 'register-button',
-                  onClick: handleRegisterModel,
-                  hidden: !isDemoMode
-                }}
+                actionButtons={[
+                  {
+                    variant: PRIMARY_BUTTON,
+                    label: 'Train model',
+                    className: 'action-button',
+                    onClick: handleTrainModel
+                  },
+                  {
+                    variant: SECONDARY_BUTTON,
+                    label: 'Register model',
+                    className: 'action-button',
+                    onClick: handleRegisterModel,
+                    hidden: !isDemoMode
+                  }
+                ]}
                 filterMenuName={MODELS_FILTERS}
                 handleRefresh={handleRefresh}
                 page={MODELS_PAGE}
@@ -168,6 +177,7 @@ ModelsView.propTypes = {
   handleExpandRow: PropTypes.func.isRequired,
   handleRefresh: PropTypes.func.isRequired,
   handleRegisterModel: PropTypes.func.isRequired,
+  handleTrainModel: PropTypes.func.isRequired,
   isDemoMode: PropTypes.bool.isRequired,
   models: PropTypes.arrayOf(PropTypes.object).isRequired,
   pageData: PropTypes.object.isRequired,
diff --git a/src/components/Project/ProjectOverview/ProjectOverview.util.js b/src/components/Project/ProjectOverview/ProjectOverview.util.js
index ae39511bb..7f5a5c28e 100644
--- a/src/components/Project/ProjectOverview/ProjectOverview.util.js
+++ b/src/components/Project/ProjectOverview/ProjectOverview.util.js
@@ -317,7 +317,7 @@ export const getInitialCards = (params, navigate, isDemoMode) => {
               component: JobWizard,
               props: {
                 params,
-                isTrainModel: true,
+                isTrain: true,
                 wizardTitle: 'Train model',
                 isOverview: true
               },
@@ -326,8 +326,7 @@ export const getInitialCards = (params, navigate, isDemoMode) => {
           },
           label: 'Train model',
           tooltip:
-            'Train a new model based on an input dataset. You can also define hyperparameters to execute and compare multiple models.',
-          hidden: !isDemoMode
+            'Train a new model based on an input dataset. You can also define hyperparameters to execute and compare multiple models.'
         },
         {
           id: 'createWorkflow',
diff --git a/src/scss/main.scss b/src/scss/main.scss
index 89d40b776..adaec1cc5 100644
--- a/src/scss/main.scss
+++ b/src/scss/main.scss
@@ -810,8 +810,12 @@ iframe {
     display: flex;
     margin-left: 20px;
 
-    .register-button {
+    .action-button {
       width: 160px;
+
+      &:not(:last-of-type) {
+        margin-right: 10px;
+      }
     }
   }
 }
diff --git a/src/utils/generateTemplatesCategories.js b/src/utils/generateTemplatesCategories.js
index 73c3b9e75..77ef3f188 100644
--- a/src/utils/generateTemplatesCategories.js
+++ b/src/utils/generateTemplatesCategories.js
@@ -18,6 +18,8 @@ under the Apache 2.0 license is conditioned upon your compliance with
 such restriction.
 */
 
+import { isEmpty } from 'lodash'
+
 import { functionRunKinds } from '../components/Jobs/jobs.util'
 
 const excludedFunctionNames = ['batch-inference']
@@ -77,8 +79,15 @@ export const generateCategories = functionTemplates => {
   return { templates, templatesCategories }
 }
 
-export const generateHubCategories = functionTemplates => {
+export const generateHubCategories = (functionTemplates, allowedHubFunctions = {}) => {
   const hubFunctions = functionTemplates
+    .filter(
+      template =>
+        functionRunKinds.includes(template.spec.kind) &&
+        !excludedFunctionNames.includes(template.metadata.name) &&
+        (isEmpty(allowedHubFunctions) ||
+          Object.keys(allowedHubFunctions).includes(template.metadata.name))
+    )
     .map(template => ({
       ...template,
       ui: {
@@ -90,11 +99,6 @@ export const generateHubCategories = functionTemplates => {
           .map(funcTemplate => funcTemplate.metadata.version)
       }
     }))
-    .filter(
-      template =>
-        functionRunKinds.includes(template.spec.kind) &&
-        !excludedFunctionNames.includes(template.metadata.name)
-    )
 
   const hubFunctionsCategories = []
 
diff --git a/src/utils/panelPathScheme.js b/src/utils/panelPathScheme.js
index b424e2237..7c5568625 100644
--- a/src/utils/panelPathScheme.js
+++ b/src/utils/panelPathScheme.js
@@ -44,7 +44,7 @@ export const generateProjectsList = (projectsList, currentProject) =>
 export const generateArtifactsList = artifacts => {
   const generatedArtifacts = artifacts
     .map(artifact => {
-      const key = artifact.link_iteration ? artifact.link_iteration.db_key : artifact.key ?? ''
+      const key = artifact.link_iteration?.db_key || artifact.db_key || artifact.key || ''
       return {
         label: key,
         id: key