6  Secure code development for AI Systems

TipLearning outcomes

After completing this chapter, you will:

  • Understand the principles of secure software development, including the Secure Software Development Lifecycle (SDLC) and privacy by design in practice.
  • Explore best practices for maintaining secure code, such as version control, code reviews, and continuous integration.
  • Identify common security vulnerabilities in popular machine learning tools and understand how to ensure reproducibility through practices like containerization.

After covering the fundamental principles of personal data protection, AI, cybersecurity, data preparation, and privacy-preserving machine learning, this chapter focuses on essential practices of secure AI systems implementation with the Secure Software Development Lifecycle (SDLC) and practical applications of privacy by design during code development.

This chapter also covers concepts like version control, continuous integration, and code reviews. While common to software development, these practices are essential here to ensure the transparency and accountability of an AI system trained on personal data.

The chapter further examines the security of common machine learning frameworks and touches on the topic of “reproducibility” to ensure that AI models remain secure and consistent across different platforms, and to ensure the possibility of auditing past versions of an AI system.

6.1 Secure MLOps: experimentation.

If we go back to our secMLOps pipeline figure, we are now going to cover the stages of “experimentation” and “production”, although the production stage only focuses on the training of the final ML model, without the deployment phase.

Figure 6.1: The data engineering stage of the secMLOps workflow Please refer to the text, for a detailed explanation on the figure.

This is the typical scenario of a team of developers working on training an AI model from scratch:

  1. Prepare the code and the related dependencies to perform the actual training.
  2. Test the code with a small dataset by running the “model training” and “model validation” stages while running monitoring tools to assess the performance of the code.
  3. Repeat steps 1 and 2 with larger amounts of data and larger amounts of computing nodes. Depending on the computational resources (local cluster, shared HPC cluster, cloud computing) different strategies might need to be adopted to bring data and code to each computational node
  4. Trained models are exported and versioned along with versions of the code.

6.2 Secure Software Development Lifecycle in AI Systems

Designing AI systems, especially those trained on personal data, demands a comprehensive integration of security throughout the Software Development Lifecycle (SDLC). The Secure Software Development Lifecycle (SSDLC, although it is almost always written SDLC) enhances the traditional SDLC by embedding proactive and reactive security measures at every phase. This section is based on Fujdiak et al. (2019).

Figure 6.2: The secure software development lifecycle Different stages of the AI lifecycle can be mapped into the software development lifecycle. It should also be noted that with sensitive data, often the “proactive approach” stages are happening in secure computing environments, while the stages related to reactive approach are happening in the live/production environment which has a larger attack surface

6.2.1 Proactive vs. Reactive Security Approaches

The SSDLC incorporates both proactive and reactive approaches:

  1. Proactive Security: Focuses on preventing vulnerabilities during early SDLC stages. Measures include security training, threat modeling, and defining security requirements before coding begins. Proactive efforts minimize resource costs by addressing issues before they escalate.

  2. Reactive Security: Addresses security post-development or during deployment through activities like penetration testing, dynamic analysis, and incident response planning. While critical, reactive approaches often add complexity and resource demands if not preceded by proactive measures.

6.2.2 Blind analysis: separating development teams?

When developing AI systems trained on personal data, the SSDLC must account for the risks associated related the sensitive data being used: unauthorized access, data breaches, and non-compliance with GDPR. From the software developer perspective, these are risks related to data access, rather than to software vulnerabilities per se. If access to the personal data is strictly controlled, a possible alternative is for the developers teams to split in two: one team develops the software using only synthetic data that mimics the properties and structure of the real data. A second team – with access to the sensitive data – obtains the software from the first team and run it in secure processing environments. Separating the two development teams further minimises potential risks associated with the sensitive nature of the data. This approach is sometimes called Role-Based Access Control, and mimics what is sometimes done in science to remove potential biases from those who analyse scientific data (MacCoun and Perlmutter 2015).

6.3 Version control, testing, continuous integration, code coverage

Version control, testing, continuous integration, and code coverage are essential practices in modern software development that help teams work efficiently and ensure the quality of their code.

  • Version control is a system that keeps track of changes made to files, allowing developers to collaborate, track history, and revert to previous versions if needed; Git is a popular tool for this.
  • Testing involves running checks on the code to ensure it works as expected and avoids bugs; this can include small tests for specific pieces (unit tests) or tests for the whole system.
  • Continuous integration (CI) is a practice where code changes are automatically tested and integrated into a shared codebase frequently, ensuring new changes don’t break the system.
  • Code coverage measures how much of the code is tested by your tests, showing areas that may need more attention. Together, these practices make software development smoother, more reliable, and less error-prone. In the next chapter we will look at another type of testing, not the code unit tests to ensure the validity of the code, but the actual testing of the AI model and its resilience to potential attacks.

