diff --git a/engine/docker/docker.go b/engine/docker/docker.go
index 5fb48189..8b3c66ba 100644
--- a/engine/docker/docker.go
+++ b/engine/docker/docker.go
@@ -205,7 +205,7 @@ func (e *Docker) destroyContainers(
 			e.softStop(ctx, ctr.ID)
 		} else {
 			if err := e.client.ContainerKill(ctx, ctr.ID, "9"); err != nil {
-				logrus.WithField("container", ctr.ID).WithField("error", err).Warnln("failed to kill container")
+				logrus.WithContext(ctx).WithField("container", ctr.ID).WithField("error", err).Warnln("failed to kill container")
 			}
 		}
 	}
@@ -213,7 +213,7 @@ func (e *Docker) destroyContainers(
 	// cleanup all containers
 	for _, ctr := range containers {
 		if err := e.client.ContainerRemove(ctx, ctr.ID, removeOpts); err != nil {
-			logrus.WithField("container", ctr.ID).WithField("error", err).Warnln("failed to remove container")
+			logrus.WithContext(ctx).WithField("container", ctr.ID).WithField("error", err).Warnln("failed to remove container")
 		}
 	}
 
@@ -228,13 +228,13 @@ func (e *Docker) destroyContainers(
 			continue
 		}
 		if err := e.client.VolumeRemove(ctx, vol.EmptyDir.ID, true); err != nil {
-			logrus.WithField("volume", vol.EmptyDir.ID).WithField("error", err).Warnln("failed to remove volume")
+			logrus.WithContext(ctx).WithField("volume", vol.EmptyDir.ID).WithField("error", err).Warnln("failed to remove volume")
 		}
 	}
 
 	// cleanup the network
 	if err := e.client.NetworkRemove(ctx, pipelineConfig.Network.ID); err != nil {
-		logrus.WithField("network", pipelineConfig.Network.ID).WithField("error", err).Warnln("failed to remove network")
+		logrus.WithContext(ctx).WithField("network", pipelineConfig.Network.ID).WithField("error", err).Warnln("failed to remove network")
 	}
 
 	// notice that we never collect or return any errors.
@@ -359,7 +359,7 @@ func (e *Docker) create(ctx context.Context, pipelineConfig *spec.PipelineConfig
 		step.ID,
 	)
 	if err == nil {
-		logrus.WithField("step", step.Name).WithField("body", containerCreateBody).Infoln("Created container for the step")
+		logrus.WithContext(ctx).WithField("step", step.Name).WithField("body", containerCreateBody).Infoln("Created container for the step")
 	}
 
 	// automatically pull and try to re-create the image if the
