two version of R2R are here HEAD master

7 files changed, 1485 insertions, 0 deletions

diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/__init__.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/__init__.py
new file mode 100644
index 00000000..fdf8caba
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/__init__.py
@@ -0,0 +1,5 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/early_termination_policy.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/early_termination_policy.py
new file mode 100644
index 00000000..b1b928fc
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/early_termination_policy.py
@@ -0,0 +1,191 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+from abc import ABC
+from typing import Any, Optional, cast
+
+from azure.ai.ml._restclient.v2023_04_01_preview.models import BanditPolicy as RestBanditPolicy
+from azure.ai.ml._restclient.v2023_04_01_preview.models import EarlyTerminationPolicy as RestEarlyTerminationPolicy
+from azure.ai.ml._restclient.v2023_04_01_preview.models import EarlyTerminationPolicyType
+from azure.ai.ml._restclient.v2023_04_01_preview.models import MedianStoppingPolicy as RestMedianStoppingPolicy
+from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+    TruncationSelectionPolicy as RestTruncationSelectionPolicy,
+)
+from azure.ai.ml._utils.utils import camel_to_snake
+from azure.ai.ml.entities._mixins import RestTranslatableMixin
+
+
+class EarlyTerminationPolicy(ABC, RestTranslatableMixin):
+    def __init__(
+        self,
+        *,
+        delay_evaluation: int,
+        evaluation_interval: int,
+    ):
+        self.type = None
+        self.delay_evaluation = delay_evaluation
+        self.evaluation_interval = evaluation_interval
+
+    @classmethod
+    def _from_rest_object(cls, obj: RestEarlyTerminationPolicy) -> Optional["EarlyTerminationPolicy"]:
+        if not obj:
+            return None
+
+        policy: Any = None
+        if obj.policy_type == EarlyTerminationPolicyType.BANDIT:
+            policy = BanditPolicy._from_rest_object(obj)  # pylint: disable=protected-access
+
+        if obj.policy_type == EarlyTerminationPolicyType.MEDIAN_STOPPING:
+            policy = MedianStoppingPolicy._from_rest_object(obj)  # pylint: disable=protected-access
+
+        if obj.policy_type == EarlyTerminationPolicyType.TRUNCATION_SELECTION:
+            policy = TruncationSelectionPolicy._from_rest_object(obj)  # pylint: disable=protected-access
+
+        return cast(Optional["EarlyTerminationPolicy"], policy)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, EarlyTerminationPolicy):
+            raise NotImplementedError
+        res: bool = self._to_rest_object() == other._to_rest_object()
+        return res
+
+
+class BanditPolicy(EarlyTerminationPolicy):
+    """Defines an early termination policy based on slack criteria and a frequency and delay interval for evaluation.
+
+    :keyword delay_evaluation: Number of intervals by which to delay the first evaluation. Defaults to 0.
+    :paramtype delay_evaluation: int
+    :keyword evaluation_interval: Interval (number of runs) between policy evaluations. Defaults to 0.
+    :paramtype evaluation_interval: int
+    :keyword slack_amount: Absolute distance allowed from the best performing run. Defaults to 0.
+    :paramtype slack_amount: float
+    :keyword slack_factor: Ratio of the allowed distance from the best performing run. Defaults to 0.
+    :paramtype slack_factor: float
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_bandit_policy]
+            :end-before: [END configure_sweep_job_bandit_policy]
+            :language: python
+            :dedent: 8
+            :caption: Configuring BanditPolicy early termination of a hyperparameter sweep on a Command job.
+    """
+
+    def __init__(
+        self,
+        *,
+        delay_evaluation: int = 0,
+        evaluation_interval: int = 0,
+        slack_amount: float = 0,
+        slack_factor: float = 0,
+    ) -> None:
+        super().__init__(delay_evaluation=delay_evaluation, evaluation_interval=evaluation_interval)
+        self.type = EarlyTerminationPolicyType.BANDIT.lower()
+        self.slack_factor = slack_factor
+        self.slack_amount = slack_amount
+
+    def _to_rest_object(self) -> RestBanditPolicy:
+        return RestBanditPolicy(
+            delay_evaluation=self.delay_evaluation,
+            evaluation_interval=self.evaluation_interval,
+            slack_factor=self.slack_factor,
+            slack_amount=self.slack_amount,
+        )
+
+    @classmethod
+    def _from_rest_object(cls, obj: RestBanditPolicy) -> "BanditPolicy":
+        return cls(
+            delay_evaluation=obj.delay_evaluation,
+            evaluation_interval=obj.evaluation_interval,
+            slack_factor=obj.slack_factor,
+            slack_amount=obj.slack_amount,
+        )
+
+
+class MedianStoppingPolicy(EarlyTerminationPolicy):
+    """Defines an early termination policy based on a running average of the primary metric of all runs.
+
+    :keyword delay_evaluation: Number of intervals by which to delay the first evaluation. Defaults to 0.
+    :paramtype delay_evaluation: int
+    :keyword evaluation_interval: Interval (number of runs) between policy evaluations. Defaults to 1.
+    :paramtype evaluation_interval: int
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_median_stopping_policy]
+            :end-before: [END configure_sweep_job_median_stopping_policy]
+            :language: python
+            :dedent: 8
+            :caption: Configuring an early termination policy for a hyperparameter sweep job using MedianStoppingPolicy
+    """
+
+    def __init__(
+        self,
+        *,
+        delay_evaluation: int = 0,
+        evaluation_interval: int = 1,
+    ) -> None:
+        super().__init__(delay_evaluation=delay_evaluation, evaluation_interval=evaluation_interval)
+        self.type = camel_to_snake(EarlyTerminationPolicyType.MEDIAN_STOPPING)
+
+    def _to_rest_object(self) -> RestMedianStoppingPolicy:
+        return RestMedianStoppingPolicy(
+            delay_evaluation=self.delay_evaluation, evaluation_interval=self.evaluation_interval
+        )
+
+    @classmethod
+    def _from_rest_object(cls, obj: RestMedianStoppingPolicy) -> "MedianStoppingPolicy":
+        return cls(
+            delay_evaluation=obj.delay_evaluation,
+            evaluation_interval=obj.evaluation_interval,
+        )
+
+
+class TruncationSelectionPolicy(EarlyTerminationPolicy):
+    """Defines an early termination policy that cancels a given percentage of runs at each evaluation interval.
+
+    :keyword delay_evaluation: Number of intervals by which to delay the first evaluation. Defaults to 0.
+    :paramtype delay_evaluation: int
+    :keyword evaluation_interval: Interval (number of runs) between policy evaluations. Defaults to 0.
+    :paramtype evaluation_interval: int
+    :keyword truncation_percentage: The percentage of runs to cancel at each evaluation interval. Defaults to 0.
+    :paramtype truncation_percentage: int
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_truncation_selection_policy]
+            :end-before: [END configure_sweep_job_truncation_selection_policy]
+            :language: python
+            :dedent: 8
+            :caption: Configuring an early termination policy for a hyperparameter sweep job
+                using TruncationStoppingPolicy
+    """
+
+    def __init__(
+        self,
+        *,
+        delay_evaluation: int = 0,
+        evaluation_interval: int = 0,
+        truncation_percentage: int = 0,
+    ) -> None:
+        super().__init__(delay_evaluation=delay_evaluation, evaluation_interval=evaluation_interval)
+        self.type = camel_to_snake(EarlyTerminationPolicyType.TRUNCATION_SELECTION)
+        self.truncation_percentage = truncation_percentage
+
+    def _to_rest_object(self) -> RestTruncationSelectionPolicy:
+        return RestTruncationSelectionPolicy(
+            delay_evaluation=self.delay_evaluation,
+            evaluation_interval=self.evaluation_interval,
+            truncation_percentage=self.truncation_percentage,
+        )
+
+    @classmethod
+    def _from_rest_object(cls, obj: RestTruncationSelectionPolicy) -> "TruncationSelectionPolicy":
+        return cls(
+            delay_evaluation=obj.delay_evaluation,
+            evaluation_interval=obj.evaluation_interval,
+            truncation_percentage=obj.truncation_percentage,
+        )
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/objective.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/objective.py
new file mode 100644
index 00000000..45e13332
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/objective.py
@@ -0,0 +1,53 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+from typing import Optional
+
+from azure.ai.ml._restclient.v2023_08_01_preview.models import Objective as RestObjective
+from azure.ai.ml.entities._mixins import RestTranslatableMixin
+
+
+class Objective(RestTranslatableMixin):
+    """Optimization objective.
+
+    :param goal: Defines supported metric goals for hyperparameter tuning. Accepted values
+     are: "minimize", "maximize".
+    :type goal: str
+    :param primary_metric: The name of the metric to optimize.
+    :type primary_metric: str
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_bayesian_sampling_algorithm]
+            :end-before: [END configure_sweep_job_bayesian_sampling_algorithm]
+            :language: python
+            :dedent: 8
+            :caption: Assigning an objective to a SweepJob.
+    """
+
+    def __init__(self, goal: Optional[str], primary_metric: Optional[str] = None) -> None:
+        """Optimization objective.
+
+        :param goal: Defines supported metric goals for hyperparameter tuning. Acceptable values
+            are: "minimize" or "maximize".
+        :type goal: str
+        :param primary_metric: The name of the metric to optimize.
+        :type primary_metric: str
+        """
+        if goal is not None:
+            self.goal = goal.lower()
+        self.primary_metric = primary_metric
+
+    def _to_rest_object(self) -> RestObjective:
+        return RestObjective(
+            goal=self.goal,
+            primary_metric=self.primary_metric,
+        )
+
+    @classmethod
+    def _from_rest_object(cls, obj: RestObjective) -> Optional["Objective"]:
+        if not obj:
+            return None
+
+        return cls(goal=obj.goal, primary_metric=obj.primary_metric)
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/parameterized_sweep.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/parameterized_sweep.py
new file mode 100644
index 00000000..5d69201f
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/parameterized_sweep.py
@@ -0,0 +1,341 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+from typing import Any, Dict, List, Optional, Type, Union
+
+from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException
+
+from ..job_limits import SweepJobLimits
+from ..job_resource_configuration import JobResourceConfiguration
+from ..queue_settings import QueueSettings
+from .early_termination_policy import (
+    BanditPolicy,
+    EarlyTerminationPolicy,
+    EarlyTerminationPolicyType,
+    MedianStoppingPolicy,
+    TruncationSelectionPolicy,
+)
+from .objective import Objective
+from .sampling_algorithm import (
+    BayesianSamplingAlgorithm,
+    GridSamplingAlgorithm,
+    RandomSamplingAlgorithm,
+    RestBayesianSamplingAlgorithm,
+    RestGridSamplingAlgorithm,
+    RestRandomSamplingAlgorithm,
+    RestSamplingAlgorithm,
+    SamplingAlgorithm,
+    SamplingAlgorithmType,
+)
+
+SAMPLING_ALGORITHM_TO_REST_CONSTRUCTOR: Dict[SamplingAlgorithmType, Type[RestSamplingAlgorithm]] = {
+    SamplingAlgorithmType.RANDOM: RestRandomSamplingAlgorithm,
+    SamplingAlgorithmType.GRID: RestGridSamplingAlgorithm,
+    SamplingAlgorithmType.BAYESIAN: RestBayesianSamplingAlgorithm,
+}
+
+SAMPLING_ALGORITHM_CONSTRUCTOR: Dict[SamplingAlgorithmType, Type[SamplingAlgorithm]] = {
+    SamplingAlgorithmType.RANDOM: RandomSamplingAlgorithm,
+    SamplingAlgorithmType.GRID: GridSamplingAlgorithm,
+    SamplingAlgorithmType.BAYESIAN: BayesianSamplingAlgorithm,
+}
+
+
+class ParameterizedSweep:  # pylint:disable=too-many-instance-attributes
+    """Shared logic for standalone and pipeline sweep job."""
+
+    def __init__(
+        self,
+        limits: Optional[SweepJobLimits] = None,
+        sampling_algorithm: Optional[Union[str, SamplingAlgorithm]] = None,
+        objective: Optional[Union[Dict, Objective]] = None,
+        early_termination: Optional[Any] = None,
+        search_space: Optional[Dict] = None,
+        queue_settings: Optional[QueueSettings] = None,
+        resources: Optional[Union[dict, JobResourceConfiguration]] = None,
+    ) -> None:
+        """
+        :param limits: Limits for sweep job.
+        :type limits: ~azure.ai.ml.sweep.SweepJobLimits
+        :param sampling_algorithm: Sampling algorithm for sweep job.
+        :type sampling_algorithm: ~azure.ai.ml.sweep.SamplingAlgorithm
+        :param objective: Objective for sweep job.
+        :type objective: ~azure.ai.ml.sweep.Objective
+        :param early_termination: Early termination policy for sweep job.
+        :type early_termination: ~azure.ai.ml.entities._job.sweep.early_termination_policy.EarlyTerminationPolicy
+        :param search_space: Search space for sweep job.
+        :type search_space: Dict[str, Union[
+            ~azure.ai.ml.sweep.Choice,
+            ~azure.ai.ml.sweep.LogNormal,
+            ~azure.ai.ml.sweep.LogUniform,
+            ~azure.ai.ml.sweep.Normal,
+            ~azure.ai.ml.sweep.QLogNormal,
+            ~azure.ai.ml.sweep.QLogUniform,
+            ~azure.ai.ml.sweep.QNormal,
+            ~azure.ai.ml.sweep.QUniform,
+            ~azure.ai.ml.sweep.Randint,
+            ~azure.ai.ml.sweep.Uniform
+
+            ]]
+        :param queue_settings: Queue settings for sweep job.
+        :type queue_settings: ~azure.ai.ml.entities.QueueSettings
+        :param resources: Compute Resource configuration for the job.
+        :type resources: ~azure.ai.ml.entities.ResourceConfiguration
+        """
+        self.sampling_algorithm = sampling_algorithm
+        self.early_termination = early_termination  # type: ignore[assignment]
+        self._limits = limits
+        self.search_space = search_space
+        self.queue_settings = queue_settings
+        self.objective: Optional[Objective] = None
+        self.resources = resources
+
+        if isinstance(objective, Dict):
+            self.objective = Objective(**objective)
+        else:
+            self.objective = objective
+
+    @property
+    def resources(self) -> Optional[Union[dict, JobResourceConfiguration]]:
+        """Resources for sweep job.
+
+        :returns: Resources for sweep job.
+        :rtype: ~azure.ai.ml.entities.ResourceConfiguration
+        """
+        return self._resources
+
+    @resources.setter
+    def resources(self, value: Optional[Union[dict, JobResourceConfiguration]]) -> None:
+        """Set Resources for sweep job.
+
+        :param value: Compute Resource configuration for the job.
+        :type value: ~azure.ai.ml.entities.ResourceConfiguration
+        """
+        if isinstance(value, dict):
+            value = JobResourceConfiguration(**value)
+        self._resources = value
+
+    @property
+    def limits(self) -> Optional[SweepJobLimits]:
+        """Limits for sweep job.
+
+        :returns: Limits for sweep job.
+        :rtype: ~azure.ai.ml.sweep.SweepJobLimits
+        """
+        return self._limits
+
+    @limits.setter
+    def limits(self, value: SweepJobLimits) -> None:
+        """Set limits for sweep job.
+
+        :param value: Limits for sweep job.
+        :type value: ~azure.ai.ml.sweep.SweepJobLimits
+        """
+        if not isinstance(value, SweepJobLimits):
+            msg = f"limits must be SweepJobLimits but get {type(value)} instead"
+            raise ValidationException(
+                message=msg,
+                no_personal_data_message=msg,
+                target=ErrorTarget.SWEEP_JOB,
+                error_category=ErrorCategory.USER_ERROR,
+                error_type=ValidationErrorType.INVALID_VALUE,
+            )
+        self._limits = value
+
+    def set_resources(
+        self,
+        *,
+        instance_type: Optional[Union[str, List[str]]] = None,
+        instance_count: Optional[int] = None,
+        locations: Optional[List[str]] = None,
+        properties: Optional[Dict] = None,
+        docker_args: Optional[str] = None,
+        shm_size: Optional[str] = None,
+    ) -> None:
+        """Set resources for Sweep.
+
+        :keyword instance_type: The instance type to use for the job.
+        :paramtype instance_type: Optional[Union[str, List[str]]]
+        :keyword instance_count: The number of instances to use for the job.
+        :paramtype instance_count: Optional[int]
+        :keyword locations: The locations to use for the job.
+        :paramtype locations: Optional[List[str]]
+        :keyword properties: The properties for the job.
+        :paramtype properties: Optional[Dict]
+        :keyword docker_args: The docker arguments for the job.
+        :paramtype docker_args: Optional[str]
+        :keyword shm_size: The shared memory size for the job.
+        :paramtype shm_size: Optional[str]
+        """
+        if self.resources is None:
+            self.resources = JobResourceConfiguration()
+
+        if not isinstance(self.resources, dict):
+            if locations is not None:
+                self.resources.locations = locations
+            if instance_type is not None:
+                self.resources.instance_type = instance_type
+            if instance_count is not None:
+                self.resources.instance_count = instance_count
+            if properties is not None:
+                self.resources.properties = properties
+            if docker_args is not None:
+                self.resources.docker_args = docker_args
+            if shm_size is not None:
+                self.resources.shm_size = shm_size
+
+    def set_limits(
+        self,
+        *,
+        max_concurrent_trials: Optional[int] = None,
+        max_total_trials: Optional[int] = None,
+        timeout: Optional[int] = None,
+        trial_timeout: Optional[int] = None,
+    ) -> None:
+        """Set limits for Sweep node. Leave parameters as None if you don't want to update corresponding values.
+
+        :keyword max_concurrent_trials: maximum concurrent trial number.
+        :paramtype max_concurrent_trials: int
+        :keyword max_total_trials: maximum total trial number.
+        :paramtype max_total_trials: int
+        :keyword timeout: total timeout in seconds for sweep node
+        :paramtype timeout: int
+        :keyword trial_timeout: timeout in seconds for each trial
+        :paramtype trial_timeout: int
+        """
+        # Looks related to https://github.com/pylint-dev/pylint/issues/3502, still an open issue
+        # pylint:disable=attribute-defined-outside-init
+        if self._limits is None:
+            self._limits = SweepJobLimits(
+                max_concurrent_trials=max_concurrent_trials,
+                max_total_trials=max_total_trials,
+                timeout=timeout,
+                trial_timeout=trial_timeout,
+            )
+        else:
+            if self.limits is not None:
+                if max_concurrent_trials is not None:
+                    self.limits.max_concurrent_trials = max_concurrent_trials
+                if max_total_trials is not None:
+                    self.limits.max_total_trials = max_total_trials
+                if timeout is not None:
+                    self.limits.timeout = timeout
+                if trial_timeout is not None:
+                    self.limits.trial_timeout = trial_timeout
+
+    def set_objective(self, *, goal: Optional[str] = None, primary_metric: Optional[str] = None) -> None:
+        """Set the sweep object.. Leave parameters as None if you don't want to update corresponding values.
+
+        :keyword goal: Defines supported metric goals for hyperparameter tuning. Acceptable values are:
+            "minimize" and "maximize".
+        :paramtype goal: str
+        :keyword primary_metric: Name of the metric to optimize.
+        :paramtype primary_metric: str
+        """
+
+        if self.objective is not None:
+            if goal:
+                self.objective.goal = goal
+            if primary_metric:
+                self.objective.primary_metric = primary_metric
+        else:
+            self.objective = Objective(goal=goal, primary_metric=primary_metric)
+
+    @property
+    def sampling_algorithm(self) -> Optional[Union[str, SamplingAlgorithm]]:
+        """Sampling algorithm for sweep job.
+
+        :returns: Sampling algorithm for sweep job.
+        :rtype: ~azure.ai.ml.sweep.SamplingAlgorithm
+        """
+        return self._sampling_algorithm
+
+    @sampling_algorithm.setter
+    def sampling_algorithm(self, value: Optional[Union[SamplingAlgorithm, str]] = None) -> None:
+        """Set sampling algorithm for sweep job.
+
+        :param value: Sampling algorithm for sweep job.
+        :type value: ~azure.ai.ml.sweep.SamplingAlgorithm
+        """
+        if value is None:
+            self._sampling_algorithm = None
+        elif isinstance(value, SamplingAlgorithm) or (
+            isinstance(value, str) and value.lower().capitalize() in SAMPLING_ALGORITHM_CONSTRUCTOR
+        ):
+            self._sampling_algorithm = value
+        else:
+            msg = f"unsupported sampling algorithm: {value}"
+            raise ValidationException(
+                message=msg,
+                no_personal_data_message=msg,
+                target=ErrorTarget.SWEEP_JOB,
+                error_category=ErrorCategory.USER_ERROR,
+                error_type=ValidationErrorType.INVALID_VALUE,
+            )
+
+    def _get_rest_sampling_algorithm(self) -> RestSamplingAlgorithm:
+        # TODO: self.sampling_algorithm will always return SamplingAlgorithm
+        if isinstance(self.sampling_algorithm, SamplingAlgorithm):
+            return self.sampling_algorithm._to_rest_object()  # pylint: disable=protected-access
+
+        if isinstance(self.sampling_algorithm, str):
+            return SAMPLING_ALGORITHM_CONSTRUCTOR[  # pylint: disable=protected-access
+                SamplingAlgorithmType(self.sampling_algorithm.lower().capitalize())
+            ]()._to_rest_object()
+
+        msg = f"Received unsupported value {self._sampling_algorithm} as the sampling algorithm"
+        raise ValidationException(
+            message=msg,
+            no_personal_data_message=msg,
+            target=ErrorTarget.SWEEP_JOB,
+            error_category=ErrorCategory.USER_ERROR,
+            error_type=ValidationErrorType.INVALID_VALUE,
+        )
+
+    @property
+    def early_termination(self) -> Optional[Union[str, EarlyTerminationPolicy]]:
+        """Early termination policy for sweep job.
+
+        :returns: Early termination policy for sweep job.
+        :rtype: ~azure.ai.ml.entities._job.sweep.early_termination_policy.EarlyTerminationPolicy
+        """
+        return self._early_termination
+
+    @early_termination.setter
+    def early_termination(self, value: Any) -> None:
+        """Set early termination policy for sweep job.
+
+        :param value: Early termination policy for sweep job.
+        :type value: ~azure.ai.ml.entities._job.sweep.early_termination_policy.EarlyTerminationPolicy
+        """
+        self._early_termination: Optional[Union[str, EarlyTerminationPolicy]]
+        if value is None:
+            self._early_termination = None
+        elif isinstance(value, EarlyTerminationPolicy):
+            self._early_termination = value
+        elif isinstance(value, str):
+            value = value.lower().capitalize()
+            if value == EarlyTerminationPolicyType.BANDIT:
+                self._early_termination = BanditPolicy()
+            elif value == EarlyTerminationPolicyType.MEDIAN_STOPPING:
+                self._early_termination = MedianStoppingPolicy()
+            elif value == EarlyTerminationPolicyType.TRUNCATION_SELECTION:
+                self._early_termination = TruncationSelectionPolicy()
+            else:
+                msg = f"Received unsupported value {value} as the early termination policy"
+                raise ValidationException(
+                    message=msg,
+                    no_personal_data_message=msg,
+                    target=ErrorTarget.SWEEP_JOB,
+                    error_category=ErrorCategory.USER_ERROR,
+                    error_type=ValidationErrorType.INVALID_VALUE,
+                )
+        else:
+            msg = f"Received unsupported value of type {type(value)} as the early termination policy"
+            raise ValidationException(
+                message=msg,
+                no_personal_data_message=msg,
+                target=ErrorTarget.SWEEP_JOB,
+                error_category=ErrorCategory.USER_ERROR,
+                error_type=ValidationErrorType.INVALID_VALUE,
+            )
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/sampling_algorithm.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/sampling_algorithm.py
new file mode 100644
index 00000000..d0bf795d
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/sampling_algorithm.py
@@ -0,0 +1,141 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+from abc import ABC
+from typing import Any, Optional, Union, cast
+
+from azure.ai.ml._restclient.v2023_08_01_preview.models import (
+    BayesianSamplingAlgorithm as RestBayesianSamplingAlgorithm,
+)
+from azure.ai.ml._restclient.v2023_08_01_preview.models import GridSamplingAlgorithm as RestGridSamplingAlgorithm
+from azure.ai.ml._restclient.v2023_08_01_preview.models import RandomSamplingAlgorithm as RestRandomSamplingAlgorithm
+from azure.ai.ml._restclient.v2023_08_01_preview.models import SamplingAlgorithm as RestSamplingAlgorithm
+from azure.ai.ml._restclient.v2023_08_01_preview.models import SamplingAlgorithmType
+from azure.ai.ml.entities._mixins import RestTranslatableMixin
+
+
+class SamplingAlgorithm(ABC, RestTranslatableMixin):
+    """Base class for sampling algorithms.
+
+    This class should not be instantiated directly. Instead, use one of its subclasses.
+    """
+
+    def __init__(self) -> None:
+        self.type = None
+
+    @classmethod
+    def _from_rest_object(cls, obj: RestSamplingAlgorithm) -> Optional["SamplingAlgorithm"]:
+        if not obj:
+            return None
+
+        sampling_algorithm: Any = None
+        if obj.sampling_algorithm_type == SamplingAlgorithmType.RANDOM:
+            sampling_algorithm = RandomSamplingAlgorithm._from_rest_object(obj)  # pylint: disable=protected-access
+
+        if obj.sampling_algorithm_type == SamplingAlgorithmType.GRID:
+            sampling_algorithm = GridSamplingAlgorithm._from_rest_object(obj)  # pylint: disable=protected-access
+
+        if obj.sampling_algorithm_type == SamplingAlgorithmType.BAYESIAN:
+            sampling_algorithm = BayesianSamplingAlgorithm._from_rest_object(obj)  # pylint: disable=protected-access
+
+        return cast(Optional["SamplingAlgorithm"], sampling_algorithm)
+
+
+class RandomSamplingAlgorithm(SamplingAlgorithm):
+    """Random Sampling Algorithm.
+
+    :keyword rule: The specific type of random algorithm. Accepted values are: "random" and "sobol".
+    :type rule: str
+    :keyword seed: The seed for random number generation.
+    :paramtype seed: int
+    :keyword logbase: A positive number or the number "e" in string format to be used as the base for log
+        based random sampling.
+    :paramtype logbase: Union[float, str]
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_random_sampling_algorithm]
+            :end-before: [END configure_sweep_job_random_sampling_algorithm]
+            :language: python
+            :dedent: 8
+            :caption: Assigning a random sampling algorithm for a SweepJob
+    """
+
+    def __init__(
+        self,
+        *,
+        rule: Optional[str] = None,
+        seed: Optional[int] = None,
+        logbase: Optional[Union[float, str]] = None,
+    ) -> None:
+        super().__init__()
+        self.type = SamplingAlgorithmType.RANDOM.lower()
+        self.rule = rule
+        self.seed = seed
+        self.logbase = logbase
+
+    def _to_rest_object(self) -> RestRandomSamplingAlgorithm:
+        return RestRandomSamplingAlgorithm(
+            rule=self.rule,
+            seed=self.seed,
+            logbase=self.logbase,
+        )
+
+    @classmethod
+    def _from_rest_object(cls, obj: RestRandomSamplingAlgorithm) -> "RandomSamplingAlgorithm":
+        return cls(
+            rule=obj.rule,
+            seed=obj.seed,
+            logbase=obj.logbase,
+        )
+
+
+class GridSamplingAlgorithm(SamplingAlgorithm):
+    """Grid Sampling Algorithm.
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_grid_sampling_algorithm]
+            :end-before: [END configure_sweep_job_grid_sampling_algorithm]
+            :language: python
+            :dedent: 8
+            :caption: Assigning a grid sampling algorithm for a SweepJob
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.type = SamplingAlgorithmType.GRID.lower()
+
+    def _to_rest_object(self) -> RestGridSamplingAlgorithm:
+        return RestGridSamplingAlgorithm()
+
+    @classmethod
+    def _from_rest_object(cls, obj: RestGridSamplingAlgorithm) -> "GridSamplingAlgorithm":
+        return cls()
+
+
+class BayesianSamplingAlgorithm(SamplingAlgorithm):
+    """Bayesian Sampling Algorithm.
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_bayesian_sampling_algorithm]
+            :end-before: [END configure_sweep_job_bayesian_sampling_algorithm]
+            :language: python
+            :dedent: 8
+            :caption: Assigning a Bayesian sampling algorithm for a SweepJob
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.type = SamplingAlgorithmType.BAYESIAN.lower()
+
+    def _to_rest_object(self) -> RestBayesianSamplingAlgorithm:
+        return RestBayesianSamplingAlgorithm()
+
+    @classmethod
+    def _from_rest_object(cls, obj: RestBayesianSamplingAlgorithm) -> "BayesianSamplingAlgorithm":
+        return cls()
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/search_space.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/search_space.py
new file mode 100644
index 00000000..bbc08d98
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/search_space.py
@@ -0,0 +1,393 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# pylint: disable=protected-access
+
+from abc import ABC
+from typing import Any, List, Optional, Union
+
+from azure.ai.ml.constants._common import TYPE
+from azure.ai.ml.constants._job.sweep import SearchSpace
+from azure.ai.ml.entities._mixins import RestTranslatableMixin
+from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, JobException
+
+
+class SweepDistribution(ABC, RestTranslatableMixin):
+    """Base class for sweep distribution configuration.
+
+    This class should not be instantiated directly. Instead, use one of its subclasses.
+
+    :keyword type: Type of distribution.
+    :paramtype type: str
+    """
+
+    def __init__(self, *, type: Optional[str] = None) -> None:  # pylint: disable=redefined-builtin
+        self.type = type
+
+    @classmethod
+    def _from_rest_object(cls, obj: List) -> "SweepDistribution":
+        mapping = {
+            SearchSpace.CHOICE: Choice,
+            SearchSpace.NORMAL: Normal,
+            SearchSpace.LOGNORMAL: LogNormal,
+            SearchSpace.QNORMAL: QNormal,
+            SearchSpace.QLOGNORMAL: QLogNormal,
+            SearchSpace.RANDINT: Randint,
+            SearchSpace.UNIFORM: Uniform,
+            SearchSpace.QUNIFORM: QUniform,
+            SearchSpace.LOGUNIFORM: LogUniform,
+            SearchSpace.QLOGUNIFORM: QLogUniform,
+        }
+
+        ss_class: Any = mapping.get(obj[0], None)
+        if ss_class:
+            res: SweepDistribution = ss_class._from_rest_object(obj)
+            return res
+
+        msg = f"Unknown search space type: {obj[0]}"
+        raise JobException(
+            message=msg,
+            no_personal_data_message=msg,
+            target=ErrorTarget.SWEEP_JOB,
+            error_category=ErrorCategory.SYSTEM_ERROR,
+        )
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, SweepDistribution):
+            return NotImplemented
+        res: bool = self._to_rest_object() == other._to_rest_object()
+        return res
+
+
+class Choice(SweepDistribution):
+    """Choice distribution configuration.
+
+    :param values: List of values to choose from.
+    :type values: list[Union[float, str, dict]]
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_choice_loguniform]
+            :end-before: [END configure_sweep_job_choice_loguniform]
+            :language: python
+            :dedent: 8
+            :caption: Using Choice distribution to set values for a hyperparameter sweep
+    """
+
+    def __init__(self, values: Optional[List[Union[float, str, dict]]] = None, **kwargs: Any) -> None:
+        kwargs.setdefault(TYPE, SearchSpace.CHOICE)
+        super().__init__(**kwargs)
+        self.values = values
+
+    def _to_rest_object(self) -> List:
+        items: List = []
+        if self.values is not None:
+            for value in self.values:
+                if isinstance(value, dict):
+                    rest_dict = {}
+                    for k, v in value.items():
+                        if isinstance(v, SweepDistribution):
+                            rest_dict[k] = v._to_rest_object()
+                        else:
+                            rest_dict[k] = v
+                    items.append(rest_dict)
+                else:
+                    items.append(value)
+        return [self.type, [items]]
+
+    @classmethod
+    def _from_rest_object(cls, obj: List) -> "Choice":
+        rest_values = obj[1][0]
+        from_rest_values = []
+        for rest_value in rest_values:
+            if isinstance(rest_value, dict):
+                from_rest_dict = {}
+                for k, v in rest_value.items():
+                    try:
+                        # first assume that any dictionary value is a valid distribution (i.e. normal, uniform, etc)
+                        # and try to deserialize it into a the correct SDK distribution object
+                        from_rest_dict[k] = SweepDistribution._from_rest_object(v)
+                    except Exception:  # pylint: disable=W0718
+                        # if an exception is raised, assume that the value was not a valid distribution and use the
+                        # value as it is for deserialization
+                        from_rest_dict[k] = v
+                from_rest_values.append(from_rest_dict)
+            else:
+                from_rest_values.append(rest_value)
+        return Choice(values=from_rest_values)  # type: ignore[arg-type]
+
+
+class Normal(SweepDistribution):
+    """Normal distribution configuration.
+
+    :param mu: Mean of the distribution.
+    :type mu: float
+    :param sigma: Standard deviation of the distribution.
+    :type sigma: float
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_randint_normal]
+            :end-before: [END configure_sweep_job_randint_normal]
+            :language: python
+            :dedent: 8
+            :caption: Configuring Normal distributions for a hyperparameter sweep on a Command job.
+    """
+
+    def __init__(self, mu: Optional[float] = None, sigma: Optional[float] = None, **kwargs: Any) -> None:
+        kwargs.setdefault(TYPE, SearchSpace.NORMAL)
+        super().__init__(**kwargs)
+        self.mu = mu
+        self.sigma = sigma
+
+    def _to_rest_object(self) -> List:
+        return [self.type, [self.mu, self.sigma]]
+
+    @classmethod
+    def _from_rest_object(cls, obj: List) -> "Normal":
+        return cls(mu=obj[1][0], sigma=obj[1][1])
+
+
+class LogNormal(Normal):
+    """LogNormal distribution configuration.
+
+    :param mu: Mean of the log of the distribution.
+    :type mu: float
+    :param sigma: Standard deviation of the log of the distribution.
+    :type sigma: float
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_lognormal_qlognormal]
+            :end-before: [END configure_sweep_job_lognormal_qlognormal]
+            :language: python
+            :dedent: 8
+            :caption: Configuring LogNormal distributions for a hyperparameter sweep on a Command job.
+    """
+
+    def __init__(self, mu: Optional[float] = None, sigma: Optional[float] = None, **kwargs: Any) -> None:
+        kwargs.setdefault(TYPE, SearchSpace.LOGNORMAL)
+        super().__init__(mu=mu, sigma=sigma, **kwargs)
+
+
+class QNormal(Normal):
+    """QNormal distribution configuration.
+
+    :param mu: Mean of the distribution.
+    :type mu: float
+    :param sigma: Standard deviation of the distribution.
+    :type sigma: float
+    :param q: Quantization factor.
+    :type q: int
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_qloguniform_qnormal]
+            :end-before: [END configure_sweep_job_qloguniform_qnormal]
+            :language: python
+            :dedent: 8
+            :caption: Configuring QNormal distributions for a hyperparameter sweep on a Command job.
+    """
+
+    def __init__(
+        self, mu: Optional[float] = None, sigma: Optional[float] = None, q: Optional[int] = None, **kwargs: Any
+    ) -> None:
+        kwargs.setdefault(TYPE, SearchSpace.QNORMAL)
+        super().__init__(mu=mu, sigma=sigma, **kwargs)
+        self.q = q
+
+    def _to_rest_object(self) -> List:
+        return [self.type, [self.mu, self.sigma, self.q]]
+
+    @classmethod
+    def _from_rest_object(cls, obj: List) -> "QNormal":
+        return cls(mu=obj[1][0], sigma=obj[1][1], q=obj[1][2])
+
+
+class QLogNormal(QNormal):
+    """QLogNormal distribution configuration.
+
+    :param mu: Mean of the log of the distribution.
+    :type mu: Optional[float]
+    :param sigma: Standard deviation of the log of the distribution.
+    :type sigma: Optional[float]
+    :param q: Quantization factor.
+    :type q: Optional[int]
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_lognormal_qlognormal]
+            :end-before: [END configure_sweep_job_lognormal_qlognormal]
+            :language: python
+            :dedent: 8
+            :caption: Configuring QLogNormal distributions for a hyperparameter sweep on a Command job.
+    """
+
+    def __init__(
+        self, mu: Optional[float] = None, sigma: Optional[float] = None, q: Optional[int] = None, **kwargs: Any
+    ) -> None:
+        kwargs.setdefault(TYPE, SearchSpace.QLOGNORMAL)
+        super().__init__(mu=mu, sigma=sigma, q=q, **kwargs)
+
+
+class Randint(SweepDistribution):
+    """Randint distribution configuration.
+
+    :param upper: Upper bound of the distribution.
+    :type upper: int
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_randint_normal]
+            :end-before: [END configure_sweep_job_randint_normal]
+            :language: python
+            :dedent: 8
+            :caption: Configuring Randint distributions for a hyperparameter sweep on a Command job.
+    """
+
+    def __init__(self, upper: Optional[int] = None, **kwargs: Any) -> None:
+        kwargs.setdefault(TYPE, SearchSpace.RANDINT)
+        super().__init__(**kwargs)
+        self.upper = upper
+
+    def _to_rest_object(self) -> List:
+        return [self.type, [self.upper]]
+
+    @classmethod
+    def _from_rest_object(cls, obj: List) -> "Randint":
+        return cls(upper=obj[1][0])
+
+
+class Uniform(SweepDistribution):
+    """
+
+    Uniform distribution configuration.
+
+    :param min_value: Minimum value of the distribution.
+    :type min_value: float
+    :param max_value: Maximum value of the distribution.
+    :type max_value: float
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_uniform]
+            :end-before: [END configure_sweep_job_uniform]
+            :language: python
+            :dedent: 8
+            :caption: Configuring Uniform distributions for learning rates and momentum
+                during a hyperparameter sweep on a Command job.
+    """
+
+    def __init__(self, min_value: Optional[float] = None, max_value: Optional[float] = None, **kwargs: Any) -> None:
+        kwargs.setdefault(TYPE, SearchSpace.UNIFORM)
+        super().__init__(**kwargs)
+        self.min_value = min_value
+        self.max_value = max_value
+
+    def _to_rest_object(self) -> List:
+        return [self.type, [self.min_value, self.max_value]]
+
+    @classmethod
+    def _from_rest_object(cls, obj: List) -> "Uniform":
+        return cls(min_value=obj[1][0], max_value=obj[1][1])
+
+
+class LogUniform(Uniform):
+    """LogUniform distribution configuration.
+
+    :param min_value: Minimum value of the log of the distribution.
+    :type min_value: float
+    :param max_value: Maximum value of the log of the distribution.
+    :type max_value: float
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_choice_loguniform]
+            :end-before: [END configure_sweep_job_choice_loguniform]
+            :language: python
+            :dedent: 8
+            :caption: Configuring a LogUniform distribution for a hyperparameter sweep job learning rate
+    """
+
+    def __init__(self, min_value: Optional[float] = None, max_value: Optional[float] = None, **kwargs: Any) -> None:
+        kwargs.setdefault(TYPE, SearchSpace.LOGUNIFORM)
+        super().__init__(min_value=min_value, max_value=max_value, **kwargs)
+
+
+class QUniform(Uniform):
+    """QUniform distribution configuration.
+
+    :param min_value: Minimum value of the distribution.
+    :type min_value: Optional[Union[int, float]]
+    :param max_value: Maximum value of the distribution.
+    :type max_value: Optional[Union[int, float]]
+    :param q: Quantization factor.
+    :type q: Optional[int]
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_truncation_selection_policy]
+            :end-before: [END configure_sweep_job_truncation_selection_policy]
+            :language: python
+            :dedent: 8
+            :caption: Configuring QUniform distributions for a hyperparameter sweep on a Command job.
+    """
+
+    def __init__(
+        self,
+        min_value: Optional[Union[int, float]] = None,
+        max_value: Optional[Union[int, float]] = None,
+        q: Optional[int] = None,
+        **kwargs: Any,
+    ) -> None:
+        kwargs.setdefault(TYPE, SearchSpace.QUNIFORM)
+        super().__init__(min_value=min_value, max_value=max_value, **kwargs)
+        self.q = q
+
+    def _to_rest_object(self) -> List:
+        return [self.type, [self.min_value, self.max_value, self.q]]
+
+    @classmethod
+    def _from_rest_object(cls, obj: List) -> "QUniform":
+        return cls(min_value=obj[1][0], max_value=obj[1][1], q=obj[1][2])
+
+
+class QLogUniform(QUniform):
+    """QLogUniform distribution configuration.
+
+    :param min_value: Minimum value of the log of the distribution.
+    :type min_value: Optional[float]
+    :param max_value: Maximum value of the log of the distribution.
+    :type max_value: Optional[float]
+    :param q: Quantization factor.
+    :type q: Optional[int]
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_qloguniform_qnormal]
+            :end-before: [END configure_sweep_job_qloguniform_qnormal]
+            :language: python
+            :dedent: 8
+            :caption: Configuring QLogUniform distributions for a hyperparameter sweep on a Command job.
+    """
+
+    def __init__(
+        self,
+        min_value: Optional[float] = None,
+        max_value: Optional[float] = None,
+        q: Optional[int] = None,
+        **kwargs: Any,
+    ) -> None:
+        kwargs.setdefault(TYPE, SearchSpace.QLOGUNIFORM)
+        super().__init__(min_value=min_value, max_value=max_value, q=q, **kwargs)
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/sweep_job.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/sweep_job.py
new file mode 100644
index 00000000..0a99bb39
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/entities/_job/sweep/sweep_job.py
@@ -0,0 +1,361 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# pylint: disable=protected-access
+
+import logging
+from typing import Any, Dict, NoReturn, Optional, Union
+
+from azure.ai.ml._restclient.v2023_08_01_preview.models import JobBase
+from azure.ai.ml._restclient.v2023_08_01_preview.models import SweepJob as RestSweepJob
+from azure.ai.ml._restclient.v2023_08_01_preview.models import TrialComponent
+from azure.ai.ml._schema._sweep.sweep_job import SweepJobSchema
+from azure.ai.ml._utils.utils import map_single_brackets_and_warn
+from azure.ai.ml.constants import JobType
+from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, TYPE
+from azure.ai.ml.entities._component.command_component import CommandComponent
+from azure.ai.ml.entities._credentials import (
+    AmlTokenConfiguration,
+    ManagedIdentityConfiguration,
+    UserIdentityConfiguration,
+    _BaseJobIdentityConfiguration,
+)
+from azure.ai.ml.entities._inputs_outputs import Input
+from azure.ai.ml.entities._job._input_output_helpers import (
+    from_rest_data_outputs,
+    from_rest_inputs_to_dataset_literal,
+    to_rest_data_outputs,
+    to_rest_dataset_literal_inputs,
+    validate_inputs_for_command,
+    validate_key_contains_allowed_characters,
+)
+from azure.ai.ml.entities._job.command_job import CommandJob
+from azure.ai.ml.entities._job.job import Job
+from azure.ai.ml.entities._job.job_io_mixin import JobIOMixin
+from azure.ai.ml.entities._job.job_resource_configuration import JobResourceConfiguration
+from azure.ai.ml.entities._job.sweep.sampling_algorithm import SamplingAlgorithm
+from azure.ai.ml.entities._system_data import SystemData
+from azure.ai.ml.entities._util import load_from_dict
+from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, JobException
+
+# from ..identity import AmlToken, Identity, ManagedIdentity, UserIdentity
+from ..job_limits import SweepJobLimits
+from ..parameterized_command import ParameterizedCommand
+from ..queue_settings import QueueSettings
+from .early_termination_policy import (
+    BanditPolicy,
+    EarlyTerminationPolicy,
+    MedianStoppingPolicy,
+    TruncationSelectionPolicy,
+)
+from .objective import Objective
+from .parameterized_sweep import ParameterizedSweep
+from .search_space import (
+    Choice,
+    LogNormal,
+    LogUniform,
+    Normal,
+    QLogNormal,
+    QLogUniform,
+    QNormal,
+    QUniform,
+    Randint,
+    SweepDistribution,
+    Uniform,
+)
+
+module_logger = logging.getLogger(__name__)
+
+
+class SweepJob(Job, ParameterizedSweep, JobIOMixin):
+    """Sweep job for hyperparameter tuning.
+
+    .. note::
+        For sweep jobs, inputs, outputs, and parameters are accessible as environment variables using the prefix
+        ``AZUREML_SWEEP_``. For example, if you have a parameter named "learning_rate", you can access it as
+        ``AZUREML_SWEEP_learning_rate``.
+
+    :keyword name: Name of the job.
+    :paramtype name: str
+    :keyword display_name: Display name of the job.
+    :paramtype display_name: str
+    :keyword description: Description of the job.
+    :paramtype description: str
+    :keyword tags: Tag dictionary. Tags can be added, removed, and updated.
+    :paramtype tags: dict[str, str]
+    :keyword properties: The asset property dictionary.
+    :paramtype properties: dict[str, str]
+    :keyword experiment_name:  Name of the experiment the job will be created under. If None is provided,
+        job will be created under experiment 'Default'.
+    :paramtype experiment_name: str
+    :keyword identity: Identity that the training job will use while running on compute.
+    :paramtype identity: Union[
+        ~azure.ai.ml.ManagedIdentityConfiguration,
+        ~azure.ai.ml.AmlTokenConfiguration,
+        ~azure.ai.ml.UserIdentityConfiguration
+
+    ]
+
+    :keyword inputs: Inputs to the command.
+    :paramtype inputs: dict
+    :keyword outputs: Mapping of output data bindings used in the job.
+    :paramtype outputs: dict[str, ~azure.ai.ml.Output]
+    :keyword sampling_algorithm: The hyperparameter sampling algorithm to use over the `search_space`. Defaults to
+        "random".
+
+    :paramtype sampling_algorithm: str
+    :keyword search_space: Dictionary of the hyperparameter search space. The key is the name of the hyperparameter
+        and the value is the parameter expression.
+
+    :paramtype search_space: Dict
+    :keyword objective: Metric to optimize for.
+    :paramtype objective: Objective
+    :keyword compute: The compute target the job runs on.
+    :paramtype compute: str
+    :keyword trial: The job configuration for each trial. Each trial will be provided with a different combination
+        of hyperparameter values that the system samples from the search_space.
+
+    :paramtype trial: Union[
+        ~azure.ai.ml.entities.CommandJob,
+        ~azure.ai.ml.entities.CommandComponent
+
+    ]
+
+    :keyword early_termination: The early termination policy to use. A trial job is canceled
+        when the criteria of the specified policy are met. If omitted, no early termination policy will be applied.
+
+    :paramtype early_termination:  Union[
+        ~azure.mgmt.machinelearningservices.models.BanditPolicy,
+        ~azure.mgmt.machinelearningservices.models.MedianStoppingPolicy,
+        ~azure.mgmt.machinelearningservices.models.TruncationSelectionPolicy
+
+    ]
+
+    :keyword limits: Limits for the sweep job.
+    :paramtype limits: ~azure.ai.ml.entities.SweepJobLimits
+    :keyword queue_settings: Queue settings for the job.
+    :paramtype queue_settings: ~azure.ai.ml.entities.QueueSettings
+    :keyword resources: Compute Resource configuration for the job.
+    :paramtype resources: Optional[Union[~azure.ai.ml.entities.ResourceConfiguration]
+    :keyword kwargs: A dictionary of additional configuration parameters.
+    :paramtype kwargs: dict
+
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/ml_samples_sweep_configurations.py
+            :start-after: [START configure_sweep_job_bayesian_sampling_algorithm]
+            :end-before: [END configure_sweep_job_bayesian_sampling_algorithm]
+            :language: python
+            :dedent: 8
+            :caption: Creating a SweepJob
+    """
+
+    def __init__(
+        self,
+        *,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        tags: Optional[Dict] = None,
+        display_name: Optional[str] = None,
+        experiment_name: Optional[str] = None,
+        identity: Optional[
+            Union[ManagedIdentityConfiguration, AmlTokenConfiguration, UserIdentityConfiguration]
+        ] = None,
+        inputs: Optional[Dict[str, Union[Input, str, bool, int, float]]] = None,
+        outputs: Optional[Dict] = None,
+        compute: Optional[str] = None,
+        limits: Optional[SweepJobLimits] = None,
+        sampling_algorithm: Optional[Union[str, SamplingAlgorithm]] = None,
+        search_space: Optional[
+            Dict[
+                str,
+                Union[
+                    Choice, LogNormal, LogUniform, Normal, QLogNormal, QLogUniform, QNormal, QUniform, Randint, Uniform
+                ],
+            ]
+        ] = None,
+        objective: Optional[Objective] = None,
+        trial: Optional[Union[CommandJob, CommandComponent]] = None,
+        early_termination: Optional[
+            Union[EarlyTerminationPolicy, BanditPolicy, MedianStoppingPolicy, TruncationSelectionPolicy]
+        ] = None,
+        queue_settings: Optional[QueueSettings] = None,
+        resources: Optional[Union[dict, JobResourceConfiguration]] = None,
+        **kwargs: Any,
+    ) -> None:
+        kwargs[TYPE] = JobType.SWEEP
+
+        Job.__init__(
+            self,
+            name=name,
+            description=description,
+            tags=tags,
+            display_name=display_name,
+            experiment_name=experiment_name,
+            compute=compute,
+            **kwargs,
+        )
+        self.inputs = inputs  # type: ignore[assignment]
+        self.outputs = outputs  # type: ignore[assignment]
+        self.trial = trial
+        self.identity = identity
+
+        ParameterizedSweep.__init__(
+            self,
+            limits=limits,
+            sampling_algorithm=sampling_algorithm,
+            objective=objective,
+            early_termination=early_termination,
+            search_space=search_space,
+            queue_settings=queue_settings,
+            resources=resources,
+        )
+
+    def _to_dict(self) -> Dict:
+        res: dict = SweepJobSchema(context={BASE_PATH_CONTEXT_KEY: "./"}).dump(self)
+        return res
+
+    def _to_rest_object(self) -> JobBase:
+        self._override_missing_properties_from_trial()
+        if self.trial is not None:
+            self.trial.command = map_single_brackets_and_warn(self.trial.command)
+
+        if self.search_space is not None:
+            search_space = {param: space._to_rest_object() for (param, space) in self.search_space.items()}
+
+        if self.trial is not None:
+            validate_inputs_for_command(self.trial.command, self.inputs)
+        for key in search_space.keys():  # pylint: disable=possibly-used-before-assignment
+            validate_key_contains_allowed_characters(key)
+
+        if self.trial is not None:
+            trial_component = TrialComponent(
+                code_id=self.trial.code,
+                distribution=(
+                    self.trial.distribution._to_rest_object()
+                    if self.trial.distribution and not isinstance(self.trial.distribution, Dict)
+                    else None
+                ),
+                environment_id=self.trial.environment,
+                command=self.trial.command,
+                environment_variables=self.trial.environment_variables,
+                resources=(
+                    self.trial.resources._to_rest_object()
+                    if self.trial.resources and not isinstance(self.trial.resources, Dict)
+                    else None
+                ),
+            )
+
+        sweep_job = RestSweepJob(
+            display_name=self.display_name,
+            description=self.description,
+            experiment_name=self.experiment_name,
+            search_space=search_space,
+            sampling_algorithm=self._get_rest_sampling_algorithm() if self.sampling_algorithm else None,
+            limits=self.limits._to_rest_object() if self.limits else None,
+            early_termination=(
+                self.early_termination._to_rest_object()
+                if self.early_termination and not isinstance(self.early_termination, str)
+                else None
+            ),
+            properties=self.properties,
+            compute_id=self.compute,
+            objective=self.objective._to_rest_object() if self.objective else None,
+            trial=trial_component,  # pylint: disable=possibly-used-before-assignment
+            tags=self.tags,
+            inputs=to_rest_dataset_literal_inputs(self.inputs, job_type=self.type),
+            outputs=to_rest_data_outputs(self.outputs),
+            identity=self.identity._to_job_rest_object() if self.identity else None,
+            queue_settings=self.queue_settings._to_rest_object() if self.queue_settings else None,
+            resources=(
+                self.resources._to_rest_object() if self.resources and not isinstance(self.resources, dict) else None
+            ),
+        )
+
+        if not sweep_job.resources and sweep_job.trial.resources:
+            sweep_job.resources = sweep_job.trial.resources
+
+        sweep_job_resource = JobBase(properties=sweep_job)
+        sweep_job_resource.name = self.name
+        return sweep_job_resource
+
+    def _to_component(self, context: Optional[Dict] = None, **kwargs: Any) -> NoReturn:
+        msg = "no sweep component entity"
+        raise JobException(
+            message=msg,
+            no_personal_data_message=msg,
+            target=ErrorTarget.SWEEP_JOB,
+            error_category=ErrorCategory.USER_ERROR,
+        )
+
+    @classmethod
+    def _load_from_dict(cls, data: Dict, context: Dict, additional_message: str, **kwargs: Any) -> "SweepJob":
+        loaded_schema = load_from_dict(SweepJobSchema, data, context, additional_message, **kwargs)
+        loaded_schema["trial"] = ParameterizedCommand(**(loaded_schema["trial"]))
+        sweep_job = SweepJob(base_path=context[BASE_PATH_CONTEXT_KEY], **loaded_schema)
+        return sweep_job
+
+    @classmethod
+    def _load_from_rest(cls, obj: JobBase) -> "SweepJob":
+        properties: RestSweepJob = obj.properties
+
+        # Unpack termination schema
+        early_termination = EarlyTerminationPolicy._from_rest_object(properties.early_termination)
+
+        # Unpack sampling algorithm
+        sampling_algorithm = SamplingAlgorithm._from_rest_object(properties.sampling_algorithm)
+
+        trial = ParameterizedCommand._load_from_sweep_job(obj.properties)
+        # Compute also appears in both layers of the yaml, but only one of the REST.
+        # This should be a required field in one place, but cannot be if its optional in two
+
+        _search_space = {}
+        for param, dist in properties.search_space.items():
+            _search_space[param] = SweepDistribution._from_rest_object(dist)
+
+        return SweepJob(
+            name=obj.name,
+            id=obj.id,
+            display_name=properties.display_name,
+            description=properties.description,
+            properties=properties.properties,
+            tags=properties.tags,
+            experiment_name=properties.experiment_name,
+            services=properties.services,
+            status=properties.status,
+            creation_context=SystemData._from_rest_object(obj.system_data) if obj.system_data else None,
+            trial=trial,  # type: ignore[arg-type]
+            compute=properties.compute_id,
+            sampling_algorithm=sampling_algorithm,
+            search_space=_search_space,  # type: ignore[arg-type]
+            limits=SweepJobLimits._from_rest_object(properties.limits),
+            early_termination=early_termination,
+            objective=Objective._from_rest_object(properties.objective) if properties.objective else None,
+            inputs=from_rest_inputs_to_dataset_literal(properties.inputs),
+            outputs=from_rest_data_outputs(properties.outputs),
+            identity=(
+                _BaseJobIdentityConfiguration._from_rest_object(properties.identity) if properties.identity else None
+            ),
+            queue_settings=properties.queue_settings,
+            resources=properties.resources if hasattr(properties, "resources") else None,
+        )
+
+    def _override_missing_properties_from_trial(self) -> None:
+        if not isinstance(self.trial, CommandJob):
+            return
+
+        if not self.compute:
+            self.compute = self.trial.compute
+        if not self.inputs:
+            self.inputs = self.trial.inputs
+        if not self.outputs:
+            self.outputs = self.trial.outputs
+
+        has_trial_limits_timeout = self.trial.limits and self.trial.limits.timeout
+        if has_trial_limits_timeout and not self.limits:
+            time_out = self.trial.limits.timeout if self.trial.limits is not None else None
+            self.limits = SweepJobLimits(trial_timeout=time_out)
+        elif has_trial_limits_timeout and self.limits is not None and not self.limits.trial_timeout:
+            self.limits.trial_timeout = self.trial.limits.timeout if self.trial.limits is not None else None