6.4 Other Good Practices for Software Security in AI Systems Development

In general, as in any software development scenario, it is critical to remember that a larger attack surface increases the likelihood of encountering risks due to weak software security. This is especially relevant in the development of AI systems, where supply chain vulnerabilities must be carefully addressed. These vulnerabilities not only encompass risks related to the libraries used for training and deploying AI models but also extend to the reuse of pre-trained machine learning models and serialization libraries.

6.4.1 Libraries for ML Training and How to Check for Vulnerabilities

Some of the most commonly used libraries for machine learning training include TensorFlow, PyTorch, Scikit-learn, and Keras. While these libraries are widely adopted and regularly updated, they may still contain vulnerabilities that could be exploited.

Mitigation strategies:

  1. Check for Known Vulnerabilities:
    Use tools such as Dependency-Check, Safety, or Snyk to scan your ML environment for known security issues in dependencies.

  2. Stay Updated
    Regularly update libraries to the latest stable versions to include security patches. Monitoring GitHub releases or PyPI feeds for your dependencies helps avoid lagging behind on critical fixes.

  3. Verify Source Integrity
    Always obtain libraries and models from official repositories or trusted sources. For extra assurance, verify checksums or digital signatures when provided.

  4. Monitor Security Advisories
    Subscribe to mailing lists, GitHub security advisories, or RSS feeds of the ML libraries you use. Projects like TensorFlow maintain detailed CVEs and release notes that flag patched vulnerabilities.

  5. Monitor Dependencies and Supply Chain ML pipelines often involve dozens of transitive dependencies. Use tools like pipdeptree or pip-audit to map and assess these. Consider pinning dependency versions in requirements.txt or pyproject.toml for reproducibility, but weigh the trade-off against missing out on important security patches.

  6. Inspect Network Traffic Some ML libraries or training pipelines may download additional data (e.g. weights, datasets) automatically or have telemetry enabled. Use tools like netstat, tcpdump, or a firewall to monitor unexpected outbound requests during training and inference phases.

  7. Run Software in Sandboxed or Isolated Environments Use containers (e.g., Docker or Singularity) or virtual environments to isolate the ML environment from your host system. Containers also help manage dependencies securely and ensure reproducibility.

6.4.2 Risks of Reusing Pre-trained Machine Learning Models

Pre-trained models, such as those available in repositories like Hugging Face, TensorFlow Hub, or ONNX Model Zoo, are valuable for saving time and computational resources. However, they introduce significant risks that may compromise the security and integrity of AI systems (for a detailed overview, please see Goldblum et al. (2022) and Wang et al. (2022)):

  1. Backdoors and Clean-Label Poisoning Attackers can embed backdoors in models during training, which are triggered only under specific input conditions. In clean-label attacks, the poisoned data appears legitimate and retains correct labels, making it difficult to detect.
  2. Transfer Learning Exploits Many real-world attacks exploit transfer learning by crafting poisoned data or model components that persist even after fine-tuning. These attacks rely on the reuse of common model architectures and pre-trained feature extractors.
  3. Unknown Provenance and Poisoned Pre-training Data Lack of transparency in the datasets or methods used during pre-training may introduce subtle biases, vulnerabilities, or targeted behaviors that are hard to detect post-hoc.
  4. Tampered Models and Model-Reuse Attacks Downloaded models may be directly tampered with to include malicious logic. Attackers may also poison the training of models expected to be reused by others — anticipating how those models will later be adapted.

Mitigation Strategies:

  • Use pre-trained models from trusted sources with clear documentation and version history.
  • Perform behavioral testing (e.g. outlier detection, gradient inspection) to uncover unexpected behaviors.
  • Fine-tune models with robust methods (e.g. differential privacy, certified defenses) and validate with multiple surrogate datasets.
  • Check digital signatures or hashes when distributing or ingesting model files.

6.5 Risks of Serialization Libraries

Serialization libraries such as Pickle and Joblib in Python are convenient for saving and loading models. However, they carry significant risks, particularly when handling models from external or untrusted sources:

  1. Arbitrary Code Execution: Pickle can execute arbitrary Python code embedded within serialized objects. If a malicious actor crafts the file, simply unpickling it may compromise the system.
  2. Tampering and Hidden Logic: Serialized models may be manipulated to subtly change predictions, trigger specific outputs, or install unauthorized functionality.
  3. Lack of Integrity Checks: Formats like Pickle do not include built-in mechanisms for detecting file corruption or unauthorized modification.