@@ -386,7 +386,7 @@ func (e *Docker) create(ctx context.Context, pipelineConfig *spec.PipelineConfig
 			step.ID,
 		)
 		if err == nil {
-			logrus.WithField("step", step.Name).WithField("body", containerCreateBody).Infoln("Created container for the step")
+			logrus.WithContext(ctx).WithField("step", step.Name).WithField("body", containerCreateBody).Infoln("Created container for the step")
 		}
 	}
 	if err != nil {
@@ -513,12 +513,12 @@ func (e *Docker) pullImage(ctx context.Context, image string, pullOpts types.Ima
 
 	if e.hidePull {
 		if _, cerr := io.Copy(io.Discard, rc); cerr != nil {
-			logrus.WithField("error", cerr).Warnln("failed to discard image pull logs")
+			logrus.WithContext(ctx).WithField("error", cerr).Warnln("failed to discard image pull logs")
 			return cerr
 		}
 	} else {
 		if cerr := jsonmessage.Copy(rc, output); cerr != nil {
-			logrus.WithField("error", cerr).Warnln("failed to copy image pull logs to output")
+			logrus.WithContext(ctx).WithField("error", cerr).Warnln("failed to copy image pull logs to output")
 			return cerr
 		}
 	}
@@ -534,7 +534,7 @@ func (e *Docker) pullImageWithRetries(ctx context.Context, image string,
 		if err == nil {
 			return nil
 		}
-		logrus.WithError(err).
+		logrus.WithContext(ctx).WithError(err).
 			WithField("image", image).
 			Warnln("failed to pull image")
 
@@ -548,7 +548,7 @@ func (e *Docker) pullImageWithRetries(ctx context.Context, image string,
 			return err
 		default:
 			if i < imageMaxRetries {
-				logrus.WithField("image", image).Infoln("retrying image pull")
+				logrus.WithContext(ctx).WithField("image", image).Infoln("retrying image pull")
 			}
 		}
 		time.Sleep(time.Millisecond * imageRetrySleepDuration)
@@ -633,7 +633,7 @@ func (e *Docker) setProxyInDockerDaemon(ctx context.Context, pipelineConfig *spe
 func (e *Docker) softStop(ctx context.Context, name string) {
 	timeout := 30 * time.Second
 	if err := e.client.ContainerStop(ctx, name, &timeout); err != nil {
-		logrus.WithField("container", name).WithField("error", err).Warnln("failed to stop the container")
+		logrus.WithContext(ctx).WithField("container", name).WithField("error", err).Warnln("failed to stop the container")
 	}
 
 	// Before removing the container we want to be sure that it's in a healthy state to be removed.
@@ -645,7 +645,7 @@ func (e *Docker) softStop(ctx context.Context, name string) {
 		time.Sleep(1 * time.Second)
 		containerStatus, err := e.client.ContainerInspect(ctx, name)
 		if err != nil {
-			logrus.WithField("container", name).WithField("error", err).Warnln("failed to retrieve container stats")
+			logrus.WithContext(ctx).WithField("container", name).WithField("error", err).Warnln("failed to retrieve container stats")
 			continue
 		}
 		if containerStatus.State.Status == removing || containerStatus.State.Status == running {
diff --git a/go.mod b/go.mod
index 6f1cdfb8..48ac3cf4 100644
--- a/go.mod
+++ b/go.mod
@@ -23,8 +23,8 @@ require (
 	github.com/mattn/go-zglob v0.0.4
 	github.com/mholt/archiver/v3 v3.5.1
 	github.com/pkg/errors v0.9.1
-	github.com/sirupsen/logrus v1.9.0
-	github.com/stretchr/testify v1.8.4
+	github.com/sirupsen/logrus v1.9.3
+	github.com/stretchr/testify v1.9.0
 	github.com/t-tomalak/logrus-easy-formatter v0.0.0-20190827215021-c074f06c5816
 	golang.org/x/sync v0.3.0
 	gopkg.in/alecthomas/kingpin.v2 v2.2.6
@@ -59,6 +59,7 @@ require (
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang/protobuf v1.5.3 // indirect
 	github.com/golang/snappy v0.0.4 // indirect
+	github.com/google/go-cmp v0.6.0 // indirect
 	github.com/google/uuid v1.3.0 // indirect
 	github.com/gorilla/mux v1.8.0 // indirect
 	github.com/hashicorp/errwrap v1.1.0 // indirect
diff --git a/go.sum b/go.sum
index fd0fd5a8..f4f69004 100644
--- a/go.sum
+++ b/go.sum
@@ -70,8 +70,9 @@ github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5a
 github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
-github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
 github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
+github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
 github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/gorilla/mux v1.8.0 h1:i40aqfkR1h2SlN9hojwV5ZA91wcXFOvkdNIeFDP5koI=
@@ -149,8 +150,8 @@ github.com/shoenig/go-m1cpu v0.1.6/go.mod h1:1JJMcUBvfNwpq05QDQVAnx3gUHr9IYF7GNg
 github.com/shoenig/test v0.6.4 h1:kVTaSd7WLz5WZ2IaoM0RSzRsUD+m8wRR+5qvntpn4LU=
 github.com/shoenig/test v0.6.4/go.mod h1:byHiCGXqrVaflBLAMq/srcZIHynQPQgeyvkvXnjqq0k=
 github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6MwdIuYE2rE=
-github.com/sirupsen/logrus v1.9.0 h1:trlNQbNUG3OdDrDil03MCb1H2o9nJ1x4/5LYw7byDE0=
-github.com/sirupsen/logrus v1.9.0/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
+github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
+github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
@@ -162,8 +163,8 @@ github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 github.com/stretchr/testify v1.7.5/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.3/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
-github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
-github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/t-tomalak/logrus-easy-formatter v0.0.0-20190827215021-c074f06c5816 h1:J6v8awz+me+xeb/cUTotKgceAYouhIB3pjzgRd6IlGk=
 github.com/t-tomalak/logrus-easy-formatter v0.0.0-20190827215021-c074f06c5816/go.mod h1:tzym/CEb5jnFI+Q0k4Qq3+LvRF4gO3E2pxS8fHP8jcA=
 github.com/tklauser/go-sysconf v0.3.11 h1:89WgdJhk5SNwJfu+GKyYveZ4IaJ7xAkecBo+KdJV0CM=
diff --git a/pipeline/runtime/run.go b/pipeline/runtime/run.go
index 10926690..26c53b58 100644
--- a/pipeline/runtime/run.go
+++ b/pipeline/runtime/run.go
@@ -99,7 +99,7 @@ func executeRunStep(ctx context.Context, f RunFunc, r *api.StartStepRequest, out
 
 	reportStart := time.Now()
 	if rerr := report.ParseAndUploadTests(ctx, r.TestReport, r.WorkingDir, step.Name, log, reportStart, tiConfig, r.Envs); rerr != nil {
-		logrus.WithError(rerr).WithField("step", step.Name).Errorln("failed to upload report")
+		logrus.WithContext(ctx).WithError(rerr).WithField("step", step.Name).Errorln("failed to upload report")
 		log.Errorf("Failed to upload report. Time taken: %s", time.Since(reportStart))
 	}
 
diff --git a/pipeline/runtime/step_executor.go b/pipeline/runtime/step_executor.go
index 22d8cfae..beba4cbe 100644
--- a/pipeline/runtime/step_executor.go
+++ b/pipeline/runtime/step_executor.go
@@ -88,7 +88,7 @@ func (e *StepExecutor) StartStep(ctx context.Context, r *api.StartStepRequest) e
 
 	go func() {
 		wr := getLogStreamWriter(r)
-		state, outputs, envs, artifact, outputV2, optimizationState, stepErr := e.executeStep(r, wr)
+		state, outputs, envs, artifact, outputV2, optimizationState, stepErr := e.executeStep(ctx, r, wr)
 		status := StepStatus{Status: Complete, State: state, StepErr: stepErr, Outputs: outputs, Envs: envs,
 			Artifact: artifact, OutputV2: outputV2, OptimizationState: optimizationState}
 		e.mu.Lock()
@@ -118,7 +118,7 @@ func (e *StepExecutor) StartStepWithStatusUpdate(ctx context.Context, r *api.Sta
 				setPrevStepExportEnvs(r)
 			}
 			wr = getLogStreamWriter(r)
-			state, outputs, envs, artifact, outputV2, optimizationState, stepErr := e.executeStep(r, wr)
+			state, outputs, envs, artifact, outputV2, optimizationState, stepErr := e.executeStep(ctx, r, wr)
 			status := StepStatus{Status: Complete, State: state, StepErr: stepErr, Outputs: outputs, Envs: envs,
 				Artifact: artifact, OutputV2: outputV2, OptimizationState: optimizationState}
 			pollResponse := convertStatus(status)
@@ -244,7 +244,7 @@ func (e *StepExecutor) executeStepDrone(r *api.StartStepRequest) (*runtime.State
 
 	stepLog := NewStepLog(ctx) // step output will terminate when the ctx is canceled
 
-	logr := logrus.
+	logr := logrus.WithContext(ctx).
 		WithField("id", r.ID).
 		WithField("step", r.Name)
 
@@ -288,7 +288,7 @@ func (e *StepExecutor) executeStepDrone(r *api.StartStepRequest) (*runtime.State
 	return runStep()
 }
 
-func (e *StepExecutor) executeStep(r *api.StartStepRequest, wr logstream.Writer) (*runtime.State, map[string]string, //nolint:gocritic
+func (e *StepExecutor) executeStep(ctx context.Context, r *api.StartStepRequest, wr logstream.Writer) (*runtime.State, map[string]string, //nolint:gocritic
 	map[string]string, []byte, []*api.OutputV2, string, error) {
 	if r.LogDrone {
 		state, err := e.executeStepDrone(r)
@@ -302,13 +302,14 @@ func (e *StepExecutor) executeStep(r *api.StartStepRequest, wr logstream.Writer)
 	} else {
 		tiConfig = pipeline.GetState().GetTIConfig()
 	}
-	return executeStepHelper(r, e.engine.Run, wr, tiConfig)
+	return executeStepHelper(ctx, r, e.engine.Run, wr, tiConfig)
 }
 
 // executeStepHelper is a helper function which is used both by this step executor as well as the
 // stateless step executor. This is done so as to not duplicate logic across multiple implementations.
 // Eventually, we should deprecate this step executor in favor of the stateless executor.
 func executeStepHelper( //nolint:gocritic
+	ctx context.Context,
 	r *api.StartStepRequest,
 	f RunFunc,
 	wr logstream.Writer,
@@ -319,7 +320,6 @@ func executeStepHelper( //nolint:gocritic
 	// We do here only for non-container step.
 	if r.Detach && r.Image == "" {
 		go func() {
-			ctx := context.Background()
 			var cancel context.CancelFunc
 			if r.Timeout > 0 {
 				ctx, cancel = context.WithTimeout(ctx, time.Second*time.Duration(r.Timeout))
@@ -333,7 +333,6 @@ func executeStepHelper( //nolint:gocritic
 
 	var result error
 
-	ctx := context.Background()
 	var cancel context.CancelFunc
 	if r.Timeout > 0 {
 		ctx, cancel = context.WithTimeout(ctx, time.Second*time.Duration(r.Timeout))
@@ -370,9 +369,9 @@ func executeStepHelper( //nolint:gocritic
 		}
 
 		if exited.OOMKilled {
-			logrus.WithField("id", r.ID).Infoln("received oom kill.")
+			logrus.WithContext(ctx).WithField("id", r.ID).Infoln("received oom kill.")
 		} else {
-			logrus.WithField("id", r.ID).Infof("received exit code %d\n", exited.ExitCode)
+			logrus.WithContext(ctx).WithField("id", r.ID).Infof("received exit code %d\n", exited.ExitCode)
 		}
 	}
 	return exited, outputs, envs, artifact, outputV2, optimizationState, result
diff --git a/pipeline/runtime/step_executor_stateless.go b/pipeline/runtime/step_executor_stateless.go
index 407eb5c6..cbb18d1c 100644
--- a/pipeline/runtime/step_executor_stateless.go
+++ b/pipeline/runtime/step_executor_stateless.go
@@ -45,7 +45,7 @@ func (e *StepExecutorStateless) Run(
 
 	e.stepStatus = StepStatus{Status: Running}
 
-	state, outputs, envs, artifact, outputV2, optimizationState, stepErr := e.executeStep(r, cfg, writer)
+	state, outputs, envs, artifact, outputV2, optimizationState, stepErr := e.executeStep(ctx, r, cfg, writer)
 	e.stepStatus = StepStatus{Status: Complete, State: state, StepErr: stepErr, Outputs: outputs, Envs: envs,
 		Artifact: artifact, OutputV2: outputV2, OptimizationState: optimizationState}
 	pollResponse := convertStatus(e.stepStatus)
@@ -53,6 +53,7 @@ func (e *StepExecutorStateless) Run(
 }
 
 func (e *StepExecutorStateless) executeStep( //nolint:gocritic
+	ctx context.Context,
 	r *api.StartStepRequest,
 	cfg *spec.PipelineConfig,
 	writer logstream.Writer,
@@ -64,7 +65,7 @@ func (e *StepExecutorStateless) executeStep( //nolint:gocritic
 	// Temporary: this should be removed once we have a better way of handling test intelligence.
 	tiConfig := getTiCfg(&r.TIConfig)
 
-	return executeStepHelper(r, runFunc, writer, &tiConfig)
+	return executeStepHelper(ctx, r, runFunc, writer, &tiConfig)
 }
 
 func getTiCfg(t *api.TIConfig) tiCfg.Cfg {