Safer Alternatives and Best Practices:

  • Avoid using insecure serialization formats like Pickle for untrusted or external files.
  • Prefer ONNX, TorchScript, or JSON-based formats with stronger validation mechanisms.
  • Always validate the content and structure of deserialized models before use.
  • Use cryptographic hashes or digital signatures to verify the integrity and authenticity of serialized files.

6.6 Additional Considerations

Beyond these risks, there are also a few potential risks that should also be considered when developing AI models:

  • Federated Learning Vulnerabilities: Poisoned data or model updates in decentralized training can lead to widespread model corruption, especially when adversaries control edge devices.
  • Broader Attack Objectives: Poisoning can be used to subvert fairness, reduce overall accuracy, or specifically degrade performance for certain subgroups or individuals.
  • Advanced Techniques: Methods such as bilevel optimization, feature collision, and influence functions allow for more effective and targeted poisoning, including attacks that remain invisible after standard model auditing.
  • Scalable and Automated Attacks: Generative approaches are being explored to automate poisoning across large datasets or at industrial scale.
NoteVulnerability Story: the case of ShellTorch

TorchServe, an open-source model-serving library for deploying PyTorch models, plays a critical role in AI infrastructure for some of the world’s largest organizations. However, researchers discovered severe vulnerabilities collectively named “ShellTorch,” which exposed thousands of TorchServe servers to Remote Code Execution (RCE) attacks. These vulnerabilities arose from issues like unauthenticated management API access, misconfigured Server-Side Request Forgery (SSRF), insecure deserialization of YAML files via the SnakeYAML library, and the infamous ZipSlip directory traversal flaw. Attackers could exploit these flaws to gain complete control over the server, upload malicious models, extract sensitive data, and even alter AI model behavior. The severity of these vulnerabilities (CVSS scores as high as 9.9) made them a major security concern, especially since some Fortune 500 companies were affected. The vulnerabilities showed how misconfigurations and insecure practices in AI infrastructure can lead to devastating consequences.

To mitigate such risks, the key lessons from ShellTorch are about the importance of secure configurations, proper access controls, and cautious use of third-party libraries. First, users must ensure management APIs are not exposed to external networks without authentication. Second, inputs like YAML files should always be validated using secure parsers like SafeConstructor in SnakeYAML. Third, enabling features like domain allow-lists for model registration can prevent SSRF attacks. Additionally, regularly updating dependencies to patched versions (e.g., TorchServe version 0.8.2 or higher) and using vulnerability scanners are crucial steps. The ShellTorch incident demonstrates the critical need for combining AI innovation with robust security practices to prevent malicious exploitation of AI infrastructure (Kaplan, Katz, and Lumelsky 2024).

6.7 Reproducibility in AI Systems

Reproducibility in AI systems, particularly those trained on personal data, is critical for maintaining transparency, ensuring compliance, and enabling updates to models when data changes. Reproducing the exact code, dependencies, and configurations allows us to recreate AI models faithfully. This capability is fundamental for scenarios such as responding to data subject requests, where specific personal data may need to be removed. With reproducible pipelines, the same model can be retrained on an updated dataset, ensuring that the AI system continues to function as intended without reliance on outdated or removed data.

In Python, tools like conda and pyenv are commonly used to manage dependencies and environments for reproducible code execution. However, containers like Docker and Singularity/Apptainer are increasingly the standard for AI development and deployment, especially in environments like HPC clusters (using Slurm) or cloud-based orchestration systems (e.g., Kubernetes). Containers encapsulate the entire runtime environment, including code, libraries, and dependencies, making it easy to deploy and share reproducible pipelines across diverse systems. While Docker is widely used, it poses risks due to its reliance on root access, which can be problematic in shared computing environments. In such cases, container systems like Singularity or Apptainer are preferred, as they are designed with security in mind, avoiding the root access issue while maintaining portability and ease of use.

6.8 Secure Processing Environments for Training AI Systems with Personal Data

We have covered various aspects of secure and responsible code development in the context of ML and personal data, but secure code is tightly coupled with the computing infrastructure where the code is going to run. Secure processing environments, often referred to as Trusted Research Environments (TREs), are controlled workspaces designed to ensure that the personal data used for training AI systems is processed securely. These environments implement strict access controls, encryption, and monitoring to minimize the risk of data breaches, unauthorized access, or misuse. While TREs are essential for compliance with regulations like GDPR, they also bring significant challenges when handling distributed learning scenarios.

6.8.1 Distributed Learning in Secure Environments

What is distributed learning? In distributed learning (Huyen 2022), training an AI model often requires splitting the workload across multiple computing nodes. This approach is necessary when: 1) data is too large to fit into one node’s memory; 2) the model is too large to fit in a single node; or 3) when adopting pipeline parallelism to optimize computations.

Whether we use data parallelism (the dataset is split across multiple nodes, and each node trains a local copy of the model on its subset of data) or model parallelism (the model is too large to fit into a single computing node, so the model itself is split across multiple nodes) or a mix of both, the training process can become more complex in TREs due to the increased security in the communication between nodes and the potential cost of reserving a large computing cluster isolated from the internet only for the current training process, to avoid sharing the infrastructure with other users like it is common with HPC clusters. The cost and complexity of setting up and maintaining a distributed TRE can be significant, especially for smaller organizations or research teams. The scalability of TREs may also be a challenge for large-scale AI training requiring extensive computational resources.

6.9 Beyond security: how to choose the right ML model?

To choose the best ML model, it is important to start with the simplest model, avoid the state-of-the-art trap, consider human biases, evaluate performance trade-offs, understand model assumptions, and consider if an ensemble model might be appropriate. Chapter 6 from Huyen (2022) is a great starting point to learn more about what to consider when developing an AI model from scratch. In the context of personal data and AI models (see European Data Protection Board (EDPB) (2024)), it is important to consider the guiding principles of the GDPR – especially transparency and accountability – when choosing the best ML architecture. Explainable AI methods might be more favorable than black-box algorithms, but this might result in a compromise when it comes to the performance of the model. Like we saw with anonymisation techniques, it is always a balance between protecting the privacy of the data subjects in the training set, and the performance of the AI model that is developed. For more information on the topic, see Appendix “Machine Learning algorithms and their explainability” which summarises Leslie et al. (2024).

6.10 Beyond the AI model: developing the AI system

While the main goal of MLOps workflow is to train the best performing AI model, it is also fundamental to consider the secure development of the AI system as a whole. If the AI model is like a strong powerful car engine, it won’t be able to do anything unless we build the rest of the car around it. This means that also the inference part of the AI system (how the model is queried and how predictions are made) need to be developed securely, as well as the actual interface part: how the inputs will be gathered and turned into queries for the model, and how outputs will be processed.

When it comes to personal data at the level of the AI system, inputs might bring new personal data that we do not necessarily want to use to query the model. And similarly outputs might contain unwanted personal data. While it is difficult to provide a general solution that works for all AI systems, the reader can already think of a few ways to filter and monitor unwanted inputs/outputs in the context of personal data (including attacks from malicious users, which we will cover next).

While this book is mostly targeted towards developers and deployers of AI models, many SME might rely on third party AI models/systems to form the core of the company managed AI system. With third party AI systems all the careful minimisation techniques and best practices for secure processing are delegated to the third party. Was the training of the AI model performed by the third party lawful according to the GDPR? This scenario is covered in the recent European Data Protection Board (EDPB) (2024), but the guidelines recommend a case-by-case assessment.

6.11 Conclusions and where to learn more about this

This section has presented multiple tools and practices that should be part of the daily skillset of an AI engineer. In the stages of the AI lifecycle (ISO 5338, ISO 22989), development happens at the “Design and Development” stage. As we saw in this chapter, the development stage involves the training data and the validation data to optimise the model hyperparameters. While we described the importance of code testing and code coverage, the actual testing of the AI system happens on the next “Verification & Validation” stage of the AI lifecycle. Many of the tools presented in this chapter would require their own dedicated course. We collect here an overview of existing open-source courses that can help the student to learn deeper some of the tools presented here.

CautionHint for instructors

Depending how the instructor is developing the course, this module is the one that could be expanded with some of the lessons described below. Once the instructor has decided what type of infrastructure they can teach with (are the learners just using their laptops? Do they have access to a slurm or kubernetes cluster? what are the sizes of the data used for training and evaluation?) they can then decide if it’s important to focus on some ML tools like scikit-learn, or MLOps tools like MLflow, or good coding practices beyond ML like git version control and continuous integration. There is a lot to learn/teach to beginners, so we do hope the instructor does not feel discouraged.

Version control with Git

Testing, continuous integration, and code coverage

Machine learning frameworks: sci-kit learn, pytorch, etc

Containers, HPC

Other useful tools for ensuring reproducibility of ML workflows

  • DVC (Data Version Control) https://dcv.org/ is a tool for versioning and managing machine learning projects, enabling efficient tracking of datasets, models, and experiments within version control systems like Git. Start with their DVC tutorial
  • Weights & Biases (W&B) https://wandb.ai/ is a platform for tracking, visualizing, and managing machine learning experiments for model development and collaboration. W&B server can also be run locally, to ensure the confidentiality of the development process. Start with the W&B Tutorial
  • MLflow https://mlflow.org/ is an open-source platform for managing the machine learning lifecycle, including experimentation, reproducibility, deployment, and model registry. MLflow tutorial

There are many more courses from commercial learning platforms using commercial tools provided by cloud infrastructure companies like Amazon, Google, Microsoft, Snowflake. The learner can explore what best fits their learning goals and decide to opt in into a paid course or certification from these providers.

6.12 Summary

In this chapter we briefly covered the most important concepts related to (secure) software development, with focus on good practices in machine learning. While each of the described practices could have a chapter (or a course!) of its own, the goal here was to motivate the learner to understand what might be missing in their skillset and what practices should be adopted. The next module we will focus on testing and validation, two necessary steps before finally deploying the AI model/system.

CautionExercise 6.1: Multiple choice questions
Question Options
1. What is the primary focus of the Secure Software Development Lifecycle (SDLC) in AI systems? 1) To ensure faster code deployment.
2) To embed security measures throughout the development process.
3) To reduce computational costs.
4) To replace traditional SDLC approaches.
2. Which of the following is an example of proactive security in the SDLC? 1) Penetration testing.
2) Threat modeling before coding begins.
3) Incident response planning.
4) Fixing vulnerabilities after deployment.
3. What is a key benefit of splitting development teams when working with sensitive personal data? 1) Faster code iteration.
2) Minimizing risks associated with unauthorized data access.
3) Ensuring compliance with code review practices.
4) Reducing computational resource needs.
4. Which tool is commonly used for version control in software development? 1) TensorFlow.
2) Git.
3) Docker.
4) PyTorch.
5. What is the main purpose of continuous integration (CI)? 1) Track changes to files over time.
2) Automatically test and integrate code changes into a shared codebase.
3) Reduce the attack surface of software.
4) Improve serialization processes.
6. What is a major risk of reusing pre-trained machine learning models? 1) Increased training time.
2) Hidden backdoors or poisoning during inference.
3) Lack of scalability.
4) High memory requirements.
7. How can developers reduce risks associated with serialization libraries like Pickle? 1) Use formats with stricter validation mechanisms like ONNX or JSON.
2) Avoid version control for serialized files.
3) Skip integrity checks for serialized files.
4) Always use older versions of serialization libraries.
8. What was a critical vulnerability in the “ShellTorch” case? 1) Unauthenticated management API access.
2) Issues with federated learning.
3) Lack of synthetic data for testing.
4) Misconfigured continuous integration pipelines.
9. What is the primary advantage of using containers like Singularity over Docker in shared environments? 1) Faster deployment times.
2) Avoiding root access issues.
3) Better GPU support.
4) Automatic code testing.
10. What is the purpose of Trusted Research Environments (TREs)? 1) To enhance code reproducibility.
2) To ensure sensitive data is processed in secure and controlled workspaces.
3) To reduce computational overhead in AI training.
4) To simplify the SDLC process.

Click to reveal solutions

  1. Answer: 2) To embed security measures throughout the development process.

    Explanation: The SDLC integrates security measures proactively and reactively to ensure the secure development of AI systems.

  2. Answer: 2) Threat modeling before coding begins.

    Explanation: Proactive security involves identifying and mitigating potential threats during early stages of development.

  3. Answer: 2) Minimizing risks associated with unauthorized data access.

    Explanation: Splitting development teams reduces risks by controlling access to sensitive data.

  4. Answer: 2) Git.

    Explanation: Git is a widely used tool for tracking changes in code and enabling collaboration.

  5. Answer: 2) Automatically test and integrate code changes into a shared codebase.

    Explanation: Continuous integration helps ensure that changes do not break the shared codebase.

  6. Answer: 2) Hidden backdoors or poisoning during inference.

    Explanation: Pre-trained models can be compromised with hidden vulnerabilities like backdoors or data poisoning.

  7. Answer: 1) Use formats with stricter validation mechanisms like ONNX or JSON.

    Explanation: Using secure serialization formats and adding validation checks reduces risks of malicious payloads.

  8. Answer: 1) Unauthenticated management API access.

    Explanation: ShellTorch vulnerabilities included unauthenticated API access, among others, which allowed remote exploitation.

  9. Answer: 2) Avoiding root access issues.

    Explanation: Singularity avoids root access, making it safer for use in shared environments compared to Docker.

  10. Answer: 2) To ensure sensitive data is processed in secure and controlled workspaces.

    Explanation: TREs are designed to secure sensitive data during processing and comply with privacy regulations.